diff --git a/dimos/control/tasks/trajectory_task/trajectory_task.py b/dimos/control/tasks/trajectory_task/trajectory_task.py index a3eb23dec3..e3deca4f11 100644 --- a/dimos/control/tasks/trajectory_task/trajectory_task.py +++ b/dimos/control/tasks/trajectory_task/trajectory_task.py @@ -23,6 +23,7 @@ from __future__ import annotations from dataclasses import dataclass +import math from typing import Any from dimos.control.task import ( @@ -127,7 +128,7 @@ def compute(self, state: CoordinatorState) -> JointCommandOutput | None: Returns: JointCommandOutput with positions, or None if not executing """ - if self._trajectory is None: + if self._trajectory is None or not self._trajectory.joint_names: return None # Set start time on first compute() for consistent timing @@ -143,8 +144,10 @@ def compute(self, state: CoordinatorState) -> JointCommandOutput | None: logger.info(f"Trajectory {self._name} completed after {t_elapsed:.3f}s") # Return final position to hold at goal q_ref, _ = self._trajectory.sample(self._trajectory.duration) + final_names = list(self._trajectory.joint_names) + self._clear_active_trajectory() return JointCommandOutput( - joint_names=self._joint_names_list, + joint_names=final_names, positions=list(q_ref), mode=ControlMode.SERVO_POSITION, ) @@ -153,7 +156,7 @@ def compute(self, state: CoordinatorState) -> JointCommandOutput | None: q_ref, _ = self._trajectory.sample(t_elapsed) return JointCommandOutput( - joint_names=self._joint_names_list, + joint_names=list(self._trajectory.joint_names), positions=list(q_ref), mode=ControlMode.SERVO_POSITION, ) @@ -169,6 +172,60 @@ def on_preempted(self, by_task: str, joints: frozenset[str]) -> None: # Abort if any of our joints were preempted if joints & self._joint_names: self._state = TrajectoryState.ABORTED + self._clear_active_trajectory() + + def _clear_active_trajectory(self) -> None: + """Clear stored trajectory-specific execution state.""" + self._trajectory = None + self._pending_start = False + self._start_time = 0.0 + + def _validate_trajectory(self, trajectory: JointTrajectory) -> bool: + """Validate a trajectory before execution.""" + joint_names = list(trajectory.joint_names) + if not joint_names: + logger.warning("Trajectory for %s has empty joint names", self._name) + return False + if len(set(joint_names)) != len(joint_names): + logger.warning("Trajectory for %s has duplicate joint names", self._name) + return False + unknown = [name for name in joint_names if name not in self._joint_names] + if unknown: + logger.warning("Trajectory for %s has unknown joints: %s", self._name, unknown) + return False + if not trajectory.points: + logger.warning("Empty trajectory for %s", self._name) + return False + width = len(joint_names) + previous_time: float | None = None + for index, point in enumerate(trajectory.points): + if len(point.positions) != width or len(point.velocities) != width: + logger.warning("Trajectory point %d for %s has invalid width", index, self._name) + return False + if not all(math.isfinite(value) for value in point.positions): + logger.warning( + "Trajectory point %d for %s has non-finite positions", index, self._name + ) + return False + if not all(math.isfinite(value) for value in point.velocities): + logger.warning( + "Trajectory point %d for %s has non-finite velocities", index, self._name + ) + return False + if not math.isfinite(point.time_from_start): + logger.warning("Trajectory point %d for %s has non-finite time", index, self._name) + return False + if index == 0 and point.time_from_start != 0.0: + logger.warning("Trajectory for %s must start at t=0", self._name) + return False + if previous_time is not None and point.time_from_start <= previous_time: + logger.warning("Trajectory for %s has non-increasing timestamps", self._name) + return False + previous_time = point.time_from_start + if trajectory.duration <= 0.0: + logger.warning("Trajectory for %s has nonpositive duration", self._name) + return False + return True def execute(self, trajectory: JointTrajectory) -> bool: """Start executing a trajectory. @@ -183,17 +240,17 @@ def execute(self, trajectory: JointTrajectory) -> bool: logger.warning(f"Cannot execute: {self._name} in FAULT state") return False - if trajectory is None or trajectory.duration <= 0: + if trajectory is None: logger.warning(f"Invalid trajectory for {self._name}") return False - if not trajectory.points: - logger.warning(f"Empty trajectory for {self._name}") + if not self._validate_trajectory(trajectory): return False # Preempt any active trajectory if self._state == TrajectoryState.EXECUTING: logger.info(f"Preempting active trajectory on {self._name}") + self._clear_active_trajectory() self._trajectory = trajectory self._pending_start = True # Start time set on first compute() @@ -214,6 +271,7 @@ def cancel(self) -> bool: if self._state != TrajectoryState.EXECUTING: return False self._state = TrajectoryState.ABORTED + self._clear_active_trajectory() logger.info(f"Trajectory {self._name} cancelled") return True @@ -227,7 +285,7 @@ def reset(self) -> bool: logger.warning(f"Cannot reset {self._name} while executing") return False self._state = TrajectoryState.IDLE - self._trajectory = None + self._clear_active_trajectory() logger.info(f"Trajectory {self._name} reset to IDLE") return True diff --git a/dimos/control/test_control.py b/dimos/control/test_control.py index d6f42ec3f1..29ca63ba8e 100644 --- a/dimos/control/test_control.py +++ b/dimos/control/test_control.py @@ -428,6 +428,141 @@ def test_execute_trajectory(self, trajectory_task, simple_trajectory): assert trajectory_task.is_active() assert trajectory_task.get_state() == TrajectoryState.EXECUTING + def test_execute_partial_subset_and_claims_full_configuration(self, trajectory_task): + trajectory = JointTrajectory( + joint_names=["arm/joint2", "arm/joint3"], + points=[ + TrajectoryPoint(positions=[0.0, 0.0], velocities=[0.0, 0.0], time_from_start=0.0), + TrajectoryPoint(positions=[0.5, 1.0], velocities=[0.0, 0.0], time_from_start=1.0), + ], + ) + + assert trajectory_task.execute(trajectory) is True + assert trajectory_task.claim().joints == frozenset( + {"arm/joint1", "arm/joint2", "arm/joint3"} + ) + + @pytest.mark.parametrize( + "trajectory", + [ + JointTrajectory( + joint_names=[], + points=[TrajectoryPoint(time_from_start=0.0, positions=[], velocities=[])], + ), + JointTrajectory( + joint_names=["arm/joint1", "arm/joint1"], + points=[ + TrajectoryPoint( + time_from_start=0.0, positions=[0.0, 0.0], velocities=[0.0, 0.0] + ) + ], + ), + JointTrajectory( + joint_names=["arm/missing"], + points=[TrajectoryPoint(time_from_start=0.0, positions=[0.0], velocities=[0.0])], + ), + JointTrajectory(joint_names=["arm/joint1"], points=[]), + JointTrajectory( + joint_names=["arm/joint1"], + points=[TrajectoryPoint(time_from_start=0.0, positions=[], velocities=[0.0])], + ), + JointTrajectory( + joint_names=["arm/joint1"], + points=[TrajectoryPoint(time_from_start=0.0, positions=[0.0], velocities=[0.0])], + ), + JointTrajectory( + joint_names=["arm/joint1", "arm/joint2", "arm/joint3"], + points=[ + TrajectoryPoint( + time_from_start=0.0, + positions=[0.0, 0.0, 0.0], + velocities=[0.0, 0.0, 0.0], + ) + ], + ), + JointTrajectory( + joint_names=["arm/joint1"], + points=[ + TrajectoryPoint(time_from_start=0.0, positions=[float("nan")], velocities=[0.0]) + ], + ), + JointTrajectory( + joint_names=["arm/joint1"], + points=[TrajectoryPoint(time_from_start=0.1, positions=[0.0], velocities=[0.0])], + ), + JointTrajectory( + joint_names=["arm/joint1"], + points=[ + TrajectoryPoint(time_from_start=0.0, positions=[0.0], velocities=[0.0]), + TrajectoryPoint(time_from_start=0.0, positions=[1.0], velocities=[0.0]), + ], + ), + ], + ) + def test_invalid_partial_inputs_reject_before_state_changes(self, trajectory_task, trajectory): + assert trajectory_task.get_state() == TrajectoryState.IDLE + assert trajectory_task.execute(trajectory) is False + assert trajectory_task.get_state() == TrajectoryState.IDLE + assert ( + trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=0.0, dt=0.01)) + is None + ) + + def test_compute_emits_active_subset_only_and_clears_on_completion(self, trajectory_task): + trajectory = JointTrajectory( + joint_names=["arm/joint2"], + points=[ + TrajectoryPoint(positions=[0.0], velocities=[0.0], time_from_start=0.0), + TrajectoryPoint(positions=[1.0], velocities=[0.0], time_from_start=1.0), + ], + ) + assert trajectory_task.execute(trajectory) is True + trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=10.0, dt=0.01)) + output = trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=10.5, dt=0.01)) + assert output is not None + assert output.joint_names == ["arm/joint2"] + assert output.positions == [pytest.approx(0.5)] + + final = trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=11.5, dt=0.01)) + assert final is not None + assert final.joint_names == ["arm/joint2"] + assert trajectory_task.get_state() == TrajectoryState.COMPLETED + assert ( + trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=12.0, dt=0.01)) + is None + ) + + def test_replacement_reset_and_cancel_clear_active_subset(self, trajectory_task): + first = JointTrajectory( + joint_names=["arm/joint1"], + points=[ + TrajectoryPoint(positions=[0.0], velocities=[0.0], time_from_start=0.0), + TrajectoryPoint(positions=[1.0], velocities=[0.0], time_from_start=1.0), + ], + ) + second = JointTrajectory( + joint_names=["arm/joint3"], + points=[ + TrajectoryPoint(positions=[2.0], velocities=[0.0], time_from_start=0.0), + TrajectoryPoint(positions=[3.0], velocities=[0.0], time_from_start=1.0), + ], + ) + assert trajectory_task.execute(first) is True + assert trajectory_task.execute(second) is True + trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=1.0, dt=0.01)) + output = trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=1.5, dt=0.01)) + assert output is not None + assert output.joint_names == ["arm/joint3"] + assert trajectory_task.cancel() is True + assert ( + trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=2.0, dt=0.01)) + is None + ) + assert trajectory_task.reset() is True + assert trajectory_task.claim().joints == frozenset( + {"arm/joint1", "arm/joint2", "arm/joint3"} + ) + def test_compute_during_trajectory(self, trajectory_task, simple_trajectory, coordinator_state): t_start = time.perf_counter() trajectory_task.execute(simple_trajectory) diff --git a/dimos/e2e_tests/test_manipulation_planning_groups.py b/dimos/e2e_tests/test_manipulation_planning_groups.py new file mode 100644 index 0000000000..67cf767177 --- /dev/null +++ b/dimos/e2e_tests/test_manipulation_planning_groups.py @@ -0,0 +1,216 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Large E2E tests for manipulation planning groups with a coordinator. + +These tests launch a real ManipulationModule + ControlCoordinator blueprint and +exercise the public planning RPCs over LCM, matching the self-hosted large-test +style used by the navigation stack. +""" + +from __future__ import annotations + +from collections.abc import Callable +import time +from typing import Any + +import pytest + +from dimos.control.coordinator import ControlCoordinator +from dimos.core.rpc_client import RPCClient +from dimos.e2e_tests.dimos_cli_call import DimosCliCall +from dimos.e2e_tests.lcm_spy import LcmSpy +from dimos.manipulation.manipulation_module import ManipulationModule +from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState + +pytestmark = [pytest.mark.self_hosted_large] + +JOINT_STATE_TOPIC = "/coordinator_joint_state#sensor_msgs.JointState" +BLUEPRINT = "openarm-mock-planner-coordinator" + + +def _wait_for_robot_info( + client: RPCClient, + robot_name: str, + *, + timeout: float = 120.0, +) -> dict[str, Any]: + deadline = time.time() + timeout + last_error: BaseException | None = None + while time.time() < deadline: + try: + info = client.get_robot_info(robot_name) + if info and info.get("planning_groups"): + return info + except BaseException as exc: + last_error = exc + time.sleep(0.5) + raise TimeoutError(f"Timed out waiting for {robot_name!r} robot info") from last_error + + +def _wait_for_trajectory_completion( + client: RPCClient, + robot_name: str, + *, + timeout: float = 10.0, +) -> None: + deadline = time.time() + timeout + last_status: dict[str, Any] | None = None + while time.time() < deadline: + last_status = client.get_trajectory_status(robot_name) + if last_status is not None and last_status.get("state") == TrajectoryState.COMPLETED: + return + time.sleep(0.1) + raise TimeoutError(f"{robot_name!r} trajectory did not complete; last={last_status}") + + +def _wait_for_manipulation_state( + client: RPCClient, + state_name: str, + *, + timeout: float = 10.0, +) -> None: + deadline = time.time() + timeout + last_state: str | None = None + while time.time() < deadline: + last_state = client.get_state() + if last_state == state_name: + return + time.sleep(0.1) + raise TimeoutError(f"ManipulationModule did not reach {state_name}; last={last_state}") + + +def _wait_for_current_joints( + client: RPCClient, + robot_names: tuple[str, ...], + *, + timeout: float = 10.0, +) -> None: + deadline = time.time() + timeout + missing = robot_names + while time.time() < deadline: + missing = tuple( + robot_name + for robot_name in robot_names + if client.get_current_joints(robot_name) is None + ) + if not missing: + return + time.sleep(0.1) + raise TimeoutError(f"Timed out waiting for current joints from {missing}") + + +def _prepare_for_planning(client: RPCClient, robot_names: tuple[str, ...]) -> None: + client.reset() + _wait_for_manipulation_state(client, "IDLE") + _wait_for_current_joints(client, robot_names) + # Robot info and joint-state topics can become available just before the + # manipulation module finishes finalizing world monitors. Require a stable + # ready state after joint state is flowing to avoid command-readiness flakes. + time.sleep(0.25) + _wait_for_manipulation_state(client, "IDLE") + + +def _planning_group_id(info: dict[str, Any]) -> str: + groups = info["planning_groups"] + assert len(groups) == 1 + group = groups[0] + if isinstance(group, PlanningGroup): + return group.id + group_id = group["id"] + assert isinstance(group_id, str) + return group_id + + +def _offset_target(client: RPCClient, robot_name: str, delta: float) -> JointState: + current = client.get_current_joints(robot_name) + assert current is not None + return JointState(position=[position + delta for position in current]) + + +def _start_openarm_mock_planner( + start_blueprint: Callable[..., DimosCliCall], lcm_spy: LcmSpy +) -> None: + lcm_spy.save_topic(JOINT_STATE_TOPIC) + start_blueprint(BLUEPRINT) + lcm_spy.wait_for_saved_topic(JOINT_STATE_TOPIC, timeout=120.0) + + +def test_single_arm_plans_and_executes_through_control_coordinator( + lcm_spy: LcmSpy, + start_blueprint: Callable[..., DimosCliCall], +) -> None: + """Plan with one arm and execute through its trajectory task.""" + _start_openarm_mock_planner(start_blueprint, lcm_spy) + + client = RPCClient(None, ManipulationModule) + coordinator_client = RPCClient(None, ControlCoordinator) + try: + left_info = _wait_for_robot_info(client, "left_arm") + left_id = _planning_group_id(left_info) + + tasks = coordinator_client.list_tasks() + assert left_info["coordinator_task_name"] in tasks + + _prepare_for_planning(client, ("left_arm",)) + + planned = client.plan_to_joint_targets({left_id: _offset_target(client, "left_arm", 0.02)}) + assert planned, client.get_error() + assert client.has_planned_path() + assert client.execute_plan() + + _wait_for_trajectory_completion(client, "left_arm") + finally: + coordinator_client.stop_rpc_client() + client.stop_rpc_client() + + +def test_dual_arm_plans_and_dispatches_both_arms_through_control_coordinator( + lcm_spy: LcmSpy, + start_blueprint: Callable[..., DimosCliCall], +) -> None: + """Plan one generated plan over both arms and dispatch both JTC tasks.""" + _start_openarm_mock_planner(start_blueprint, lcm_spy) + + client = RPCClient(None, ManipulationModule) + coordinator_client = RPCClient(None, ControlCoordinator) + try: + left_info = _wait_for_robot_info(client, "left_arm") + right_info = _wait_for_robot_info(client, "right_arm") + left_id = _planning_group_id(left_info) + right_id = _planning_group_id(right_info) + + tasks = coordinator_client.list_tasks() + assert left_info["coordinator_task_name"] in tasks + assert right_info["coordinator_task_name"] in tasks + + _prepare_for_planning(client, ("left_arm", "right_arm")) + + planned = client.plan_to_joint_targets( + { + left_id: _offset_target(client, "left_arm", 0.02), + right_id: _offset_target(client, "right_arm", -0.02), + } + ) + assert planned, client.get_error() + assert client.has_planned_path() + assert client.execute_plan() + + _wait_for_trajectory_completion(client, "left_arm") + _wait_for_trajectory_completion(client, "right_arm") + finally: + coordinator_client.stop_rpc_client() + client.stop_rpc_client() diff --git a/dimos/manipulation/_test_manipulation_helpers.py b/dimos/manipulation/_test_manipulation_helpers.py new file mode 100644 index 0000000000..006fa3bf0b --- /dev/null +++ b/dimos/manipulation/_test_manipulation_helpers.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. + +"""Shared lightweight test harnesses for manipulation module tests.""" + +from unittest.mock import MagicMock + +from dimos.manipulation.execution_runtime import ( + ExecutionRuntime, + Outcome, + prepare_generated_plan, +) +from dimos.manipulation.manipulation_module import ManipulationModule +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.models import GeneratedPlan + +_TEST_RUNTIMES: list[ExecutionRuntime] = [] + + +class FakeCoordinatorGateway: + """Deterministic coordinator gateway used by module-side tests.""" + + def __init__( + self, + execute_outcome: Outcome = Outcome.ACCEPTED, + status_outcome: Outcome = Outcome.INACTIVE, + status_outcomes: list[Outcome] | None = None, + cancel_outcome: Outcome = Outcome.CANCELLED, + ) -> None: + self.execute_outcome = execute_outcome + self.status_outcome = status_outcome + self.status_outcomes = status_outcomes or [] + self.cancel_outcome = cancel_outcome + self.execute_calls: list[tuple[str, object]] = [] + self.cancel_calls: list[str] = [] + self.stopped = False + + def execute(self, task_name: str, request: object) -> Outcome: + self.execute_calls.append((task_name, request)) + return self.execute_outcome + + def cancel(self, task_name: str) -> Outcome: + self.cancel_calls.append(task_name) + return self.cancel_outcome + + def status(self, task_name: str) -> Outcome: + del task_name + if self.status_outcomes: + return self.status_outcomes.pop(0) + return self.status_outcome + + def reset(self, task_name: str) -> Outcome: + del task_name + return Outcome.INACTIVE + + def set_gripper_position(self, hardware_id: str, position: float) -> Outcome: + del hardware_id, position + return Outcome.ACCEPTED + + def get_gripper_position(self, hardware_id: str) -> float | None: + del hardware_id + return 0.0 + + def stop(self) -> None: + self.stopped = True + + +class ManipulationModuleHarness(ManipulationModule): + """Manipulation module initialized only with state needed by unit tests.""" + + def __init__(self) -> None: + self._robots = {} + self._world_monitor = None + self._planner = None + self._kinematics = None + self._execution_runtime = None + self._execution_topology = None + self._init_joints = {} + self.config = MagicMock(planning_timeout=10.0, physical_operation_timeout=60.0) + + +def make_module() -> ManipulationModule: + """Create a lightweight ManipulationModule harness for behavior tests.""" + return ManipulationModuleHarness() + + +def install_runtime( + module: ManipulationModule, + robots: list[RobotModelConfig], + gateway: FakeCoordinatorGateway | None = None, +) -> FakeCoordinatorGateway: + """Install the production runtime with a test-owned fake gateway.""" + gateway = gateway or FakeCoordinatorGateway() + from dimos.manipulation.execution_runtime import ExecutionTopology + + module._execution_topology = ExecutionTopology.from_robot_configs(robots) + module._execution_runtime = ExecutionRuntime( + lambda: gateway, + topology=module._execution_topology, + action_timeout=module.config.planning_timeout, + physical_operation_timeout=module.config.physical_operation_timeout, + poll_interval=0.01, + ) + _TEST_RUNTIMES.append(module._execution_runtime) + return gateway + + +def install_ready_plan(module: ManipulationModule, plan: GeneratedPlan) -> None: + """Publish a prepared plan through the production runtime API.""" + assert module._execution_runtime is not None + assert module._execution_topology is not None + token = module._execution_runtime.start_planning() + assert isinstance(token, str) + prepared = prepare_generated_plan(plan, module._execution_topology) + result = module._execution_runtime.complete_planning(token, prepared) + assert result.accepted + + +def close_test_runtimes() -> None: + """Close runtimes created by the lightweight test harness.""" + while _TEST_RUNTIMES: + runtime = _TEST_RUNTIMES.pop() + operation = runtime.snapshot().operation + if operation is not None: + runtime.wait_for_terminal(operation.handle, timeout=10.0) + runtime.shutdown(timeout=10.0) diff --git a/dimos/manipulation/control/coordinator_client.py b/dimos/manipulation/control/coordinator_client.py index ed552e0846..6139b7b67c 100644 --- a/dimos/manipulation/control/coordinator_client.py +++ b/dimos/manipulation/control/coordinator_client.py @@ -42,10 +42,12 @@ from __future__ import annotations +import importlib +import importlib.util import math import sys import time -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Protocol, cast from dimos.control.components import split_joint_name from dimos.control.coordinator import ControlCoordinator @@ -58,6 +60,12 @@ from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +class _IPythonModule(Protocol): + """Typed subset of the optional IPython module used by this client.""" + + def start_ipython(self, *, argv: list[str], user_ns: dict[str, Any]) -> None: ... + + class CoordinatorClient: """ RPC client for the ControlCoordinator. @@ -560,7 +568,9 @@ def accel(self, value: float | None = None) -> None: def interactive_mode(client: CoordinatorClient, initial_task: str) -> None: """Start IPython interactive mode.""" - import IPython + if importlib.util.find_spec("IPython") is None: + raise RuntimeError("IPython is required for interactive mode") + ipython = cast("_IPythonModule", importlib.import_module("IPython")) shell = CoordinatorShell(client, initial_task) @@ -571,7 +581,7 @@ def interactive_mode(client: CoordinatorClient, initial_task: str) -> None: print("\nType help() for available commands") print("=" * 60 + "\n") - IPython.start_ipython( # type: ignore[no-untyped-call] + ipython.start_ipython( argv=[], user_ns={ "help": shell.help, diff --git a/dimos/manipulation/execution_auxiliary.py b/dimos/manipulation/execution_auxiliary.py new file mode 100644 index 0000000000..650e0e5f71 --- /dev/null +++ b/dimos/manipulation/execution_auxiliary.py @@ -0,0 +1,70 @@ +"""Passive auxiliary-call bookkeeping for the execution runtime.""" + +from dataclasses import dataclass +from typing import Any + +from dimos.manipulation.execution_effects import AuxiliaryDone + + +@dataclass(frozen=True) +class AuxiliaryTicket: + action_id: str + setter: bool + + +class AuxiliaryCallBook: + """Store auxiliary correlations and results without lifecycle decisions.""" + + def __init__(self) -> None: + self._pending: dict[str, float] = {} + self._inflight: set[str] = set() + self._results: dict[str, tuple[bool, Any, str]] = {} + + def register(self, action_id: str, deadline: float, setter: bool) -> AuxiliaryTicket: + self._inflight.add(action_id) + self._pending[action_id] = deadline + return AuxiliaryTicket(action_id, setter) + + def complete(self, done: AuxiliaryDone) -> bool: + pending = done.action_id in self._pending + inflight = done.action_id in self._inflight + self._pending.pop(done.action_id, None) + self._inflight.discard(done.action_id) + if not pending and not inflight: + return False + if done.action_id in self._results: + return True + if done.error is None: + result = (True, done.value, "") + else: + result = (False, None, str(done.error)) + self._results[done.action_id] = result + if len(self._results) > 16: + self._results.pop(next(iter(self._results))) + return True + + def take_result(self, action_id: str) -> tuple[bool, Any, str] | None: + return self._results.pop(action_id, None) + + def has_pending(self) -> bool: + return bool(self._pending) + + def has_inflight(self) -> bool: + return bool(self._inflight) + + def deadlines(self) -> tuple[float, ...]: + return tuple(self._pending.values()) + + def has_unsettled(self) -> bool: + return bool(self._pending or self._inflight) + + def expire(self, now: float, diagnostic: str) -> bool: + changed = False + for action_id, deadline in tuple(self._pending.items()): + if deadline <= now and action_id not in self._results: + self._pending.pop(action_id, None) + self._results[action_id] = (False, None, diagnostic) + changed = True + if len(self._results) > 16: + self._results.pop(next(iter(self._results))) + return changed diff --git a/dimos/manipulation/execution_clock.py b/dimos/manipulation/execution_clock.py new file mode 100644 index 0000000000..1ac6d4a587 --- /dev/null +++ b/dimos/manipulation/execution_clock.py @@ -0,0 +1,32 @@ +"""Validated monotonic clock dependency for execution deadlines.""" + +from collections.abc import Callable +import math +import threading + + +class InvalidMonotonicClock(ValueError): # noqa: N818 + """Raised when a clock violates the execution clock contract.""" + + +class ValidatedMonotonicClock: + """Serialize sampling and validate finite, nondecreasing seconds.""" + + def __init__(self, source: Callable[[], float]) -> None: + self._source = source + self._lock = threading.Lock() + with self._lock: + initial = self._source() + if not math.isfinite(initial): + raise InvalidMonotonicClock("monotonic clock returned a non-finite value") + self._last = initial + + def now(self) -> float: + with self._lock: + sample = self._source() + if not math.isfinite(sample): + raise InvalidMonotonicClock("monotonic clock returned a non-finite value") + if sample < self._last: + raise InvalidMonotonicClock("monotonic clock moved backwards") + self._last = sample + return sample diff --git a/dimos/manipulation/execution_effects.py b/dimos/manipulation/execution_effects.py new file mode 100644 index 0000000000..18ccaf8d98 --- /dev/null +++ b/dimos/manipulation/execution_effects.py @@ -0,0 +1,113 @@ +"""Daemon worker plumbing for execution-runtime effects.""" + +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from concurrent.futures.thread import _worker +from dataclasses import dataclass +import threading +from typing import Any, cast +import weakref + + +@dataclass(frozen=True) +class EffectDone: + action_id: str + outcome: Any + + +@dataclass(frozen=True) +class AuxiliaryDone: + action_id: str + value: Any + error: BaseException | None = None + + +@dataclass(frozen=True) +class StopDone: + success: bool + diagnostic: str = "" + + +class _DaemonThreadPoolExecutor(ThreadPoolExecutor): + def _adjust_thread_count(self) -> None: + if self._idle_semaphore.acquire(timeout=0): + return + if len(self._threads) >= self._max_workers: + return + + def weakref_cb(_ref: Any, q: Any = self._work_queue) -> None: + q.put(None) + + thread = threading.Thread( + name=f"{self._thread_name_prefix}_{len(self._threads)}", + target=_worker, + args=( + weakref.ref(self, weakref_cb), + self._work_queue, + self._initializer, # type: ignore[attr-defined] + self._initargs, # type: ignore[attr-defined] + ), + daemon=True, + ) + thread.start() + cast("Any", self._threads).add(thread) + + +class ExecutionEffectRunner: + """Execute already-admitted effects and post immutable completion events.""" + + def __init__(self) -> None: + self._executor = _DaemonThreadPoolExecutor( + max_workers=4, thread_name_prefix="execution-rpc" + ) + + def submit_action( + self, action_id: str, effect: Callable[[], Any], enqueue: Callable[[Any], None] + ) -> None: + future = self._executor.submit(effect) + + def done_callback(future: Any) -> None: + try: + outcome = future.exception() + if outcome is None: + outcome = future.result() + except BaseException as exc: + outcome = exc + enqueue(EffectDone(action_id, outcome)) + + future.add_done_callback(done_callback) + + def submit_auxiliary( + self, action_id: str, effect: Callable[[], Any], enqueue: Callable[[Any], None] + ) -> None: + future = self._executor.submit(effect) + + def done_callback(future: Any) -> None: + try: + error = future.exception() + enqueue( + AuxiliaryDone( + action_id, + None if error is not None else future.result(), + error, + ) + ) + except BaseException as exc: + enqueue(AuxiliaryDone(action_id, None, exc)) + + future.add_done_callback(done_callback) + + def submit_stop(self, effect: Callable[[], Any], enqueue: Callable[[Any], None]) -> None: + future = self._executor.submit(effect) + + def done_callback(future: Any) -> None: + try: + error = future.exception() + enqueue(StopDone(error is None, "" if error is None else str(error))) + except BaseException as exc: + enqueue(StopDone(False, str(exc))) + + future.add_done_callback(done_callback) + + def shutdown(self) -> None: + self._executor.shutdown(wait=False, cancel_futures=False) diff --git a/dimos/manipulation/execution_gateway.py b/dimos/manipulation/execution_gateway.py new file mode 100644 index 0000000000..b439d26859 --- /dev/null +++ b/dimos/manipulation/execution_gateway.py @@ -0,0 +1,105 @@ +# Copyright 2025-2026 Dimensional Inc. +# Licensed under the Apache License, Version 2.0 (the "License"). +"""Coordinator RPC conversion seam for the manipulation execution runtime.""" + +from __future__ import annotations + +from typing import Any + +from dimos.manipulation.execution_models import Outcome +from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState + + +class ControlCoordinatorGateway: + """Convert coordinator RPC requests and responses to normalized outcomes.""" + + def __init__(self, client: Any) -> None: + self._client = client + + def _invoke(self, task: str, method: str, kwargs: dict[str, Any]) -> Any: + return self._client.task_invoke(task, method, kwargs) + + def execute(self, task_name: str, request: Any) -> Outcome: + try: + value = self._invoke( + task_name, + "execute", + request if isinstance(request, dict) else {"trajectory": request}, + ) + return ( + Outcome.ACCEPTED + if value is True + else Outcome.REJECTED + if value is False + else Outcome.UNKNOWN + ) + except Exception: + return Outcome.UNKNOWN + + def cancel(self, task_name: str) -> Outcome: + try: + value = self._invoke(task_name, "cancel", {}) + return ( + Outcome.CANCELLED + if value is True + else Outcome.INACTIVE + if value is False + else Outcome.UNKNOWN + ) + except Exception: + return Outcome.UNKNOWN + + def status(self, task_name: str) -> Outcome: + try: + value = self._invoke(task_name, "get_state", {}) + value = getattr(value, "state", value) + state = TrajectoryState(value) + except (TypeError, ValueError, KeyError): + return Outcome.UNKNOWN + except Exception: + return Outcome.UNKNOWN + return { + TrajectoryState.IDLE: Outcome.INACTIVE, + TrajectoryState.EXECUTING: Outcome.RUNNING, + TrajectoryState.COMPLETED: Outcome.COMPLETED, + TrajectoryState.ABORTED: Outcome.CANCELLED, + TrajectoryState.FAULT: Outcome.FAILED, + }[state] + + def reset(self, task_name: str) -> Outcome: + try: + value = self._invoke(task_name, "reset", {}) + return ( + Outcome.INACTIVE + if value is True + else Outcome.RUNNING + if value is False + else Outcome.UNKNOWN + ) + except Exception: + return Outcome.UNKNOWN + + def set_gripper_position(self, hardware_id: str, position: float) -> Outcome: + try: + value = self._client.set_gripper_position(hardware_id, position) + return ( + Outcome.ACCEPTED + if value is True + else Outcome.REJECTED + if value is False + else Outcome.UNKNOWN + ) + except Exception: + return Outcome.UNKNOWN + + def get_gripper_position(self, hardware_id: str) -> float | None: + try: + value = self._client.get_gripper_position(hardware_id) + return float(value) if isinstance(value, (int, float)) else None + except Exception: + return None + + def stop(self) -> None: + stop = getattr(self._client, "stop_rpc_client", None) + if callable(stop): + stop() diff --git a/dimos/manipulation/execution_models.py b/dimos/manipulation/execution_models.py new file mode 100644 index 0000000000..254e4cf31b --- /dev/null +++ b/dimos/manipulation/execution_models.py @@ -0,0 +1,188 @@ +# Copyright 2025-2026 Dimensional Inc. +# Licensed under the Apache License, Version 2.0 (the "License"). +"""Pure public records and protocol models for execution runtime state.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Any, Generic, Protocol, TypeVar + +from dimos.manipulation.execution_topology import ExecutionPlan, PreparedPlan, TaskEntry + + +class LifecycleState(str, Enum): + IDLE = "IDLE" + PLANNING = "PLANNING" + READY = "READY" + DISPATCHING = "DISPATCHING" + RUNNING = "RUNNING" + CANCELLING = "CANCELLING" + FAULT = "FAULT" + + +class TaskActivity(str, Enum): + NOT_STARTED = "not_started" + EXECUTE_UNRESOLVED = "execute_unresolved" + ACTIVE = "active" + COMPLETED = "completed" + CANCELLED = "cancelled" + INACTIVE = "inactive" + REMOTE_FAULT = "remote_fault" + UNKNOWN = "unknown" + + +class ActionMethod(str, Enum): + EXECUTE = "execute" + CANCEL = "cancel" + STATUS = "get_state" + RESET = "reset" + GRIPPER_SET = "set_gripper_position" + GRIPPER_GET = "get_gripper_position" + + +class Outcome(str, Enum): + ACCEPTED = "accepted" + RUNNING = "running" + COMPLETED = "completed" + CANCELLED = "cancelled" + INACTIVE = "inactive" + REJECTED = "rejected" + FAILED = "failed" + UNKNOWN = "unknown" + + +class ShutdownState(str, Enum): + OPEN = "open" + CLOSING = "closing" + CLOSED = "closed" + + +@dataclass(frozen=True) +class OperationHandle: + plan_id: str + operation_id: str + attempt_id: str + + +ExecutionHandle = OperationHandle + + +@dataclass(frozen=True) +class ActionRecord: + action_id: str + method: ActionMethod + started: float + deadline: float + deadline_reported: bool = False + reset_id: str | None = None + + +@dataclass(frozen=True) +class TaskRecord: + task_id: str + task_name: str + entry: TaskEntry + activity: TaskActivity + action: ActionRecord | None = None + cancel_required: bool = False + reset_required: bool = False + + +PlanInput = ExecutionPlan | PreparedPlan + + +@dataclass(frozen=True) +class Operation: + handle: OperationHandle + plan: PlanInput + tasks: tuple[TaskRecord, ...] + next_index: int = 0 + cancel_requested: bool = False + uncertain: bool = False + rejected: bool = False + failed: bool = False + diagnostic: str = "" + + +@dataclass(frozen=True) +class ExecutionResult: + handle: OperationHandle + outcome: Outcome + diagnostic: str = "" + + +@dataclass(frozen=True) +class ResetHandle: + reset_id: str + + +@dataclass(frozen=True) +class ResetResult: + handle: ResetHandle + success: bool + diagnostic: str = "" + + +@dataclass(frozen=True) +class ShutdownResult: + success: bool + diagnostic: str = "" + + +@dataclass +class RuntimeContext: + state: LifecycleState = LifecycleState.IDLE + ready_plan: PlanInput | None = None + ready_plan_id: str | None = None + planning_token: str | None = None + active: Operation | None = None + fault: str | None = None + diagnostic: str | None = None + shutdown: ShutdownState = ShutdownState.OPEN + reset_handle: ResetHandle | None = None + dispatch_results: tuple[tuple[OperationHandle, ExecutionResult], ...] = () + terminal_results: tuple[tuple[OperationHandle, ExecutionResult], ...] = () + reset_results: tuple[tuple[ResetHandle, ResetResult], ...] = () + reset_proven_inactive: frozenset[str] = frozenset() + reset_completed_tasks: frozenset[str] = frozenset() + reset_deadline: float | None = None + shutdown_result: ShutdownResult | None = None + shutdown_deadline: float | None = None + physical_operation_deadline: float | None = None + revision: int = 0 + + +@dataclass(frozen=True) +class RuntimeSnapshot: + state: LifecycleState + ready_plan: PlanInput | None + ready_plan_id: str | None + planning_token: str | None + operation: Operation | None + fault: str | None + diagnostic: str | None + shutdown: ShutdownState + shutdown_result: ShutdownResult | None + revision: int + + +T = TypeVar("T") + + +@dataclass(frozen=True) +class CommandResult(Generic[T]): + accepted: bool + value: T | None = None + diagnostic: str = "" + snapshot: RuntimeSnapshot | None = None + + +class CoordinatorGateway(Protocol): + def execute(self, task_name: str, request: Any) -> Outcome: ... + def cancel(self, task_name: str) -> Outcome: ... + def status(self, task_name: str) -> Outcome: ... + def reset(self, task_name: str) -> Outcome: ... + def set_gripper_position(self, hardware_id: str, position: float) -> Outcome: ... + def get_gripper_position(self, hardware_id: str) -> float | None: ... + def stop(self) -> None: ... diff --git a/dimos/manipulation/execution_policy.py b/dimos/manipulation/execution_policy.py new file mode 100644 index 0000000000..4d08b79b7f --- /dev/null +++ b/dimos/manipulation/execution_policy.py @@ -0,0 +1,134 @@ +"""Pure policy decisions for execution-runtime reconciliation.""" + +from dataclasses import dataclass + +from dimos.manipulation.execution_models import ( + ActionMethod, + Outcome, + TaskActivity, +) + + +@dataclass(frozen=True) +class ResetCompletionDecision: + """Passive result of reconciling one reset-related coordinator call.""" + + activity: TaskActivity | None = None + cancel_required: bool | None = None + reset_required: bool | None = None + prove_inactive: bool = False + complete_reset: bool = False + advance_reset: bool = False + reset_success: bool | None = None + request_cancel: bool = False + emergency_cancel: bool = False + diagnostic: str | None = None + + +@dataclass(frozen=True) +class CompletionDecision: + """Pure normalized decision for a non-reset coordinator completion.""" + + activity: TaskActivity | None = None + fault: bool = False + rejected: bool = False + accepted: bool = False + diagnostic: str | None = None + + +_SAFE = frozenset((TaskActivity.INACTIVE, TaskActivity.CANCELLED, TaskActivity.COMPLETED)) + + +def reconcile_execute_completion(outcome: Outcome) -> CompletionDecision: + if outcome == Outcome.ACCEPTED: + return CompletionDecision(activity=TaskActivity.ACTIVE, accepted=True) + if outcome == Outcome.REJECTED: + return CompletionDecision(activity=TaskActivity.INACTIVE, rejected=True) + return CompletionDecision(fault=True, diagnostic=f"execute outcome={outcome.value}") + + +def reconcile_cancel_completion(outcome: Outcome) -> CompletionDecision: + if outcome in (Outcome.CANCELLED, Outcome.INACTIVE, Outcome.ACCEPTED): + return CompletionDecision(activity=TaskActivity.CANCELLED) + return CompletionDecision( + fault=True, + diagnostic=( + "cancellation is uncertain" + if outcome == Outcome.UNKNOWN + else "cancellation failed" + if outcome == Outcome.FAILED + else "malformed cancel outcome" + ), + ) + + +def reconcile_status_completion(outcome: Outcome) -> CompletionDecision: + activity = { + Outcome.COMPLETED: TaskActivity.COMPLETED, + Outcome.INACTIVE: TaskActivity.INACTIVE, + Outcome.CANCELLED: TaskActivity.CANCELLED, + Outcome.ACCEPTED: TaskActivity.ACTIVE, + Outcome.RUNNING: TaskActivity.ACTIVE, + }.get(outcome) + if activity is not None: + return CompletionDecision(activity=activity) + return CompletionDecision(fault=True, diagnostic=f"status outcome={outcome.value}") + + +def reconcile_reset_completion( + *, + method: ActionMethod, + outcome: Outcome, + activity: TaskActivity, + gate_sealed: bool, + clock_failed: bool, +) -> ResetCompletionDecision: + """Classify a reset completion without mutating runtime state.""" + if method == ActionMethod.STATUS: + if outcome in (Outcome.RUNNING, Outcome.ACCEPTED): + return ResetCompletionDecision( + activity=TaskActivity.ACTIVE, + cancel_required=True, + request_cancel=True, + emergency_cancel=gate_sealed or clock_failed, + ) + if outcome in (Outcome.INACTIVE, Outcome.CANCELLED, Outcome.COMPLETED): + resolved = { + Outcome.INACTIVE: TaskActivity.INACTIVE, + Outcome.CANCELLED: TaskActivity.CANCELLED, + Outcome.COMPLETED: TaskActivity.COMPLETED, + }[outcome] + return ResetCompletionDecision( + activity=resolved, + cancel_required=False, + prove_inactive=True, + advance_reset=True, + ) + safe = activity in _SAFE + return ResetCompletionDecision( + activity=activity if safe else TaskActivity.UNKNOWN, + cancel_required=not safe, + request_cancel=not safe, + emergency_cancel=not safe and (gate_sealed or clock_failed), + reset_success=False, + diagnostic="status reconciliation is unsafe", + ) + if method == ActionMethod.CANCEL: + if outcome in (Outcome.CANCELLED, Outcome.INACTIVE, Outcome.ACCEPTED): + return ResetCompletionDecision( + activity=TaskActivity.ACTIVE, + cancel_required=False, + advance_reset=True, + ) + return ResetCompletionDecision( + cancel_required=False, + reset_success=False, + diagnostic=None, + ) + if method == ActionMethod.RESET and outcome == Outcome.INACTIVE: + return ResetCompletionDecision( + reset_required=False, + complete_reset=True, + advance_reset=True, + ) + return ResetCompletionDecision(reset_success=False) diff --git a/dimos/manipulation/execution_runtime.py b/dimos/manipulation/execution_runtime.py new file mode 100644 index 0000000000..802a6709e4 --- /dev/null +++ b/dimos/manipulation/execution_runtime.py @@ -0,0 +1,1474 @@ +# Copyright 2025-2026 Dimensional Inc. +# Licensed under the Apache License, Version 2.0 (the "License"). +"""Small, fail-closed owner-thread execution runtime. + +The owner is the only thread which changes :class:`RuntimeContext`. Gateway +calls are submitted to a fixed executor and return as owner events; in +particular, an RPC can never hold the owner while it is in flight. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, replace +import queue +import threading +import time +from typing import Any, TypeVar, cast +import uuid + +from dimos.manipulation.execution_auxiliary import AuxiliaryCallBook, AuxiliaryTicket +from dimos.manipulation.execution_clock import InvalidMonotonicClock, ValidatedMonotonicClock +from dimos.manipulation.execution_effects import ( + AuxiliaryDone, + EffectDone, + ExecutionEffectRunner, + StopDone, +) +from dimos.manipulation.execution_gateway import ( + ControlCoordinatorGateway as ControlCoordinatorGateway, +) +from dimos.manipulation.execution_models import ( + ActionMethod as ActionMethod, + ActionRecord as ActionRecord, + CommandResult as CommandResult, + CoordinatorGateway as CoordinatorGateway, + ExecutionHandle as ExecutionHandle, + ExecutionResult as ExecutionResult, + LifecycleState as LifecycleState, + Operation as Operation, + OperationHandle as OperationHandle, + Outcome as Outcome, + PlanInput as PlanInput, + ResetHandle as ResetHandle, + ResetResult as ResetResult, + RuntimeContext as RuntimeContext, + RuntimeSnapshot as RuntimeSnapshot, + ShutdownResult as ShutdownResult, + ShutdownState as ShutdownState, + TaskActivity as TaskActivity, + TaskRecord as TaskRecord, +) +from dimos.manipulation.execution_policy import ( + reconcile_cancel_completion, + reconcile_execute_completion, + reconcile_reset_completion, + reconcile_status_completion, +) +from dimos.manipulation.execution_topology import ( + ExecutionPlan as ExecutionPlan, + ExecutionTopology as ExecutionTopology, + PreparedPlan as PreparedPlan, + TaskEntry as TaskEntry, + materialize_prepared_plan as materialize_prepared_plan, + prepare_execution_plan as prepare_execution_plan, + prepare_generated_plan as prepare_generated_plan, + prepare_plan as prepare_plan, +) +from dimos.manipulation.planning.spec.models import GeneratedPlan + +T = TypeVar("T") + + +class _Call: + def __init__(self, fn: Callable[[], Any]) -> None: + self.fn = fn + self.answer: queue.Queue[Any] = queue.Queue(1) + + +@dataclass(frozen=True) +class _PendingAction: + handle: OperationHandle + task_name: str + task: TaskRecord + method: ActionMethod + deadline: float + + +class ExecutionRuntime: + """Serialized plan executor with a sticky, fail-closed fault state.""" + + def __init__( + self, + gateway_factory: Callable[[], CoordinatorGateway], + *, + topology: ExecutionTopology | None = None, + plan_validator: Callable[[PlanInput], bool] | None = None, + action_timeout: float = 1.0, + physical_operation_timeout: float = 60.0, + poll_interval: float = 0.1, + monotonic_clock: Callable[[], float] = time.monotonic, + ) -> None: + self._gateway = gateway_factory() + self._topology = topology + self._validator = plan_validator + self._action_timeout = max(0.0, action_timeout) + self._physical_timeout = max(0.0, physical_operation_timeout) + self._poll_interval = max(0.001, poll_interval) + self._clock = ValidatedMonotonicClock(monotonic_clock) + self._context = RuntimeContext() + self._lock = threading.RLock() + self._condition = threading.Condition(self._lock) + self._events: queue.Queue[_Call | EffectDone | StopDone | AuxiliaryDone | None] = ( + queue.Queue() + ) + self._effects = ExecutionEffectRunner() + self._closed = False + self._gate = "open" + self._pending: dict[str, _PendingAction] = {} + self._fault_terminal_results: dict[OperationHandle, str] = {} + self._next_poll: float | None = None + self._shutdown_lock = threading.Lock() + self._shutdown_initiated = False + self._shutdown_stop_started = False + self._shutdown_stop_finished = False + self._shutdown_drain_started = False + self._clock_invalid = False + self._auxiliary = AuxiliaryCallBook() + self._owner = threading.Thread(target=self._run, daemon=True, name="execution-owner") + self._owner.start() + + def _snapshot(self) -> RuntimeSnapshot: + c = self._context + return RuntimeSnapshot( + c.state, + c.ready_plan, + c.ready_plan_id, + c.planning_token, + c.active, + c.fault, + c.diagnostic, + c.shutdown, + c.shutdown_result, + c.revision, + ) + + def snapshot(self) -> RuntimeSnapshot: + with self._lock: + return self._snapshot() + + def _commit(self, **kwargs: Any) -> None: + with self._lock: + self._context = replace(self._context, **kwargs, revision=self._context.revision + 1) + self._condition.notify_all() + + def _submit(self, fn: Callable[[], T], *, closing: bool = False) -> T: + with self._lock: + if self._closed or ( + self._shutdown_drain_started + or (self._shutdown_initiated and not self._owner.is_alive()) + or ( + not closing + and (self._context.shutdown != ShutdownState.OPEN or self._gate != "open") + ) + ): + return cast( + "T", + CommandResult( + False, diagnostic="runtime is closing", snapshot=self._snapshot() + ), + ) + call = _Call(fn) + self._events.put(call) + result = call.answer.get() + if isinstance(result, BaseException): + raise result + return cast("T", result) + + def _run(self) -> None: + while True: + timeout = self._wait_timeout() + try: + event = self._events.get(timeout=timeout) + except queue.Empty: + with self._lock: + self._timers() + continue + if event is None: + return + try: + with self._lock: + if isinstance(event, _Call): + self._timers() + event.answer.put(event.fn()) + elif isinstance(event, EffectDone): + self._done(event) + elif isinstance(event, AuxiliaryDone): + self._aux_done(event) + else: + self._stop_done(event) + self._timers() + except BaseException as exc: + if isinstance(event, _Call): + event.answer.put(exc) + + def _wait_timeout(self) -> float: + if self._clock_invalid: + return 3600.0 + deadlines = [ + action.deadline + for action in self._pending.values() + if action.task.action is None or not action.task.action.deadline_reported + ] + deadlines.extend(self._auxiliary.deadlines()) + if self._next_poll is not None: + deadlines.append(self._next_poll) + if self._context.reset_deadline is not None: + deadlines.append(self._context.reset_deadline) + if self._context.shutdown_deadline is not None: + deadlines.append(self._context.shutdown_deadline) + if self._context.physical_operation_deadline is not None: + deadlines.append(self._context.physical_operation_deadline) + if not deadlines: + return 3600.0 + now = self._clock_now() + if now is None: + return 3600.0 + return max(0.0, min(deadlines) - now) + + def _clock_failure_owner(self, diagnostic: str) -> None: + self._gate = "sealed" + self._shutdown_initiated = True + op = self._context.active + if op is not None: + potentially_active = { + TaskActivity.EXECUTE_UNRESOLVED, + TaskActivity.ACTIVE, + TaskActivity.UNKNOWN, + TaskActivity.REMOTE_FAULT, + } + op = replace( + op, + tasks=tuple( + replace( + task, + cancel_required=task.cancel_required or task.activity in potentially_active, + reset_required=(task.reset_required or task.activity in potentially_active), + ) + for task in op.tasks + ), + uncertain=True, + diagnostic=diagnostic, + ) + self._commit( + state=LifecycleState.FAULT + if self._context.state == LifecycleState.FAULT + else LifecycleState.CANCELLING, + shutdown=ShutdownState.CLOSING, + shutdown_result=self._context.shutdown_result or ShutdownResult(False, diagnostic), + active=op, + diagnostic=diagnostic, + ) + if op is not None: + self._schedule_cancels(op, emergency=True) + self._shutdown_progress() + + def _clock_now(self) -> float | None: + if self._clock_invalid: + return None + try: + return self._clock.now() + except InvalidMonotonicClock: + self._clock_invalid = True + if threading.current_thread() is self._owner: + self._clock_failure_owner("invalid monotonic clock") + else: + with self._lock: + if not self._shutdown_drain_started and self._owner.is_alive(): + self._events.put( + _Call(lambda: self._clock_failure_owner("invalid monotonic clock")) + ) + return None + + def _timers(self) -> None: + now = self._clock_now() + if now is None: + return + for action_id, item in tuple(self._pending.items()): + if item.deadline <= now: + handle, task_name, task, method = ( + item.handle, + item.task_name, + item.task, + item.method, + ) + op = self._context.active + report = False + if op is not None and op.handle == handle: + current = next((t for t in op.tasks if t.task_id == task.task_id), None) + if ( + current is not None + and current.action is not None + and not current.action.deadline_reported + ): + report = True + action = replace(current.action, deadline_reported=True) + tasks = tuple( + replace( + t, + action=action, + activity=TaskActivity.UNKNOWN, + cancel_required=True, + ) + if t.task_id == task.task_id + else t + for t in op.tasks + ) + self._commit(active=replace(op, tasks=tasks, uncertain=True)) + self._pending[action_id] = replace( + item, + task=replace(task, action=action), + ) + if report: + self._fault( + f"coordinator {method.value} deadline for {task_name}", + handle, + task, + ) + if self._auxiliary.expire(now, "coordinator auxiliary action deadline exceeded"): + self._condition.notify_all() + physical_deadline = self._context.physical_operation_deadline + if ( + physical_deadline is not None + and now >= physical_deadline + and self._context.active is not None + and self._context.state == LifecycleState.RUNNING + ): + self._commit(physical_operation_deadline=None) + self._begin_cancellation("physical operation deadline exceeded") + if self._next_poll is not None and now >= self._next_poll: + self._next_poll = now + self._poll_interval + self._poll_active() + if self._context.reset_deadline is not None and now >= self._context.reset_deadline: + if self._context.reset_handle and self._context.reset_handle not in dict( + self._context.reset_results + ): + self._finish_reset(False, "reset deadline exceeded") + if self._context.shutdown_deadline is not None and now >= self._context.shutdown_deadline: + if self._context.shutdown_result is None: + self._finish_shutdown(False, "shutdown safety deadline exceeded") + + def _valid_plan(self, plan: PlanInput) -> bool: + if not isinstance(plan, (ExecutionPlan, PreparedPlan)) or not plan.entries: + return False + if isinstance(plan, ExecutionPlan) and isinstance(plan.generated_plan, GeneratedPlan): + return False + if isinstance(plan, PreparedPlan) and (not plan._canonical or plan.topology is None): + return False + groups = tuple(getattr(plan.generated_plan, "group_ids", ())) + names = tuple(e.planning_group for e in plan.entries) + if len(set(names)) != len(names) or len({e.task_name for e in plan.entries}) != len(names): + return False + if groups and groups != names: + return False + topology = plan.topology if isinstance(plan, PreparedPlan) else self._topology + if topology is not None: + routes = {g: (r, t) for g, r, t in topology.routes} + for entry in plan.entries: + route = routes.get(entry.planning_group) + if ( + route is None + or route[1] != entry.task_name + or (entry.robot_name and route[0] != entry.robot_name) + ): + return False + return True + + def _valid(self, plan: PlanInput) -> bool: + try: + return bool(self._validator(plan) if self._validator else self._valid_plan(plan)) + except Exception: + return False + + def _start_action( + self, + op: Operation, + task: TaskRecord, + method: ActionMethod, + *, + emergency: bool = False, + ) -> None: + if task.action is not None or (self._gate == "sealed" and not emergency): + return + if self._gate != "open" and not ( + method == ActionMethod.CANCEL + and task.cancel_required + and task.activity != TaskActivity.NOT_STARTED + and (emergency or self._gate == "cancel_only") + ): + return + if method not in (ActionMethod.STATUS, ActionMethod.RESET) and task.activity in ( + TaskActivity.INACTIVE, + TaskActivity.COMPLETED, + TaskActivity.CANCELLED, + ): + return + now = self._clock_now() + if now is None and method != ActionMethod.CANCEL: + return + action = ActionRecord( + str(uuid.uuid4()), + method, + now or 0.0, + float("inf") if now is None else now + self._action_timeout, + ) + tasks = tuple( + replace( + t, + action=action, + activity=TaskActivity.EXECUTE_UNRESOLVED + if method == ActionMethod.EXECUTE + else t.activity, + ) + if t.task_id == task.task_id + else t + for t in op.tasks + ) + op = replace(op, tasks=tasks) + self._commit(active=op) + self._pending[action.action_id] = _PendingAction( + handle=op.handle, + task_name=task.task_name, + task=next(t for t in tasks if t.task_id == task.task_id), + method=method, + deadline=action.deadline, + ) + fn: Callable[[], Any] + if method == ActionMethod.EXECUTE: + fn = lambda: self._gateway.execute(task.task_name, task.entry.request) # noqa: E731 + elif method == ActionMethod.CANCEL: + fn = lambda: self._gateway.cancel(task.task_name) # noqa: E731 + elif method == ActionMethod.STATUS: + fn = lambda: self._gateway.status(task.task_name) # noqa: E731 + else: + fn = lambda: self._gateway.reset(task.task_name) # noqa: E731 + self._effects.submit_action(action.action_id, fn, self._events.put) + + @dataclass(frozen=True) + class _RetiredAction: + operation: Operation + index: int + task: TaskRecord + method: ActionMethod + + def _retire_action(self, done: EffectDone) -> _RetiredAction | None: + item = self._pending.get(done.action_id) + active = self._context.active + if item is None or active is None: + return None + if active.handle != item.handle: + return None + index = next( + (i for i, t in enumerate(active.tasks) if t.task_id == item.task.task_id), None + ) + if index is None: + return None + current_task = active.tasks[index] + if current_task.action is None or current_task.action.action_id != done.action_id: + # The registered operation/task identity is sufficient to consume this + # physical completion, but it must not retire a newer action. + self._pending.pop(done.action_id, None) + return None + self._pending.pop(done.action_id, None) + task = replace(active.tasks[index], action=None) + operation = replace(active, tasks=(*active.tasks[:index], task, *active.tasks[index + 1 :])) + self._commit(active=operation) + return self._RetiredAction(operation, index, task, item.method) + + def _done(self, done: EffectDone) -> None: + retired = self._retire_action(done) + if retired is None: + return + op, index, task, method = ( + retired.operation, + retired.index, + retired.task, + retired.method, + ) + outcome = done.outcome if isinstance(done.outcome, Outcome) else Outcome.UNKNOWN + if self._context.reset_handle is not None and method in ( + ActionMethod.STATUS, + ActionMethod.CANCEL, + ActionMethod.RESET, + ): + self._done_reset(op, index, task, method, outcome) + return + if method == ActionMethod.EXECUTE: + self._done_execute(op, index, task, outcome) + return + if method == ActionMethod.CANCEL: + self._done_cancel(op, index, task, outcome) + return + if method == ActionMethod.STATUS: + self._done_status(op, index, task, outcome) + return + + def _aux_done(self, done: AuxiliaryDone) -> None: + self._auxiliary.complete(done) + self._condition.notify_all() + self._shutdown_progress() + + def _reset_next(self, op: Operation) -> None: + reset = self._context.reset_handle + if reset is None or reset in dict(self._context.reset_results): + return + if self._gate == "sealed": + self._finish_reset(False, "reset unavailable while runtime is closing") + self._shutdown_progress() + return + proven = self._context.reset_proven_inactive + completed = self._context.reset_completed_tasks + for task in op.tasks: + if task.action is not None or task.activity == TaskActivity.NOT_STARTED: + continue + if task.task_id in completed: + continue + if task.task_id in proven: + if task.reset_required: + self._start_action(op, task, ActionMethod.RESET) + return + self._commit(reset_completed_tasks=frozenset((*completed, task.task_id))) + completed = self._context.reset_completed_tasks + continue + self._start_action(op, task, ActionMethod.STATUS) + return + if all(t.task_id in completed or t.activity == TaskActivity.NOT_STARTED for t in op.tasks): + self._finish_reset(True) + + def _done_reset( + self, + op: Operation, + index: int, + task: TaskRecord, + method: ActionMethod, + outcome: Outcome, + ) -> None: + current = op.tasks[index] + decision = reconcile_reset_completion( + method=method, + outcome=outcome, + activity=current.activity, + gate_sealed=self._gate == "sealed", + clock_failed=self._clock_invalid, + ) + current = replace( + current, + activity=decision.activity or current.activity, + cancel_required=( + current.cancel_required + if decision.cancel_required is None + else decision.cancel_required + ), + reset_required=( + current.reset_required + if decision.reset_required is None + else decision.reset_required + ), + ) + self._commit(active=replace(op, tasks=(*op.tasks[:index], current, *op.tasks[index + 1 :]))) + if decision.prove_inactive: + self._commit( + reset_proven_inactive=frozenset( + (*self._context.reset_proven_inactive, current.task_id) + ) + ) + if decision.complete_reset: + self._commit( + reset_completed_tasks=frozenset( + (*self._context.reset_completed_tasks, current.task_id) + ) + ) + diagnostic = decision.diagnostic or ( + f"task={current.task_name} action={method.value} outcome={outcome.value}" + ) + if decision.reset_success is not None: + self._finish_reset(decision.reset_success, diagnostic) + if decision.request_cancel: + self._schedule_cancels( + self._context.active or op, + emergency=decision.emergency_cancel, + ) + if decision.advance_reset: + self._reset_next(self._context.active or op) + self._shutdown_progress() + + def _record_dispatch(self, op: Operation, outcome: Outcome, diagnostic: str = "") -> None: + if dict(self._context.dispatch_results).get(op.handle) is None: + self._commit( + dispatch_results=( + *self._context.dispatch_results, + (op.handle, ExecutionResult(op.handle, outcome, diagnostic)), + )[-16:] + ) + + def _begin_cancellation(self, diagnostic: str = "") -> None: + op = self._context.active + if op is None: + return + tasks = tuple( + replace( + t, + cancel_required=( + t.cancel_required + or t.activity + in ( + TaskActivity.EXECUTE_UNRESOLVED, + TaskActivity.ACTIVE, + TaskActivity.UNKNOWN, + TaskActivity.REMOTE_FAULT, + ) + ), + ) + for t in op.tasks + ) + op = replace( + op, + tasks=tasks, + cancel_requested=True, + failed=op.failed or diagnostic == "physical operation deadline exceeded", + diagnostic=diagnostic or op.diagnostic, + ) + self._commit(state=LifecycleState.CANCELLING, active=op) + self._schedule_cancels(op) + + def _done_execute(self, op: Operation, index: int, task: TaskRecord, outcome: Outcome) -> None: + if self._context.shutdown == ShutdownState.CLOSING: + op = replace(op, cancel_requested=True) + decision = reconcile_execute_completion(outcome) + if decision.fault: + diagnostic = f"task={task.task_name} action=execute outcome={outcome.value}" + self._record_dispatch(op, Outcome.UNKNOWN, diagnostic) + self._fault( + diagnostic, + op.handle, + replace(task, action=None), + ) + return + if decision.rejected: + tasks = tuple( + replace( + t, + action=None, + activity=TaskActivity.INACTIVE if i >= index else t.activity, + cancel_required=(t.cancel_required or i < index) if i < index else False, + ) + for i, t in enumerate(op.tasks) + ) + op = replace(op, tasks=tasks, rejected=True, cancel_requested=True) + diagnostic = f"task={task.task_name} action=execute outcome={outcome.value}" + op = replace(op, diagnostic=diagnostic) + if self._context.state == LifecycleState.FAULT or op.uncertain: + self._commit(active=op) + else: + self._commit(state=LifecycleState.CANCELLING, active=op) + self._record_dispatch(op, Outcome.REJECTED, diagnostic) + self._schedule_cancels(op) + self._finish_if_terminal(op) + return + activity = TaskActivity.ACTIVE + tasks = tuple( + replace( + t, + action=None, + activity=activity, + cancel_required=( + t.cancel_required + or op.cancel_requested + or self._context.state == LifecycleState.FAULT + ), + ) + if t.task_id == task.task_id + else t + for t in op.tasks + ) + op = replace(op, tasks=tasks) + self._commit(active=op) + if op.cancel_requested or self._context.state in ( + LifecycleState.CANCELLING, + LifecycleState.FAULT, + ): + self._record_dispatch(op, Outcome.CANCELLED) + self._schedule_cancels(op) + self._shutdown_progress() + return + if index != op.next_index: + self._fault("out-of-order execute completion", op.handle, task) + return + if index + 1 < len(op.tasks): + op = replace(op, next_index=index + 1) + self._commit(active=op) + self._start_action(op, op.tasks[index + 1], ActionMethod.EXECUTE) + else: + self._record_dispatch(op, Outcome.ACCEPTED) + now = self._clock_now() + if now is None: + return + self._commit( + state=LifecycleState.RUNNING, + physical_operation_deadline=now + self._physical_timeout, + ) + self._next_poll = now + self._poll_interval + + def _done_cancel(self, op: Operation, index: int, task: TaskRecord, outcome: Outcome) -> None: + decision = reconcile_cancel_completion(outcome) + if decision.fault and outcome in (Outcome.UNKNOWN, Outcome.FAILED): + diagnostic = decision.diagnostic or "cancellation failed" + self._fault( + diagnostic, + op.handle, + replace(task, action=None), + schedule_cancels=False, + mark_peers=False, + ) + current = self._context.active + if current is not None and current.handle == op.handle: + self._commit( + active=replace( + current, + tasks=tuple( + replace( + item, + action=None, + cancel_required=False, + reset_required=True, + ) + if item.task_id == task.task_id + else item + for item in current.tasks + ), + ) + ) + current = self._context.active + if current is not None: + if self._context.shutdown == ShutdownState.CLOSING: + self._finish_shutdown(False, diagnostic) + self._schedule_cancels(current) + self._shutdown_progress() + return + if decision.fault: + diagnostic = f"malformed cancel outcome for {task.task_name}" + self._fault( + diagnostic, + op.handle, + replace(task, action=None), + schedule_cancels=False, + mark_peers=False, + ) + current = self._context.active + if current is not None and current.handle == op.handle: + current = replace( + current, + tasks=tuple( + replace(item, action=None, cancel_required=False) + if item.task_id == task.task_id + else item + for item in current.tasks + ), + ) + self._commit(active=current) + if self._context.shutdown == ShutdownState.CLOSING: + self._finish_shutdown(False, diagnostic) + self._schedule_cancels(current) + self._shutdown_progress() + return + tasks = tuple( + replace(t, action=None, activity=TaskActivity.CANCELLED, cancel_required=False) + if t.task_id == task.task_id + else t + for t in op.tasks + ) + op = replace(op, tasks=tasks) + self._commit(active=op) + self._schedule_cancels(op) + self._finish_if_terminal(op) + self._shutdown_progress() + + def _done_status(self, op: Operation, index: int, task: TaskRecord, outcome: Outcome) -> None: + decision = reconcile_status_completion(outcome) + if decision.fault: + self._fault( + f"status is unsafe: task={task.task_name} action=get_state outcome={outcome.value}", + op.handle, + replace(task, action=None), + ) + return + activity = cast("TaskActivity", decision.activity) + tasks = tuple( + replace( + t, + action=None, + activity=activity, + cancel_required=False + if t.task_id == task.task_id + and activity + in (TaskActivity.INACTIVE, TaskActivity.CANCELLED, TaskActivity.COMPLETED) + else t.cancel_required, + ) + if t.task_id == task.task_id + else t + for t in op.tasks + ) + op = replace(op, tasks=tasks) + if activity in (TaskActivity.INACTIVE, TaskActivity.CANCELLED): + op = replace( + op, + cancel_requested=True, + failed=True, + diagnostic=( + f"task={task.task_name} outcome={outcome.value}" + if activity in (TaskActivity.INACTIVE, TaskActivity.CANCELLED) + else op.diagnostic + ), + tasks=tuple( + replace( + t, + cancel_required=( + False + if t.task_id == task.task_id + else t.cancel_required + or t.activity + in ( + TaskActivity.EXECUTE_UNRESOLVED, + TaskActivity.ACTIVE, + TaskActivity.UNKNOWN, + TaskActivity.REMOTE_FAULT, + ) + ), + ) + for t in op.tasks + ), + ) + self._commit(state=LifecycleState.CANCELLING, active=op) + self._schedule_cancels(op) + self._finish_if_terminal(op) + return + self._commit(active=op) + self._finish_if_terminal(op) + if self._context.state in (LifecycleState.CANCELLING, LifecycleState.FAULT): + self._schedule_cancels(op) + + def _poll_active(self) -> None: + op = self._context.active + if op is not None and self._context.state == LifecycleState.RUNNING: + task = next( + (t for t in op.tasks if t.activity == TaskActivity.ACTIVE and t.action is None), + None, + ) + if task: + self._start_action(op, task, ActionMethod.STATUS) + + def _schedule_cancels(self, op: Operation, *, emergency: bool = False) -> None: + for task in op.tasks: + current = self._context.active + if current is None or current.handle != op.handle: + return + task = next(t for t in current.tasks if t.task_id == task.task_id) + if task.action is None and task.cancel_required: + self._start_action( + current, task, ActionMethod.CANCEL, emergency=emergency or self._clock_invalid + ) + + def _finish_if_terminal(self, op: Operation) -> None: + if any(t.action is not None for t in op.tasks): + return + if self._context.state == LifecycleState.CANCELLING: + if all( + t.activity + in ( + TaskActivity.NOT_STARTED, + TaskActivity.CANCELLED, + TaskActivity.INACTIVE, + TaskActivity.COMPLETED, + ) + and not t.cancel_required + for t in op.tasks + ): + if self._context.shutdown == ShutdownState.CLOSING: + self._shutdown_progress() + else: + self._terminal( + Outcome.REJECTED + if op.rejected and not op.failed + else Outcome.FAILED + if op.failed + else Outcome.CANCELLED + ) + elif all(t.activity == TaskActivity.COMPLETED for t in op.tasks): + self._terminal(Outcome.COMPLETED) + + def _terminal(self, outcome: Outcome) -> None: + op = self._context.active + if op is None: + return + result = ExecutionResult(op.handle, outcome, op.diagnostic) + self._commit( + state=LifecycleState.IDLE, + active=None, + diagnostic=op.diagnostic, + terminal_results=(*self._context.terminal_results, (op.handle, result))[-16:], + ) + + def _fault( + self, + diagnostic: str, + handle: OperationHandle | None = None, + task: TaskRecord | None = None, + *, + schedule_cancels: bool = True, + mark_peers: bool = True, + ) -> None: + op = self._context.active + if op is not None: + tasks = tuple( + replace( + t, + action=None if task is not None and task.action is None else t.action, + activity=TaskActivity.UNKNOWN, + cancel_required=( + t.cancel_required + or t.activity + in ( + TaskActivity.EXECUTE_UNRESOLVED, + TaskActivity.ACTIVE, + TaskActivity.UNKNOWN, + TaskActivity.REMOTE_FAULT, + ) + ), + reset_required=( + t.reset_required + or t.activity + in ( + TaskActivity.EXECUTE_UNRESOLVED, + TaskActivity.ACTIVE, + TaskActivity.UNKNOWN, + TaskActivity.REMOTE_FAULT, + ) + ), + ) + if ( + mark_peers + and t.activity + in ( + TaskActivity.EXECUTE_UNRESOLVED, + TaskActivity.ACTIVE, + TaskActivity.UNKNOWN, + TaskActivity.REMOTE_FAULT, + ) + ) + or (task is not None and t.task_id == task.task_id) + else t + for t in op.tasks + ) + op = replace(op, tasks=tasks, uncertain=True, diagnostic=diagnostic) + self._commit( + state=LifecycleState.FAULT, + fault=self._context.fault or diagnostic, + diagnostic=diagnostic, + active=op, + ) + fault_handle = handle or (op.handle if op is not None else None) + if fault_handle is not None: + self._fault_terminal_results[fault_handle] = diagnostic + if len(self._fault_terminal_results) > 8: + oldest = next(iter(self._fault_terminal_results)) + del self._fault_terminal_results[oldest] + if handle and not dict(self._context.dispatch_results).get(handle): + self._commit( + dispatch_results=( + *self._context.dispatch_results, + (handle, ExecutionResult(handle, Outcome.UNKNOWN, diagnostic)), + )[-16:] + ) + if op and schedule_cancels: + self._schedule_cancels(op) + + def _stop_done(self, done: StopDone) -> None: + self._shutdown_stop_finished = True + if self._context.shutdown_result is not None: + self._shutdown_drain_owner() + return + if not done.success: + self._finish_shutdown( + False, + f"gateway close failed: {done.diagnostic or 'stop failed'}", + ) + self._shutdown_drain_owner() + return + deadline = self._context.shutdown_deadline + now = self._clock_now() + if now is None: + return + if deadline is None or now >= deadline: + self._finish_shutdown(False, "shutdown safety deadline exceeded") + self._shutdown_drain_owner() + return + self._commit( + state=LifecycleState.IDLE, + shutdown=ShutdownState.CLOSED, + shutdown_result=ShutdownResult(True, done.diagnostic), + shutdown_deadline=None, + ) + self._shutdown_drain_owner() + + def _has_unsettled_shutdown_work(self) -> bool: + active = self._context.active + return bool( + self._pending + or self._auxiliary.has_unsettled() + or ( + active is not None + and any(t.action is not None or t.cancel_required for t in active.tasks) + ) + ) + + def _shutdown_progress(self) -> None: + if self._context.shutdown != ShutdownState.CLOSING: + return + if self._has_unsettled_shutdown_work(): + return + if self._shutdown_stop_started: + return + self._shutdown_stop_started = True + self._gate = "sealed" + self._effects.submit_stop(self._gateway.stop, self._events.put) + + def _shutdown_drain_owner(self) -> None: + if self._shutdown_drain_started: + return + if self._has_unsettled_shutdown_work(): + return + self._shutdown_drain_started = True + self._effects.shutdown() + self._events.put(None) + + def _finish_reset(self, success: bool, diagnostic: str = "") -> None: + h = self._context.reset_handle + if h is None or h in dict(self._context.reset_results): + return + result = ResetResult(h, success, diagnostic) + self._commit( + state=LifecycleState.IDLE if success else LifecycleState.FAULT, + reset_results=(*self._context.reset_results, (h, result))[-16:], + reset_handle=None if success else h, + reset_deadline=None, + fault=None if success else self._context.fault or diagnostic, + active=None if success else self._context.active, + physical_operation_deadline=None + if success + else self._context.physical_operation_deadline, + ready_plan=None if success else self._context.ready_plan, + ready_plan_id=None if success else self._context.ready_plan_id, + diagnostic=diagnostic, + ) + + def _finish_shutdown(self, success: bool, diagnostic: str = "") -> None: + if self._context.shutdown_result is None: + self._commit( + shutdown_result=ShutdownResult(success, diagnostic), + shutdown_deadline=None if not success else self._context.shutdown_deadline, + ) + + def start_planning(self) -> str | None: + def fn() -> str | None: + if self._context.state not in (LifecycleState.IDLE, LifecycleState.READY): + return None + token = str(uuid.uuid4()) + self._commit( + state=LifecycleState.PLANNING, + planning_token=token, + ready_plan=None, + ready_plan_id=None, + diagnostic=None, + ) + return token + + return self._submit(fn) + + def complete_planning(self, token: str, plan: PlanInput) -> CommandResult[str]: + def fn() -> CommandResult[str]: + if token != self._context.planning_token: + return CommandResult( + False, diagnostic="planning token mismatch", snapshot=self._snapshot() + ) + if not self._valid(plan): + self._commit( + state=LifecycleState.IDLE, + planning_token=None, + diagnostic="invalid prepared plan", + ) + return CommandResult( + False, diagnostic="invalid prepared plan", snapshot=self._snapshot() + ) + pid = str(uuid.uuid4()) + self._commit( + state=LifecycleState.READY, ready_plan=plan, ready_plan_id=pid, planning_token=None + ) + return CommandResult(True, pid, snapshot=self._snapshot()) + + return self._submit(fn) + + def fail_planning(self, token: str, diagnostic: str) -> CommandResult[None]: + def fn() -> CommandResult[None]: + ok = token == self._context.planning_token + if ok: + self._commit(state=LifecycleState.IDLE, planning_token=None, diagnostic=diagnostic) + return CommandResult( + ok, + diagnostic=diagnostic if ok else "planning token mismatch", + snapshot=self._snapshot(), + ) + + return self._submit(fn) + + def cancel_planning(self) -> CommandResult[None]: + return self._submit(lambda: self._planning_command(LifecycleState.PLANNING, "cancelled")) + + def _planning_command(self, state: LifecycleState, diagnostic: str) -> CommandResult[None]: + ok = self._context.state == state + if ok: + self._commit(state=LifecycleState.IDLE, planning_token=None, diagnostic=diagnostic) + return CommandResult(ok, snapshot=self._snapshot()) + + def clear_ready_plan(self) -> CommandResult[None]: + def fn() -> CommandResult[None]: + if self._context.state != LifecycleState.READY: + return CommandResult(False, snapshot=self._snapshot()) + self._commit( + state=LifecycleState.IDLE, + ready_plan=None, + ready_plan_id=None, + diagnostic="", + ) + return CommandResult(True, snapshot=self._snapshot()) + + return self._submit(fn) + + def _execute( + self, plan: PlanInput, plan_id: str | None = None + ) -> CommandResult[OperationHandle]: + if self._context.state not in ( + LifecycleState.IDLE, + LifecycleState.READY, + ) or not self._valid(plan): + return CommandResult( + False, diagnostic="execution unavailable or invalid plan", snapshot=self._snapshot() + ) + h = OperationHandle(plan_id or str(uuid.uuid4()), str(uuid.uuid4()), str(uuid.uuid4())) + op = Operation( + h, + plan, + tuple( + TaskRecord(str(uuid.uuid4()), e.task_name, e, TaskActivity.NOT_STARTED) + for e in plan.entries + ), + ) + self._commit( + state=LifecycleState.DISPATCHING, + ready_plan=None, + ready_plan_id=None, + active=op, + diagnostic=None, + ) + self._start_action(op, op.tasks[0], ActionMethod.EXECUTE) + return CommandResult(True, h, snapshot=self._snapshot()) + + def execute_explicit(self, prepared_plan: PlanInput) -> CommandResult[OperationHandle]: + return self._submit(lambda: self._execute(prepared_plan)) + + def execute_ready(self) -> CommandResult[OperationHandle]: + return self._submit( + lambda: self._execute(self._context.ready_plan, self._context.ready_plan_id) + if self._context.ready_plan is not None + else CommandResult(False, diagnostic="no ready plan", snapshot=self._snapshot()) + ) + + def cancel(self) -> CommandResult[None]: + return self._submit(self._cancel_owner, closing=True) + + def _cancel_owner(self) -> CommandResult[None]: + op = self._context.active + if op is None or self._context.state not in ( + LifecycleState.DISPATCHING, + LifecycleState.RUNNING, + ): + return CommandResult(False, snapshot=self._snapshot()) + self._begin_cancellation() + return CommandResult(True, snapshot=self._snapshot()) + + def cancel_if_current(self, handle: OperationHandle) -> CommandResult[None]: + return self._submit( + lambda: self._cancel_owner() + if self._context.active and self._context.active.handle == handle + else CommandResult( + False, diagnostic="operation is no longer current", snapshot=self._snapshot() + ), + closing=True, + ) + + def poll(self) -> CommandResult[None]: + def fn() -> CommandResult[None]: + self._poll_active() + return CommandResult( + self._context.state == LifecycleState.RUNNING, + snapshot=self._snapshot(), + ) + + return self._submit(fn) + + def set_gripper_position(self, hardware_id: str, position: float) -> CommandResult[Outcome]: + result = self._submit(lambda: self._gripper(True, hardware_id, position)) + if isinstance(result, CommandResult): + return cast("CommandResult[Outcome]", result) + return cast("CommandResult[Outcome]", self._wait_aux(result)) + + def get_gripper_position(self, hardware_id: str) -> CommandResult[float]: + result = self._submit(lambda: self._gripper(False, hardware_id, None)) + if isinstance(result, CommandResult): + return cast("CommandResult[float]", result) + return cast("CommandResult[float]", self._wait_aux(result)) + + def _gripper( + self, setter: bool, hardware: str, position: float | None + ) -> AuxiliaryTicket | CommandResult[Any]: + if self._gate != "open": + return CommandResult(False, diagnostic="gateway is closing", snapshot=self._snapshot()) + if setter and position is None: + return CommandResult( + False, diagnostic="position is required", snapshot=self._snapshot() + ) + now = self._clock_now() + if now is None: + return CommandResult( + False, diagnostic="invalid monotonic clock", snapshot=self._snapshot() + ) + action_id = str(uuid.uuid4()) + ticket = self._auxiliary.register(action_id, now + self._action_timeout, setter) + if setter: + assert position is not None + fn: Callable[[], Any] = lambda: self._gateway.set_gripper_position(hardware, position) # noqa: E731 + else: + fn = lambda: self._gateway.get_gripper_position(hardware) # noqa: E731 + self._effects.submit_auxiliary(action_id, fn, self._events.put) + return ticket + + def _wait_aux(self, ticket: AuxiliaryTicket) -> CommandResult[Any]: + start = self._clock_now() + if start is None: + return CommandResult( + False, diagnostic="invalid monotonic clock", snapshot=self._snapshot() + ) + end = start + self._action_timeout + with self._condition: + while True: + result = self._auxiliary.take_result(ticket.action_id) + if result is not None: + accepted, value, diagnostic = result + if ticket.setter: + return CommandResult( + accepted and isinstance(value, Outcome), + value if isinstance(value, Outcome) else None, + diagnostic=diagnostic, + snapshot=self._snapshot(), + ) + valid = isinstance(value, (int, float)) and not isinstance(value, bool) + return CommandResult( + accepted and valid, + float(value) if valid else None, + diagnostic=diagnostic, + snapshot=self._snapshot(), + ) + now = self._clock_now() + if now is None: + return CommandResult( + False, diagnostic="invalid monotonic clock", snapshot=self._snapshot() + ) + if now >= end: + return CommandResult( + False, + diagnostic="coordinator auxiliary action timeout", + snapshot=self._snapshot(), + ) + self._condition.wait(max(0.0, end - now)) + + def _wait(self, handle: Any, kind: str, timeout: float) -> CommandResult[Any]: + end = time.monotonic() + max(0.0, timeout) + with self._condition: + while time.monotonic() < end: + values = ( + self._context.dispatch_results + if kind == "dispatch" + else self._context.terminal_results + if kind == "terminal" + else self._context.reset_results + ) + found = dict(values).get(handle) + if kind == "terminal": + fault_diagnostic = self._fault_terminal_results.get(handle) + if fault_diagnostic is not None: + return CommandResult( + False, + diagnostic=fault_diagnostic, + snapshot=self._snapshot(), + ) + if found is not None: + return CommandResult(True, found, snapshot=self._snapshot()) + if kind == "terminal": + active = self._context.active + if ( + active is not None + and active.handle == handle + and self._context.state == LifecycleState.FAULT + ): + return CommandResult( + False, + diagnostic=self._context.diagnostic or "runtime fault", + snapshot=self._snapshot(), + ) + remaining = end - time.monotonic() + self._condition.wait(min(0.1, max(0.0, remaining))) + return CommandResult(False, diagnostic="timeout", snapshot=self._snapshot()) + + def wait_for_dispatch( + self, handle: OperationHandle, timeout: float = 1.0 + ) -> CommandResult[ExecutionResult]: + return self._wait(handle, "dispatch", timeout) + + def wait_for_terminal( + self, handle: OperationHandle, timeout: float = 1.0 + ) -> CommandResult[ExecutionResult]: + return self._wait(handle, "terminal", timeout) + + def reset(self) -> CommandResult[ResetHandle]: + def fn() -> CommandResult[ResetHandle]: + if self._gate != "open": + return CommandResult( + False, + diagnostic="reset unavailable while runtime is closing", + snapshot=self._snapshot(), + ) + existing = self._context.reset_handle + if existing is not None: + prior = dict(self._context.reset_results).get(existing) + if prior is not None: + return CommandResult( + False, + existing, + diagnostic=prior.diagnostic, + snapshot=self._snapshot(), + ) + return CommandResult(True, existing, snapshot=self._snapshot()) + if self._context.state != LifecycleState.FAULT: + return CommandResult( + False, diagnostic="reset requires FAULT", snapshot=self._snapshot() + ) + if self._pending: + return CommandResult( + False, + diagnostic="reset blocked by unresolved coordinator RPC", + snapshot=self._snapshot(), + ) + now = self._clock_now() + if now is None: + return CommandResult( + False, diagnostic="invalid monotonic clock", snapshot=self._snapshot() + ) + h = ResetHandle(str(uuid.uuid4())) + self._commit( + reset_handle=h, + reset_deadline=now + self._action_timeout, + reset_proven_inactive=frozenset(), + reset_completed_tasks=frozenset(), + ) + op = self._context.active + if op: + self._reset_next(op) + else: + self._finish_reset(True) + return CommandResult(True, h, snapshot=self._snapshot()) + + return self._submit(fn, closing=True) + + def wait_for_reset( + self, handle: ResetHandle, timeout: float = 1.0 + ) -> CommandResult[ResetResult]: + return self._wait(handle, "reset", timeout) + + def shutdown(self, timeout: float = 1.0) -> CommandResult[ShutdownResult]: + with self._shutdown_lock: + with self._lock: + if self._context.shutdown_result is not None: + return CommandResult( + self._context.shutdown_result.success, + self._context.shutdown_result, + snapshot=self._snapshot(), + ) + if self._shutdown_initiated: + return CommandResult( + False, + diagnostic="shutdown already initiated", + snapshot=self._snapshot(), + ) + self._shutdown_initiated = True + self._gate = "cancel_only" + return self._shutdown_wait(timeout) + + def _shutdown_wait(self, timeout: float) -> CommandResult[ShutdownResult]: + start = self._clock_now() + if start is None: + end = time.monotonic() + 1.0 + with self._condition: + while self._context.shutdown_result is None and time.monotonic() < end: + self._condition.wait(0.01) + result = self._context.shutdown_result + if result is None: + return CommandResult( + False, diagnostic="invalid monotonic clock", snapshot=self.snapshot() + ) + return CommandResult(False, result, snapshot=self.snapshot()) + deadline = start + max(0.0, timeout) + self._submit(lambda: self._shutdown_owner(deadline), closing=True) + while True: + with self._condition: + if self._context.shutdown_result is not None: + break + now = self._clock_now() + if now is None or now >= deadline: + break + with self._condition: + self._condition.wait(min(0.1, max(0.0, deadline - now))) + if self._context.shutdown_result is None: + self._submit( + lambda: self._finish_shutdown(False, "shutdown safety deadline exceeded"), + closing=True, + ) + result = self._context.shutdown_result + if result is None: + return CommandResult( + False, + diagnostic="shutdown safety deadline exceeded", + snapshot=self.snapshot(), + ) + if not result.success: + return CommandResult(False, result, snapshot=self.snapshot()) + self._owner.join() + if self._owner.is_alive(): + return CommandResult( + False, + diagnostic="owner thread did not stop", + snapshot=self.snapshot(), + ) + self._closed = True + return CommandResult(True, result, snapshot=self.snapshot()) + + close = shutdown + + def _shutdown_owner(self, deadline: float) -> CommandResult[None]: + self._commit( + shutdown=ShutdownState.CLOSING, + shutdown_deadline=deadline, + state=LifecycleState.CANCELLING if self._context.active else LifecycleState.IDLE, + ) + if self._context.active: + self._begin_cancellation("shutdown cancellation requested") + self._shutdown_progress() + return CommandResult(True, snapshot=self._snapshot()) diff --git a/dimos/manipulation/execution_topology.py b/dimos/manipulation/execution_topology.py new file mode 100644 index 0000000000..d3a528c30a --- /dev/null +++ b/dimos/manipulation/execution_topology.py @@ -0,0 +1,273 @@ +# Copyright 2025-2026 Dimensional Inc. +# Licensed under the Apache License, Version 2.0 (the "License"). +"""Pure topology binding and generated-plan materialization helpers.""" + +from __future__ import annotations + +from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass, field +from typing import Any + +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.models import GeneratedPlan +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint + + +@dataclass(frozen=True) +class ExecutionTopology: + """Canonical group -> robot -> coordinator-task routing.""" + + routes: tuple[tuple[str, str, str], ...] + robot_configs: tuple[RobotModelConfig, ...] = () + inverse_joint_mappings: tuple[tuple[str, tuple[tuple[str, str], ...]], ...] = () + + @classmethod + def from_robot_configs( + cls, + robot_configs: Iterable[RobotModelConfig], + group_routes: Mapping[str, tuple[str, str]] | Iterable[tuple[str, str, str]] | None = None, + group_resolver: Callable[[str], tuple[str, str]] | None = None, + ) -> ExecutionTopology: + """Build canonical routes while retaining the supplied config references.""" + configs = tuple(robot_configs) + by_robot: dict[str, RobotModelConfig] = {} + for config in configs: + if config.name in by_robot: + raise ValueError(f"duplicate robot config: {config.name}") + by_robot[config.name] = config + explicit: dict[str, tuple[str, str]] = {} + if isinstance(group_routes, Mapping): + for group, route in group_routes.items(): + if not isinstance(group, str) or not isinstance(route, tuple) or len(route) != 2: + raise ValueError("invalid group route mapping") + explicit[group] = (str(route[0]), str(route[1])) + elif group_routes is not None: + for group, robot, task in group_routes: + if group in explicit: + raise ValueError(f"duplicate group route: {group}") + explicit[group] = (robot, task) + routes: list[tuple[str, str, str]] = [] + inverse: list[tuple[str, tuple[tuple[str, str], ...]]] = [] + for config in configs: + local_names = set(config.joint_names) + local_to_coordinator: dict[str, str] = {} + for coordinator_name, local_name in config.joint_name_mapping.items(): + if local_name not in local_names: + raise ValueError( + f"mapped local joint is not configured: {config.name}/{local_name}" + ) + previous = local_to_coordinator.get(local_name) + if previous is not None and previous != coordinator_name: + raise ValueError( + f"duplicate local joint mapping for {config.name}/{local_name}" + ) + local_to_coordinator[local_name] = coordinator_name + coordinator_names: set[str] = set() + for local_name in config.joint_names: + coordinator_name = local_to_coordinator.get(local_name, local_name) + if coordinator_name in coordinator_names: + raise ValueError( + f"duplicate coordinator joint mapping for {config.name}/{coordinator_name}" + ) + coordinator_names.add(coordinator_name) + inverse.append( + ( + config.name, + tuple( + (local, local_to_coordinator.get(local, local)) + for local in config.joint_names + ), + ) + ) + for definition in config.planning_groups: + if not set(definition.joint_names).issubset(local_names): + raise ValueError( + f"planning group joints are not configured: {config.name}/{definition.name}" + ) + group_id = f"{config.name}/{definition.name}" + if group_resolver is not None: + route_robot, route_task = group_resolver(group_id) + else: + if not config.coordinator_task_name and group_id not in explicit: + continue + route_robot, route_task = explicit.get( + group_id, (config.name, config.coordinator_task_name or "") + ) + if route_robot != config.name or route_task != config.coordinator_task_name: + raise ValueError(f"route/config mismatch for {group_id}") + routes.append((group_id, route_robot, route_task)) + if explicit and set(explicit) != {group for group, _, _ in routes}: + raise ValueError("route set does not match configured planning groups") + return cls(tuple(routes), configs, tuple(inverse)) + + from_configs = from_robot_configs + from_robot_model_configs = from_robot_configs + + def config_for_robot(self, robot_name: str) -> RobotModelConfig: + for config in self.robot_configs: + if config.name == robot_name: + return config + raise KeyError(f"unknown robot: {robot_name}") + + def coordinator_joint_name(self, robot_name: str, local_name: str) -> str: + for name, mappings in self.inverse_joint_mappings: + if name == robot_name: + for local, coordinator in mappings: + if local == local_name: + return coordinator + break + return local_name + + def route_for_group(self, group_id: str) -> tuple[str, str]: + routes = {group: (robot, task) for group, robot, task in self.routes} + try: + return routes[group_id] + except KeyError as exc: + raise ValueError(f"missing route for planning group {group_id}") from exc + + def affected_robots(self, group_ids: Iterable[str]) -> tuple[str, ...]: + robots: list[str] = [] + for group_id in group_ids: + robot_name, _task_name = self.route_for_group(group_id) + if robot_name in robots: + raise ValueError(f"multiple selected groups resolve to robot {robot_name}") + robots.append(robot_name) + return tuple(robots) + + @property + def groups(self) -> tuple[str, ...]: + return tuple(group for group, _, _ in self.routes) + + +@dataclass(frozen=True) +class TaskEntry: + planning_group: str + task_name: str + request: Any + robot_name: str | None = None + + +@dataclass(frozen=True) +class ExecutionPlan: + generated_plan: Any + entries: tuple[TaskEntry, ...] + + +@dataclass(frozen=True) +class PreparedPlan: + generated_plan: GeneratedPlan + entries: tuple[TaskEntry, ...] + topology: ExecutionTopology | None = None + _canonical: bool = field(default=False, repr=False, compare=False) + + +def prepare_generated_plan( + generated_plan: GeneratedPlan, topology: ExecutionTopology +) -> PreparedPlan: + """Purely bind a generated multi-robot plan to canonical coordinator tasks.""" + if not isinstance(generated_plan, GeneratedPlan): + raise TypeError("expected GeneratedPlan") + group_ids = tuple(generated_plan.group_ids) + if not group_ids or len(set(group_ids)) != len(group_ids): + raise ValueError("generated plan must contain unique planning groups") + route_by_group = {group: (robot, task) for group, robot, task in topology.routes} + if len(route_by_group) != len(topology.routes): + raise ValueError("topology contains duplicate group routes") + selected_routes: list[tuple[str, str, str, RobotModelConfig, tuple[str, ...]]] = [] + affected: set[str] = set() + for group_id in group_ids: + route = route_by_group.get(group_id) + if route is None: + robot_name = group_id.split("/", 1)[0] + try: + if not topology.config_for_robot(robot_name).coordinator_task_name: + raise ValueError(f"no coordinator task for selected robot {robot_name}") + except KeyError: + pass + raise ValueError(f"missing route for planning group {group_id}") + robot_name, task_name = route + if not task_name: + raise ValueError(f"missing coordinator task for {group_id}") + if robot_name in affected: + raise ValueError(f"multiple selected groups resolve to robot {robot_name}") + config = topology.config_for_robot(robot_name) + prefix, _, definition_name = group_id.partition("/") + if prefix != robot_name or not definition_name: + raise ValueError(f"invalid canonical planning group {group_id}") + definitions = [ + definition + for definition in config.planning_groups + if definition.name == definition_name + ] + if len(definitions) != 1: + raise ValueError(f"missing config group for {group_id}") + expected_local = tuple(definitions[0].joint_names) + if len(set(expected_local)) != len(expected_local): + raise ValueError(f"duplicate intended joints for {group_id}") + affected.add(robot_name) + selected_routes.append((group_id, robot_name, task_name, config, expected_local)) + expected_by_robot = {robot: set(expected) for _, robot, _, _, expected in selected_routes} + + trajectory = generated_plan.trajectory + names = tuple(trajectory.joint_names) + if not names: + raise ValueError("generated trajectory has no joint names") + if len(set(names)) != len(names): + raise ValueError("duplicate global trajectory joint assignment") + columns: dict[str, dict[str, int]] = {robot: {} for robot in affected} + known_robots = {config.name for config in topology.robot_configs} + for index, global_name in enumerate(names): + if global_name.count("/") != 1: + raise ValueError(f"malformed global joint name: {global_name}") + robot_name, local_name = global_name.split("/", 1) + if not robot_name or not local_name or robot_name not in known_robots: + raise ValueError(f"unknown global joint assignment: {global_name}") + if robot_name in columns: + if local_name not in expected_by_robot[robot_name]: + raise ValueError(f"additional joint column for selected robot: {global_name}") + if local_name in columns[robot_name]: + raise ValueError(f"duplicate robot joint assignment: {global_name}") + columns[robot_name][local_name] = index + points = tuple(trajectory.points) + for point in points: + if len(point.positions) != len(names): + raise ValueError("trajectory point dimension does not match joint names") + if point.velocities and len(point.velocities) != len(names): + raise ValueError("trajectory velocity dimension does not match joint names") + + entries: list[TaskEntry] = [] + for group_id, robot_name, task_name, _config, expected_local in selected_routes: + robot_columns = columns[robot_name] + missing = [local for local in expected_local if local not in robot_columns] + if missing: + raise ValueError(f"missing robot joints for {robot_name}: {missing}") + intended = set(expected_local) + local_order = tuple(local for local in robot_columns if local in intended) + coordinator_names = tuple( + topology.coordinator_joint_name(robot_name, local) for local in local_order + ) + if len(set(coordinator_names)) != len(coordinator_names): + raise ValueError(f"duplicate coordinator joint names for {robot_name}") + local_trajectory = JointTrajectory( + joint_names=list(coordinator_names), + points=[ + TrajectoryPoint( + time_from_start=point.time_from_start, + positions=[point.positions[robot_columns[local]] for local in local_order], + velocities=( + [point.velocities[robot_columns[local]] for local in local_order] + if point.velocities + else None + ), + ) + for point in points + ], + ) + entries.append(TaskEntry(group_id, task_name, {"trajectory": local_trajectory}, robot_name)) + return PreparedPlan(generated_plan, tuple(entries), topology, True) + + +prepare_plan = prepare_generated_plan +prepare_execution_plan = prepare_generated_plan +materialize_prepared_plan = prepare_generated_plan diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index 20a8d167c0..c80b31e072 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -24,12 +24,13 @@ from __future__ import annotations -from enum import Enum +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +import math import threading import time -from typing import TYPE_CHECKING, Any, TypeAlias +from typing import Any, TypeAlias -import numpy as np from pydantic import Field from dimos.agents.annotation import skill @@ -38,6 +39,18 @@ from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In +from dimos.manipulation.execution_runtime import ( + ControlCoordinatorGateway, + ExecutionRuntime, + ExecutionTopology, + LifecycleState, + OperationHandle, + Outcome, + PreparedPlan, + RuntimeSnapshot, + ShutdownState, + prepare_generated_plan, +) from dimos.manipulation.planning.factory import ( KinematicsName, PlannerName, @@ -45,6 +58,12 @@ create_planning_specs, create_world, ) +from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.manipulation.planning.groups.utils import ( + filter_joint_state_to_selected_joints, + joint_target_to_global_names, + planning_group_id_from_selector, +) from dimos.manipulation.planning.kinematics.config import ( ManipulationKinematicsConfig, PinkKinematicsConfig, @@ -53,9 +72,10 @@ from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import IKStatus, ObstacleType from dimos.manipulation.planning.spec.models import ( + GeneratedPlan, IKResult, - JointPath, Obstacle, + PlanningGroupID, RobotName, WorldRobotID, ) @@ -69,17 +89,16 @@ NoManipulationVisualizationConfig, ) from dimos.manipulation.visualization.factory import create_manipulation_visualization +from dimos.manipulation.visualization.operator import ManipulationOperator from dimos.manipulation.visualization.types import TargetEvaluation 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.Vector3 import Vector3 from dimos.msgs.sensor_msgs.JointState import JointState from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory from dimos.utils.logging_config import setup_logger -if TYPE_CHECKING: - from dimos.core.rpc_client import RPCClient - logger = setup_logger() # Composite type aliases for readability (using semantic IDs from planning.spec) @@ -89,21 +108,20 @@ RobotRegistry: TypeAlias = dict[RobotName, RobotEntry] """Maps robot_name -> RobotEntry""" -PlannedPaths: TypeAlias = dict[RobotName, JointPath] -"""Maps robot_name -> planned joint path""" - -PlannedTrajectories: TypeAlias = dict[RobotName, JointTrajectory] -"""Maps robot_name -> planned trajectory""" +RobotInfoValue: TypeAlias = ( + str | bool | float | list[str] | list[float] | list[PlanningGroup] | None +) +RobotInfoPayload: TypeAlias = dict[str, RobotInfoValue] -class ManipulationState(Enum): - """State machine for manipulation module.""" +@dataclass(frozen=True) +class ManipulationExecutionSnapshot: + """Atomic execution state exposed to manipulation operators and UIs.""" - IDLE = 0 - PLANNING = 1 - EXECUTING = 2 - COMPLETED = 3 - FAULT = 4 + state: str + diagnostic: str + ready_plan_id: str | None + has_ready_plan: bool class ManipulationModuleConfig(ModuleConfig): @@ -111,6 +129,7 @@ class ManipulationModuleConfig(ModuleConfig): robots: list[RobotModelConfig] = Field(default_factory=list) planning_timeout: float = 10.0 + physical_operation_timeout: float = Field(default=60.0, gt=0, allow_inf_nan=False) world_backend: WorldBackend = "drake" visualization: ManipulationVisualizationConfig = Field( default_factory=NoManipulationVisualizationConfig @@ -135,6 +154,7 @@ class ManipulationModule(Module): """ config: ManipulationModuleConfig + _runtime_factory = ExecutionRuntime # Input: Joint state from coordinator (for world sync) coordinator_joint_state: In[JointState] @@ -142,11 +162,7 @@ class ManipulationModule(Module): def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) - # State machine - self._state = ManipulationState.IDLE self._lock = threading.Lock() - self._error_message = "" - self._planning_epoch = 0 # Planning components (initialized in start()) self._world_monitor: WorldMonitor | None = None @@ -156,12 +172,8 @@ def __init__(self, **kwargs: Any) -> None: # Robot registry: maps robot_name -> (world_robot_id, config, trajectory_gen) self._robots: RobotRegistry = {} - # Stored path for plan/preview/execute workflow (per robot) - self._planned_paths: PlannedPaths = {} - self._planned_trajectories: PlannedTrajectories = {} - - # Coordinator integration (lazy initialized) - self._coordinator_client: RPCClient | None = None + self._execution_runtime: ExecutionRuntime | None = None + self._execution_topology: ExecutionTopology | None = None # Init joints: captured from first joint state per robot, used by go_init self._init_joints: dict[RobotName, JointState] = {} @@ -177,16 +189,54 @@ def start(self) -> None: """Start the manipulation module.""" super().start() - # Initialize planning stack - self._initialize_planning() - - # Subscribe to joint state via port - if self.coordinator_joint_state is not None: - self.coordinator_joint_state.subscribe(self._on_joint_state) - logger.info("Subscribed to coordinator_joint_state port") + try: + # Keep subscription and all resource creation in one transaction. A + # subscription can fail just as readily as planner construction. + self._initialize_planning() + if self.coordinator_joint_state is not None: + self.coordinator_joint_state.subscribe(self._on_joint_state) + logger.info("Subscribed to coordinator_joint_state port") + except Exception: + logger.exception("Failed to initialize manipulation planning") + self._unwind_startup() + raise logger.info("ManipulationModule started") + def _unwind_startup(self) -> None: + """Best-effort cleanup for a partially completed start transaction.""" + try: + self._stop_tf() + except Exception: + logger.exception("Failed to stop TF during startup unwind") + try: + if self._world_monitor is not None: + self._world_monitor.stop_all_monitors() + except Exception: + logger.exception("Failed to stop world monitor during startup unwind") + try: + runtime = self._execution_runtime + if runtime is not None: + shutdown = runtime.shutdown(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) + if ( + shutdown.snapshot is not None + and shutdown.snapshot.shutdown == ShutdownState.CLOSED + ): + self._execution_runtime = None + except Exception: + logger.exception("Failed to shut down runtime during startup unwind") + finally: + try: + super().stop() + except Exception: + logger.exception("Failed to stop superclass during startup unwind") + + def _stop_tf(self) -> None: + self._tf_stop_event.set() + if self._tf_thread is not None: + self._tf_thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) + self._tf_thread = None + def _initialize_planning(self) -> None: """Initialize world, planner, and trajectory generator.""" if not self.config.robots: @@ -225,6 +275,21 @@ def _initialize_planning(self) -> None: self._world_monitor.finalize() + self._execution_topology = ExecutionTopology.from_robot_configs(self.config.robots) + + def gateway_factory() -> ControlCoordinatorGateway: + from dimos.control.coordinator import ControlCoordinator + from dimos.core.rpc_client import RPCClient + + return ControlCoordinatorGateway(RPCClient(None, ControlCoordinator)) + + self._execution_runtime = self._runtime_factory( + gateway_factory, + topology=self._execution_topology, + action_timeout=self.config.planning_timeout, + physical_operation_timeout=self.config.physical_operation_timeout, + ) + # Add floor obstacle to prevent trajectories below the table surface if self.config.floor_z is not None: fz = self.config.floor_z @@ -246,7 +311,8 @@ def _initialize_planning(self) -> None: self._world_monitor.start_state_monitor(robot_id) self._world_monitor.set_visualization(visualization) - self._world_monitor.sync_visualization_scene() + operator = ManipulationOperator(self, self._world_monitor) + self._world_monitor.initialize_visualization(operator=operator) if self._world_monitor.visualization is not None: self._world_monitor.start_visualization_thread(rate_hz=10.0) @@ -332,7 +398,7 @@ def _on_joint_state(self, msg: JointState) -> None: # Route to specific monitor self._world_monitor.on_joint_state(sub_msg, robot_id=robot_id) - # Capture per-robot init joints on first receipt + # Capture per-robot init joints on first update if robot_name not in self._init_joints: self._init_joints[robot_name] = sub_msg logger.info( @@ -382,7 +448,33 @@ def _tf_publish_loop(self) -> None: @rpc def get_state(self) -> str: """Get current manipulation state name.""" - return self._state.name + snapshot = self._execution_snapshot() + if snapshot is None: + return LifecycleState.IDLE.name + return snapshot.state.name + + def _execution_snapshot(self) -> RuntimeSnapshot | None: + """Read the runtime-owned execution projection exactly once.""" + runtime = self._execution_runtime + return runtime.snapshot() if runtime is not None else None + + @rpc + def get_execution_snapshot(self) -> ManipulationExecutionSnapshot: + """Return one atomic runtime execution snapshot for observers.""" + snapshot = self._execution_snapshot() + if snapshot is None: + return ManipulationExecutionSnapshot( + state=LifecycleState.IDLE.name, + diagnostic="", + ready_plan_id=None, + has_ready_plan=False, + ) + return ManipulationExecutionSnapshot( + state=snapshot.state.name, + diagnostic=snapshot.diagnostic or "", + ready_plan_id=snapshot.ready_plan_id, + has_ready_plan=snapshot.ready_plan is not None, + ) @rpc def get_error(self) -> str: @@ -391,21 +483,25 @@ def get_error(self) -> str: Returns: Error message or empty string """ - return self._error_message + snapshot = self._execution_snapshot() + if snapshot is None: + return "" + return snapshot.diagnostic or "" @rpc def cancel(self) -> bool: """Cancel current motion or invalidate an in-progress plan.""" - if self._state == ManipulationState.PLANNING: - self._planning_epoch += 1 - self._state = ManipulationState.IDLE - logger.info("Planning cancelled") - return True - if self._state != ManipulationState.EXECUTING: + if self._execution_runtime is None: return False - self._state = ManipulationState.IDLE - logger.info("Motion cancelled") - return True + snapshot = self._execution_runtime.snapshot() + if snapshot.state == LifecycleState.PLANNING: + return self._execution_runtime.cancel_planning().accepted + if snapshot.state == LifecycleState.READY: + result = self._execution_runtime.clear_ready_plan() + if result.accepted and isinstance(snapshot.ready_plan, PreparedPlan): + self._dismiss_preview(snapshot.ready_plan.generated_plan.group_ids) + return result.accepted + return self._execution_runtime.cancel().accepted @rpc @skill @@ -415,15 +511,39 @@ def reset(self) -> SkillResult[ManipulationSkillError]: Use this after an error or fault to allow new commands. Cannot reset while a motion is executing — cancel first. """ - if self._state == ManipulationState.EXECUTING: - return SkillResult.fail( - "INVALID_STATE", - "Cannot reset while executing — cancel the motion first", + if self._execution_runtime is None: + return SkillResult.fail("INVALID_STATE", "Execution runtime is not started") + snapshot = self._execution_runtime.snapshot() + if snapshot.state == LifecycleState.PLANNING: + result = self._execution_runtime.cancel_planning() + return ( + SkillResult.ok("Reset to IDLE — ready for new commands") + if result.accepted + else SkillResult.fail("INVALID_STATE", result.diagnostic) ) - if self._state == ManipulationState.PLANNING: - self._planning_epoch += 1 - self._state = ManipulationState.IDLE - self._error_message = "" + if snapshot.state == LifecycleState.READY: + result = self._execution_runtime.clear_ready_plan() + if not result.accepted: + return SkillResult.fail("INVALID_STATE", result.diagnostic) + if isinstance(snapshot.ready_plan, PreparedPlan): + self._dismiss_preview(snapshot.ready_plan.generated_plan.group_ids) + return SkillResult.ok("Reset to IDLE — ready for new commands") + if snapshot.state in ( + LifecycleState.DISPATCHING, + LifecycleState.RUNNING, + LifecycleState.CANCELLING, + ): + return SkillResult.fail("INVALID_STATE", "Cannot reset while execution is active") + if snapshot.state != LifecycleState.FAULT: + return SkillResult.ok("Reset to IDLE — ready for new commands") + reset_result = self._execution_runtime.reset() + if not reset_result.accepted or reset_result.value is None: + return SkillResult.fail("INVALID_STATE", reset_result.diagnostic) + reconciled = self._execution_runtime.wait_for_reset( + reset_result.value, timeout=float("inf") + ) + if not reconciled.accepted or reconciled.value is None or not reconciled.value.success: + return SkillResult.fail("RESET_FAILED", reconciled.diagnostic or "Reset failed") return SkillResult.ok("Reset to IDLE — ready for new commands") @rpc @@ -447,7 +567,11 @@ def get_ee_pose(self, robot_name: RobotName | None = None) -> Pose | None: robot_name: Robot to query (required if multiple robots configured) """ if (robot := self._get_robot(robot_name)) and self._world_monitor: - return self._world_monitor.get_ee_pose(robot[1], joint_state=None) + try: + return self._world_monitor.get_ee_pose(robot[1], joint_state=None) + except ValueError as exc: + logger.warning("End-effector pose unavailable: %s", exc) + return None return None @rpc @@ -464,40 +588,215 @@ def is_collision_free(self, joints: list[float], robot_name: RobotName | None = return self._world_monitor.is_state_valid(robot_id, joint_state) return False - def _begin_planning( - self, robot_name: RobotName | None = None - ) -> tuple[RobotName, WorldRobotID] | None: - """Check state and begin planning. Returns (robot_name, robot_id) or None. - - Args: - robot_name: Robot to plan for (required if multiple robots configured) - """ - if self._world_monitor is None: + def _begin_group_planning(self) -> str | None: + """Check state and begin planning for explicit planning-group APIs.""" + if self._world_monitor is None or self._execution_runtime is None: logger.error("Planning not initialized") return None - if (robot := self._get_robot(robot_name)) is None: + ready = self._execution_runtime.snapshot().ready_plan + previous = ready.generated_plan if isinstance(ready, PreparedPlan) else None + token = self._execution_runtime.start_planning() + if isinstance(token, str) and previous is not None: + self._dismiss_preview(previous.group_ids) + return token if isinstance(token, str) else None + + def _require_unique_pose_group_id_for_robot(self, robot_name: RobotName) -> PlanningGroupID: + """Return the unique pose-targetable group or raise if it is ambiguous.""" + if self._world_monitor is None: + raise ValueError("Planning not initialized") + group_id = self._world_monitor.planning_groups.primary_pose_group_id_for_robot(robot_name) + if group_id is None: + raise ValueError( + f"Robot '{robot_name}' has no pose-targetable planning group; " + "use an explicit planning group ID" + ) + return group_id + + @staticmethod + def _assert_finite_sequence(values: Sequence[float], label: str) -> None: + for value in values: + if not math.isfinite(value): + raise ValueError(f"{label} contains non-finite value") + + def _limits_for_global_joints( + self, joint_names: Sequence[str] + ) -> tuple[list[float], list[float]]: + velocities: list[float] = [] + accelerations: list[float] = [] + for global_name in joint_names: + if "/" not in global_name: + raise ValueError(f"Joint '{global_name}' is not globally named") + robot_name, local_name = global_name.split("/", 1) + robot = self._get_robot(robot_name) + if robot is None: + raise ValueError(f"Unknown robot for joint '{global_name}'") + _, _, config, _ = robot + if local_name not in config.joint_names: + raise ValueError(f"Unknown local joint '{global_name}'") + velocity = float(config.max_velocity) + acceleration = float(config.max_acceleration) + if not math.isfinite(velocity) or velocity <= 0.0: + raise ValueError(f"Invalid velocity limit for '{global_name}'") + if not math.isfinite(acceleration) or acceleration <= 0.0: + raise ValueError(f"Invalid acceleration limit for '{global_name}'") + velocities.append(velocity) + accelerations.append(acceleration) + return velocities, accelerations + + def _validate_selected_path( + self, path: Sequence[JointState], expected_names: Sequence[str] + ) -> list[list[float]]: + if len(path) < 2: + raise ValueError("Planner returned fewer than two waypoints") + expected = list(expected_names) + waypoints: list[list[float]] = [] + for waypoint_index, state in enumerate(path): + if list(state.name) != expected: + raise ValueError( + f"Waypoint {waypoint_index} joint names do not match selected order" + ) + positions = list(state.position) + if len(positions) != len(expected): + raise ValueError(f"Waypoint {waypoint_index} position dimension mismatch") + self._assert_finite_sequence(positions, f"Waypoint {waypoint_index} positions") + waypoints.append(positions) + return waypoints + + def _validate_generated_trajectory( + self, + trajectory: JointTrajectory, + expected_names: Sequence[str], + waypoints: Sequence[Sequence[float]], + ) -> None: + expected = list(expected_names) + if list(trajectory.joint_names) != expected: + raise ValueError("Generated trajectory joint names do not match selected order") + if not trajectory.points: + raise ValueError("Generated trajectory has no points") + previous_time: float | None = None + for point_index, point in enumerate(trajectory.points): + if len(point.positions) != len(expected) or len(point.velocities) != len(expected): + raise ValueError(f"Generated point {point_index} dimension mismatch") + self._assert_finite_sequence( + point.positions, f"Generated point {point_index} positions" + ) + self._assert_finite_sequence( + point.velocities, f"Generated point {point_index} velocities" + ) + if not math.isfinite(point.time_from_start): + raise ValueError(f"Generated point {point_index} time is non-finite") + if point_index == 0 and point.time_from_start != 0.0: + raise ValueError("Generated trajectory must start at time 0") + if previous_time is not None and point.time_from_start <= previous_time: + raise ValueError("Generated trajectory times must be strictly increasing") + previous_time = point.time_from_start + non_noop = any(list(waypoint) != list(waypoints[0]) for waypoint in waypoints[1:]) + if non_noop and trajectory.duration <= 0.0: + raise ValueError("Generated trajectory duration must be positive") + waypoint_index = 0 + for point in trajectory.points: + if list(point.positions) == list(waypoints[waypoint_index]): + waypoint_index += 1 + if waypoint_index == len(waypoints): + break + if waypoint_index != len(waypoints): + raise ValueError("Generated trajectory does not contain ordered waypoint boundaries") + + def _materialize_generated_plan( + self, group_ids: tuple[PlanningGroupID, ...], result_path: Sequence[JointState] + ) -> tuple[list[JointState], JointTrajectory]: + assert self._world_monitor is not None + selection = self._world_monitor.planning_groups.select(group_ids) + expected_names = list(selection.joint_names) + path = [JointState(state) for state in result_path] + waypoints = self._validate_selected_path(path, expected_names) + velocities, accelerations = self._limits_for_global_joints(expected_names) + generator = JointTrajectoryGenerator( + num_joints=len(expected_names), + max_velocity=velocities, + max_acceleration=accelerations, + ) + generated = generator.generate(waypoints) + trajectory = JointTrajectory( + joint_names=expected_names, + points=generated.points, + timestamp=generated.timestamp, + ) + self._validate_generated_trajectory(trajectory, expected_names, waypoints) + return path, trajectory + + def _plan_selected_path( + self, + group_ids: tuple[PlanningGroupID, ...], + start: JointState, + goal: JointState, + planning_epoch: str, + ) -> GeneratedPlan | None: + """Plan over explicit planning groups and store the resulting plan.""" + assert self._world_monitor and self._planner + result = self._planner.plan_selected_joint_path( + world=self._world_monitor.world, + selection=self._world_monitor.planning_groups.select(group_ids), + start=start, + goal=goal, + timeout=self.config.planning_timeout, + ) + if not result.is_success(): + self._fail_planning_epoch(planning_epoch, f"Planning failed: {result.status.name}") + return None + + logger.info("Path: %d waypoints, groups=%s", len(result.path), group_ids) + try: + path, trajectory = self._materialize_generated_plan(group_ids, result.path) + except Exception as exc: + self._fail_planning_epoch(planning_epoch, f"Failed to materialize plan: {exc}") return None - with self._lock: - if self._state not in (ManipulationState.IDLE, ManipulationState.COMPLETED): - logger.warning(f"Cannot plan: state is {self._state.name}") + plan = GeneratedPlan( + group_ids=group_ids, + trajectory=trajectory, + path=path, + status=result.status, + planning_time=result.planning_time, + path_length=result.path_length, + iterations=result.iterations, + message=result.message, + ) + if self._execution_runtime is None or self._execution_topology is None: + return None + try: + prepared = prepare_generated_plan(plan, self._execution_topology) + if not self._execution_runtime.complete_planning(planning_epoch, prepared).accepted: return None - self._planning_epoch += 1 - self._state = ManipulationState.PLANNING - return robot[0], robot[1] + except Exception as exc: + self._fail_planning_epoch(planning_epoch, f"Failed to prepare plan: {exc}") + return None + return plan - def _fail(self, msg: str) -> bool: - """Set FAULT state with error message.""" + def _fail_planning_epoch(self, planning_epoch: str, msg: str) -> bool: + """Fault only the still-current planning operation.""" logger.warning(msg) - self._state = ManipulationState.FAULT - self._error_message = msg - return False + return ( + self._execution_runtime is not None + and self._execution_runtime.fail_planning(planning_epoch, msg).accepted + ) - def _dismiss_preview(self, robot_id: WorldRobotID) -> None: + def _dismiss_preview(self, group_ids: Sequence[PlanningGroupID]) -> None: """Hide the preview ghost if the world supports it.""" if self._world_monitor is None: return - self._world_monitor.hide_preview(robot_id) - self._world_monitor.publish_visualization() + try: + robot_names = self._world_monitor.planning_groups.select(tuple(group_ids)).robot_names + robot_ids = tuple( + robot_id + for robot_name in robot_names + if (robot_id := self.robot_id_for_name(robot_name)) is not None + ) + except (KeyError, ValueError): + robot_ids = () + if robot_ids: + self._world_monitor.cancel_preview_animation(robot_ids=robot_ids) + else: + self._world_monitor.cancel_preview_animation() def _solve_ik_for_pose( self, @@ -509,9 +808,6 @@ def _solve_ik_for_pose( """Run the configured kinematics backend for a world-frame pose.""" assert self._world_monitor and self._kinematics - # Convert Pose to PoseStamped for the IK solver - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - target_pose = PoseStamped( frame_id="world", position=pose.position, @@ -527,133 +823,246 @@ def _solve_ik_for_pose( ) @rpc - def solve_ik( + def inverse_kinematics( + self, + pose_targets: Mapping[PlanningGroupID, PoseStamped], + auxiliary_group_ids: Sequence[PlanningGroupID] = (), + seed: JointState | None = None, + check_collision: bool = True, + ) -> IKResult: + """Solve planning-group pose targets without planning a joint path.""" + if self._kinematics is None or self._world_monitor is None: + return IKResult(status=IKStatus.NO_SOLUTION, message="Planning not initialized") + if not pose_targets: + return IKResult( + status=IKStatus.NO_SOLUTION, message="At least one pose target is required" + ) + + try: + stamped_targets = dict(pose_targets) + auxiliary_ids = tuple(auxiliary_group_ids) + group_ids = tuple(dict.fromkeys((*stamped_targets.keys(), *auxiliary_ids))) + target_groups = { + self._world_monitor.planning_groups.get(group_id): pose + for group_id, pose in stamped_targets.items() + } + auxiliary_groups = tuple( + self._world_monitor.planning_groups.get(group_id) for group_id in auxiliary_ids + ) + seed_state = seed + if seed_state is None: + selection = self._world_monitor.planning_groups.select(group_ids) + current = self._world_monitor.current_global_joint_state() + if not current.name and not current.position: + return IKResult(status=IKStatus.NO_SOLUTION, message="No joint state") + seed_state = filter_joint_state_to_selected_joints(current, selection.joint_names) + except (KeyError, ValueError) as exc: + return IKResult(status=IKStatus.NO_SOLUTION, message=str(exc)) + if seed_state is None: + return IKResult(status=IKStatus.NO_SOLUTION, message="No joint state") + return self._kinematics.solve_pose_targets( + world=self._world_monitor.world, + pose_targets=target_groups, + auxiliary_groups=auxiliary_groups, + seed=seed_state, + check_collision=check_collision, + ) + + @rpc + def inverse_kinematics_single( self, pose: Pose, robot_name: RobotName | None = None, - check_collision: bool = True, seed: JointState | None = None, + check_collision: bool = True, ) -> IKResult: - """Solve IK for a pose without planning a joint path. - - Args: - pose: Target end-effector pose - robot_name: Robot to solve for (required if multiple robots configured) - check_collision: Whether to reject IK candidates in collision - seed: Optional joint state to initialize local IK. Uses current state when omitted. - """ - if self._kinematics is None or self._world_monitor is None: + """Solve IK for one robot's unique pose-targetable planning group.""" + if self._world_monitor is None: return IKResult(status=IKStatus.NO_SOLUTION, message="Planning not initialized") robot = self._get_robot(robot_name) if robot is None: return IKResult(status=IKStatus.NO_SOLUTION, message="Robot not found") + selected_robot_name, _, _, _ = robot + try: + group_id = self._require_unique_pose_group_id_for_robot(selected_robot_name) + except ValueError as exc: + return IKResult(status=IKStatus.NO_SOLUTION, message=str(exc)) + target_pose = PoseStamped( + frame_id="world", + position=pose.position, + orientation=pose.orientation, + ) + return self.inverse_kinematics( + {group_id: target_pose}, seed=seed, check_collision=check_collision + ) - with self._lock: - if self._state not in (ManipulationState.IDLE, ManipulationState.COMPLETED): - return IKResult( - status=IKStatus.NO_SOLUTION, - message=f"Cannot solve IK while state is {self._state.name}", - ) - self._state = ManipulationState.PLANNING - - _, robot_id, _, _ = robot - seed_state = seed or self._world_monitor.get_current_joint_state(robot_id) - if seed_state is None: - self._state = ManipulationState.IDLE - return IKResult(status=IKStatus.NO_SOLUTION, message="No joint state") - - result = self._solve_ik_for_pose(robot_id, pose, seed_state, check_collision) - self._state = ManipulationState.COMPLETED if result.is_success() else ManipulationState.IDLE + @rpc + def solve_ik( + self, + pose: Pose, + robot_name: RobotName | None = None, + check_collision: bool = True, + seed: JointState | None = None, + ) -> IKResult: + """Compatibility wrapper for inverse_kinematics_single().""" + result = self.inverse_kinematics_single( + pose, + robot_name=robot_name, + seed=seed, + check_collision=check_collision, + ) if result.is_success(): logger.info(f"IK solved, error: {result.position_error:.4f}m") return result @rpc def plan_to_pose(self, pose: Pose, robot_name: RobotName | None = None) -> bool: - """Plan motion to pose. Use preview_path() then execute(). + """Plan motion to pose. Use preview_plan() then execute(). Args: pose: Target end-effector pose robot_name: Robot to plan for (required if multiple robots configured) """ - if self._kinematics is None or (r := self._begin_planning(robot_name)) is None: + if self._kinematics is None or self._world_monitor is None: return False - robot_name, robot_id = r - planning_epoch = self._planning_epoch - assert self._world_monitor # guaranteed by _begin_planning - - current = self._world_monitor.get_current_joint_state(robot_id) - if current is None: - return self._fail("No joint state") - - ik = self._solve_ik_for_pose(robot_id, pose, current, check_collision=True) - if not ik.is_success() or ik.joint_state is None: - return self._fail(f"IK failed: {ik.status.name}") + robot = self._get_robot(robot_name) + if robot is None: + return False + selected_robot_name, _, _, _ = robot + try: + group_id = self._require_unique_pose_group_id_for_robot(selected_robot_name) + except ValueError as exc: + logger.warning("Pose planning unavailable: %s", exc) + return False + return self.plan_to_pose_targets({group_id: pose}) - logger.info(f"IK solved, error: {ik.position_error:.4f}m") - return self._plan_path_only(robot_name, robot_id, ik.joint_state, planning_epoch) + @rpc + def plan_to_pose_targets( + self, + pose_targets: Mapping[PlanningGroupID | PlanningGroup, Pose], + auxiliary_groups: Sequence[PlanningGroupID | PlanningGroup] = (), + ) -> bool: + """Plan to one or more group pose targets with optional auxiliary groups.""" + return self.generate_plan_to_pose_targets(pose_targets, auxiliary_groups) is not None @rpc def plan_to_joints(self, joints: JointState, robot_name: RobotName | None = None) -> bool: - """Plan motion to joint config. Use preview_path() then execute(). + """Plan motion to joint config. Use preview_plan() then execute(). Args: joints: Target joint state (names + positions) robot_name: Robot to plan for (required if multiple robots configured) """ - if (r := self._begin_planning(robot_name)) is None: + robot = self._get_robot(robot_name) + if robot is None: + return False + selected_robot_name, _, _, _ = robot + logger.info( + f"Planning to joints for {selected_robot_name}: {[f'{j:.3f}' for j in joints.position]}" + ) + if self._world_monitor is None: + return False + group_id = self._world_monitor.planning_groups.default_group_id_for_robot( + selected_robot_name + ) + if group_id is None: + logger.error( + "Robot '%s' has no unique default planning group; use explicit group APIs", + selected_robot_name, + ) return False - robot_name, robot_id = r - planning_epoch = self._planning_epoch - logger.info(f"Planning to joints for {robot_name}: {[f'{j:.3f}' for j in joints.position]}") - return self._plan_path_only(robot_name, robot_id, joints, planning_epoch) + return self.plan_to_joint_targets({group_id: joints}) - def _plan_path_only( - self, - robot_name: RobotName, - robot_id: WorldRobotID, - goal: JointState, - planning_epoch: int, + @rpc + def plan_to_joint_targets( + self, joint_targets: Mapping[PlanningGroupID | PlanningGroup, JointState] ) -> bool: - """Plan path from current position to goal, store result.""" - assert self._world_monitor and self._planner # guaranteed by _begin_planning - self._dismiss_preview(robot_id) - start = self._world_monitor.get_current_joint_state(robot_id) - if start is None: - return self._fail("No joint state") - - # Trim goal to planner DOF (e.g. strip gripper joint from coordinator state) - planner_dof = len(start.position) - if len(goal.position) > planner_dof: - goal = JointState( - name=list(goal.name[:planner_dof]) if goal.name else [], - position=list(goal.position[:planner_dof]), - ) + """Plan to joint targets keyed by planning group.""" + return self.generate_plan_to_joint_targets(joint_targets) is not None + + def generate_plan_to_joint_targets( + self, joint_targets: Mapping[PlanningGroupID | PlanningGroup, JointState] + ) -> GeneratedPlan | None: + """Plan to joint targets and return the exact stored GeneratedPlan.""" + if self._world_monitor is None or self._planner is None: + return None + if not joint_targets: + logger.warning("At least one joint target is required") + return None - result = self._planner.plan_joint_path( - world=self._world_monitor.world, - robot_id=robot_id, - start=start, - goal=goal, - timeout=self.config.planning_timeout, + group_ids = tuple( + dict.fromkeys(planning_group_id_from_selector(group) for group in joint_targets) ) - if self._state != ManipulationState.PLANNING or planning_epoch != self._planning_epoch: - logger.info("Discarding cancelled planning result") - return False - if not result.is_success(): - return self._fail(f"Planning failed: {result.status.name}") + planning_epoch = self._begin_group_planning() + if planning_epoch is None: + return None - logger.info(f"Path: {len(result.path)} waypoints") - self._planned_paths[robot_name] = result.path + try: + selection = self._world_monitor.planning_groups.select(group_ids) + current = self._world_monitor.current_global_joint_state() + start = filter_joint_state_to_selected_joints(current, selection.joint_names) + except Exception as exc: + self._fail_planning_epoch(planning_epoch, f"Failed to resolve planning groups: {exc}") + return None - _, _, traj_gen = self._robots[robot_name] - # Convert JointState path to list of position lists for trajectory generator - traj = traj_gen.generate([list(state.position) for state in result.path]) - self._planned_trajectories[robot_name] = traj - logger.info(f"Trajectory: {traj.duration:.3f}s") + goal_names: list[str] = [] + goal_positions: list[float] = [] + for group, target in joint_targets.items(): + group_id = planning_group_id_from_selector(group) + try: + target_group = self._world_monitor.planning_groups.get(group_id) + target_global = joint_target_to_global_names(target_group, target) + except (KeyError, ValueError) as exc: + logger.error(str(exc)) + self._fail_planning_epoch(planning_epoch, f"Invalid joint target for '{group_id}'") + return None + goal_names.extend(target_global.name) + goal_positions.extend(target_global.position) - self._state = ManipulationState.COMPLETED - return True + goal = JointState(name=goal_names, position=goal_positions) + return self._plan_selected_path(group_ids, start, goal, planning_epoch) + + def generate_plan_to_pose_targets( + self, + pose_targets: Mapping[PlanningGroupID | PlanningGroup, Pose], + auxiliary_groups: Sequence[PlanningGroupID | PlanningGroup] = (), + ) -> GeneratedPlan | None: + """Plan to pose targets and return the exact stored GeneratedPlan.""" + if self._world_monitor is None or self._kinematics is None: + return None + if not pose_targets: + logger.warning("At least one pose target is required") + return None + stamped_targets = { + planning_group_id_from_selector(group): PoseStamped( + frame_id="world", + position=pose.position, + orientation=pose.orientation, + ) + for group, pose in pose_targets.items() + } + auxiliary_ids = tuple(planning_group_id_from_selector(group) for group in auxiliary_groups) + group_ids = tuple(dict.fromkeys((*stamped_targets.keys(), *auxiliary_ids))) + planning_epoch = self._begin_group_planning() + if planning_epoch is None: + return None + try: + selection = self._world_monitor.planning_groups.select(group_ids) + current = self._world_monitor.current_global_joint_state() + start = filter_joint_state_to_selected_joints(current, selection.joint_names) + except Exception as exc: + self._fail_planning_epoch(planning_epoch, f"Failed to resolve planning groups: {exc}") + return None + ik = self.inverse_kinematics( + pose_targets=stamped_targets, + auxiliary_group_ids=auxiliary_ids, + seed=start, + ) + if not ik.is_success() or ik.joint_state is None: + self._fail_planning_epoch(planning_epoch, f"IK failed: {ik.status.name}") + return None + logger.info(f"IK solved, error: {ik.position_error:.4f}m") + return self._plan_selected_path(group_ids, start, ik.joint_state, planning_epoch) @rpc def preview_path( @@ -662,54 +1071,43 @@ def preview_path( robot_name: RobotName | None = None, target_fps: float = 30.0, ) -> bool: - """Preview the planned path in the visualizer. + """Compatibility wrapper for preview_plan(). Args: - duration: Total animation duration in seconds. Uses trajectory duration if None. - robot_name: Robot to preview (required if multiple robots configured) - target_fps: Nominal preview update rate. Set <= 0 to use planned waypoints directly. + duration: Total animation duration in seconds. Defaults to one second. + robot_name: Compatibility affected-robot validation; does not filter the preview. + target_fps: Deprecated compatibility argument; shared-clock previews use plan waypoints. """ - if self._world_monitor is None: - return False + return self.preview_plan(None, duration, robot_name, target_fps) - robot = self._get_robot(robot_name) - if robot is None: + @rpc + def preview_plan( + self, + plan: GeneratedPlan | None = None, + duration: float | None = None, + robot_name: RobotName | None = None, + target_fps: float = 30.0, + ) -> bool: + """Preview a complete generated plan in the visualizer.""" + if plan is None and self._execution_runtime is not None: + ready = self._execution_runtime.snapshot().ready_plan + plan = ready.generated_plan if isinstance(ready, PreparedPlan) else None + if plan is None or not plan.path: + logger.warning("No generated plan to preview") return False - robot_name, robot_id, _, _ = robot - - planned_path = self._planned_paths.get(robot_name) - if planned_path is None or len(planned_path) == 0: - logger.warning(f"No planned path to preview for {robot_name}") + try: + assert self._world_monitor is not None + affected = self._world_monitor.planning_groups.select(plan.group_ids).robot_names + except Exception as exc: + logger.error("Generated plan cannot be resolved: %s", exc) return False - - if duration is None: - trajectory = self._planned_trajectories.get(robot_name) - animation_duration = trajectory.duration if trajectory is not None else 3.0 - else: - trajectory = self._planned_trajectories.get(robot_name) - animation_duration = duration - - interpolated = list(planned_path) - if trajectory is not None and target_fps > 0 and animation_duration > 0: - times = np.array( - [point.time_from_start for point in trajectory.points], dtype=np.float64 - ) - positions = np.array([point.positions for point in trajectory.points], dtype=np.float64) - if len(times) > 1 and positions.ndim == 2 and times[-1] > times[0]: - frame_count = int(np.ceil(animation_duration * target_fps)) + 1 - sample_times = np.linspace(times[0], times[-1], frame_count) - joint_names = trajectory.joint_names or planned_path[0].name - sampled_positions = np.column_stack( - [ - np.interp(sample_times, times, positions[:, joint]) - for joint in range(positions.shape[1]) - ] - ) - interpolated = [ - JointState(name=joint_names, position=position.tolist()) - for position in sampled_positions - ] - self._world_monitor.animate_path(robot_id, interpolated, animation_duration) + if robot_name is not None: + if robot_name not in affected: + logger.error("Generated plan does not affect robot '%s'", robot_name) + return False + if self._world_monitor is None: + return False + self._world_monitor.animate_trajectory(plan.trajectory, duration) return True @rpc @@ -719,13 +1117,11 @@ def has_planned_path(self) -> bool: Returns: True if a path is planned and ready """ - robot = self._get_robot() - if robot is None: + snapshot = self._execution_snapshot() + if snapshot is None: return False - robot_name, _, _, _ = robot - - path = self._planned_paths.get(robot_name) - return path is not None and len(path) > 0 + ready = snapshot.ready_plan + return isinstance(ready, PreparedPlan) and bool(ready.generated_plan.path) @rpc def get_visualization_url(self) -> str | None: @@ -745,15 +1141,20 @@ def clear_planned_path(self) -> bool: Returns: True if cleared """ - if self._world_monitor is None: - return False - robot = self._get_robot() - if robot is None: - return False - robot_name, _, _, _ = robot - - self._planned_paths.pop(robot_name, None) - self._planned_trajectories.pop(robot_name, None) + plan = None + if self._execution_runtime is not None: + snapshot = self._execution_runtime.snapshot() + if snapshot.state == LifecycleState.PLANNING: + return self._execution_runtime.cancel_planning().accepted + ready = snapshot.ready_plan + plan = ready.generated_plan if isinstance(ready, PreparedPlan) else None + result = self._execution_runtime.clear_ready_plan() + if not result.accepted: + return False + if plan is not None: + # Preserve the group selection until the public visualization + # transaction has invalidated and hidden its preview. + self._dismiss_preview(plan.group_ids) return True @rpc @@ -766,7 +1167,23 @@ def list_robots(self) -> list[str]: return list(self._robots.keys()) @rpc - def get_robot_info(self, robot_name: RobotName | None = None) -> dict[str, Any] | None: + def list_planning_groups(self) -> list[PlanningGroup]: + """Return all configured planning groups.""" + if self._world_monitor is None: + return [] + return list(self._world_monitor.planning_groups.list()) + + def get_current_joint_state(self, robot_name: RobotName) -> JointState | None: + """Return the named robot's current local joint state with names.""" + if self._world_monitor is None: + return None + robot_id = self.robot_id_for_name(robot_name) + if robot_id is None: + return None + return self._world_monitor.get_current_joint_state(robot_id) + + @rpc + def get_robot_info(self, robot_name: RobotName | None = None) -> RobotInfoPayload | None: """Get information about a robot. Args: @@ -780,12 +1197,22 @@ def get_robot_info(self, robot_name: RobotName | None = None) -> dict[str, Any] return None robot_name, robot_id, config, _ = robot + planning_groups = ( + list(self._world_monitor.planning_groups.groups_for_robot(robot_name)) + if self._world_monitor is not None + else [] + ) + try: + end_effector_link = config.end_effector_link + except ValueError: + end_effector_link = None return { "name": config.name, "world_robot_id": robot_id, "joint_names": config.joint_names, - "end_effector_link": config.end_effector_link, + "planning_groups": planning_groups, + "end_effector_link": end_effector_link, "base_link": config.base_link, "max_velocity": config.max_velocity, "max_acceleration": config.max_acceleration, @@ -908,20 +1335,6 @@ def evaluate_pose_target(self, pose: Pose, robot_name: RobotName) -> TargetEvalu "collision_free": collision_free, } - def get_planned_path(self, robot_name: RobotName) -> JointPath | None: - """Return a copy of the stored planned path for visualization.""" - path = self._planned_paths.get(robot_name) - if path is None: - return None - return [JointState(point) for point in path] - - def get_planned_trajectory_duration(self, robot_name: RobotName) -> float | None: - """Return the stored planned trajectory duration for visualization.""" - trajectory = self._planned_trajectories.get(robot_name) - if trajectory is None: - return None - return float(trajectory.duration) - @rpc def set_init_joints(self, joint_state: JointState, robot_name: RobotName | None = None) -> bool: """Set the init joint state. @@ -964,97 +1377,81 @@ def set_init_joints_to_current(self, robot_name: RobotName | None = None) -> boo ) return True - def _get_coordinator_client(self) -> RPCClient | None: - """Get or create coordinator RPC client (lazy init).""" - if not any( - c.coordinator_task_name or c.gripper_hardware_id for _, c, _ in self._robots.values() - ): - return None - if self._coordinator_client is None: - from dimos.control.coordinator import ControlCoordinator - from dimos.core.rpc_client import RPCClient - - self._coordinator_client = RPCClient(None, ControlCoordinator) - return self._coordinator_client - - def _translate_trajectory_to_coordinator( - self, - trajectory: JointTrajectory, - robot_config: RobotModelConfig, - ) -> JointTrajectory: - """Translate trajectory joint names from URDF to coordinator namespace. - - Args: - trajectory: Trajectory with URDF joint names - robot_config: Robot config with joint name mapping - - Returns: - Trajectory with coordinator joint names - """ - if not robot_config.joint_name_mapping: - return trajectory # No translation needed - - # Translate joint names - coordinator_names = [ - robot_config.get_coordinator_joint_name(j) for j in trajectory.joint_names - ] - - # Create new trajectory with translated names - # Note: duration is computed automatically from points in JointTrajectory.__init__ - return JointTrajectory( - joint_names=coordinator_names, - points=trajectory.points, - timestamp=trajectory.timestamp, - ) - @rpc - def execute(self, robot_name: RobotName | None = None) -> bool: - """Execute planned trajectory via ControlCoordinator.""" - if (robot := self._get_robot(robot_name)) is None: + def execute(self, plan: GeneratedPlan | None = None) -> bool: + """Execute a complete plan, or the runtime-owned READY plan when omitted.""" + handle = self._submit_execution(plan) + if handle is None or self._execution_runtime is None: return False - robot_name, _, config, _ = robot - - if (traj := self._planned_trajectories.get(robot_name)) is None: - logger.warning("No planned trajectory") + dispatch = self._execution_runtime.wait_for_dispatch(handle, timeout=float("inf")) + if not dispatch.accepted or dispatch.value is None: return False - if not config.coordinator_task_name: - logger.error(f"No coordinator_task_name for '{robot_name}'") + return bool(dispatch.value.outcome == Outcome.ACCEPTED) + + @rpc + def execute_plan(self, plan: GeneratedPlan | None = None) -> bool: + """Execute a complete generated plan through the runtime owner.""" + handle = self._submit_execution(plan) + if handle is None or self._execution_runtime is None: return False - if (client := self._get_coordinator_client()) is None: - logger.error("No coordinator client") + dispatch = self._execution_runtime.wait_for_dispatch(handle, timeout=float("inf")) + if not dispatch.accepted or dispatch.value is None: return False + return bool(dispatch.value.outcome == Outcome.ACCEPTED) - translated = self._translate_trajectory_to_coordinator(traj, config) - logger.info( - f"Executing: task='{config.coordinator_task_name}', {len(translated.points)} pts, {translated.duration:.2f}s" - ) - - self._state = ManipulationState.EXECUTING - result = client.task_invoke( - config.coordinator_task_name, "execute", {"trajectory": translated} - ) - if result: - logger.info("Trajectory accepted") - self._state = ManipulationState.COMPLETED - return True + def _submit_execution(self, plan: GeneratedPlan | None) -> OperationHandle | None: + """Submit one explicit or exactly-ready generated plan to the runtime.""" + runtime = self._execution_runtime + topology = self._execution_topology + if runtime is None or topology is None: + return None + if plan is None: + result = runtime.execute_ready() else: - return self._fail("Coordinator rejected trajectory") + try: + prepared = prepare_generated_plan(plan, topology) + except Exception as exc: + logger.warning("Failed to prepare execution: %s", exc) + return None + result = runtime.execute_explicit(prepared) + return result.value if result.accepted else None @rpc def get_trajectory_status(self, robot_name: RobotName | None = None) -> dict[str, Any] | None: - """Get trajectory execution status via coordinator task_invoke.""" - if (robot := self._get_robot(robot_name)) is None: - return None - _, _, config, _ = robot - if not config.coordinator_task_name or (client := self._get_coordinator_client()) is None: - return None - try: - state = client.task_invoke(config.coordinator_task_name, "get_state", {}) - if state is not None: - return {"state": int(state), "task": config.coordinator_task_name} - return None - except Exception: + """Return a snapshot-only runtime status, optionally scoped to a robot.""" + snapshot = self._execution_snapshot() + if snapshot is None: return None + operation = snapshot.operation + tasks: list[dict[str, Any]] = [] + if operation is not None: + for task in operation.tasks: + task_robot = task.entry.robot_name + if task_robot is None and self._execution_topology is not None: + task_robot, _ = self._execution_topology.route_for_group( + task.entry.planning_group + ) + if robot_name is None or task_robot == robot_name: + tasks.append( + { + "task_id": task.task_id, + "task_name": task.task_name, + "robot_name": task_robot, + "activity": task.activity.value, + "cancel_required": task.cancel_required, + "reset_required": task.reset_required, + } + ) + return { + "state": snapshot.state.name, + "diagnostic": snapshot.diagnostic, + "fault": snapshot.fault, + "ready_plan_id": snapshot.ready_plan_id, + "operation_id": ( + snapshot.operation.handle.operation_id if snapshot.operation else None + ), + "tasks": tasks, + } @property def world_monitor(self) -> WorldMonitor | None: @@ -1126,11 +1523,10 @@ def _set_gripper_position(self, position: float, robot_name: RobotName | None = hw_id = self._get_gripper_hardware_id(robot_name) if hw_id is None: return False - client = self._get_coordinator_client() - if client is None: - logger.error("No coordinator client for gripper control") + if self._execution_runtime is None: return False - return bool(client.set_gripper_position(hw_id, position)) + result = self._execution_runtime.set_gripper_position(hw_id, position) + return bool(result.accepted and result.value == Outcome.ACCEPTED) @rpc def get_gripper(self, robot_name: RobotName | None = None) -> float | None: @@ -1142,11 +1538,10 @@ def get_gripper(self, robot_name: RobotName | None = None) -> float | None: hw_id = self._get_gripper_hardware_id(robot_name) if hw_id is None: return None - client = self._get_coordinator_client() - if client is None: + if self._execution_runtime is None: return None - result = client.get_gripper_position(hw_id) - return float(result) if result is not None else None + result = self._execution_runtime.get_gripper_position(hw_id) + return result.value if result.accepted else None @skill def set_gripper( @@ -1184,67 +1579,6 @@ def close_gripper(self, robot_name: str | None = None) -> SkillResult[Manipulati return SkillResult.ok("Gripper closed") return SkillResult.fail("GRIPPER_FAILED", "Failed to close gripper") - def _wait_for_trajectory_completion( - self, robot_name: RobotName | None = None, timeout: float = 60.0, poll_interval: float = 0.2 - ) -> bool: - """Wait for trajectory execution to complete. - - Polls the coordinator task state via task_invoke. Falls back to waiting - for the trajectory duration if the coordinator is unavailable. - - Args: - robot_name: Robot to monitor - timeout: Maximum wait time in seconds - poll_interval: Time between status checks - - Returns: - True if trajectory completed successfully - """ - robot = self._get_robot(robot_name) - if robot is None: - return True - rname, _, config, _ = robot - client = self._get_coordinator_client() - - if client is None or not config.coordinator_task_name: - # No coordinator — wait for trajectory duration as fallback - traj = self._planned_trajectories.get(rname) - if traj is not None: - logger.info(f"No coordinator status — waiting {traj.duration:.1f}s for trajectory") - time.sleep(traj.duration + 0.5) - return True - - # Poll task state via task_invoke - start = time.time() - while (time.time() - start) < timeout: - try: - state = client.task_invoke(config.coordinator_task_name, "get_state", {}) - # TrajectoryState is an IntEnum: IDLE=0, EXECUTING=1, COMPLETED=2, ABORTED=3, FAULT=4 - if state is not None: - state_val = int(state) - if state_val in (0, 2): # IDLE or COMPLETED - return True - if state_val in (3, 4): # ABORTED or FAULT - logger.warning(f"Trajectory failed: state={state}") - return False - # state_val == 1 means EXECUTING, keep polling - else: - # task_invoke returned None — task not found, assume done - return True - except Exception: - # Fallback: wait for trajectory duration - traj = self._planned_trajectories.get(rname) - if traj is not None: - remaining = traj.duration - (time.time() - start) - if remaining > 0: - logger.info(f"Status poll failed — waiting {remaining:.1f}s for trajectory") - time.sleep(remaining + 0.5) - return True - time.sleep(poll_interval) - - logger.warning(f"Trajectory execution timed out after {timeout}s") - return False - def _lift_if_low( self, robot_name: RobotName | None = None, min_z: float = 0.05 ) -> SkillResult[ManipulationSkillError]: @@ -1276,13 +1610,50 @@ def _preview_execute_wait( self.preview_path(preview_duration, robot_name) logger.info("Executing trajectory...") - if not self.execute(robot_name): + if self._execution_runtime is None: + return SkillResult.fail("EXECUTION_FAILED", "Execution runtime is not started") + ready = self._execution_runtime.snapshot().ready_plan + plan = ready.generated_plan if isinstance(ready, PreparedPlan) else None + handle = self._submit_execution(plan) + if handle is None: return SkillResult.fail("EXECUTION_FAILED", "Trajectory execution failed") - if not self._wait_for_trajectory_completion(robot_name): - return SkillResult.fail("EXECUTION_TIMEOUT", "Trajectory execution timed out") - - return SkillResult.ok() + dispatch = self._execution_runtime.wait_for_dispatch(handle, timeout=float("inf")) + if not dispatch.accepted or dispatch.value is None: + return SkillResult.fail( + "EXECUTION_DISPATCH_FAILED", + dispatch.diagnostic or "Execution dispatch failed", + ) + if dispatch.value.outcome == Outcome.REJECTED: + return SkillResult.fail( + "EXECUTION_REJECTED", dispatch.value.diagnostic or "Trajectory rejected" + ) + if dispatch.value.outcome != Outcome.ACCEPTED: + return SkillResult.fail( + "EXECUTION_DISPATCH_FAILED", + dispatch.value.diagnostic or "Unexpected dispatch outcome", + ) + terminal = self._execution_runtime.wait_for_terminal(handle, timeout=float("inf")) + if not terminal.accepted or terminal.value is None: + return SkillResult.fail( + "EXECUTION_FAILED", terminal.diagnostic or "Execution terminal result unavailable" + ) + result = terminal.value + if result.outcome == Outcome.COMPLETED: + return SkillResult.ok() + if result.outcome == Outcome.CANCELLED: + return SkillResult.fail( + "EXECUTION_CANCELLED", result.diagnostic or "Trajectory cancelled" + ) + if result.outcome == Outcome.FAILED: + return SkillResult.fail("EXECUTION_FAILED", result.diagnostic or "Trajectory failed") + if result.outcome == Outcome.REJECTED: + return SkillResult.fail( + "EXECUTION_REJECTED", result.diagnostic or "Trajectory rejected" + ) + return SkillResult.fail( + "EXECUTION_FAILED", result.diagnostic or "Execution terminal result was unresolved" + ) @skill def get_robot_state(self, robot_name: str | None = None) -> SkillResult[ManipulationSkillError]: @@ -1524,14 +1895,31 @@ def stop(self) -> None: """Stop the manipulation module.""" logger.info("Stopping ManipulationModule") - # Stop TF thread - if self._tf_thread is not None: - self._tf_stop_event.set() - self._tf_thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) - self._tf_thread = None - - # Stop world monitor (includes visualization thread) - if self._world_monitor is not None: - self._world_monitor.stop_all_monitors() - - super().stop() + try: + try: + runtime = self._execution_runtime + if runtime is not None: + shutdown = runtime.shutdown(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) + if ( + shutdown.snapshot is not None + and shutdown.snapshot.shutdown == ShutdownState.CLOSED + ): + self._execution_runtime = None + except Exception: + logger.exception("Failed to shut down execution runtime") + finally: + try: + self._stop_tf() + except Exception: + logger.exception("Failed to stop TF thread") + finally: + try: + if self._world_monitor is not None: + self._world_monitor.stop_all_monitors() + except Exception: + logger.exception("Failed to stop world monitor") + finally: + try: + super().stop() + except Exception: + logger.exception("Failed to stop superclass") diff --git a/dimos/manipulation/planning/examples/manipulation_client.py b/dimos/manipulation/planning/examples/manipulation_client.py index 1185f28f21..5c9718ecb1 100644 --- a/dimos/manipulation/planning/examples/manipulation_client.py +++ b/dimos/manipulation/planning/examples/manipulation_client.py @@ -25,10 +25,14 @@ Available functions: joints() Get current joint positions ee() Get end-effector pose - state() Get module state (IDLE, PLANNING, EXECUTING, ...) + groups() List explicit planning groups + state() Get module state (IDLE, PLANNING, RUNNING, ...) ik_pose(x,y,z, seed_joints=None) Solve IK only, without path planning + ik_group_pose(group_id,x,y,z) Solve IK for an explicit planning group plan(joints) Plan to joint configuration, e.g. plan([0.1]*7) + plan_group(group_id,joints) Plan to an explicit planning-group joint target plan_pose(x,y,z) Plan to Cartesian pose + plan_group_pose(group_id,x,y,z) Plan to an explicit planning-group pose target preview(duration=None) Preview planned path in Meshcat execute() Execute planned trajectory via coordinator home() Move to home position @@ -50,8 +54,10 @@ from dimos.core.rpc_client import RPCClient from dimos.manipulation.manipulation_module import ManipulationModule +from dimos.manipulation.planning.groups.models import PlanningGroup from dimos.manipulation.planning.spec.models import IKResult 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.Vector3 import Vector3 from dimos.msgs.sensor_msgs.JointState import JointState @@ -80,6 +86,21 @@ def plan(target_joints: list[float], robot_name: str | None = None) -> bool: return _client.plan_to_joints(js, robot_name) +def groups() -> list[PlanningGroup]: + """List explicit planning groups available for group APIs.""" + return _client.list_planning_groups() + + +def plan_group(group_id: str, target_joints: list[float] | JointState) -> bool: + """Plan to a joint target for an explicit planning group.""" + target = ( + target_joints + if isinstance(target_joints, JointState) + else JointState(position=target_joints) + ) + return _client.plan_to_joint_targets({group_id: target}) + + def _make_target_pose( x: float, y: float, @@ -124,7 +145,6 @@ def ik_pose( pitch: float | None = None, yaw: float | None = None, robot_name: str | None = None, - check_collision: bool = True, seed_joints: list[float] | JointState | None = None, ) -> IKResult: """Solve IK for a Cartesian pose without path planning. @@ -137,13 +157,32 @@ def ik_pose( pitch: Optional target pitch. Preserves current orientation if omitted. yaw: Optional target yaw. Preserves current orientation if omitted. robot_name: Robot to solve for when multiple robots are configured. - check_collision: Whether to reject IK candidates in collision. seed_joints: Optional initial joint configuration for local IK. Pass either a list of joint positions in robot joint order or a named JointState. """ target = _make_target_pose(x, y, z, roll, pitch, yaw, robot_name) seed = _make_seed_joint_state(seed_joints, robot_name) - return _client.solve_ik(target, robot_name, check_collision, seed) + return _client.inverse_kinematics_single(target, robot_name, seed) + + +def ik_group_pose( + group_id: str, + x: float, + y: float, + z: float, + roll: float | None = None, + pitch: float | None = None, + yaw: float | None = None, + seed: JointState | None = None, +) -> IKResult: + """Solve IK for an explicit planning group pose target.""" + target = _make_target_pose(x, y, z, roll, pitch, yaw) + stamped = PoseStamped( + frame_id="world", + position=target.position, + orientation=target.orientation, + ) + return _client.inverse_kinematics({group_id: stamped}, seed=seed) def plan_pose( @@ -160,18 +199,31 @@ def plan_pose( return _client.plan_to_pose(target, robot_name) +def plan_group_pose( + group_id: str, + x: float, + y: float, + z: float, + roll: float | None = None, + pitch: float | None = None, + yaw: float | None = None, +) -> bool: + """Plan to a Cartesian pose for an explicit planning group.""" + target = _make_target_pose(x, y, z, roll, pitch, yaw) + return _client.plan_to_pose_targets({group_id: target}) + + def preview( duration: float | None = None, robot_name: str | None = None, - target_fps: float = 30.0, ) -> bool: - """Preview planned path in Meshcat.""" - return _client.preview_path(duration, robot_name, target_fps) + """Preview the last generated plan in the visualizer.""" + return _client.preview_plan(None, duration, robot_name) -def execute(robot_name: str | None = None) -> bool: +def execute() -> bool: """Execute planned trajectory via coordinator.""" - return _client.execute(robot_name) + return _client.execute() def home(robot_name: str | None = None) -> bool: @@ -181,7 +233,7 @@ def home(robot_name: str | None = None) -> bool: home_joints = _client.get_robot_info(robot_name).get("home_joints", [0.0] * 7) success = _client.plan_to_joints(JointState(position=home_joints), robot_name) if success: - return _client.execute(robot_name) + return _client.execute() return False diff --git a/dimos/manipulation/planning/groups/discovery.py b/dimos/manipulation/planning/groups/discovery.py new file mode 100644 index 0000000000..49cad388ff --- /dev/null +++ b/dimos/manipulation/planning/groups/discovery.py @@ -0,0 +1,352 @@ +# 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. + +"""Planning group discovery from SRDF or conservative model fallback.""" + +from __future__ import annotations + +import itertools +from pathlib import Path +import xml.etree.ElementTree as ET + +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition +from dimos.robot.model_parser import JointDescription, ModelDescription +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +FALLBACK_PLANNING_GROUP_NAME = "manipulator" + + +class PlanningGroupDiscoveryError(ValueError): + """Raised when planning groups cannot be discovered for a model.""" + + +def discover_planning_group_definitions( + *, + robot_name: str, + model_path: Path, + model: ModelDescription, + controllable_joint_names: list[str], + srdf_path: Path | None = None, +) -> list[PlanningGroupDefinition]: + """Discover planning groups from SRDF or fallback generation. + + Precedence is explicit SRDF path, conservative auto-discovery with warning, + then fallback generation from the controllable joint set. + """ + resolved_srdf_path = _resolve_srdf_path(model_path, srdf_path) + if resolved_srdf_path is not None: + groups = parse_srdf_planning_groups( + resolved_srdf_path, + model=model, + controllable_joint_names=controllable_joint_names, + ) + if groups: + return groups + logger.warning( + f"No supported planning groups found in SRDF {resolved_srdf_path} " + f"for robot {robot_name}; trying fallback generation" + ) + + return [ + generate_fallback_planning_group( + model=model, + controllable_joint_names=controllable_joint_names, + ) + ] + + +def parse_srdf_planning_groups( + srdf_path: Path, + *, + model: ModelDescription, + controllable_joint_names: list[str], +) -> list[PlanningGroupDefinition]: + """Extract supported SRDF planning group definitions. + + Supported forms are a single ```` + child or an ordered list of ```` children. Other forms, + including SRDF ```` metadata, are ignored for planning group + extraction. This is intentionally a minimal SRDF group extractor rather + than a full SRDF parser; adopting a ROS/MoveIt parser such as srdfdom would + add substantial dependency overhead for this narrow subset. + """ + root = ET.parse(srdf_path).getroot() + groups: list[PlanningGroupDefinition] = [] + for group_elem in root.findall("group"): + group_name = group_elem.get("name") + if not group_name: + logger.warning(f"Skipping SRDF group without a name in {srdf_path}") + continue + + children = [child for child in list(group_elem) if isinstance(child.tag, str)] + chain_children = [child for child in children if child.tag == "chain"] + joint_children = [child for child in children if child.tag == "joint"] + unsupported_children = [child for child in children if child.tag not in {"chain", "joint"}] + + if len(chain_children) == 1 and not joint_children and not unsupported_children: + definition = _parse_chain_group( + group_name, + chain_children[0], + model=model, + controllable_joint_names=controllable_joint_names, + srdf_path=srdf_path, + ) + elif joint_children and len(joint_children) == len(children): + definition = _parse_joint_list_group( + group_name, + joint_children, + model=model, + controllable_joint_names=controllable_joint_names, + srdf_path=srdf_path, + ) + else: + child_tags = [child.tag for child in children] + logger.warning( + f"Skipping unsupported SRDF planning group {group_name} in " + f"{srdf_path} with child tags {child_tags}" + ) + definition = None + + if definition is not None: + groups.append(definition) + + return groups + + +def generate_fallback_planning_group( + *, + model: ModelDescription, + controllable_joint_names: list[str], +) -> PlanningGroupDefinition: + """Generate one conservative fallback planning group named ``manipulator``.""" + ordered_joints = _validate_and_order_serial_joints(model, controllable_joint_names) + while ordered_joints and ordered_joints[-1].type == "prismatic": + removed = ordered_joints.pop() + logger.warning( + f"Excluding terminal prismatic joint {removed.name} from " + f"fallback planning group {FALLBACK_PLANNING_GROUP_NAME}" + ) + + if not ordered_joints: + raise PlanningGroupDiscoveryError( + "Fallback planning group generation removed all candidate joints; provide SRDF" + ) + + return PlanningGroupDefinition( + name=FALLBACK_PLANNING_GROUP_NAME, + joint_names=tuple(joint.name for joint in ordered_joints), + base_link=ordered_joints[0].parent_link, + tip_link=ordered_joints[-1].child_link, + source="fallback", + ) + + +def _resolve_srdf_path(model_path: Path, srdf_path: Path | None) -> Path | None: + if srdf_path is not None: + if srdf_path.exists(): + return srdf_path + raise FileNotFoundError(f"SRDF file not found: {srdf_path}") + + for candidate in _srdf_auto_discovery_candidates(model_path): + if candidate.exists(): + logger.warning(f"Auto-discovered SRDF at {candidate}") + return candidate + return None + + +def _srdf_auto_discovery_candidates(model_path: Path) -> list[Path]: + candidates: list[Path] = [] + name = model_path.name + if name.endswith(".urdf.xacro"): + candidates.append(model_path.with_name(name.removesuffix(".urdf.xacro") + ".srdf")) + elif model_path.suffix: + candidates.append(model_path.with_suffix(".srdf")) + candidates.append(model_path.parent / "config" / "robot.srdf") + candidates.append(model_path.parent.parent / "config" / "robot.srdf") + return list(dict.fromkeys(candidates)) + + +def _parse_chain_group( + group_name: str, + chain_elem: ET.Element, + *, + model: ModelDescription, + controllable_joint_names: list[str], + srdf_path: Path, +) -> PlanningGroupDefinition | None: + base_link = chain_elem.get("base_link") + tip_link = chain_elem.get("tip_link") + if not base_link or not tip_link: + logger.warning( + f"Skipping SRDF chain group {group_name} in {srdf_path} because " + "base_link or tip_link is missing" + ) + return None + + try: + ordered_joints = _ordered_joints_between_links(model, base_link, tip_link) + controlled_joints = [joint for joint in ordered_joints if joint.type != "fixed"] + _validate_controllable(group_name, controlled_joints, controllable_joint_names) + except PlanningGroupDiscoveryError as exc: + logger.warning(f"Skipping SRDF chain group {group_name} in {srdf_path}: {exc}") + return None + + return PlanningGroupDefinition( + name=group_name, + joint_names=tuple(joint.name for joint in controlled_joints), + base_link=base_link, + tip_link=tip_link, + ) + + +def _parse_joint_list_group( + group_name: str, + joint_children: list[ET.Element], + *, + model: ModelDescription, + controllable_joint_names: list[str], + srdf_path: Path, +) -> PlanningGroupDefinition | None: + joint_names = [child.get("name", "") for child in joint_children] + if any(not name for name in joint_names): + logger.warning( + f"Skipping SRDF joint-list group {group_name} in {srdf_path} with empty joint name" + ) + return None + try: + ordered_joints = _validate_ordered_serial_joints(model, joint_names) + _validate_controllable(group_name, ordered_joints, controllable_joint_names) + except PlanningGroupDiscoveryError as exc: + logger.warning(f"Skipping SRDF joint-list group {group_name} in {srdf_path}: {exc}") + return None + + return PlanningGroupDefinition( + name=group_name, + joint_names=tuple(joint.name for joint in ordered_joints), + base_link=ordered_joints[0].parent_link, + tip_link=ordered_joints[-1].child_link, + ) + + +def _ordered_joints_between_links( + model: ModelDescription, + base_link: str, + tip_link: str, +) -> list[JointDescription]: + joints_by_parent: dict[str, list[JointDescription]] = {} + for joint in model.joints: + joints_by_parent.setdefault(joint.parent_link, []).append(joint) + + ordered_joints: list[JointDescription] = [] + current_link = base_link + visited_links = {base_link} + while current_link != tip_link: + children = joints_by_parent.get(current_link, []) + if len(children) != 1: + raise PlanningGroupDiscoveryError( + f"chain from {base_link} to {tip_link} is branching or disconnected at {current_link}" + ) + joint = children[0] + ordered_joints.append(joint) + current_link = joint.child_link + if current_link in visited_links: + raise PlanningGroupDiscoveryError("chain contains a cycle") + visited_links.add(current_link) + + return ordered_joints + + +def _validate_ordered_serial_joints( + model: ModelDescription, + joint_names: list[str], +) -> list[JointDescription]: + ordered_joints: list[JointDescription] = [] + for joint_name in joint_names: + joint = model.get_joint(joint_name) + if joint is None: + raise PlanningGroupDiscoveryError(f"joint {joint_name} does not exist in model") + if joint.type == "fixed": + raise PlanningGroupDiscoveryError(f"joint {joint_name} is fixed") + ordered_joints.append(joint) + + if not ordered_joints: + raise PlanningGroupDiscoveryError("planning group contains no joints") + + for previous, current in itertools.pairwise(ordered_joints): + if previous.child_link != current.parent_link: + raise PlanningGroupDiscoveryError( + f"joints {previous.name} and {current.name} are not adjacent in a serial chain" + ) + return ordered_joints + + +def _validate_and_order_serial_joints( + model: ModelDescription, + joint_names: list[str], +) -> list[JointDescription]: + if not joint_names: + raise PlanningGroupDiscoveryError("fallback requires at least one controllable joint") + + joints: list[JointDescription] = [] + for joint_name in joint_names: + joint = model.get_joint(joint_name) + if joint is None: + raise PlanningGroupDiscoveryError(f"joint {joint_name} does not exist in model") + if joint.type == "fixed": + raise PlanningGroupDiscoveryError(f"joint {joint_name} is fixed") + joints.append(joint) + + by_parent = {joint.parent_link: joint for joint in joints} + by_child = {joint.child_link: joint for joint in joints} + if len(by_parent) != len(joints) or len(by_child) != len(joints): + raise PlanningGroupDiscoveryError("controllable joints branch or merge; provide SRDF") + + starts = [joint for joint in joints if joint.parent_link not in by_child] + ends = [joint for joint in joints if joint.child_link not in by_parent] + if len(starts) != 1 or len(ends) != 1: + raise PlanningGroupDiscoveryError( + "controllable joints are disconnected or cyclic; provide SRDF" + ) + + ordered_joints: list[JointDescription] = [] + current = starts[0] + while True: + ordered_joints.append(current) + next_joint = by_parent.get(current.child_link) + if next_joint is None: + break + current = next_joint + + if len(ordered_joints) != len(joints): + raise PlanningGroupDiscoveryError("controllable joints are disconnected; provide SRDF") + return ordered_joints + + +def _validate_controllable( + group_name: str, + joints: list[JointDescription], + controllable_joint_names: list[str], +) -> None: + if not joints: + raise PlanningGroupDiscoveryError( + f"planning group {group_name} contains no controllable joints" + ) + controllable = set(controllable_joint_names) + missing = [joint.name for joint in joints if joint.name not in controllable] + if missing: + raise PlanningGroupDiscoveryError( + f"planning group {group_name} includes joints outside controllable set: {missing}" + ) diff --git a/dimos/manipulation/planning/groups/identifiers.py b/dimos/manipulation/planning/groups/identifiers.py new file mode 100644 index 0000000000..db9d990081 --- /dev/null +++ b/dimos/manipulation/planning/groups/identifiers.py @@ -0,0 +1,112 @@ +# 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. + +"""Planning-group and global-joint identifier helpers.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from dimos.manipulation.planning.spec.models import ( + GlobalJointName, + LocalModelJointName, + PlanningGroupID, + RobotName, +) + + +def assert_valid_robot_name(robot_name: RobotName) -> None: + """Validate a robot name for delimiter-based public IDs.""" + if not robot_name or "/" in robot_name: + raise ValueError(f"Invalid robot name: {robot_name!r}") + + +def assert_valid_local_joint_name(local_joint_name: LocalModelJointName) -> None: + """Validate a local model joint name for delimiter-based global joint names.""" + if not local_joint_name or "/" in local_joint_name: + raise ValueError(f"Invalid local joint name: {local_joint_name!r}") + + +def assert_local_joint_names(names: Sequence[LocalModelJointName]) -> None: + """Validate that names are local model joint names, not global joint names.""" + for name in names: + assert_valid_local_joint_name(name) + + +def make_planning_group_id(robot_name: RobotName, group_name: str) -> PlanningGroupID: + """Build a public planning group ID.""" + assert_valid_robot_name(robot_name) + if not group_name or "/" in group_name: + raise ValueError(f"Invalid planning group name: {group_name!r}") + return f"{robot_name}/{group_name}" + + +def parse_planning_group_id(group_id: PlanningGroupID) -> tuple[RobotName, str]: + """Split and validate a planning group ID.""" + parts = group_id.split("/", maxsplit=1) + if len(parts) != 2 or not parts[0] or not parts[1] or "/" in parts[1]: + raise ValueError( + f"Invalid planning group ID {group_id!r}; expected '{{robot_name}}/{{group_name}}'" + ) + return parts[0], parts[1] + + +def make_global_joint_name( + robot_name: RobotName, + local_joint_name: LocalModelJointName, +) -> GlobalJointName: + """Convert a local model joint name to a public global joint name.""" + assert_valid_robot_name(robot_name) + assert_valid_local_joint_name(local_joint_name) + return f"{robot_name}/{local_joint_name}" + + +def make_global_joint_names( + robot_name: RobotName, + local_joint_names: list[LocalModelJointName] | tuple[LocalModelJointName, ...], +) -> list[GlobalJointName]: + """Convert local model joint names to public global joint names.""" + return [make_global_joint_name(robot_name, name) for name in local_joint_names] + + +def is_global_joint_name(name: str) -> bool: + """Return whether name has the exact global joint-name shape.""" + parts = name.split("/") + return len(parts) == 2 and bool(parts[0]) and bool(parts[1]) + + +def assert_global_joint_names(names: Sequence[GlobalJointName]) -> None: + """Validate that names are global joint names.""" + invalid = [name for name in names if not is_global_joint_name(name)] + if invalid: + raise ValueError(f"Expected global joint names; got invalid names: {invalid}") + + +def local_joint_name_from_global( + robot_name: RobotName, + global_joint_name: GlobalJointName, +) -> LocalModelJointName: + """Validate and strip a global joint name for backend internals.""" + assert_valid_robot_name(robot_name) + prefix = f"{robot_name}/" + if not global_joint_name.startswith(prefix): + raise ValueError( + f"Global joint name {global_joint_name!r} does not belong to robot {robot_name!r}" + ) + local_name = global_joint_name[len(prefix) :] + try: + assert_valid_local_joint_name(local_name) + except ValueError as exc: + raise ValueError(f"Invalid global joint name: {global_joint_name!r}") from exc + return local_name diff --git a/dimos/manipulation/planning/groups/models.py b/dimos/manipulation/planning/groups/models.py new file mode 100644 index 0000000000..c08b1bd4b5 --- /dev/null +++ b/dimos/manipulation/planning/groups/models.py @@ -0,0 +1,112 @@ +# 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. + +"""Backend-independent planning-group domain models.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal, TypeAlias + +from dimos.manipulation.planning.spec.models import ( + GlobalJointName, + LocalModelJointName, + PlanningGroupID, + RobotName, +) + +PlanningGroupSource: TypeAlias = Literal["srdf", "fallback"] + + +@dataclass(frozen=True) +class PlanningGroupDefinition: + """Model-level declaration of a planning group. + + Joint names are local model names. The definition is safe to store on + ``RobotModelConfig`` and is not bound to any runtime world robot ID. + """ + + name: str + joint_names: tuple[LocalModelJointName, ...] + base_link: str + tip_link: str | None = None + source: PlanningGroupSource = "srdf" + + @property + def has_pose_target(self) -> bool: + """Whether this group has a valid pose target frame.""" + return self.tip_link is not None + + +@dataclass(frozen=True) +class PlanningGroup: + """Public backend-independent planning group. + + A planning group exposes stable public IDs and global joint names for + planning APIs. It intentionally does not include backend runtime robot IDs. + """ + + id: PlanningGroupID + robot_name: RobotName + group_name: str + joint_names: tuple[GlobalJointName, ...] + local_joint_names: tuple[LocalModelJointName, ...] + base_link: str + tip_link: str | None = None + source: PlanningGroupSource = "srdf" + + @property + def has_pose_target(self) -> bool: + """Whether this group can be directly pose-targeted.""" + return self.tip_link is not None + + +@dataclass(frozen=True) +class PlanningGroupSelection: + """Validated ordered selection of planning groups. + + Selection validates ID existence and selected-joint overlap outside any + world backend. Requested group order is preserved. + """ + + groups: tuple[PlanningGroup, ...] + group_ids: tuple[PlanningGroupID, ...] + joint_names: tuple[GlobalJointName, ...] + robot_names: tuple[RobotName, ...] + + @classmethod + def from_groups(cls, groups: tuple[PlanningGroup, ...]) -> PlanningGroupSelection: + """Build a selection, rejecting overlapping selected global joints.""" + seen_joints: dict[GlobalJointName, PlanningGroupID] = {} + joint_names: list[GlobalJointName] = [] + robot_names: list[RobotName] = [] + for group in groups: + if group.robot_name not in robot_names: + robot_names.append(group.robot_name) + for joint_name in group.joint_names: + previous_group_id = seen_joints.get(joint_name) + if previous_group_id is not None: + raise ValueError( + "Selected planning groups overlap on global joint " + f"{joint_name}: {previous_group_id} and {group.id}" + ) + seen_joints[joint_name] = group.id + joint_names.append(joint_name) + + return cls( + groups=groups, + group_ids=tuple(group.id for group in groups), + joint_names=tuple(joint_names), + robot_names=tuple(robot_names), + ) diff --git a/dimos/manipulation/planning/groups/registry.py b/dimos/manipulation/planning/groups/registry.py new file mode 100644 index 0000000000..fb982b3562 --- /dev/null +++ b/dimos/manipulation/planning/groups/registry.py @@ -0,0 +1,117 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Backend-independent planning-group registry.""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import TYPE_CHECKING + +from dimos.manipulation.planning.groups.discovery import FALLBACK_PLANNING_GROUP_NAME +from dimos.manipulation.planning.groups.identifiers import ( + make_global_joint_names, + make_planning_group_id, +) +from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupSelection +from dimos.manipulation.planning.spec.models import PlanningGroupID, RobotName + +if TYPE_CHECKING: + from dimos.manipulation.planning.spec.config import RobotModelConfig + + +class PlanningGroupRegistry: + """Registry of public planning groups derived from robot configs.""" + + def __init__(self, robot_configs: Iterable[RobotModelConfig] = ()) -> None: + self._groups: dict[PlanningGroupID, PlanningGroup] = {} + self._groups_by_robot: dict[RobotName, list[PlanningGroup]] = {} + for config in robot_configs: + self.add_robot(config) + + def add_robot(self, config: RobotModelConfig) -> None: + """Register all planning groups declared by one robot config.""" + if config.name in self._groups_by_robot: + raise ValueError(f"Robot '{config.name}' is already registered") + + robot_groups: list[PlanningGroup] = [] + for definition in config.planning_groups: + group_id = make_planning_group_id(config.name, definition.name) + if group_id in self._groups: + raise ValueError(f"Planning group '{group_id}' is already registered") + group = PlanningGroup( + id=group_id, + robot_name=config.name, + group_name=definition.name, + joint_names=tuple(make_global_joint_names(config.name, definition.joint_names)), + local_joint_names=definition.joint_names, + base_link=definition.base_link, + tip_link=definition.tip_link, + source=definition.source, + ) + self._groups[group_id] = group + robot_groups.append(group) + self._groups_by_robot[config.name] = robot_groups + + def list(self) -> tuple[PlanningGroup, ...]: + """List planning groups in robot registration order.""" + groups: list[PlanningGroup] = [] + for robot_groups in self._groups_by_robot.values(): + groups.extend(robot_groups) + return tuple(groups) + + def get(self, group_id: PlanningGroupID) -> PlanningGroup: + """Return one planning group by public ID.""" + try: + return self._groups[group_id] + except KeyError as exc: + raise KeyError(f"Unknown planning group ID: {group_id}") from exc + + def select(self, group_ids: Iterable[PlanningGroupID]) -> PlanningGroupSelection: + """Validate and return an ordered planning-group selection.""" + return PlanningGroupSelection.from_groups( + tuple(self.get(group_id) for group_id in group_ids) + ) + + def groups_for_robot(self, robot_name: RobotName) -> tuple[PlanningGroup, ...]: + """Return planning groups for one robot.""" + return tuple(self._groups_by_robot.get(robot_name, ())) + + def default_group_id_for_robot(self, robot_name: RobotName) -> PlanningGroupID | None: + """Return the group ID used by robot-scoped joint wrappers. + + Prefer the generated whole-robot fallback group. If a robot only has one + configured planning group, use that group as the unambiguous fallback. + """ + group_id = make_planning_group_id(robot_name, FALLBACK_PLANNING_GROUP_NAME) + if group_id in self._groups: + return group_id + robot_groups = self.groups_for_robot(robot_name) + if len(robot_groups) == 1: + return robot_groups[0].id + return None + + def primary_pose_group_id_for_robot(self, robot_name: RobotName) -> PlanningGroupID | None: + """Return the unique pose-targetable group ID for robot-scoped wrappers.""" + pose_groups = [ + group for group in self.groups_for_robot(robot_name) if group.has_pose_target + ] + if not pose_groups: + return None + if len(pose_groups) > 1: + raise ValueError( + f"Robot '{robot_name}' has {len(pose_groups)} pose-targetable planning groups; " + "use an explicit planning group ID" + ) + return pose_groups[0].id diff --git a/dimos/manipulation/planning/groups/test_planning_groups.py b/dimos/manipulation/planning/groups/test_planning_groups.py new file mode 100644 index 0000000000..88757ec254 --- /dev/null +++ b/dimos/manipulation/planning/groups/test_planning_groups.py @@ -0,0 +1,728 @@ +# 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. + +"""Tests for planning groups.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from dimos.manipulation.planning.groups.discovery import ( + FALLBACK_PLANNING_GROUP_NAME, + PlanningGroupDiscoveryError, + discover_planning_group_definitions, + generate_fallback_planning_group, + parse_srdf_planning_groups, +) +from dimos.manipulation.planning.groups.identifiers import local_joint_name_from_global +from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupDefinition +from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry +from dimos.manipulation.planning.groups.utils import ( + filter_joint_state_to_selected_joints, + joint_state_to_ordered_positions, + joint_target_to_global_names, + matching_global_joint_name, + planning_group_id_from_selector, + project_global_joint_path_to_robot, +) +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.robot.model_parser import JointDescription, ModelDescription + + +def _serial_model(*joint_types: str) -> ModelDescription: + joints = [ + JointDescription( + name=f"joint{i + 1}", + type=joint_type, + parent_link=f"link{i}", + child_link=f"link{i + 1}", + ) + for i, joint_type in enumerate(joint_types) + ] + return ModelDescription( + joints=joints, + root_link="link0", + links=[f"link{i}" for i in range(len(joint_types) + 1)], + ) + + +def _branching_model() -> ModelDescription: + return ModelDescription( + joints=[ + JointDescription( + name="left_joint", + type="revolute", + parent_link="base", + child_link="left_link", + ), + JointDescription( + name="right_joint", + type="revolute", + parent_link="base", + child_link="right_link", + ), + ], + root_link="base", + links=["base", "left_link", "right_link"], + ) + + +def _write_srdf(tmp_path: Path, body: str) -> Path: + srdf_path = tmp_path / "robot.srdf" + srdf_path.write_text(f"{body}") + return srdf_path + + +def _make_group() -> PlanningGroup: + return PlanningGroup( + id="left/arm", + robot_name="left", + group_name="arm", + joint_names=("left/j1", "left/j2", "left/j3"), + local_joint_names=("j1", "j2", "j3"), + base_link="base", + tip_link="ee", + ) + + +def _robot_config( + name: str = "robot", + planning_groups: list[PlanningGroupDefinition] | None = None, +) -> RobotModelConfig: + return RobotModelConfig( + name=name, + model_path=Path("/tmp/robot.urdf"), + base_pose=PoseStamped(), + joint_names=["joint1", "joint2", "joint3"], + planning_groups=planning_groups + if planning_groups is not None + else [ + PlanningGroupDefinition( + name=FALLBACK_PLANNING_GROUP_NAME, + joint_names=("joint1", "joint2"), + base_link="base", + tip_link="tool", + ) + ], + ) + + +def test_parse_srdf_chain_group(tmp_path: Path) -> None: + model = _serial_model("revolute", "revolute", "revolute") + srdf_path = _write_srdf( + tmp_path, + "", + ) + + groups = parse_srdf_planning_groups( + srdf_path, + model=model, + controllable_joint_names=["joint1", "joint2", "joint3"], + ) + + assert len(groups) == 1 + assert groups[0].name == "arm" + assert groups[0].joint_names == ("joint1", "joint2", "joint3") + assert groups[0].base_link == "link0" + assert groups[0].tip_link == "link3" + assert groups[0].source == "srdf" + + +def test_parse_srdf_ordered_joint_list_group(tmp_path: Path) -> None: + model = _serial_model("revolute", "prismatic", "revolute") + srdf_path = _write_srdf( + tmp_path, + """ + + + + + + """, + ) + + groups = parse_srdf_planning_groups( + srdf_path, + model=model, + controllable_joint_names=["joint1", "joint2", "joint3"], + ) + + assert len(groups) == 1 + assert groups[0].joint_names == ("joint1", "joint2", "joint3") + assert groups[0].base_link == "link0" + assert groups[0].tip_link == "link3" + + +def test_parse_srdf_skips_unsupported_groups_and_ignores_end_effector( + tmp_path: Path, +) -> None: + model = _serial_model("revolute", "revolute") + srdf_path = _write_srdf( + tmp_path, + """ + + + + + """, + ) + + groups = parse_srdf_planning_groups( + srdf_path, + model=model, + controllable_joint_names=["joint1", "joint2"], + ) + + assert [group.name for group in groups] == ["arm"] + + +def test_fallback_generates_manipulator_for_unambiguous_serial_chain() -> None: + model = _serial_model("revolute", "prismatic", "revolute") + + group = generate_fallback_planning_group( + model=model, + controllable_joint_names=["joint2", "joint1", "joint3"], + ) + + assert group.name == FALLBACK_PLANNING_GROUP_NAME + assert group.joint_names == ("joint1", "joint2", "joint3") + assert group.base_link == "link0" + assert group.tip_link == "link3" + assert group.source == "fallback" + + +def test_fallback_strips_terminal_prismatic_joints() -> None: + model = _serial_model("revolute", "revolute", "prismatic") + + group = generate_fallback_planning_group( + model=model, + controllable_joint_names=["joint1", "joint2", "joint3"], + ) + + assert group.joint_names == ("joint1", "joint2") + assert group.tip_link == "link2" + assert group.source == "fallback" + + +def test_fallback_rejects_branching_model() -> None: + with pytest.raises(PlanningGroupDiscoveryError, match="branch"): + generate_fallback_planning_group( + model=_branching_model(), + controllable_joint_names=["left_joint", "right_joint"], + ) + + +def test_fallback_rejects_all_terminal_prismatic_candidates() -> None: + with pytest.raises(PlanningGroupDiscoveryError, match="removed all candidate joints"): + generate_fallback_planning_group( + model=_serial_model("prismatic", "prismatic"), + controllable_joint_names=["joint1", "joint2"], + ) + + +def test_parse_srdf_skips_invalid_groups_and_keeps_valid_group(tmp_path: Path) -> None: + model = _serial_model("revolute", "revolute") + srdf_path = _write_srdf( + tmp_path, + """ + + + + + """, + ) + + groups = parse_srdf_planning_groups( + srdf_path, + model=model, + controllable_joint_names=["joint1", "joint2"], + ) + + assert [group.name for group in groups] == ["arm"] + + +def test_discovery_rejects_missing_explicit_srdf(tmp_path: Path) -> None: + with pytest.raises(FileNotFoundError, match="SRDF file not found"): + discover_planning_group_definitions( + robot_name="robot", + model_path=tmp_path / "robot.urdf", + model=_serial_model("revolute"), + controllable_joint_names=["joint1"], + srdf_path=tmp_path / "missing.srdf", + ) + + +def test_discovery_falls_back_when_srdf_has_no_supported_groups(tmp_path: Path) -> None: + model_path = tmp_path / "robot.urdf.xacro" + model_path.write_text("") + (tmp_path / "robot.srdf").write_text( + "" + ) + + groups = discover_planning_group_definitions( + robot_name="robot", + model_path=model_path, + model=_serial_model("revolute"), + controllable_joint_names=["joint1"], + ) + + assert [group.name for group in groups] == [FALLBACK_PLANNING_GROUP_NAME] + assert [group.source for group in groups] == ["fallback"] + + +def test_discovery_prefers_explicit_srdf_over_fallback(tmp_path: Path) -> None: + model = _serial_model("revolute", "revolute") + model_path = tmp_path / "robot.urdf" + model_path.write_text("") + srdf_path = _write_srdf( + tmp_path, + "", + ) + + groups = discover_planning_group_definitions( + robot_name="robot", + model_path=model_path, + model=model, + controllable_joint_names=["joint1", "joint2"], + srdf_path=srdf_path, + ) + + assert [group.name for group in groups] == ["srdf_arm"] + + +def test_discovery_auto_discovers_srdf(tmp_path: Path) -> None: + model = _serial_model("revolute") + model_path = tmp_path / "robot.urdf" + model_path.write_text("") + _write_srdf( + tmp_path, + "", + ) + + groups = discover_planning_group_definitions( + robot_name="robot", + model_path=model_path, + model=model, + controllable_joint_names=["joint1"], + ) + + assert [group.name for group in groups] == ["auto_arm"] + + +def test_primary_pose_group_id_for_robot_raises_when_ambiguous() -> None: + registry = PlanningGroupRegistry( + [ + RobotModelConfig( + name="robot", + model_path=Path("/tmp/robot.urdf"), + base_pose=PoseStamped(), + joint_names=["joint1", "joint2"], + planning_groups=[ + PlanningGroupDefinition( + name="left", + joint_names=("joint1",), + base_link="base", + tip_link="left_tool", + ), + PlanningGroupDefinition( + name="right", + joint_names=("joint2",), + base_link="base", + tip_link="right_tool", + ), + ], + ) + ] + ) + + with pytest.raises(ValueError, match="multiple|2 pose-targetable|explicit planning group"): + registry.primary_pose_group_id_for_robot("robot") + + +def test_registry_preserves_order_and_exposes_defaults() -> None: + registry = PlanningGroupRegistry([_robot_config("left"), _robot_config("right")]) + + assert [group.id for group in registry.list()] == ["left/manipulator", "right/manipulator"] + assert registry.default_group_id_for_robot("left") == "left/manipulator" + assert registry.primary_pose_group_id_for_robot("right") == "right/manipulator" + assert registry.get("left/manipulator").source == "srdf" + assert registry.groups_for_robot("missing") == () + assert registry.default_group_id_for_robot("missing") is None + + +def test_registry_uses_single_group_as_robot_scoped_default() -> None: + registry = PlanningGroupRegistry( + [ + _robot_config( + "solo", + planning_groups=[PlanningGroupDefinition("arm", ("joint1",), "base", "tool")], + ), + _robot_config( + "multi", + planning_groups=[ + PlanningGroupDefinition("arm", ("joint1",), "base", "tool"), + PlanningGroupDefinition("gripper", ("joint2",), "tool"), + ], + ), + ] + ) + + assert registry.default_group_id_for_robot("solo") == "solo/arm" + assert registry.default_group_id_for_robot("multi") is None + + +def test_project_global_joint_path_to_robot_overlays_selected_joints() -> None: + path = [ + JointState(name=["robot/joint1", "robot/joint3"], position=[0.1, 0.3]), + JointState(name=["robot/joint1", "robot/joint3"], position=[0.2, 0.4]), + ] + current = JointState(name=["joint1", "joint2", "joint3"], position=[0.0, 0.5, 0.0]) + + projected = project_global_joint_path_to_robot( + path, + robot_name="robot", + local_joint_names=("joint1", "joint2", "joint3"), + current_joint_state=current, + ) + + assert [point.name for point in projected] == [ + ["joint1", "joint2", "joint3"], + ["joint1", "joint2", "joint3"], + ] + assert [point.position for point in projected] == [[0.1, 0.5, 0.3], [0.2, 0.5, 0.4]] + + +def test_project_global_joint_path_to_robot_rejects_inconsistent_path() -> None: + path = [ + JointState(name=["robot/joint1"], position=[0.1]), + JointState(name=["robot/joint2"], position=[0.2]), + ] + + with pytest.raises(ValueError, match="inconsistent waypoint joint names"): + project_global_joint_path_to_robot( + path, + robot_name="robot", + local_joint_names=("joint1", "joint2"), + current_joint_state=JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), + ) + + +def test_project_global_joint_path_to_robot_requires_current_non_selected_joints() -> None: + path = [JointState(name=["robot/joint1"], position=[0.1])] + + with pytest.raises(ValueError, match="missing joint 'joint2'"): + project_global_joint_path_to_robot( + path, + robot_name="robot", + local_joint_names=("joint1", "joint2"), + current_joint_state=JointState(name=["joint1"], position=[0.0]), + ) + + +def test_registry_rejects_duplicate_robot_and_unknown_group() -> None: + registry = PlanningGroupRegistry([_robot_config()]) + + with pytest.raises(ValueError, match="already registered"): + registry.add_robot(_robot_config()) + with pytest.raises(KeyError, match="Unknown planning group ID"): + registry.get("robot/missing") + + +def test_selection_preserves_group_order_and_rejects_overlapping_joints() -> None: + registry = PlanningGroupRegistry( + [ + _robot_config( + planning_groups=[ + PlanningGroupDefinition("arm", ("joint1", "joint2"), "base", "tool"), + PlanningGroupDefinition("gripper", ("joint3",), "tool"), + ] + ) + ] + ) + + selection = registry.select(["robot/gripper", "robot/arm"]) + + assert selection.group_ids == ("robot/gripper", "robot/arm") + assert selection.joint_names == ("robot/joint3", "robot/joint1", "robot/joint2") + assert selection.robot_names == ("robot",) + + overlapping = ( + PlanningGroup("robot/first", "robot", "first", ("robot/joint1",), ("joint1",), "base"), + PlanningGroup("robot/second", "robot", "second", ("robot/joint1",), ("joint1",), "base"), + ) + with pytest.raises(ValueError, match="overlap"): + type(selection).from_groups(overlapping) + + +def test_joint_target_to_global_names_accepts_named_global_targets_in_group_order() -> None: + group = _make_group() + target = JointState({"name": ["left/j3", "left/j1", "left/j2"], "position": [3.0, 1.0, 2.0]}) + + normalized = joint_target_to_global_names(group, target) + + assert normalized.name == ["left/j1", "left/j2", "left/j3"] + assert normalized.position == [1.0, 2.0, 3.0] + + +def test_joint_target_to_global_names_accepts_named_local_targets_in_group_order() -> None: + group = _make_group() + target = JointState({"name": ["j2", "j3", "j1"], "position": [2.0, 3.0, 1.0]}) + + normalized = joint_target_to_global_names(group, target) + + assert normalized.name == ["left/j1", "left/j2", "left/j3"] + assert normalized.position == [1.0, 2.0, 3.0] + + +def test_joint_target_to_global_names_rejects_mixed_global_and_local_target_names() -> None: + group = _make_group() + target = JointState({"name": ["left/j1", "j2", "left/j3"], "position": [1.0, 2.0, 3.0]}) + + with pytest.raises(ValueError, match="mixes global and local joint names"): + joint_target_to_global_names(group, target) + + +def test_joint_target_to_global_names_rejects_bad_counts_missing_and_extra() -> None: + group = _make_group() + + with pytest.raises(ValueError, match="2 positions, expected 3"): + joint_target_to_global_names(group, JointState({"position": [1.0, 2.0]})) + with pytest.raises(ValueError, match="2 names but 3 positions"): + joint_target_to_global_names( + group, JointState({"name": ["j1", "j2"], "position": [1.0, 2.0, 3.0]}) + ) + with pytest.raises(ValueError, match="missing joints"): + joint_target_to_global_names( + group, JointState({"name": ["j1", "j2"], "position": [1.0, 2.0]}) + ) + with pytest.raises(ValueError, match="extra joints"): + joint_target_to_global_names( + group, JointState({"name": ["j1", "j2", "j3", "j4"], "position": [1.0, 2.0, 3.0, 4.0]}) + ) + + +def test_filter_joint_state_to_selected_joints_uses_local_fallbacks() -> None: + joint_state = JointState({"name": ["j1", "robot/j2"], "position": [1.0, 2.0]}) + + filtered = filter_joint_state_to_selected_joints( + joint_state, + ["robot/j1", "robot/j2"], + ["j1", "j2"], + ) + + assert filtered.name == ["robot/j1", "robot/j2"] + assert filtered.position == [1.0, 2.0] + + +def test_filter_joint_state_to_selected_joints_rejects_mismatched_and_missing_names() -> None: + joint_state = JointState({"name": ["robot/j1"], "position": [1.0]}) + + with pytest.raises(ValueError, match="same length"): + filter_joint_state_to_selected_joints(joint_state, ["robot/j1", "robot/j2"], ["j1"]) + with pytest.raises(ValueError, match="missing selected joints"): + filter_joint_state_to_selected_joints(joint_state, ["robot/j1", "robot/j2"]) + + +def test_matching_global_joint_name_requires_unique_suffix_match() -> None: + assert matching_global_joint_name({"left/j1": 1.0, "right/j2": 2.0}, "j1") == "left/j1" + assert matching_global_joint_name({"left/j1": 1.0, "right/j1": 2.0}, "j1") is None + assert matching_global_joint_name({"left/j1": 1.0}, "j2") is None + + +def test_filter_joint_state_to_selected_joints_uses_local_fallbacks() -> None: + state = JointState(name=["j2", "arm/j1"], position=[2.0, 1.0]) + + filtered = filter_joint_state_to_selected_joints( + state, + ["arm/j1", "arm/j2"], + ["j1", "j2"], + ) + + assert filtered.name == ["arm/j1", "arm/j2"] + assert filtered.position == [1.0, 2.0] + + +def test_joint_target_to_global_names_accepts_unnamed_positions_in_group_order() -> None: + target = joint_target_to_global_names( + PlanningGroup( + id="left/arm", + robot_name="left", + group_name="arm", + joint_names=("left/j2", "left/j1"), + local_joint_names=("j2", "j1"), + base_link="base", + tip_link="ee", + ), + JointState(name=[], position=[2.0, 1.0]), + ) + + assert target.name == ["left/j2", "left/j1"] + assert target.position == [2.0, 1.0] + + +def test_planning_group_id_from_selector_accepts_id_or_group() -> None: + group = _make_group() + + assert planning_group_id_from_selector(group) == "left/arm" + assert planning_group_id_from_selector("left/arm") == "left/arm" + + +def test_local_joint_name_from_global_validates_robot_prefix_and_local_shape() -> None: + assert local_joint_name_from_global("robot", "robot/j1") == "j1" + with pytest.raises(ValueError, match="does not belong"): + local_joint_name_from_global("robot", "other/j1") + with pytest.raises(ValueError, match="Invalid global joint name"): + local_joint_name_from_global("robot", "robot/") + + +def test_robot_model_config_derives_legacy_end_effector_link_from_pose_group() -> None: + config = RobotModelConfig( + name="arm", + model_path=Path("robot.urdf"), + joint_names=["j1", "j2"], + joint_name_mapping={"hw_j1": "j1", "hw_j2": "j2"}, + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=("j1", "j2"), + base_link="base", + tip_link="tool", + ) + ], + ) + + assert config.end_effector_link == "tool" + assert config.get_urdf_joint_name("hw_j1") == "j1" + assert config.get_coordinator_joint_name("j2") == "hw_j2" + assert config.get_coordinator_joint_names() == ["hw_j1", "hw_j2"] + + +def test_robot_model_config_end_effector_link_requires_pose_group() -> None: + config = RobotModelConfig( + name="arm", + model_path=Path("robot.urdf"), + joint_names=["j1"], + planning_groups=[ + PlanningGroupDefinition( + name="joint_only", + joint_names=("j1",), + base_link="base", + ) + ], + ) + + with pytest.raises(ValueError, match="no pose-target planning group"): + assert config.end_effector_link + + +def test_robot_model_config_end_effector_link_rejects_ambiguous_pose_groups() -> None: + config = RobotModelConfig( + name="arm", + model_path=Path("robot.urdf"), + joint_names=["j1", "j2"], + planning_groups=[ + PlanningGroupDefinition( + name="left", + joint_names=("j1",), + base_link="base", + tip_link="left_tool", + ), + PlanningGroupDefinition( + name="right", + joint_names=("j2",), + base_link="base", + tip_link="right_tool", + ), + ], + ) + + with pytest.raises(ValueError, match="multiple pose-target planning groups"): + assert config.end_effector_link + + +def test_joint_state_to_ordered_positions_accepts_all_supported_name_forms() -> None: + joint_names = ["joint1", "joint2", "joint3"] + mapping = {"hw1": "joint1", "hw2": "joint2", "hw3": "joint3"} + + unnamed = joint_state_to_ordered_positions( + JointState(name=[], position=[1.0, 2.0, 3.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + local = joint_state_to_ordered_positions( + JointState(name=["joint3", "joint1", "joint2"], position=[30.0, 10.0, 20.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + coordinator = joint_state_to_ordered_positions( + JointState(name=["hw2", "hw3", "hw1"], position=[200.0, 300.0, 100.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + global_names = joint_state_to_ordered_positions( + JointState(name=["arm/joint2", "arm/joint1", "arm/joint3"], position=[2.0, 1.0, 3.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + + assert unnamed.tolist() == [1.0, 2.0, 3.0] + assert local.tolist() == [10.0, 20.0, 30.0] + assert coordinator.tolist() == [100.0, 200.0, 300.0] + assert global_names.tolist() == [1.0, 2.0, 3.0] + + +def test_joint_state_to_ordered_positions_rejects_invalid_inputs() -> None: + joint_names = ["joint1", "joint2"] + mapping = {"hw1": "joint1"} + + with pytest.raises(ValueError, match="position length"): + joint_state_to_ordered_positions( + JointState(name=[], position=[1.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + with pytest.raises(ValueError, match="name and position"): + joint_state_to_ordered_positions( + JointState(name=["joint1", "joint2"], position=[1.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + with pytest.raises(ValueError, match="duplicate"): + joint_state_to_ordered_positions( + JointState(name=["joint1", "hw1"], position=[1.0, 2.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + with pytest.raises(ValueError, match="Unknown global"): + joint_state_to_ordered_positions( + JointState(name=["arm/joint3", "joint2"], position=[1.0, 2.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + with pytest.raises(ValueError, match="missing joints"): + joint_state_to_ordered_positions( + JointState(name=["joint1"], position=[1.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + with pytest.raises(ValueError, match="Unrecognized joint name"): + joint_state_to_ordered_positions( + JointState(name=["mystery", "joint2"], position=[1.0, 2.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) diff --git a/dimos/manipulation/planning/groups/utils.py b/dimos/manipulation/planning/groups/utils.py new file mode 100644 index 0000000000..5d9161e5bd --- /dev/null +++ b/dimos/manipulation/planning/groups/utils.py @@ -0,0 +1,236 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared helpers for planning-group selectors and joint-state projection.""" + +from collections.abc import Mapping, Sequence + +import numpy as np +from numpy.typing import NDArray + +from dimos.manipulation.planning.groups.identifiers import ( + assert_global_joint_names, + assert_local_joint_names, + is_global_joint_name, + make_global_joint_names, +) +from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.manipulation.planning.spec.models import ( + GlobalJointName, + JointPath, + LocalModelJointName, + PlanningGroupID, + RobotName, +) +from dimos.msgs.sensor_msgs.JointState import JointState + + +def planning_group_id_from_selector(selector: PlanningGroupID | PlanningGroup) -> PlanningGroupID: + """Return the planning-group ID represented by a selector.""" + if isinstance(selector, PlanningGroup): + return selector.id + return selector + + +def matching_global_joint_name( + positions_by_name: Mapping[str, float], local_joint_name: LocalModelJointName +) -> GlobalJointName | None: + """Find the unique global joint name ending with a local joint name.""" + suffix = f"/{local_joint_name}" + matches = [name for name in positions_by_name if name.endswith(suffix)] + if len(matches) == 1: + return matches[0] + return None + + +def filter_joint_state_to_selected_joints( + joint_state: JointState, + global_joint_names: Sequence[GlobalJointName], + local_joint_names: Sequence[LocalModelJointName] = (), +) -> JointState: + """Project a joint state to selected global joints. + + Values are looked up by global name first. When ``local_joint_names`` is + provided, each corresponding local name is used as a fallback. + """ + if local_joint_names and len(global_joint_names) != len(local_joint_names): + raise ValueError("Global and local selected joint lists must have the same length") + + positions_by_name = dict(zip(joint_state.name, joint_state.position, strict=True)) + selected_positions: list[float] = [] + missing: list[str] = [] + for index, global_name in enumerate(global_joint_names): + if global_name in positions_by_name: + selected_positions.append(float(positions_by_name[global_name])) + continue + if local_joint_names: + local_name = local_joint_names[index] + if local_name in positions_by_name: + selected_positions.append(float(positions_by_name[local_name])) + continue + missing.append(global_name) + + if missing: + raise ValueError(f"IK result is missing selected joints: {missing}") + + return JointState({"name": list(global_joint_names), "position": selected_positions}) + + +def joint_target_to_global_names( + group: PlanningGroup, + target: JointState, +) -> JointState: + """Convert a group joint target to global joint names in group order. + + Named targets may use either the public global planning names or the + robot-local model names used by legacy robot-scoped callers, but the two + namespaces must not be mixed in one target. + """ + if not target.name: + if len(target.position) != len(group.joint_names): + raise ValueError( + f"Target for '{group.id}' has {len(target.position)} positions, " + f"expected {len(group.joint_names)}" + ) + return JointState(name=list(group.joint_names), position=list(target.position)) + + if len(target.name) != len(target.position): + raise ValueError( + f"Target for '{group.id}' has {len(target.name)} names but " + f"{len(target.position)} positions" + ) + + target_names = list(target.name) + global_flags = [is_global_joint_name(name) for name in target_names] + if any(global_flags) and not all(global_flags): + raise ValueError( + f"Target for '{group.id}' mixes global and local joint names: {target_names}" + ) + + if all(global_flags): + assert_global_joint_names(target_names) + expected_names = group.joint_names + else: + assert_local_joint_names(target_names) + expected_names = group.local_joint_names + + positions_by_name = dict(zip(target_names, target.position, strict=True)) + global_positions: list[float] = [] + missing: list[str] = [] + for expected_name in expected_names: + if expected_name in positions_by_name: + global_positions.append(positions_by_name[expected_name]) + else: + missing.append(expected_name) + if missing: + raise ValueError(f"Target for '{group.id}' is missing joints: {missing}") + + extra = set(target_names) - set(expected_names) + if extra: + raise ValueError(f"Target for '{group.id}' has extra joints: {sorted(extra)}") + return JointState(name=list(group.joint_names), position=global_positions) + + +def project_global_joint_path_to_robot( + path: Sequence[JointState], + *, + robot_name: RobotName, + local_joint_names: Sequence[LocalModelJointName], + current_joint_state: JointState | None, +) -> JointPath: + """Project a selected-global-joint path into one robot's local joint path.""" + if not path: + return [] + + selected_joint_names = tuple(path[0].name) + assert_global_joint_names(selected_joint_names) + if any( + len(waypoint.name) != len(waypoint.position) or tuple(waypoint.name) != selected_joint_names + for waypoint in path + ): + raise ValueError("inconsistent waypoint joint names") + + selected_joint_indices = dict( + zip(selected_joint_names, range(len(selected_joint_names)), strict=True) + ) + selected_joint_set = set(selected_joint_names) + waypoint_positions = [[float(position) for position in waypoint.position] for waypoint in path] + current_by_name = ( + dict(zip(current_joint_state.name, current_joint_state.position, strict=False)) + if current_joint_state is not None + else {} + ) + global_joint_names = make_global_joint_names(robot_name, tuple(local_joint_names)) + joint_pairs = list(zip(local_joint_names, global_joint_names, strict=True)) + try: + base_positions = [ + 0.0 if global_name in selected_joint_set else float(current_by_name[local_name]) + for local_name, global_name in joint_pairs + ] + except KeyError as exc: + raise ValueError(f"missing joint '{exc.args[0]}'") from exc + + overlay_indices = [ + (local_index, selected_joint_indices[global_name]) + for local_index, (_, global_name) in enumerate(joint_pairs) + if global_name in selected_joint_indices + ] + local_path: JointPath = [] + for waypoint_positions_by_joint in waypoint_positions: + projected_positions = base_positions.copy() + for local_index, selected_index in overlay_indices: + projected_positions[local_index] = waypoint_positions_by_joint[selected_index] + local_path.append(JointState(name=list(local_joint_names), position=projected_positions)) + return local_path + + +def joint_state_to_ordered_positions( + joint_state: JointState, + *, + joint_names: Sequence[str], + joint_name_mapping: Mapping[str, str], +) -> NDArray[np.float64]: + """Convert a JointState to an array ordered by local robot joint names.""" + if not joint_state.name: + if len(joint_state.position) != len(joint_names): + raise ValueError("JointState position length must match configured joint count") + return np.asarray(joint_state.position, dtype=np.float64) + + if len(joint_state.name) != len(joint_state.position): + raise ValueError("JointState name and position lengths must match") + + joint_name_set = set(joint_names) + name_to_pos: dict[str, float] = {} + for name, position in zip(joint_state.name, joint_state.position, strict=True): + if name in joint_name_set: + resolved_name = name + elif name in joint_name_mapping: + resolved_name = joint_name_mapping[name] + elif is_global_joint_name(name): + resolved_name = name.split("/", maxsplit=1)[1] + if resolved_name not in joint_name_set: + raise ValueError(f"Unknown global joint name: {name}") + else: + raise ValueError( + f"Unrecognized joint name '{name}': not a known local name, not in joint_name_mapping, and not a global name" + ) + + if resolved_name in name_to_pos: + raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") + name_to_pos[resolved_name] = float(position) + + missing = [name for name in joint_names if name not in name_to_pos] + if missing: + raise ValueError(f"JointState missing joints: {missing}") + return np.asarray([name_to_pos[name] for name in joint_names], dtype=np.float64) diff --git a/dimos/manipulation/planning/kinematics/drake_optimization_ik.py b/dimos/manipulation/planning/kinematics/drake_optimization_ik.py index abccee119d..91da25986c 100644 --- a/dimos/manipulation/planning/kinematics/drake_optimization_ik.py +++ b/dimos/manipulation/planning/kinematics/drake_optimization_ik.py @@ -16,10 +16,17 @@ from __future__ import annotations +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING import numpy as np +from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.manipulation.planning.kinematics.utils import ( + filter_result_to_group as _filter_result_to_group, + resolve_single_pose_target_request as _resolve_single_pose_target_request, + unique_pose_target_frame_for_robot as _unique_pose_target_frame_for_robot, +) from dimos.manipulation.planning.spec.enums import IKStatus from dimos.manipulation.planning.spec.models import IKResult, WorldRobotID from dimos.manipulation.planning.spec.protocols import WorldSpec @@ -28,7 +35,6 @@ from dimos.msgs.geometry_msgs.Transform import Transform from dimos.msgs.sensor_msgs.JointState import JointState from dimos.utils.logging_config import setup_logger -from dimos.utils.transform_utils import pose_to_matrix if TYPE_CHECKING: from numpy.typing import NDArray @@ -84,6 +90,13 @@ def solve( if error is not None: return error + target_frame_name = _unique_pose_target_frame_for_robot(world, robot_id) + if target_frame_name is None: + return _create_failure_result( + IKStatus.UNSUPPORTED, + "DrakeOptimizationIK requires exactly one pose-targetable planning group for legacy solve()", + ) + # Convert PoseStamped to 4x4 matrix via Transform target_matrix = Transform( translation=target_pose.position, @@ -127,6 +140,7 @@ def solve( orientation_tolerance=orientation_tolerance, lower_limits=lower_limits, upper_limits=upper_limits, + target_frame_name=target_frame_name, ) if result.is_success() and result.joint_state is not None: @@ -156,6 +170,96 @@ def solve( f"IK failed after {max_attempts} attempts", ) + def solve_pose_targets( + self, + world: WorldSpec, + pose_targets: Mapping[PlanningGroup, PoseStamped], + auxiliary_groups: Sequence[PlanningGroup] = (), + seed: JointState | None = None, + position_tolerance: float = 0.001, + orientation_tolerance: float = 0.01, + check_collision: bool = True, + max_attempts: int = 10, + ) -> IKResult: + """Solve a planning-group-scoped pose target with Drake IK.""" + error = self._validate_world(world) + if error is not None: + return error + request, request_error = _resolve_single_pose_target_request( + world, + pose_targets, + auxiliary_groups, + seed, + "DrakeOptimizationIK", + ) + if request_error is not None: + return request_error + if request is None or request.group.tip_link is None: + return _create_failure_result( + IKStatus.UNSUPPORTED, + "DrakeOptimizationIK requires a pose-targetable planning group", + ) + + lower_limits, upper_limits = world.get_joint_limits(request.robot_id) + target_matrix = Transform( + translation=request.target_pose.position, + rotation=request.target_pose.orientation, + ).to_matrix() + target_transform = RigidTransform(target_matrix) + locked_positions = { + index: float(request.seed_positions[index]) + for index in range(len(request.joint_names)) + if index not in set(request.group_indices) + } + + best_result: IKResult | None = None + best_error = float("inf") + for attempt in range(max_attempts): + if attempt == 0: + current_seed = request.seed_positions + else: + current_seed = request.seed_positions.copy() + random_group_positions = np.random.uniform( + lower_limits[request.group_indices], upper_limits[request.group_indices] + ) + current_seed[request.group_indices] = random_group_positions + + result = self._solve_single( + world=world, + robot_id=request.robot_id, + target_transform=target_transform, + seed=current_seed, + joint_names=request.joint_names, + position_tolerance=position_tolerance, + orientation_tolerance=orientation_tolerance, + lower_limits=lower_limits, + upper_limits=upper_limits, + target_frame_name=request.group.tip_link, + locked_joint_positions=locked_positions, + ) + if not result.is_success() or result.joint_state is None: + continue + if check_collision and not world.check_config_collision_free( + request.robot_id, result.joint_state + ): + continue + total_error = result.position_error + result.orientation_error + if total_error < best_error: + best_error = total_error + best_result = result + if ( + result.position_error <= position_tolerance + and result.orientation_error <= orientation_tolerance + ): + return _filter_result_to_group(result, request.group) + + if best_result is not None: + return _filter_result_to_group(best_result, request.group) + return _create_failure_result( + IKStatus.NO_SOLUTION, + f"IK failed after {max_attempts} attempts", + ) + def _solve_single( self, world: WorldSpec, @@ -167,6 +271,8 @@ def _solve_single( orientation_tolerance: float, lower_limits: NDArray[np.float64], upper_limits: NDArray[np.float64], + target_frame_name: str, + locked_joint_positions: Mapping[int, float] | None = None, ) -> IKResult: # Get robot data from world internals (Drake-specific access) robot_data = world._robots[robot_id] # type: ignore[attr-defined] @@ -175,12 +281,13 @@ def _solve_single( # Create IK problem ik = InverseKinematics(plant) - # Get end-effector frame - ee_frame = robot_data.ee_frame + target_frame = plant.GetBodyByName( + target_frame_name, robot_data.model_instance + ).body_frame() # Add position constraint ik.AddPositionConstraint( - frameB=ee_frame, + frameB=target_frame, p_BQ=np.array([0.0, 0.0, 0.0]), # type: ignore[arg-type] frameA=plant.world_frame(), p_AQ_lower=target_transform.translation() - np.array([position_tolerance] * 3), @@ -191,7 +298,7 @@ def _solve_single( ik.AddOrientationConstraint( frameAbar=plant.world_frame(), R_AbarA=target_transform.rotation(), - frameBbar=ee_frame, + frameBbar=target_frame, R_BbarB=RotationMatrix(), theta_bound=orientation_tolerance, ) @@ -200,6 +307,10 @@ def _solve_single( prog = ik.get_mutable_prog() q = ik.q() + for local_index, value in (locked_joint_positions or {}).items(): + joint_idx = robot_data.joint_indices[local_index] + prog.AddBoundingBoxConstraint(value, value, q[joint_idx]) + # Set initial guess (full positions vector) full_seed = np.zeros(plant.num_positions()) for i, joint_idx in enumerate(robot_data.joint_indices): @@ -223,13 +334,13 @@ def _solve_single( joint_solution = np.clip(joint_solution, lower_limits, upper_limits) # Compute actual error using FK - solution_state = JointState(name=joint_names, position=joint_solution.tolist()) + solution_state = JointState({"name": joint_names, "position": joint_solution.tolist()}) with world.scratch_context() as ctx: world.set_joint_state(ctx, robot_id, solution_state) - actual_pose = world.get_ee_pose(ctx, robot_id) + actual_matrix = world.get_link_pose(ctx, robot_id, target_frame_name) position_error, orientation_error = compute_pose_error( - pose_to_matrix(actual_pose), + actual_matrix, target_transform.GetAsMatrix4(), # type: ignore[arg-type] ) @@ -251,7 +362,7 @@ def _create_success_result( ) -> IKResult: return IKResult( status=IKStatus.SUCCESS, - joint_state=JointState(name=joint_names, position=joint_positions.tolist()), + joint_state=JointState({"name": joint_names, "position": joint_positions.tolist()}), position_error=position_error, orientation_error=orientation_error, iterations=iterations, diff --git a/dimos/manipulation/planning/kinematics/jacobian_ik.py b/dimos/manipulation/planning/kinematics/jacobian_ik.py index 7727b6fa0f..4c4e16207a 100644 --- a/dimos/manipulation/planning/kinematics/jacobian_ik.py +++ b/dimos/manipulation/planning/kinematics/jacobian_ik.py @@ -24,10 +24,15 @@ from __future__ import annotations +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING import numpy as np +from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.manipulation.planning.kinematics.utils import ( + resolve_single_pose_target_request as _resolve_single_pose_target_request, +) from dimos.manipulation.planning.spec.enums import IKStatus from dimos.manipulation.planning.spec.models import IKResult, WorldRobotID from dimos.manipulation.planning.spec.protocols import WorldSpec @@ -145,7 +150,9 @@ def solve( else: # Random seed within joint limits random_positions = np.random.uniform(lower_limits, upper_limits) - current_seed = JointState(name=joint_names, position=random_positions.tolist()) + current_seed = JointState( + {"name": joint_names, "position": random_positions.tolist()} + ) # Solve iterative IK result = self.solve_iterative( @@ -185,6 +192,61 @@ def solve( f"IK failed after {max_attempts} attempts", ) + def solve_pose_targets( + self, + world: WorldSpec, + pose_targets: Mapping[PlanningGroup, PoseStamped], + auxiliary_groups: Sequence[PlanningGroup] = (), + seed: JointState | None = None, + position_tolerance: float = 0.001, + orientation_tolerance: float = 0.01, + check_collision: bool = True, + max_attempts: int = 10, + ) -> IKResult: + """Solve a planning-group pose target using group FK/Jacobian.""" + if not world.is_finalized: + return _create_failure_result(IKStatus.NO_SOLUTION, "World must be finalized before IK") + request, request_error = _resolve_single_pose_target_request( + world, + pose_targets, + auxiliary_groups, + seed, + "JacobianIK", + ) + if request_error is not None: + return request_error + if request is None: + return _create_failure_result(IKStatus.NO_SOLUTION, "Invalid pose target request") + + full_seed = JointState( + {"name": request.joint_names, "position": request.seed_positions.tolist()} + ) + result = self.solve_iterative( + world=world, + robot_id=request.robot_id, + target_pose=request.target_pose, + seed=full_seed, + max_iterations=self._max_iterations * max(1, max_attempts), + position_tolerance=position_tolerance, + orientation_tolerance=orientation_tolerance, + group=request.group, + active_joint_indices=request.group_indices, + ) + if not result.is_success() or result.joint_state is None: + return result + + if check_collision: + full_positions = request.seed_positions.copy() + full_positions[request.group_indices] = np.asarray( + result.joint_state.position, dtype=np.float64 + ) + full_state = JointState( + {"name": request.joint_names, "position": full_positions.tolist()} + ) + if not world.check_config_collision_free(request.robot_id, full_state): + return _create_failure_result(IKStatus.COLLISION, "IK solution is in collision") + return result + def solve_iterative( self, world: WorldSpec, @@ -194,6 +256,8 @@ def solve_iterative( max_iterations: int = 100, position_tolerance: float = 0.001, orientation_tolerance: float = 0.01, + group: PlanningGroup | None = None, + active_joint_indices: list[int] | None = None, ) -> IKResult: """Iterative Jacobian-based IK until convergence. @@ -219,6 +283,8 @@ def solve_iterative( ).to_matrix() current_joints = np.array(seed.position, dtype=np.float64) joint_names = seed.name + active_joint_indices = active_joint_indices or list(range(len(joint_names))) + result_joint_names = list(group.joint_names) if group is not None else joint_names max_iterations = max_iterations or self._max_iterations lower_limits, upper_limits = world.get_joint_limits(robot_id) @@ -226,11 +292,15 @@ def solve_iterative( for iteration in range(max_iterations): with world.scratch_context() as ctx: # Set current position (convert to JointState for API) - current_state = JointState(name=joint_names, position=current_joints.tolist()) + current_state = JointState( + {"name": joint_names, "position": current_joints.tolist()} + ) world.set_joint_state(ctx, robot_id, current_state) - # Get current pose (as matrix for error computation) - current_pose = pose_to_matrix(world.get_ee_pose(ctx, robot_id)) + if group is None: + current_pose = pose_to_matrix(world.get_ee_pose(ctx, robot_id)) + else: + current_pose = pose_to_matrix(world.get_group_ee_pose(ctx, group.id)) # Compute error pos_error, ori_error = compute_pose_error(current_pose, target_matrix) @@ -238,28 +308,30 @@ def solve_iterative( # Check convergence if pos_error <= position_tolerance and ori_error <= orientation_tolerance: return _create_success_result( - joint_names=joint_names, - joint_positions=current_joints, + joint_names=result_joint_names, + joint_positions=current_joints[active_joint_indices], position_error=pos_error, orientation_error=ori_error, iterations=iteration + 1, ) + if group is None: + jacobian = world.get_jacobian(ctx, robot_id) + else: + jacobian = world.get_group_jacobian(ctx, group.id) + # Compute twist to reduce error twist = compute_error_twist(current_pose, target_matrix, gain=0.5) - # Get Jacobian - J = world.get_jacobian(ctx, robot_id) - # Adaptive damping near singularities - if check_singularity(J, threshold=self._singularity_threshold): + if check_singularity(jacobian, threshold=self._singularity_threshold): # Increase damping near singularity instead of failing effective_damping = self._damping * 10.0 else: effective_damping = self._damping # Compute joint velocities - J_pinv = damped_pseudoinverse(J, effective_damping) + J_pinv = damped_pseudoinverse(jacobian, effective_damping) q_dot = J_pinv @ twist # Clamp maximum joint change per iteration (like reference implementations) @@ -268,16 +340,23 @@ def solve_iterative( if max_change > max_delta: q_dot = q_dot * (max_delta / max_change) - current_joints = current_joints + q_dot + current_joints[active_joint_indices] = current_joints[active_joint_indices] + q_dot # Clip to limits - current_joints = np.clip(current_joints, lower_limits, upper_limits) + current_joints[active_joint_indices] = np.clip( + current_joints[active_joint_indices], + lower_limits[active_joint_indices], + upper_limits[active_joint_indices], + ) # Compute final error with world.scratch_context() as ctx: - final_state = JointState(name=joint_names, position=current_joints.tolist()) + final_state = JointState({"name": joint_names, "position": current_joints.tolist()}) world.set_joint_state(ctx, robot_id, final_state) - final_pose = pose_to_matrix(world.get_ee_pose(ctx, robot_id)) + if group is None: + final_pose = pose_to_matrix(world.get_ee_pose(ctx, robot_id)) + else: + final_pose = pose_to_matrix(world.get_group_ee_pose(ctx, group.id)) pos_error, ori_error = compute_pose_error(final_pose, target_matrix) return _create_failure_result( @@ -349,7 +428,7 @@ def solve_differential( if max_ratio > 1.0: q_dot = q_dot / max_ratio - return JointState(name=joint_names, velocity=q_dot.tolist()) + return JointState({"name": joint_names, "velocity": q_dot.tolist()}) def solve_differential_position_only( self, @@ -397,7 +476,7 @@ def solve_differential_position_only( # Compute joint velocities q_dot = J_pinv @ vel_array - return JointState(name=joint_names, velocity=q_dot.tolist()) + return JointState({"name": joint_names, "velocity": q_dot.tolist()}) # Result Helpers @@ -413,7 +492,7 @@ def _create_success_result( """Create a successful IK result.""" return IKResult( status=IKStatus.SUCCESS, - joint_state=JointState(name=joint_names, position=joint_positions.tolist()), + joint_state=JointState({"name": joint_names, "position": joint_positions.tolist()}), position_error=position_error, orientation_error=orientation_error, iterations=iterations, diff --git a/dimos/manipulation/planning/kinematics/pink_ik.py b/dimos/manipulation/planning/kinematics/pink_ik.py index 1245e5aea4..3456c2e27f 100644 --- a/dimos/manipulation/planning/kinematics/pink_ik.py +++ b/dimos/manipulation/planning/kinematics/pink_ik.py @@ -16,6 +16,7 @@ from __future__ import annotations +from collections.abc import Mapping, Sequence from dataclasses import dataclass import importlib from pathlib import Path @@ -24,10 +25,17 @@ import numpy as np +from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupSelection from dimos.manipulation.planning.kinematics.config import PinkKinematicsConfig +from dimos.manipulation.planning.kinematics.utils import ( + groups_by_robot as _groups_by_robot, + robot_ids_by_name as _robot_ids_by_name, + seed_positions_with_world_fallback as _seed_positions_with_world_fallback, + unique_pose_target_frame_for_robot as _unique_pose_target_frame_for_robot, +) from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import IKStatus -from dimos.manipulation.planning.spec.models import IKResult, WorldRobotID +from dimos.manipulation.planning.spec.models import IKResult, RobotName, WorldRobotID from dimos.manipulation.planning.spec.protocols import WorldSpec from dimos.manipulation.planning.utils.kinematics_utils import compute_pose_error from dimos.manipulation.planning.utils.mesh_utils import prepare_urdf_for_drake @@ -98,7 +106,7 @@ def __init__( config_values.update(overrides) self.config = PinkKinematicsConfig(**config_values) self._modules = _load_optional_dependencies(self.config.solver) - self._robot_contexts: dict[str, _PinkRobotContext] = {} + self._robot_contexts: dict[tuple[str, str], _PinkRobotContext] = {} def solve( self, @@ -115,8 +123,15 @@ def solve( if not world.is_finalized: return _failure(IKStatus.NO_SOLUTION, "World must be finalized before IK") + target_frame_name = _unique_pose_target_frame_for_robot(world, robot_id) + if target_frame_name is None: + return _failure( + IKStatus.NO_SOLUTION, + "PinkIK requires exactly one pose-targetable planning group for legacy solve()", + ) + try: - robot_context = self._get_robot_context(world, robot_id) + robot_context = self._get_robot_context(world, robot_id, target_frame_name) except (FileNotFoundError, ImportError, ValueError) as exc: return _failure(IKStatus.NO_SOLUTION, f"Pink IK model setup failed: {exc}") @@ -164,6 +179,265 @@ def solve( return _failure(IKStatus.NO_SOLUTION, f"Pink IK failed after {max_attempts} attempts") + def solve_pose_targets( + self, + world: WorldSpec, + pose_targets: Mapping[PlanningGroup, PoseStamped], + auxiliary_groups: Sequence[PlanningGroup] = (), + seed: JointState | None = None, + position_tolerance: float = 0.001, + orientation_tolerance: float = 0.01, + check_collision: bool = True, + max_attempts: int = 10, + ) -> IKResult: + """Solve planning-group-scoped pose targets with Pink IK.""" + if not world.is_finalized: + return _failure(IKStatus.NO_SOLUTION, "World must be finalized before IK") + all_groups = tuple(pose_targets.keys()) + tuple(auxiliary_groups) + if not all_groups: + return _failure( + IKStatus.NO_SOLUTION, "At least one pose target or auxiliary group is required" + ) + bad_groups = [ + group.id + for group in pose_targets + if not group.has_pose_target or group.tip_link is None + ] + if bad_groups: + return _failure( + IKStatus.UNSUPPORTED, + f"Planning groups have no pose target frame: {bad_groups}", + ) + + try: + selection = PlanningGroupSelection.from_groups(all_groups) + robot_ids_by_name = _robot_ids_by_name(world, selection.robot_names) + except ValueError as exc: + return _failure(IKStatus.NO_SOLUTION, str(exc)) + + results_by_robot: dict[RobotName, IKResult] = {} + for robot_name, groups in _groups_by_robot(all_groups).items(): + robot_id = robot_ids_by_name[robot_name] + config = world.get_robot_config(robot_id) + joint_names = list(config.joint_names) + try: + selected_indices = [ + joint_names.index(name) for group in groups for name in group.local_joint_names + ] + seed_positions = _seed_positions_with_world_fallback( + world, robot_id, config.name, joint_names, seed + ) + except ValueError as exc: + return _failure(IKStatus.NO_SOLUTION, f"Pink IK mapping failed: {exc}") + robot_pose_targets = [group for group in groups if group in pose_targets] + if not robot_pose_targets: + robot_result = _success(joint_names, seed_positions, 0.0, 0.0, 0) + results_by_robot[robot_name] = robot_result + continue + + lower_limits, upper_limits = world.get_joint_limits(robot_id) + locked_positions = { + index: float(seed_positions[index]) + for index in range(len(joint_names)) + if index not in set(selected_indices) + } + targets: list[tuple[_PinkRobotContext, NDArray[np.float64]]] = [] + try: + for group in robot_pose_targets: + if group.tip_link is None: + raise ValueError(f"Planning group '{group.id}' has no pose target frame") + targets.append( + ( + self._get_robot_context(world, robot_id, group.tip_link), + self._target_in_model_frame(config, pose_targets[group]), + ) + ) + except (FileNotFoundError, ImportError, ValueError) as exc: + return _failure(IKStatus.NO_SOLUTION, f"Pink IK model setup failed: {exc}") + + fallback_result: IKResult | None = None + for attempt in range(max_attempts): + current_positions = seed_positions.copy() + if attempt > 0: + current_positions[selected_indices] = np.random.uniform( + lower_limits[selected_indices], upper_limits[selected_indices] + ) + try: + q0 = self._q_from_dimos_positions(targets[0][0], current_positions) + if len(targets) == 1: + result = self._solve_single( + robot_context=targets[0][0], + target_model=targets[0][1], + seed_q=q0, + lower_limits=lower_limits, + upper_limits=upper_limits, + position_tolerance=position_tolerance, + orientation_tolerance=orientation_tolerance, + locked_joint_positions=locked_positions, + ) + else: + result = self._solve_multi( + targets=targets, + seed_q=q0, + lower_limits=lower_limits, + upper_limits=upper_limits, + position_tolerance=position_tolerance, + orientation_tolerance=orientation_tolerance, + locked_joint_positions=locked_positions, + ) + except ValueError as exc: + return _failure(IKStatus.NO_SOLUTION, f"Pink IK mapping failed: {exc}") + except Exception as exc: + return _failure(IKStatus.NO_SOLUTION, f"Pink IK solver failed: {exc}") + + if not result.is_success() or result.joint_state is None: + if fallback_result is None: + fallback_result = result + continue + results_by_robot[robot_name] = result + break + else: + if fallback_result is not None: + return fallback_result + return _failure( + IKStatus.NO_SOLUTION, f"Pink IK failed after {max_attempts} attempts" + ) + + positions_by_robot: dict[RobotName, dict[str, float]] = {} + max_position_error = 0.0 + max_orientation_error = 0.0 + iterations = 0 + for robot_name, result in results_by_robot.items(): + if not result.is_success() or result.joint_state is None: + return result + positions_by_robot[robot_name] = dict( + zip(result.joint_state.name, result.joint_state.position, strict=True) + ) + max_position_error = max(max_position_error, result.position_error) + max_orientation_error = max(max_orientation_error, result.orientation_error) + iterations = max(iterations, result.iterations) + + selected_names: list[str] = [] + selected_positions: list[float] = [] + for group in selection.groups: + robot_positions = positions_by_robot[group.robot_name] + for global_name, local_name in zip( + group.joint_names, + group.local_joint_names, + strict=True, + ): + if global_name in robot_positions: + position = robot_positions[global_name] + elif local_name in robot_positions: + position = robot_positions[local_name] + else: + return _failure( + IKStatus.NO_SOLUTION, + f"Pink IK result is missing selected joint '{global_name}'", + ) + selected_names.append(global_name) + selected_positions.append(float(position)) + + combined = IKResult( + status=IKStatus.SUCCESS, + joint_state=JointState( + { + "name": selected_names, + "position": selected_positions, + } + ), + position_error=max_position_error, + orientation_error=max_orientation_error, + iterations=iterations, + message="Pink IK solution found", + ) + if check_collision and not _combined_robot_results_collision_free( + world, + robot_ids_by_name, + results_by_robot, + ): + return _collision_failure(combined) + return combined + + def _solve_multi( + self, + targets: Sequence[tuple[_PinkRobotContext, NDArray[np.float64]]], + seed_q: NDArray[np.float64], + lower_limits: NDArray[np.float64], + upper_limits: NDArray[np.float64], + position_tolerance: float, + orientation_tolerance: float, + locked_joint_positions: Mapping[int, float] | None = None, + ) -> IKResult: + robot_context = targets[0][0] + pink = self._modules.pink + pinocchio = self._modules.pinocchio + configuration = pink.Configuration(robot_context.model, robot_context.data, seed_q.copy()) + tasks: list[Any] = [] + for target_context, target_model in targets: + frame_task = pink.tasks.FrameTask( + target_context.frame_name, + position_cost=self.config.position_cost, + orientation_cost=self.config.orientation_cost, + lm_damping=self.config.lm_damping, + gain=self.config.gain, + ) + frame_task.set_target(_matrix_to_se3(pinocchio, target_model)) + tasks.append(frame_task) + if self.config.posture_cost > 0.0: + posture_task = pink.tasks.PostureTask(cost=self.config.posture_cost) + posture_task.set_target_from_configuration(configuration) + tasks.append(posture_task) + final_position_error = float("inf") + final_orientation_error = float("inf") + for iteration in range(self.config.max_iterations): + errors = [ + compute_pose_error(self._current_frame_matrix(ctx, configuration.q), target_model) + for ctx, target_model in targets + ] + final_position_error = max(error[0] for error in errors) + final_orientation_error = max(error[1] for error in errors) + if ( + final_position_error <= position_tolerance + and final_orientation_error <= orientation_tolerance + ): + return _success( + robot_context.mapping.dimos_joint_names, + self._q_to_dimos_positions(robot_context, configuration.q), + final_position_error, + final_orientation_error, + iteration + 1, + ) + velocity = pink.solve_ik( + configuration, + tasks, + self.config.dt, + solver=self.config.solver, + damping=self.config.damping, + safety_break=self.config.safety_break, + ) + configuration.integrate_inplace(velocity, self.config.dt) + for local_index, value in (locked_joint_positions or {}).items(): + configuration.q[robot_context.mapping.idx_q[local_index]] = value + joint_positions = self._q_to_dimos_positions(robot_context, configuration.q) + if not _within_limits(joint_positions, lower_limits, upper_limits): + return IKResult( + status=IKStatus.JOINT_LIMITS, + joint_state=None, + position_error=final_position_error, + orientation_error=final_orientation_error, + iterations=iteration + 1, + message="Pink IK candidate violates DimOS joint limits", + ) + return IKResult( + status=IKStatus.NO_SOLUTION, + joint_state=None, + position_error=final_position_error, + orientation_error=final_orientation_error, + iterations=self.config.max_iterations, + message="Pink IK did not converge within the iteration budget", + ) + def _solve_single( self, robot_context: _PinkRobotContext, @@ -173,6 +447,7 @@ def _solve_single( upper_limits: NDArray[np.float64], position_tolerance: float, orientation_tolerance: float, + locked_joint_positions: Mapping[int, float] | None = None, ) -> IKResult: pink = self._modules.pink pinocchio = self._modules.pinocchio @@ -224,6 +499,8 @@ def _solve_single( safety_break=self.config.safety_break, ) configuration.integrate_inplace(velocity, self.config.dt) + for local_index, value in (locked_joint_positions or {}).items(): + configuration.q[robot_context.mapping.idx_q[local_index]] = value joint_positions = self._q_to_dimos_positions(robot_context, configuration.q) if not _within_limits(joint_positions, lower_limits, upper_limits): @@ -245,15 +522,20 @@ def _solve_single( message="Pink IK did not converge within the iteration budget", ) - def _get_robot_context(self, world: WorldSpec, robot_id: WorldRobotID) -> _PinkRobotContext: - cache_key = str(robot_id) + def _get_robot_context( + self, + world: WorldSpec, + robot_id: WorldRobotID, + frame_name: str, + ) -> _PinkRobotContext: + cache_key = (str(robot_id), frame_name) if cache_key not in self._robot_contexts: self._robot_contexts[cache_key] = self._build_robot_context( - world.get_robot_config(robot_id) + world.get_robot_config(robot_id), frame_name ) return self._robot_contexts[cache_key] - def _build_robot_context(self, config: RobotModelConfig) -> _PinkRobotContext: + def _build_robot_context(self, config: RobotModelConfig, frame_name: str) -> _PinkRobotContext: pinocchio = self._modules.pinocchio model_path = Path(config.model_path).resolve() if not model_path.exists(): @@ -271,13 +553,14 @@ def _build_robot_context(self, config: RobotModelConfig) -> _PinkRobotContext: model = pinocchio.buildModelFromUrdf(str(prepared_path)) data = model.createData() - frame_id = _get_frame_id(model, config.end_effector_link) + _assert_base_link_is_model_root(model, config.base_link) + frame_id = _get_frame_id(model, frame_name) mapping = _build_joint_mapping(model, config) return _PinkRobotContext( model=model, data=data, frame_id=frame_id, - frame_name=config.end_effector_link, + frame_name=frame_name, mapping=mapping, ) @@ -302,6 +585,21 @@ def _initial_q( q[idx_q] = value return q + def _q_from_dimos_positions( + self, + context: _PinkRobotContext, + positions: NDArray[np.float64], + ) -> NDArray[np.float64]: + pinocchio = self._modules.pinocchio + q = np.array(pinocchio.neutral(context.model), dtype=np.float64) + if len(positions) != len(context.mapping.idx_q): + raise ValueError( + f"Seed has {len(positions)} positions for {len(context.mapping.idx_q)} joints" + ) + for value, idx_q in zip(positions, context.mapping.idx_q, strict=True): + q[idx_q] = value + return q + def _q_to_dimos_positions( self, context: _PinkRobotContext, q: NDArray[np.float64] ) -> NDArray[np.float64]: @@ -407,6 +705,18 @@ def _get_frame_id(model: Any, frame_name: str) -> int: return frame_id +def _assert_base_link_is_model_root(model: Any, base_link: str) -> None: + """Validate that the configured base link is fixed at the Pinocchio model root.""" + frame_id = _get_frame_id(model, base_link) + frame = model.frames[frame_id] + parent_joint = int(getattr(frame, "parentJoint", 0)) + if parent_joint != 0: + raise ValueError( + f"PinkIK expects RobotModelConfig.base_link '{base_link}' to be the model root; " + f"Pinocchio frame parentJoint is {parent_joint}" + ) + + def _missing_joint_message(model: Any, joint_name: str) -> str: available = [str(name) for name in getattr(model, "names", [])] return f"Joint '{joint_name}' not found in Pinocchio model. Available joints: {available}" @@ -456,6 +766,23 @@ def _within_limits( ) +def _combined_robot_results_collision_free( + world: WorldSpec, + robot_ids_by_name: Mapping[RobotName, WorldRobotID], + results_by_robot: Mapping[RobotName, IKResult], +) -> bool: + with world.scratch_context() as ctx: + for robot_name, result in results_by_robot.items(): + if result.joint_state is None: + return False + world.set_joint_state(ctx, robot_ids_by_name[robot_name], result.joint_state) + return all( + world.is_collision_free(ctx, robot_id) + for robot_name, robot_id in robot_ids_by_name.items() + if robot_name in results_by_robot + ) + + def _success( joint_names: list[str], joint_positions: NDArray[np.float64], @@ -465,7 +792,7 @@ def _success( ) -> IKResult: return IKResult( status=IKStatus.SUCCESS, - joint_state=JointState(name=joint_names, position=joint_positions.tolist()), + joint_state=JointState({"name": joint_names, "position": joint_positions.tolist()}), position_error=position_error, orientation_error=orientation_error, iterations=iterations, diff --git a/dimos/manipulation/planning/kinematics/test_drake_optimization_ik_selection.py b/dimos/manipulation/planning/kinematics/test_drake_optimization_ik_selection.py new file mode 100644 index 0000000000..afac663e72 --- /dev/null +++ b/dimos/manipulation/planning/kinematics/test_drake_optimization_ik_selection.py @@ -0,0 +1,253 @@ +# 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 + +from contextlib import contextmanager +from pathlib import Path + +import numpy as np + +from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.manipulation.planning.kinematics import drake_optimization_ik as drake_ik +from dimos.manipulation.planning.kinematics.drake_optimization_ik import DrakeOptimizationIK +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.enums import IKStatus +from dimos.manipulation.planning.spec.models import IKResult +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.sensor_msgs.JointState import JointState + + +class FakeWorld: + def __init__(self) -> None: + self.robot_id = "robot-instance" + self.config = RobotModelConfig( + name="arm", + model_path=Path("/tmp/fake.urdf"), + joint_names=["base", "shoulder", "elbow", "wrist"], + ) + self.current_state = JointState( + {"name": ["base", "shoulder", "elbow", "wrist"], "position": [1.0, 2.0, 3.0, 4.0]} + ) + self.collision_checked_state: JointState | None = None + + def get_robot_ids(self) -> list[str]: + return [self.robot_id] + + def get_robot_config(self, robot_id: str) -> RobotModelConfig: + assert robot_id == self.robot_id + return self.config + + def get_joint_limits(self, robot_id: str) -> tuple[np.ndarray, np.ndarray]: + assert robot_id == self.robot_id + return np.array([-10.0] * 4), np.array([10.0] * 4) + + @contextmanager + def scratch_context(self): + yield object() + + def get_joint_state(self, ctx: object, robot_id: str) -> JointState: + assert robot_id == self.robot_id + return self.current_state + + def check_config_collision_free(self, robot_id: str, joint_state: JointState) -> bool: + assert robot_id == self.robot_id + self.collision_checked_state = joint_state + return True + + +def test_solve_pose_targets_uses_group_tip_locks_seed_fallback_and_filters(monkeypatch) -> None: + monkeypatch.setattr(drake_ik, "DRAKE_AVAILABLE", True) + monkeypatch.setattr(DrakeOptimizationIK, "_validate_world", lambda self, world: None) + + world = FakeWorld() + group = PlanningGroup( + id="arm/reach", + robot_name="arm", + group_name="reach", + joint_names=("arm/shoulder", "arm/wrist"), + local_joint_names=("shoulder", "wrist"), + base_link="base_link", + tip_link="group_tip_link", + ) + calls = [] + + def fake_solve_single(self, **kwargs) -> IKResult: + calls.append(kwargs) + return IKResult( + status=IKStatus.SUCCESS, + joint_state=JointState( + { + "name": ["base", "shoulder", "elbow", "wrist"], + "position": [10.0, 20.0, 30.0, 40.0], + } + ), + position_error=0.0, + orientation_error=0.0, + ) + + monkeypatch.setattr(DrakeOptimizationIK, "_solve_single", fake_solve_single) + monkeypatch.setattr(drake_ik, "RigidTransform", lambda matrix: matrix, raising=False) + + result = DrakeOptimizationIK().solve_pose_targets( + world=world, # type: ignore[arg-type] + pose_targets={group: PoseStamped()}, + seed=JointState({"name": ["shoulder"], "position": [22.0]}), + check_collision=False, + max_attempts=1, + ) + + assert result.is_success() + assert result.joint_state is not None + assert result.joint_state.name == ["arm/shoulder", "arm/wrist"] + assert result.joint_state.position == [20.0, 40.0] + assert calls[0]["target_frame_name"] == "group_tip_link" + np.testing.assert_allclose(calls[0]["seed"], [1.0, 22.0, 3.0, 4.0]) + assert calls[0]["locked_joint_positions"] == {0: 1.0, 2: 3.0} + + +class FakeRigidTransform: + def translation(self) -> np.ndarray: + return np.array([0.5, 0.6, 0.7]) + + def rotation(self) -> str: + return "target-rotation" + + def GetAsMatrix4(self) -> np.ndarray: + return np.eye(4) + + +class FakeBody: + def __init__(self, frame_name: str) -> None: + self.frame_name = frame_name + + def body_frame(self) -> str: + return f"frame:{self.frame_name}" + + +class FakePlant: + def __init__(self) -> None: + self.requested_bodies: list[tuple[str, str]] = [] + + def GetBodyByName(self, frame_name: str, model_instance: str) -> FakeBody: + self.requested_bodies.append((frame_name, model_instance)) + return FakeBody(frame_name) + + def world_frame(self) -> str: + return "world-frame" + + def num_positions(self) -> int: + return 6 + + +class FakeProgram: + def __init__(self) -> None: + self.locks: list[tuple[float, float, str]] = [] + self.initial_guess: tuple[list[str], np.ndarray] | None = None + + def AddBoundingBoxConstraint(self, lower: float, upper: float, variable: str) -> None: + self.locks.append((lower, upper, variable)) + + def SetInitialGuess(self, q: list[str], full_seed: np.ndarray) -> None: + self.initial_guess = (q, full_seed.copy()) + + +class FakeInverseKinematics: + instances: list[FakeInverseKinematics] = [] + + def __init__(self, plant: FakePlant) -> None: + self.plant = plant + self.program = FakeProgram() + self.q_vars = [f"q{i}" for i in range(6)] + self.position_constraints = [] + self.orientation_constraints = [] + self.instances.append(self) + + def AddPositionConstraint(self, **kwargs) -> None: + self.position_constraints.append(kwargs) + + def AddOrientationConstraint(self, **kwargs) -> None: + self.orientation_constraints.append(kwargs) + + def get_mutable_prog(self) -> FakeProgram: + return self.program + + def q(self) -> list[str]: + return self.q_vars + + +class FakeSolveResult: + def is_success(self) -> bool: + return True + + def GetSolution(self, q: list[str]) -> np.ndarray: + return np.array([0.0, 11.0, 0.0, 22.0, 33.0, 0.0]) + + +class FakeDrakeWorld: + def __init__(self) -> None: + self.plant = FakePlant() + self._robots = {"robot-instance": _FakeRobotData()} + self.link_pose_calls: list[tuple[str, str]] = [] + self.set_joint_state_calls: list[JointState] = [] + + @contextmanager + def scratch_context(self): + yield "ctx" + + def set_joint_state(self, ctx: str, robot_id: str, joint_state: JointState) -> None: + self.set_joint_state_calls.append(joint_state) + + def get_link_pose(self, ctx: str, robot_id: str, target_frame_name: str) -> np.ndarray: + self.link_pose_calls.append((robot_id, target_frame_name)) + return np.eye(4) + + +class _FakeRobotData: + model_instance = "model-instance" + joint_indices = [1, 3, 4] + + +def test_solve_single_uses_target_frame_for_constraints_error_and_joint_locks(monkeypatch) -> None: + FakeInverseKinematics.instances.clear() + monkeypatch.setattr(drake_ik, "DRAKE_AVAILABLE", True) + monkeypatch.setattr(drake_ik, "InverseKinematics", FakeInverseKinematics, raising=False) + monkeypatch.setattr(drake_ik, "RotationMatrix", lambda: "identity-rotation", raising=False) + monkeypatch.setattr(drake_ik, "Solve", lambda prog: FakeSolveResult(), raising=False) + monkeypatch.setattr(drake_ik, "compute_pose_error", lambda actual, target: (0.01, 0.02)) + + world = FakeDrakeWorld() + result = DrakeOptimizationIK()._solve_single( + world=world, # type: ignore[arg-type] + robot_id="robot-instance", + target_transform=FakeRigidTransform(), + seed=np.array([1.0, 2.0, 3.0]), + joint_names=["j0", "j1", "j2"], + position_tolerance=0.1, + orientation_tolerance=0.2, + lower_limits=np.array([-5.0, -5.0, -5.0]), + upper_limits=np.array([5.0, 5.0, 5.0]), + target_frame_name="selected_tip_link", + locked_joint_positions={0: 1.5, 2: 3.5}, + ) + + ik = FakeInverseKinematics.instances[0] + assert result.is_success() + assert world.plant.requested_bodies == [("selected_tip_link", "model-instance")] + assert ik.position_constraints[0]["frameB"] == "frame:selected_tip_link" + assert ik.orientation_constraints[0]["frameBbar"] == "frame:selected_tip_link" + assert world.link_pose_calls == [("robot-instance", "selected_tip_link")] + assert ik.program.locks == [(1.5, 1.5, "q1"), (3.5, 3.5, "q4")] + assert ik.program.initial_guess is not None + np.testing.assert_allclose(ik.program.initial_guess[1], [0.0, 1.0, 0.0, 2.0, 3.0, 0.0]) diff --git a/dimos/manipulation/planning/kinematics/test_jacobian_ik_selection.py b/dimos/manipulation/planning/kinematics/test_jacobian_ik_selection.py new file mode 100644 index 0000000000..15d966b7a4 --- /dev/null +++ b/dimos/manipulation/planning/kinematics/test_jacobian_ik_selection.py @@ -0,0 +1,147 @@ +# 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. + +"""Focused tests for group-aware Jacobian IK.""" + +from __future__ import annotations + +from contextlib import nullcontext +from pathlib import Path + +import numpy as np + +from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupDefinition +from dimos.manipulation.planning.kinematics.jacobian_ik import JacobianIK +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.enums import IKStatus +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.sensor_msgs.JointState import JointState + + +def _pose(x: float = 0.0) -> PoseStamped: + return PoseStamped(position=Vector3(x, 0.0, 0.0), orientation=Quaternion(0.0, 0.0, 0.0, 1.0)) + + +def _group(tip_link: str | None = "tool") -> PlanningGroup: + return PlanningGroup( + id="arm/manipulator", + robot_name="arm", + group_name="manipulator", + joint_names=("arm/joint_a", "arm/joint_b"), + local_joint_names=("joint_a", "joint_b"), + base_link="base", + tip_link=tip_link, + ) + + +class _World: + is_finalized = True + + def __init__(self) -> None: + self.group_pose_calls = 0 + self.group_jacobian_calls = 0 + self.legacy_pose_calls = 0 + self.legacy_jacobian_calls = 0 + self.config = RobotModelConfig( + name="arm", + model_path=Path("robot.urdf"), + base_pose=_pose(), + joint_names=["joint_a", "joint_b", "gripper"], + base_link="base", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=("joint_a", "joint_b"), + base_link="base", + tip_link="tool", + ) + ], + ) + + def get_robot_ids(self) -> list[str]: + return ["robot"] + + def get_robot_config(self, robot_id: str) -> RobotModelConfig: + return self.config + + def get_joint_limits(self, robot_id: str) -> tuple[np.ndarray, np.ndarray]: + return np.array([-1.0, -1.0, -1.0]), np.array([1.0, 1.0, 1.0]) + + def scratch_context(self) -> nullcontext[None]: + return nullcontext(None) + + def get_joint_state(self, ctx: object, robot_id: str) -> JointState: + return JointState({"name": ["joint_a", "joint_b", "gripper"], "position": [0.0, 0.0, 0.9]}) + + def set_joint_state(self, ctx: object, robot_id: str, joint_state: JointState) -> None: + self.last_state = joint_state + + def get_group_ee_pose(self, ctx: object, group_id: str) -> PoseStamped: + self.group_pose_calls += 1 + return _pose(0.0) + + def get_group_jacobian(self, ctx: object, group_id: str) -> np.ndarray: + self.group_jacobian_calls += 1 + return np.eye(6, 2) + + def get_ee_pose(self, ctx: object, robot_id: str) -> PoseStamped: + self.legacy_pose_calls += 1 + raise AssertionError("legacy EE pose should not be used") + + def get_jacobian(self, ctx: object, robot_id: str) -> np.ndarray: + self.legacy_jacobian_calls += 1 + raise AssertionError("legacy Jacobian should not be used") + + def check_config_collision_free(self, robot_id: str, joint_state: JointState) -> bool: + return True + + +def test_solve_pose_targets_filters_to_group_and_uses_group_world_methods() -> None: + world = _World() + result = JacobianIK(max_iterations=2).solve_pose_targets( + world=world, + pose_targets={_group(): _pose()}, + seed=JointState( + {"name": ["arm/joint_a", "arm/joint_b", "arm/gripper"], "position": [0.0, 0.0, 0.9]} + ), + max_attempts=1, + ) + + assert result.status == IKStatus.SUCCESS + assert result.joint_state is not None + assert result.joint_state.name == ["arm/joint_a", "arm/joint_b"] + assert world.group_pose_calls == 1 + assert world.group_jacobian_calls == 0 + assert world.legacy_pose_calls == 0 + assert world.legacy_jacobian_calls == 0 + + +def test_solve_pose_targets_rejects_auxiliary_groups() -> None: + result = JacobianIK().solve_pose_targets( + world=_World(), + pose_targets={_group(): _pose()}, + auxiliary_groups=[_group()], + ) + + assert result.status == IKStatus.UNSUPPORTED + assert "no auxiliary" in result.message + + +def test_solve_pose_targets_rejects_group_without_pose_target_frame() -> None: + result = JacobianIK().solve_pose_targets(world=_World(), pose_targets={_group(None): _pose()}) + + assert result.status == IKStatus.UNSUPPORTED + assert "no pose target frame" in result.message diff --git a/dimos/manipulation/planning/kinematics/test_pink_ik.py b/dimos/manipulation/planning/kinematics/test_pink_ik.py index 9bfb1abeff..d0324edb56 100644 --- a/dimos/manipulation/planning/kinematics/test_pink_ik.py +++ b/dimos/manipulation/planning/kinematics/test_pink_ik.py @@ -26,6 +26,7 @@ from pytest_mock import MockerFixture from dimos.manipulation.planning.factory import create_kinematics +from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupDefinition from dimos.manipulation.planning.kinematics.config import PinkKinematicsConfig import dimos.manipulation.planning.kinematics.pink_ik as pink_ik from dimos.manipulation.planning.kinematics.pink_ik import ( @@ -53,8 +54,9 @@ def __init__(self, idx_q: int) -> None: class _FakeFrame: - def __init__(self, name: str) -> None: + def __init__(self, name: str, parent_joint: int = 0) -> None: self.name = name + self.parentJoint = parent_joint class _FakePlacement: @@ -66,7 +68,7 @@ def __init__(self, translation: np.ndarray) -> None: class _FakeData: def __init__(self) -> None: self.q = np.zeros(3) - self.oMf = [_FakePlacement(np.zeros(3))] + self.oMf = [_FakePlacement(np.zeros(3)), _FakePlacement(np.zeros(3))] class _FakeModel: @@ -75,9 +77,9 @@ class _FakeModel: def __init__(self) -> None: self.names = ["universe", "joint_b", "joint_a", "joint_c"] self.joints = [SimpleNamespace(idx_q=-1, nq=0), _FakeJoint(0), _FakeJoint(1), _FakeJoint(2)] - self.frames = [_FakeFrame("tool")] + self.frames = [_FakeFrame("base", 0), _FakeFrame("tool", 3)] self._joint_ids = {"joint_b": 1, "joint_a": 2, "joint_c": 3} - self._frame_ids = {"tool": 0} + self._frame_ids = {"base": 0, "tool": 1} def createData(self) -> _FakeData: return _FakeData() @@ -137,7 +139,7 @@ def forward_kinematics(model: _FakeModel, data: _FakeData, q: np.ndarray) -> Non data.q = q.copy() def update_frame_placements(model: _FakeModel, data: _FakeData) -> None: - data.oMf[0] = _FakePlacement(data.q.copy()) + data.oMf[1] = _FakePlacement(data.q.copy()) pinocchio.forwardKinematics = forward_kinematics # type: ignore[attr-defined] pinocchio.updateFramePlacements = update_frame_placements # type: ignore[attr-defined] @@ -169,8 +171,15 @@ def _robot_config() -> RobotModelConfig: model_path=Path("/tmp/fake.urdf"), base_pose=PoseStamped(position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0)), joint_names=["joint_a", "joint_b", "joint_c"], - end_effector_link="tool", base_link="base", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=("joint_a", "joint_b", "joint_c"), + base_link="base", + tip_link="tool", + ) + ], ) @@ -187,7 +196,7 @@ def _context() -> _PinkRobotContext: return _PinkRobotContext( model=model, data=model.createData(), - frame_id=0, + frame_id=1, frame_name="tool", mapping=mapping, ) @@ -199,6 +208,39 @@ class _FakeWorld: def __init__(self, collision_free: bool = True) -> None: self.config = _robot_config() self.collision_free = collision_free + self.joint_state_calls = 0 + self.groups = { + "arm/manipulator": PlanningGroup( + id="arm/manipulator", + robot_name="arm", + group_name="manipulator", + joint_names=("arm/joint_a", "arm/joint_b"), + local_joint_names=("joint_a", "joint_b"), + base_link="base", + tip_link="tool", + ), + "arm/no_tip": PlanningGroup( + id="arm/no_tip", + robot_name="arm", + group_name="no_tip", + joint_names=("arm/joint_c",), + local_joint_names=("joint_c",), + base_link="base", + tip_link=None, + ), + "arm/wrist": PlanningGroup( + id="arm/wrist", + robot_name="arm", + group_name="wrist", + joint_names=("arm/joint_c",), + local_joint_names=("joint_c",), + base_link="base", + tip_link="base", + ), + } + + def get_robot_ids(self) -> list[str]: + return ["robot"] def get_robot_config(self, robot_id: str) -> RobotModelConfig: return self.config @@ -207,10 +249,8 @@ def scratch_context(self) -> nullcontext[None]: return nullcontext(None) def get_joint_state(self, ctx: object, robot_id: str) -> JointState: - return JointState( - name=["joint_b", "joint_c", "joint_a"], - position=[0.0, 0.0, 0.0], - ) + self.joint_state_calls += 1 + return JointState({"name": ["joint_b", "joint_c", "joint_a"], "position": [0.0, 0.0, 0.0]}) def get_joint_limits(self, robot_id: str) -> tuple[np.ndarray, np.ndarray]: return np.array([-1.0, -1.0, -1.0]), np.array([1.0, 1.0, 1.0]) @@ -218,6 +258,72 @@ def get_joint_limits(self, robot_id: str) -> tuple[np.ndarray, np.ndarray]: def check_config_collision_free(self, robot_id: str, joint_state: JointState) -> bool: return self.collision_free + def set_joint_state(self, ctx: object, robot_id: str, joint_state: JointState) -> None: + self.joint_state = joint_state + + def is_collision_free(self, ctx: object, robot_id: str) -> bool: + return self.collision_free + + +class _MultiRobotCollisionWorld: + is_finalized = True + + def __init__(self) -> None: + left_config = _robot_config() + left_config.name = "left" + right_config = _robot_config() + right_config.name = "right" + self.configs = {"left-id": left_config, "right-id": right_config} + self.groups = { + "left/manipulator": PlanningGroup( + id="left/manipulator", + robot_name="left", + group_name="manipulator", + joint_names=("left/joint_a", "left/joint_b"), + local_joint_names=("joint_a", "joint_b"), + base_link="base", + tip_link="tool", + ), + "right/manipulator": PlanningGroup( + id="right/manipulator", + robot_name="right", + group_name="manipulator", + joint_names=("right/joint_a", "right/joint_b"), + local_joint_names=("joint_a", "joint_b"), + base_link="base", + tip_link="tool", + ), + } + self.config_collision_checks = 0 + self.context_collision_checks = 0 + self.context_states: dict[str, JointState] = {} + + def get_robot_ids(self) -> list[str]: + return ["left-id", "right-id"] + + def get_robot_config(self, robot_id: str) -> RobotModelConfig: + return self.configs[robot_id] + + def scratch_context(self) -> nullcontext[None]: + return nullcontext(None) + + def get_joint_state(self, ctx: object, robot_id: str) -> JointState: + return JointState({"name": ["joint_a", "joint_b", "joint_c"], "position": [0.0, 0.0, 0.0]}) + + def get_joint_limits(self, robot_id: str) -> tuple[np.ndarray, np.ndarray]: + return np.array([-1.0, -1.0, -1.0]), np.array([1.0, 1.0, 1.0]) + + def check_config_collision_free(self, robot_id: str, joint_state: JointState) -> bool: + self.config_collision_checks += 1 + return True + + def set_joint_state(self, ctx: object, robot_id: str, joint_state: JointState) -> None: + self.context_states[robot_id] = joint_state + + def is_collision_free(self, ctx: object, robot_id: str) -> bool: + self.context_collision_checks += 1 + return len(self.context_states) < 2 + def test_create_kinematics_pink_missing_dependency_is_actionable( mocker: MockerFixture, @@ -341,7 +447,7 @@ def test_solve_single_reports_non_convergence(mocker: MockerFixture) -> None: def test_solve_rejects_collision_candidate(mocker: MockerFixture) -> None: ik = _pink_ik(mocker, converge=True) context = _context() - ik._robot_contexts = {"robot": context} + ik._robot_contexts = {("robot", "tool"): context} result = ik.solve( world=cast("Any", _FakeWorld(collision_free=False)), @@ -361,7 +467,7 @@ def test_solve_rejects_collision_candidate(mocker: MockerFixture) -> None: def test_solve_retries_after_joint_limit_failure(mocker: MockerFixture) -> None: ik = _pink_ik(mocker, converge=True) context = _context() - ik._robot_contexts = {"robot": context} + ik._robot_contexts = {("robot", "tool"): context} calls = 0 def fake_solve_single(**_: object) -> IKResult: @@ -399,3 +505,245 @@ def fake_solve_single(**_: object) -> IKResult: assert solve_single.call_count == 2 assert result.status == IKStatus.SUCCESS + + +def test_robot_context_cache_key_includes_tip_frame(mocker: MockerFixture, tmp_path: Path) -> None: + modules = _fake_modules() + modules.pinocchio.buildModelFromUrdf = lambda path: _FakeModel() # type: ignore[attr-defined] + mocker.patch.object(pink_ik, "_load_optional_dependencies", return_value=modules) + mocker.patch.object(pink_ik, "prepare_urdf_for_drake", return_value=tmp_path / "prepared.urdf") + model_path = tmp_path / "fake.urdf" + model_path.write_text("") + world = _FakeWorld() + world.config.model_path = model_path + ik = PinkIK(PinkIKConfig(max_iterations=1)) + + first = ik._get_robot_context(cast("Any", world), "robot", "tool") + second = ik._get_robot_context(cast("Any", world), "robot", "base") + + assert first is not second + assert set(ik._robot_contexts) == {("robot", "tool"), ("robot", "base")} + + +def test_build_robot_context_rejects_base_link_not_model_root( + mocker: MockerFixture, tmp_path: Path +) -> None: + model = _FakeModel() + model.frames[0] = _FakeFrame("base", parent_joint=1) + modules = _fake_modules() + modules.pinocchio.buildModelFromUrdf = lambda path: model # type: ignore[attr-defined] + mocker.patch.object(pink_ik, "_load_optional_dependencies", return_value=modules) + mocker.patch.object(pink_ik, "prepare_urdf_for_drake", return_value=tmp_path / "prepared.urdf") + model_path = tmp_path / "fake.urdf" + model_path.write_text("") + config = _robot_config() + config.model_path = model_path + + with pytest.raises(ValueError, match="base_link 'base'.*model root"): + PinkIK(PinkIKConfig(max_iterations=1))._build_robot_context(config, "tool") + + +def test_solve_pose_targets_uses_group_tip_and_filters_group_joints( + mocker: MockerFixture, +) -> None: + ik = _pink_ik(mocker, converge=True) + context = _context() + get_context = mocker.patch.object(ik, "_get_robot_context", return_value=context) + mocker.patch.object( + ik, + "_solve_single", + return_value=IKResult( + status=IKStatus.SUCCESS, + joint_state=JointState( + {"name": ["joint_a", "joint_b", "joint_c"], "position": [0.1, 0.2, 0.3]} + ), + ), + ) + world = _FakeWorld() + + result = ik.solve_pose_targets( + world=cast("Any", world), + pose_targets={ + world.groups["arm/manipulator"]: PoseStamped( + position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + ) + }, + seed=JointState( + {"name": ["arm/joint_a", "arm/joint_b", "arm/joint_c"], "position": [0.0, 0.0, 0.0]} + ), + max_attempts=1, + ) + + get_context.assert_called_once_with(cast("Any", world), "robot", "tool") + assert result.status == IKStatus.SUCCESS + assert result.joint_state is not None + assert result.joint_state.name == ["arm/joint_a", "arm/joint_b"] + assert result.joint_state.position == [0.1, 0.2] + assert world.joint_state_calls == 0 + + +def test_solve_pose_targets_rejects_group_without_tip(mocker: MockerFixture) -> None: + ik = _pink_ik(mocker) + world = _FakeWorld() + + result = ik.solve_pose_targets( + world=cast("Any", world), + pose_targets={ + world.groups["arm/no_tip"]: PoseStamped( + position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + ) + }, + ) + + assert result.status == IKStatus.UNSUPPORTED + assert "no pose target frame" in result.message + + +def test_solve_pose_targets_partial_seed_reads_world_state(mocker: MockerFixture) -> None: + ik = _pink_ik(mocker) + mocker.patch.object(ik, "_get_robot_context", return_value=_context()) + mocker.patch.object( + ik, + "_solve_single", + return_value=IKResult( + status=IKStatus.SUCCESS, + joint_state=JointState( + {"name": ["joint_a", "joint_b", "joint_c"], "position": [0.1, 0.2, 0.3]} + ), + ), + ) + world = _FakeWorld() + + result = ik.solve_pose_targets( + world=cast("Any", world), + pose_targets={ + world.groups["arm/manipulator"]: PoseStamped( + position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + ) + }, + seed=JointState({"name": ["arm/joint_a"], "position": [0.0]}), + max_attempts=1, + ) + + assert result.status == IKStatus.SUCCESS + assert world.joint_state_calls == 1 + + +def test_solve_pose_targets_multi_target_uses_multi_frame_solve(mocker: MockerFixture) -> None: + ik = _pink_ik(mocker) + world = _FakeWorld() + mocker.patch.object(ik, "_get_robot_context", return_value=_context()) + solve_multi = mocker.patch.object( + ik, + "_solve_multi", + return_value=IKResult( + status=IKStatus.SUCCESS, + joint_state=JointState( + {"name": ["joint_a", "joint_b", "joint_c"], "position": [0.1, 0.2, 0.3]} + ), + position_error=0.0, + orientation_error=0.0, + ), + ) + + result = ik.solve_pose_targets( + world=cast("Any", world), + pose_targets={ + world.groups["arm/manipulator"]: PoseStamped( + position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + ), + world.groups["arm/wrist"]: PoseStamped( + position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + ), + }, + seed=JointState( + {"name": ["arm/joint_a", "arm/joint_b", "arm/joint_c"], "position": [0.0, 0.0, 0.0]} + ), + max_attempts=1, + ) + + solve_multi.assert_called_once() + assert len(solve_multi.call_args.kwargs["targets"]) == 2 + assert result.joint_state is not None + assert result.joint_state.name == ["arm/joint_a", "arm/joint_b", "arm/joint_c"] + assert result.joint_state.position == [0.1, 0.2, 0.3] + + +def test_solve_pose_targets_checks_multi_robot_solution_together( + mocker: MockerFixture, +) -> None: + ik = _pink_ik(mocker) + world = _MultiRobotCollisionWorld() + mocker.patch.object(ik, "_get_robot_context", return_value=_context()) + solve_single = mocker.patch.object( + ik, + "_solve_single", + side_effect=[ + IKResult( + status=IKStatus.SUCCESS, + joint_state=JointState( + {"name": ["joint_a", "joint_b", "joint_c"], "position": [0.1, 0.2, 0.3]} + ), + ), + IKResult( + status=IKStatus.SUCCESS, + joint_state=JointState( + {"name": ["joint_a", "joint_b", "joint_c"], "position": [0.4, 0.5, 0.6]} + ), + ), + ], + ) + + result = ik.solve_pose_targets( + world=cast("Any", world), + pose_targets={ + world.groups["left/manipulator"]: PoseStamped( + position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + ), + world.groups["right/manipulator"]: PoseStamped( + position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + ), + }, + seed=JointState( + { + "name": [ + "left/joint_a", + "left/joint_b", + "left/joint_c", + "right/joint_a", + "right/joint_b", + "right/joint_c", + ], + "position": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + } + ), + max_attempts=1, + ) + + assert solve_single.call_count == 2 + assert result.status == IKStatus.COLLISION + assert world.config_collision_checks == 0 + assert world.context_collision_checks == 1 + assert set(world.context_states) == {"left-id", "right-id"} + + +def test_solve_pose_targets_auxiliary_only_retains_seed_selection_order( + mocker: MockerFixture, +) -> None: + ik = _pink_ik(mocker) + world = _FakeWorld() + + result = ik.solve_pose_targets( + world=cast("Any", world), + pose_targets={}, + auxiliary_groups=[world.groups["arm/no_tip"], world.groups["arm/manipulator"]], + seed=JointState( + {"name": ["arm/joint_a", "arm/joint_b", "arm/joint_c"], "position": [0.1, 0.2, 0.3]} + ), + ) + + assert result.status == IKStatus.SUCCESS + assert result.joint_state is not None + assert result.joint_state.name == ["arm/joint_c", "arm/joint_a", "arm/joint_b"] + assert result.joint_state.position == [0.3, 0.1, 0.2] + assert world.joint_state_calls == 0 diff --git a/dimos/manipulation/planning/kinematics/utils.py b/dimos/manipulation/planning/kinematics/utils.py new file mode 100644 index 0000000000..3d0fa0509c --- /dev/null +++ b/dimos/manipulation/planning/kinematics/utils.py @@ -0,0 +1,236 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared IK-only helpers for planning-group-scoped kinematics backends.""" + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass + +import numpy as np +from numpy.typing import NDArray + +from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupSelection +from dimos.manipulation.planning.groups.utils import filter_joint_state_to_selected_joints +from dimos.manipulation.planning.spec.enums import IKStatus +from dimos.manipulation.planning.spec.models import IKResult, RobotName, WorldRobotID +from dimos.manipulation.planning.spec.protocols import WorldSpec +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.sensor_msgs.JointState import JointState + + +@dataclass(frozen=True) +class SinglePoseTargetRequest: + group: PlanningGroup + target_pose: PoseStamped + robot_id: WorldRobotID + joint_names: list[str] + seed_positions: NDArray[np.float64] + group_indices: list[int] + + +def unique_pose_target_frame_for_robot(world: WorldSpec, robot_id: WorldRobotID) -> str | None: + config = world.get_robot_config(robot_id) + pose_target_frames = [ + group.tip_link for group in config.planning_groups if group.tip_link is not None + ] + unique_frames = list(dict.fromkeys(pose_target_frames)) + if len(unique_frames) != 1: + return None + return unique_frames[0] + + +def robot_ids_by_name( + world: WorldSpec, + robot_names: tuple[RobotName, ...], +) -> dict[RobotName, WorldRobotID]: + robot_ids_by_name: dict[RobotName, WorldRobotID] = {} + for robot_name in robot_names: + matches = [ + robot_id + for robot_id in world.get_robot_ids() + if world.get_robot_config(robot_id).name == robot_name + ] + if not matches: + raise ValueError(f"Robot '{robot_name}' not found") + if len(matches) > 1: + raise ValueError(f"Robot name '{robot_name}' is not unique in planning world") + robot_ids_by_name[robot_name] = matches[0] + return robot_ids_by_name + + +def seed_positions_with_world_fallback( + world: WorldSpec, + robot_id: WorldRobotID, + robot_name: RobotName, + local_joint_names: list[str], + seed: JointState | None, +) -> NDArray[np.float64]: + """Return full robot positions, reading world only for absent seed joints.""" + if seed is None: + with world.scratch_context() as ctx: + current = world.get_joint_state(ctx, robot_id) + return positions_by_local_name(current, robot_name, local_joint_names) + + try: + return positions_by_local_name(seed, robot_name, local_joint_names) + except ValueError: + with world.scratch_context() as ctx: + current = world.get_joint_state(ctx, robot_id) + fallback_positions = positions_by_local_name(current, robot_name, local_joint_names) + seed_positions = partial_positions_by_local_name(seed, robot_name, local_joint_names) + local_indices = {name: index for index, name in enumerate(local_joint_names)} + for local_name, position in seed_positions.items(): + fallback_positions[local_indices[local_name]] = position + return fallback_positions + + +def resolve_single_pose_target_request( + world: WorldSpec, + pose_targets: dict[PlanningGroup, PoseStamped] | Mapping[PlanningGroup, PoseStamped], + auxiliary_groups: Sequence[PlanningGroup], + seed: JointState | None, + backend_name: str, +) -> tuple[SinglePoseTargetRequest | None, IKResult | None]: + if not pose_targets: + return None, _failure(IKStatus.NO_SOLUTION, "At least one pose target is required") + if len(pose_targets) != 1 or auxiliary_groups: + return None, _failure( + IKStatus.UNSUPPORTED, + f"{backend_name} supports exactly one pose target and no auxiliary planning groups", + ) + + target_group = next(iter(pose_targets.keys())) + if not target_group.has_pose_target: + return None, _failure( + IKStatus.UNSUPPORTED, + f"Planning group '{target_group.id}' has no pose target frame", + ) + + try: + selection = PlanningGroupSelection.from_groups((target_group,)) + robot_id = robot_ids_by_name(world, selection.robot_names)[target_group.robot_name] + config = world.get_robot_config(robot_id) + joint_names = list(config.joint_names) + seed_positions = seed_positions_with_world_fallback( + world, + robot_id, + config.name, + joint_names, + seed, + ) + group_indices = [joint_names.index(name) for name in target_group.local_joint_names] + except ValueError as exc: + return None, _failure(IKStatus.NO_SOLUTION, str(exc)) + + return ( + SinglePoseTargetRequest( + group=target_group, + target_pose=pose_targets[target_group], + robot_id=robot_id, + joint_names=joint_names, + seed_positions=seed_positions, + group_indices=group_indices, + ), + None, + ) + + +def positions_by_local_name( + joint_state: JointState, + robot_name: RobotName, + local_joint_names: list[str], +) -> NDArray[np.float64]: + if not joint_state.name: + if len(joint_state.position) != len(local_joint_names): + raise ValueError( + f"JointState has {len(joint_state.position)} positions for " + f"{len(local_joint_names)} joints" + ) + return np.asarray(joint_state.position, dtype=np.float64) + + positions_by_name = dict(zip(joint_state.name, joint_state.position, strict=True)) + positions: list[float] = [] + missing: list[str] = [] + for local_name in local_joint_names: + global_name = f"{robot_name}/{local_name}" + if local_name in positions_by_name: + positions.append(float(positions_by_name[local_name])) + elif global_name in positions_by_name: + positions.append(float(positions_by_name[global_name])) + else: + missing.append(local_name) + if missing: + raise ValueError(f"JointState missing joints: {missing}") + return np.asarray(positions, dtype=np.float64) + + +def partial_positions_by_local_name( + joint_state: JointState, + robot_name: RobotName, + local_joint_names: list[str], +) -> dict[str, float]: + if len(joint_state.name) != len(joint_state.position): + raise ValueError( + f"Seed has {len(joint_state.name)} names but {len(joint_state.position)} positions" + ) + positions_by_name = dict(zip(joint_state.name, joint_state.position, strict=True)) + known_local_names = set(local_joint_names) + positions: dict[str, float] = {} + for name, position in positions_by_name.items(): + if name in known_local_names: + positions[name] = float(position) + continue + prefix = f"{robot_name}/" + if name.startswith(prefix): + local_name = name[len(prefix) :] + if local_name in known_local_names: + positions[local_name] = float(position) + continue + raise ValueError(f"Unrecognized seed joint '{name}'") + return positions + + +def filter_result_to_group(result: IKResult, group: PlanningGroup) -> IKResult: + return filter_result_to_selection(result, PlanningGroupSelection.from_groups((group,))) + + +def filter_result_to_selection(result: IKResult, selection: PlanningGroupSelection) -> IKResult: + if result.joint_state is None: + return result + local_joint_names = tuple( + local_name for group in selection.groups for local_name in group.local_joint_names + ) + return IKResult( + status=result.status, + joint_state=filter_joint_state_to_selected_joints( + result.joint_state, + selection.joint_names, + local_joint_names, + ), + position_error=result.position_error, + orientation_error=result.orientation_error, + iterations=result.iterations, + message=result.message, + ) + + +def groups_by_robot(groups: Sequence[PlanningGroup]) -> dict[RobotName, list[PlanningGroup]]: + grouped: dict[RobotName, list[PlanningGroup]] = {} + for group in groups: + grouped.setdefault(group.robot_name, []).append(group) + return grouped + + +def _failure(status: IKStatus, message: str) -> IKResult: + return IKResult(status=status, joint_state=None, message=message) diff --git a/dimos/manipulation/planning/monitor/test_world_monitor.py b/dimos/manipulation/planning/monitor/test_world_monitor.py index df75d23d03..aa868e2892 100644 --- a/dimos/manipulation/planning/monitor/test_world_monitor.py +++ b/dimos/manipulation/planning/monitor/test_world_monitor.py @@ -17,22 +17,64 @@ from pathlib import Path from typing import Any +import numpy as np +import pytest + from dimos.manipulation.planning import factory as planning_factory +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.monitor import world_monitor as world_monitor_module from dimos.manipulation.planning.spec.config import RobotModelConfig -from dimos.manipulation.planning.spec.models import PlanningSceneInfo +from dimos.manipulation.planning.spec.models import ( + PlanningSceneInfo, + VisualizationSession, + VisualizationStateFrame, +) +from dimos.manipulation.planning.spec.protocols import VisualizationSpec 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.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory + + +class _VectorLike(list[float]): + def tolist(self) -> list[float]: + return list(self) + + +class _FakeStateMonitor: + def __init__(self, positions: list[float], stale: bool = False) -> None: + self._positions = _VectorLike(positions) + self._stale = stale + + def get_current_positions(self) -> _VectorLike: + return self._positions + + def get_current_velocities(self) -> None: + return None + + def is_state_stale(self, max_age: float) -> bool: + return self._stale + + +class _ScratchContext: + def __enter__(self) -> str: + return "scratch" + + def __exit__(self, exc_type: object, exc: object, traceback: object) -> bool: + return False class FakeWorld: def __init__(self) -> None: - self.calls: list[tuple[str, Any]] = [] + self.calls: list[tuple[Any, ...]] = [] + self.configs: dict[str, RobotModelConfig] = {} def add_robot(self, config): self.calls.append(("add_robot", config)) - return "robot-1" + robot_id = f"robot-{len(self.configs) + 1}" + self.configs[robot_id] = config + return robot_id def get_robot_ids(self): return [] @@ -69,12 +111,14 @@ def get_live_context(self): return None def scratch_context(self): - return self + self.calls.append(("scratch_context", None)) + return _ScratchContext() def sync_from_joint_state(self, robot_id, joint_state) -> None: return None def set_joint_state(self, ctx, robot_id, joint_state) -> None: + self.calls.append(("set_joint_state", ctx, robot_id, joint_state)) return None def get_joint_state(self, ctx, robot_id): @@ -95,28 +139,33 @@ def check_edge_collision_free(self, robot_id, start, end, step_size: float = 0.0 def get_ee_pose(self, ctx, robot_id): return None + def get_group_ee_pose(self, ctx, group_id): + self.calls.append(("get_group_ee_pose", ctx, group_id)) + return PoseStamped(position=Vector3(1, 2, 3), orientation=Quaternion([0, 0, 0, 1])) + def get_link_pose(self, ctx, robot_id, link_name): return [] def get_jacobian(self, ctx, robot_id): return [] - def get_visualization_url(self): - return None + def get_group_jacobian(self, ctx, group_id): + self.calls.append(("get_group_jacobian", ctx, group_id)) + return np.ones((6, 2)) - def initialize_scene(self, scene: PlanningSceneInfo) -> None: + def get_visualization_url(self): return None - def publish_visualization(self, ctx=None) -> None: + def initialize(self, session: VisualizationSession) -> None: return None - def show_preview(self, robot_id) -> None: + def update_state(self, frame: VisualizationStateFrame) -> None: return None - def hide_preview(self, robot_id) -> None: + def animate_trajectory(self, trajectory, duration: float | None = None) -> None: return None - def animate_path(self, robot_id, path, duration: float = 3.0) -> None: + def cancel_preview_animation(self) -> None: return None def close(self) -> None: @@ -130,20 +179,17 @@ def __init__(self) -> None: def get_visualization_url(self): return None - def initialize_scene(self, scene: PlanningSceneInfo) -> None: - self.calls.append(("initialize_scene", scene)) + def initialize(self, session: VisualizationSession) -> None: + self.calls.append(("initialize", session)) - def publish_visualization(self, ctx=None) -> None: - return None + def update_state(self, frame: VisualizationStateFrame) -> None: + self.calls.append(("update_state", frame)) - def show_preview(self, robot_id) -> None: - self.calls.append(("show_preview", robot_id)) + def animate_trajectory(self, trajectory, duration: float | None = None) -> None: + self.calls.append(("animate_trajectory", trajectory, duration)) - def hide_preview(self, robot_id) -> None: - self.calls.append(("hide_preview", robot_id)) - - def animate_path(self, robot_id, path, duration: float = 3.0) -> None: - return None + def cancel_preview_animation(self) -> None: + self.calls.append(("cancel_preview_animation",)) def close(self) -> None: self.calls.append(("close", None)) @@ -155,8 +201,32 @@ def _robot_config() -> RobotModelConfig: model_path=Path("/tmp/arm.urdf"), base_pose=PoseStamped(position=Vector3(), orientation=Quaternion([0, 0, 0, 1])), joint_names=["j1", "j2"], - end_effector_link="ee", base_link="base", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", joint_names=("j1", "j2"), base_link="base", tip_link="ee" + ) + ], + ) + + +def _robot_config_with_groups(groups: list[PlanningGroupDefinition]) -> RobotModelConfig: + return _robot_config().model_copy(update={"planning_groups": groups}) + + +def _three_joint_reordered_group_config() -> RobotModelConfig: + return _robot_config().model_copy( + update={ + "joint_names": ["j1", "j2", "j3"], + "planning_groups": [ + PlanningGroupDefinition( + name="manipulator", + joint_names=("j2", "j1"), + base_link="base", + tip_link="ee", + ) + ], + } ) @@ -180,10 +250,28 @@ def test_world_monitor_syncs_planning_scene_to_visualization() -> None: monitor.add_robot(_robot_config()) monitor.sync_visualization_scene() - assert fake_viz.calls[0][0] == "initialize_scene" - scene = fake_viz.calls[0][1] + assert fake_viz.calls[0][0] == "initialize" + session = fake_viz.calls[0][1] + scene = session.scene assert isinstance(scene, PlanningSceneInfo) assert scene.robots["robot-1"].name == "arm" + assert scene.planning_groups[0].id == "arm/manipulator" + + +def test_world_monitor_forwards_raw_trajectory_preview_protocol() -> None: + fake_viz = FakeViz() + monitor = world_monitor_module.WorldMonitor(world=FakeWorld(), visualization=fake_viz) # type: ignore[arg-type] + trajectory = JointTrajectory(joint_names=["arm/j1"], points=[]) + + assert isinstance(fake_viz, VisualizationSpec) + monitor.cancel_preview_animation() + monitor.animate_trajectory(trajectory, 2.0) + + assert fake_viz.calls == [ + ("cancel_preview_animation",), + ("cancel_preview_animation",), + ("animate_trajectory", trajectory, 2.0), + ] def test_create_planning_specs_wraps_existing_world(monkeypatch) -> None: @@ -204,3 +292,217 @@ def test_create_planning_specs_wraps_existing_world(monkeypatch) -> None: assert planning_specs.world_monitor.visualization is None assert planning_specs.kinematics is fake_kinematics assert planning_specs.planner is fake_planner + + +def test_world_monitor_exposes_planning_groups_and_duplicate_names_do_not_mutate() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + monitor.add_robot(_robot_config()) + + assert [group.id for group in monitor.planning_groups.list()] == ["arm/manipulator"] + with pytest.raises(ValueError, match="already registered"): + monitor.add_robot(_robot_config()) + assert [call[0] for call in fake_world.calls].count("add_robot") == 1 + + +def test_world_monitor_invalid_duplicate_group_config_does_not_mutate_backend() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + invalid_config = _robot_config_with_groups( + [ + PlanningGroupDefinition( + name="manipulator", joint_names=("j1",), base_link="base", tip_link="ee" + ), + PlanningGroupDefinition( + name="manipulator", joint_names=("j2",), base_link="base", tip_link="ee" + ), + ] + ) + + with pytest.raises(ValueError, match="already registered"): + monitor.add_robot(invalid_config) + + assert [call[0] for call in fake_world.calls].count("add_robot") == 0 + + +def test_world_monitor_invalid_group_joint_name_does_not_mutate_backend() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + invalid_config = _robot_config_with_groups( + [ + PlanningGroupDefinition( + name="manipulator", + joint_names=("j1", "bad/joint"), + base_link="base", + tip_link="ee", + ) + ] + ) + + with pytest.raises(ValueError, match="Invalid local joint name"): + monitor.add_robot(invalid_config) + + assert [call[0] for call in fake_world.calls].count("add_robot") == 0 + + +def test_current_group_joint_state_uses_public_names_in_group_order() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + robot_id = monitor.add_robot(_three_joint_reordered_group_config()) + monitor._state_monitors[robot_id] = _FakeStateMonitor([0.1, 0.2, 0.3]) # type: ignore[attr-defined] + + state = monitor.current_group_joint_state("arm/manipulator") + + assert state.name == ["arm/j2", "arm/j1"] + assert state.position == [0.2, 0.1] + + +def test_current_global_joint_state_skips_stale_robots_and_preserves_state_order() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + fresh_id = monitor.add_robot(_three_joint_reordered_group_config()) + stale_id = monitor.add_robot( + RobotModelConfig( + name="arm2", + model_path=Path("/tmp/arm2.urdf"), + joint_names=["a", "b"], + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", joint_names=("a", "b"), base_link="base", tip_link="ee" + ) + ], + ) + ) + monitor._state_monitors[fresh_id] = _FakeStateMonitor([0.1, 0.2, 0.3]) # type: ignore[attr-defined] + monitor._state_monitors[stale_id] = _FakeStateMonitor([1.0, 2.0], stale=True) # type: ignore[attr-defined] + monitor.add_robot( + RobotModelConfig( + name="arm3", + model_path=Path("/tmp/arm3.urdf"), + joint_names=["x"], + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", joint_names=("x",), base_link="base", tip_link="ee" + ) + ], + ) + ) + + state = monitor.current_global_joint_state(max_age=0.5) + + assert state.name == ["arm/j1", "arm/j2", "arm/j3"] + assert state.position == [0.1, 0.2, 0.3] + + +def test_current_group_joint_state_rejects_stale_or_unavailable_state() -> None: + stale_world = FakeWorld() + stale_monitor = world_monitor_module.WorldMonitor(world=stale_world) # type: ignore[arg-type] + stale_id = stale_monitor.add_robot(_three_joint_reordered_group_config()) + stale_monitor._state_monitors[stale_id] = _FakeStateMonitor([0.1, 0.2, 0.3], stale=True) # type: ignore[attr-defined] + + with pytest.raises(ValueError, match="stale"): + stale_monitor.current_group_joint_state("arm/manipulator") + + unavailable_monitor = world_monitor_module.WorldMonitor(world=FakeWorld()) # type: ignore[arg-type] + unavailable_monitor.add_robot(_three_joint_reordered_group_config()) + with pytest.raises(ValueError, match="unavailable"): + unavailable_monitor.current_group_joint_state("arm/manipulator") + + +def test_group_ee_pose_uses_current_state_when_no_joint_state_is_provided() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + robot_id = monitor.add_robot(_three_joint_reordered_group_config()) + monitor._state_monitors[robot_id] = _FakeStateMonitor([0.1, 0.2, 0.3]) # type: ignore[attr-defined] + + pose = monitor.get_group_ee_pose("arm/manipulator") + + set_calls = [call for call in fake_world.calls if call[0] == "set_joint_state"] + assert set_calls[0][3].name == ["j1", "j2", "j3"] + assert set_calls[0][3].position == [0.1, 0.2, 0.3] + assert pose.position.x == 1 + + +def test_group_ee_pose_without_joint_state_rejects_stale_or_unavailable_state() -> None: + stale_world = FakeWorld() + stale_monitor = world_monitor_module.WorldMonitor(world=stale_world) # type: ignore[arg-type] + stale_id = stale_monitor.add_robot(_three_joint_reordered_group_config()) + stale_monitor._state_monitors[stale_id] = _FakeStateMonitor([0.1, 0.2, 0.3], stale=True) # type: ignore[attr-defined] + + with pytest.raises(ValueError, match="stale"): + stale_monitor.get_group_ee_pose("arm/manipulator") + + unavailable_monitor = world_monitor_module.WorldMonitor(world=FakeWorld()) # type: ignore[arg-type] + unavailable_monitor.add_robot(_three_joint_reordered_group_config()) + with pytest.raises(ValueError, match="unavailable"): + unavailable_monitor.get_group_ee_pose("arm/manipulator") + + +def test_group_kinematics_with_full_state_does_not_require_current_state() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + monitor.add_robot(_three_joint_reordered_group_config()) + + pose = monitor.get_group_ee_pose( + "arm/manipulator", + JointState(name=["j1", "j2", "j3"], position=[0.1, 0.2, 0.3]), + ) + + set_calls = [call for call in fake_world.calls if call[0] == "set_joint_state"] + assert set_calls[0][3].name == ["j1", "j2", "j3"] + assert set_calls[0][3].position == [0.1, 0.2, 0.3] + assert pose.position.x == 1 + + +def test_group_kinematics_route_full_state_to_backend() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + monitor.add_robot(_three_joint_reordered_group_config()) + + pose = monitor.get_group_ee_pose( + "arm/manipulator", + JointState(name=["j1", "j2", "j3"], position=[0.9, 0.8, 0.3]), + ) + jacobian = monitor.get_group_jacobian( + "arm/manipulator", + JointState(name=["j1", "j2", "j3"], position=[0.4, 0.3, 0.3]), + ) + + set_calls = [call for call in fake_world.calls if call[0] == "set_joint_state"] + assert set_calls[0][3].name == ["j1", "j2", "j3"] + assert set_calls[0][3].position == [0.9, 0.8, 0.3] + assert set_calls[1][3].name == ["j1", "j2", "j3"] + assert set_calls[1][3].position == [0.4, 0.3, 0.3] + assert pose.position.x == 1 + assert jacobian.shape == (6, 2) + assert ("get_group_ee_pose", "scratch", "arm/manipulator") in fake_world.calls + assert ("get_group_jacobian", "scratch", "arm/manipulator") in fake_world.calls + + +def test_legacy_wrappers_fail_for_no_pose_and_ambiguous_pose_groups() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + no_pose_id = monitor.add_robot( + _robot_config_with_groups( + [PlanningGroupDefinition(name="base", joint_names=("j1",), base_link="base")] + ) + ) + with pytest.raises(ValueError, match="no pose-targetable"): + monitor.get_ee_pose(no_pose_id, JointState(name=["j1", "j2"], position=[0.0, 0.0])) + + fake_world2 = FakeWorld() + monitor2 = world_monitor_module.WorldMonitor(world=fake_world2) # type: ignore[arg-type] + ambiguous_id = monitor2.add_robot( + _robot_config_with_groups( + [ + PlanningGroupDefinition( + name="a", joint_names=("j1",), base_link="base", tip_link="ee1" + ), + PlanningGroupDefinition( + name="b", joint_names=("j2",), base_link="base", tip_link="ee2" + ), + ] + ) + ) + with pytest.raises(ValueError, match="pose-targetable planning groups"): + monitor2.get_jacobian(ambiguous_id, JointState(name=["j1", "j2"], position=[0.0, 0.0])) diff --git a/dimos/manipulation/planning/monitor/world_monitor.py b/dimos/manipulation/planning/monitor/world_monitor.py index 5e12568874..ed4da16837 100644 --- a/dimos/manipulation/planning/monitor/world_monitor.py +++ b/dimos/manipulation/planning/monitor/world_monitor.py @@ -16,17 +16,29 @@ from __future__ import annotations +from collections.abc import Sequence from contextlib import contextmanager import threading from typing import TYPE_CHECKING, Any from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT +from dimos.manipulation.planning.groups.identifiers import ( + make_global_joint_names, + make_planning_group_id, +) +from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry +from dimos.manipulation.planning.groups.utils import filter_joint_state_to_selected_joints from dimos.manipulation.planning.monitor.robot_state_monitor import RobotStateMonitor from dimos.manipulation.planning.monitor.world_obstacle_monitor import WorldObstacleMonitor -from dimos.manipulation.planning.spec.models import PlanningSceneInfo +from dimos.manipulation.planning.spec.models import ( + PlanningSceneInfo, + VisualizationSession, + VisualizationStateFrame, +) from dimos.manipulation.planning.spec.protocols import VisualizationSpec, WorldSpec from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: @@ -40,6 +52,8 @@ CollisionObjectMessage, JointPath, Obstacle, + PlanningGroupID, + RobotName, WorldRobotID, ) from dimos.msgs.vision_msgs.Detection3D import Detection3D @@ -59,8 +73,13 @@ def __init__( self._world = world self._visualization = visualization self._lock = threading.RLock() + # Keep renderer mutations and periodic publishes ordered. Cancellation is + # deliberately issued outside this lock so it can interrupt an animation. + self._visualization_lock = threading.RLock() self._robot_joints: dict[WorldRobotID, list[str]] = {} self._robot_configs: dict[WorldRobotID, RobotModelConfig] = {} + self._robot_ids_by_name: dict[RobotName, WorldRobotID] = {} + self._planning_groups = PlanningGroupRegistry() self._state_monitors: dict[WorldRobotID, RobotStateMonitor] = {} self._obstacle_monitor: WorldObstacleMonitor | None = None self._viz_thread: threading.Thread | None = None @@ -72,23 +91,42 @@ def __init__( def add_robot(self, config: RobotModelConfig) -> WorldRobotID: """Add a robot. Returns robot_id.""" with self._lock: + if config.name in self._robot_ids_by_name: + raise ValueError(f"Robot name '{config.name}' is already registered") + self._validate_planning_group_config(config) robot_id = self._world.add_robot(config) self._robot_joints[robot_id] = config.joint_names self._robot_configs[robot_id] = config + self._robot_ids_by_name[config.name] = robot_id + self._planning_groups.add_robot(config) logger.info(f"Added robot '{config.name}' as '{robot_id}'") return robot_id + @property + def planning_groups(self) -> PlanningGroupRegistry: + """Registered public planning groups.""" + return self._planning_groups + def planning_scene_info(self) -> PlanningSceneInfo: """Return a stable metadata snapshot of the initialized planning scene.""" with self._lock: - return PlanningSceneInfo(robots=dict(self._robot_configs)) + return PlanningSceneInfo( + robots=dict(self._robot_configs), + planning_groups=tuple(self._planning_groups.list()), + ) - def sync_visualization_scene(self) -> None: - """Synchronize startup scene metadata to the attached visualization.""" + def initialize_visualization(self, operator: object | None = None) -> None: + """Initialize attached visualization with immutable startup metadata.""" visualization = self._visualization if visualization is None: return - visualization.initialize_scene(self.planning_scene_info()) + visualization.initialize( + VisualizationSession(scene=self.planning_scene_info(), operator=operator) + ) + + def sync_visualization_scene(self) -> None: + """Compatibility wrapper for initializing visualization metadata.""" + self.initialize_visualization() def get_robot_ids(self) -> list[WorldRobotID]: """Get all robot IDs.""" @@ -194,7 +232,10 @@ def stop_all_monitors(self) -> None: logger.info("All monitors stopped") if self._visualization is not None: - self._visualization.close() + self._visualization.cancel_preview_animation() + # Wait for cancelled animation cleanup before releasing renderer resources. + with self._visualization_lock: + self._visualization.close() # Message Handlers @@ -294,6 +335,46 @@ def get_current_joint_state(self, robot_id: WorldRobotID) -> JointState | None: ctx = self._world.get_live_context() return self._world.get_joint_state(ctx, robot_id) + def current_global_joint_state(self, max_age: float = 1.0) -> JointState: + """Return current state for all fresh robots with public global joint names.""" + names: list[str] = [] + positions: list[float] = [] + for robot_name, robot_id in self._robot_ids_by_name.items(): + if robot_id in self._state_monitors and self.is_state_stale(robot_id, max_age): + continue + state = self.get_current_joint_state(robot_id) + if state is None: + continue + for name, position in zip(state.name, state.position, strict=True): + names.append(f"{robot_name}/{name}") + positions.append(float(position)) + return JointState(name=names, position=positions) + + def current_group_joint_state( + self, group_id: PlanningGroupID, max_age: float = 1.0 + ) -> JointState: + """Return current joint state scoped and ordered for one planning group.""" + group = self._planning_groups.get(group_id) + robot_id = self._robot_ids_by_name[group.robot_name] + if robot_id in self._state_monitors and self.is_state_stale(robot_id, max_age): + raise ValueError(f"Current state for robot '{group.robot_name}' is stale") + state = self.get_current_joint_state(robot_id) + if state is None: + raise ValueError(f"Current state for robot '{group.robot_name}' is unavailable") + return filter_joint_state_to_selected_joints( + state, group.joint_names, group.local_joint_names + ) + + def _validate_planning_group_config(self, config: RobotModelConfig) -> None: + """Validate planning groups before mutating world/backend state.""" + seen_group_names: set[str] = set() + for definition in config.planning_groups: + group_id = make_planning_group_id(config.name, definition.name) + if definition.name in seen_group_names: + raise ValueError(f"Planning group '{group_id}' is already registered") + make_global_joint_names(config.name, definition.joint_names) + seen_group_names.add(definition.name) + def get_current_velocities(self, robot_id: WorldRobotID) -> JointState | None: """Get current joint velocities as JointState. Returns None if not available.""" if robot_id in self._state_monitors: @@ -367,15 +448,28 @@ def get_ee_pose( self, robot_id: WorldRobotID, joint_state: JointState | None = None ) -> PoseStamped: """Get end-effector pose. Uses current state if joint_state is None.""" + robot_name = self._robot_configs[robot_id].name + group_id = self._planning_groups.primary_pose_group_id_for_robot(robot_name) + if group_id is None: + raise ValueError(f"Robot '{robot_name}' has no pose-targetable planning group") + return self.get_group_ee_pose(group_id, joint_state) + + def get_group_ee_pose( + self, group_id: PlanningGroupID, joint_state: JointState | None = None + ) -> PoseStamped: + """Get planning-group tip pose. Uses current robot state if joint_state is None.""" + group = self._planning_groups.get(group_id) + robot_id = self._robot_ids_by_name[group.robot_name] with self._world.scratch_context() as ctx: - # If no state provided, fetch current from state monitor if joint_state is None: + if robot_id in self._state_monitors and self.is_state_stale(robot_id): + raise ValueError(f"Current state for robot '{group.robot_name}' is stale") joint_state = self.get_current_joint_state(robot_id) + if joint_state is None: + raise ValueError(f"Current state for robot '{group.robot_name}' is unavailable") + self._world.set_joint_state(ctx, robot_id, joint_state) - if joint_state is not None: - self._world.set_joint_state(ctx, robot_id, joint_state) - - return self._world.get_ee_pose(ctx, robot_id) + return self._world.get_group_ee_pose(ctx, group_id) def get_link_pose( self, robot_id: WorldRobotID, link_name: str, joint_state: JointState | None = None @@ -411,9 +505,21 @@ def get_link_pose( def get_jacobian(self, robot_id: WorldRobotID, joint_state: JointState) -> NDArray[np.float64]: """Get 6xN Jacobian matrix.""" + robot_name = self._robot_configs[robot_id].name + group_id = self._planning_groups.primary_pose_group_id_for_robot(robot_name) + if group_id is None: + raise ValueError(f"Robot '{robot_name}' has no pose-targetable planning group") + return self.get_group_jacobian(group_id, joint_state) + + def get_group_jacobian( + self, group_id: PlanningGroupID, joint_state: JointState + ) -> NDArray[np.float64]: + """Get 6xN planning-group Jacobian matrix.""" + group = self._planning_groups.get(group_id) + robot_id = self._robot_ids_by_name[group.robot_name] with self._world.scratch_context() as ctx: self._world.set_joint_state(ctx, robot_id, joint_state) - return self._world.get_jacobian(ctx, robot_id) + return self._world.get_group_jacobian(ctx, group_id) # Lifecycle @@ -437,30 +543,57 @@ def get_visualization_url(self) -> str | None: return str(url) if url else None return None - def publish_visualization(self) -> None: - """Force publish current state to visualization.""" - if self._visualization is not None: - self._visualization.publish_visualization() - - def show_preview(self, robot_id: WorldRobotID) -> None: - """Show the preview representation for a robot if visualization is available.""" + def visualization_state_frame(self) -> VisualizationStateFrame: + """Build a pushed visualization state frame without freshness policy.""" + joint_states: dict[str, JointState] = {} + with self._lock: + robot_ids = list(self._robot_configs.keys()) + for robot_id in robot_ids: + state = self.get_current_joint_state(robot_id) + if state is not None: + joint_states[robot_id] = state + return VisualizationStateFrame(joint_states=joint_states) + + def update_visualization_state(self) -> None: + """Push current state to visualization.""" if self._visualization is not None: - self._visualization.show_preview(robot_id) + with self._visualization_lock: + self._visualization.update_state(self.visualization_state_frame()) - def hide_preview(self, robot_id: WorldRobotID) -> None: - """Hide the preview representation for a robot if visualization is available.""" + def cancel_preview_animation(self, robot_ids: Sequence[WorldRobotID] | None = None) -> None: + """Cancel active visualization preview animation.""" if self._visualization is not None: - self._visualization.hide_preview(robot_id) + if robot_ids is None: + self._visualization.cancel_preview_animation() + else: + self._visualization.cancel_preview_animation(robot_ids) - def animate_path( - self, - robot_id: WorldRobotID, - path: JointPath, - duration: float = 3.0, + def animate_trajectory( + self, trajectory: JointTrajectory, duration: float | None = None ) -> None: - """Animate a path if visualization is available.""" + """Animate a raw generated-plan trajectory if visualization is available.""" if self._visualization is not None: - self._visualization.animate_path(robot_id, path, duration) + robot_ids = self.robot_ids_for_global_joints(trajectory.joint_names) + if robot_ids: + self._visualization.cancel_preview_animation(robot_ids) + else: + self._visualization.cancel_preview_animation() + with self._visualization_lock: + self._visualization.animate_trajectory(trajectory, duration) + + def robot_ids_for_global_joints(self, joint_names: Sequence[str]) -> tuple[WorldRobotID, ...]: + """Return visualization robot IDs affected by globally named trajectory joints.""" + robot_ids: list[WorldRobotID] = [] + with self._lock: + by_name = {config.name: robot_id for robot_id, config in self._robot_configs.items()} + for joint_name in joint_names: + if "/" not in joint_name: + continue + robot_name, _ = joint_name.split("/", 1) + robot_id = by_name.get(robot_name) + if robot_id is not None and robot_id not in robot_ids: + robot_ids.append(robot_id) + return tuple(robot_ids) def start_visualization_thread(self, rate_hz: float = 10.0) -> None: """Start background thread for visualization updates at given rate.""" @@ -501,7 +634,7 @@ def _visualization_loop(self) -> None: period = 1.0 / self._viz_rate_hz while not self._viz_stop_event.is_set(): try: - self.publish_visualization() + self.update_visualization_state() except Exception as e: logger.debug(f"Visualization publish failed: {e}") time.sleep(period) diff --git a/dimos/manipulation/planning/planners/rrt_planner.py b/dimos/manipulation/planning/planners/rrt_planner.py index 3ca19eb099..a98ea358e6 100644 --- a/dimos/manipulation/planning/planners/rrt_planner.py +++ b/dimos/manipulation/planning/planners/rrt_planner.py @@ -26,8 +26,17 @@ import numpy as np +from dimos.manipulation.planning.groups.models import PlanningGroupSelection +from dimos.manipulation.planning.planners.selected_joint_space import ( + SelectedJointSpace, + normalize_selection_target, +) from dimos.manipulation.planning.spec.enums import PlanningStatus -from dimos.manipulation.planning.spec.models import JointPath, PlanningResult, WorldRobotID +from dimos.manipulation.planning.spec.models import ( + JointPath, + PlanningResult, + WorldRobotID, +) from dimos.manipulation.planning.spec.protocols import WorldSpec from dimos.manipulation.planning.utils.path_utils import compute_path_length from dimos.msgs.sensor_msgs.JointState import JointState @@ -98,6 +107,9 @@ def plan_joint_path( if error is not None: return error + if world.check_edge_collision_free(robot_id, start, goal, self._collision_step_size): + return _create_success_result([start, goal], time.time() - start_time, 0) + lower, upper = world.get_joint_limits(robot_id) start_tree = [TreeNode(config=q_start.copy())] goal_tree = [TreeNode(config=q_goal.copy())] @@ -147,6 +159,183 @@ def get_name(self) -> str: """Get planner name.""" return "RRTConnect" + def plan_selected_joint_path( + self, + world: WorldSpec, + selection: PlanningGroupSelection, + start: JointState, + goal: JointState, + timeout: float = 10.0, + max_iterations: int = 5000, + ) -> PlanningResult: + """Plan over an explicit planning-group selection. + + The search space is the selected global-joint order. Collision checks project + candidates into full per-robot states, holding unselected joints at the world + current state. + """ + start_time = time.time() + if not world.is_finalized: + return _create_failure_result( + PlanningStatus.NO_SOLUTION, + "World must be finalized before planning", + ) + + if not selection.groups: + return _create_failure_result( + PlanningStatus.INVALID_GOAL, "No planning groups selected" + ) + + selected_joint_names = list(selection.joint_names) + try: + normalized_start = normalize_selection_target(selection, start, "start") + except ValueError as exc: + return _create_failure_result(PlanningStatus.INVALID_START, str(exc)) + try: + normalized_goal = normalize_selection_target(selection, goal, "goal") + except ValueError as exc: + return _create_failure_result(PlanningStatus.INVALID_GOAL, str(exc)) + try: + selected_space = SelectedJointSpace.from_world(world, selection) + q_start = np.asarray(normalized_start.position, dtype=np.float64) + q_goal = np.asarray(normalized_goal.position, dtype=np.float64) + lower, upper = selected_space.joint_limits() + except ValueError as exc: + return _create_failure_result(PlanningStatus.NO_SOLUTION, str(exc)) + + if np.any(q_start < lower) or np.any(q_start > upper): + return _create_failure_result( + PlanningStatus.INVALID_START, + "Start configuration is outside joint limits", + ) + if np.any(q_goal < lower) or np.any(q_goal > upper): + return _create_failure_result( + PlanningStatus.INVALID_GOAL, + "Goal configuration is outside joint limits", + ) + + if not selected_space.config_collision_free(world, q_start): + return _create_failure_result( + PlanningStatus.COLLISION_AT_START, + "Start configuration is in collision", + ) + if not selected_space.config_collision_free(world, q_goal): + return _create_failure_result( + PlanningStatus.COLLISION_AT_GOAL, + "Goal configuration is in collision", + ) + + if selected_space.edge_collision_free( + world, + q_start, + q_goal, + self._collision_step_size, + ): + return _create_success_result( + [normalized_start, normalized_goal], time.time() - start_time, 0 + ) + + start_tree = [TreeNode(config=q_start.copy())] + goal_tree = [TreeNode(config=q_goal.copy())] + trees_swapped = False + + for iteration in range(max_iterations): + if time.time() - start_time > timeout: + return _create_failure_result( + PlanningStatus.TIMEOUT, + f"Timeout after {iteration} iterations", + time.time() - start_time, + iteration, + ) + + sample = np.random.uniform(lower, upper) + extended = self._extend_selected_tree( + selected_space, + world, + start_tree, + sample, + self._step_size, + ) + if extended is not None: + connected = self._connect_selected_tree( + selected_space, + world, + goal_tree, + extended.config, + self._connect_step_size, + ) + if connected is not None: + path = self._extract_path(extended, connected, selected_joint_names) + if trees_swapped: + path = list(reversed(path)) + path = selected_space.simplify_path( + world, + path, + self._collision_step_size, + ) + return _create_success_result(path, time.time() - start_time, iteration + 1) + + start_tree, goal_tree = goal_tree, start_tree + trees_swapped = not trees_swapped + + return _create_failure_result( + PlanningStatus.NO_SOLUTION, + f"No path found after {max_iterations} iterations", + time.time() - start_time, + max_iterations, + ) + + def _extend_selected_tree( + self, + selected_space: SelectedJointSpace, + world: WorldSpec, + tree: list[TreeNode], + target: NDArray[np.float64], + step_size: float, + ) -> TreeNode | None: + """Extend a tree in selected-joint space.""" + nearest = min(tree, key=lambda node: float(np.linalg.norm(node.config - target))) + diff = target - nearest.config + dist = float(np.linalg.norm(diff)) + if dist <= step_size: + new_config = target.copy() + else: + new_config = nearest.config + step_size * (diff / dist) + + if selected_space.edge_collision_free( + world, + nearest.config, + new_config, + self._collision_step_size, + ): + new_node = TreeNode(config=new_config, parent=nearest) + nearest.children.append(new_node) + tree.append(new_node) + return new_node + return None + + def _connect_selected_tree( + self, + selected_space: SelectedJointSpace, + world: WorldSpec, + tree: list[TreeNode], + target: NDArray[np.float64], + step_size: float, + ) -> TreeNode | None: + """Try to connect a selected-joint tree to a target.""" + while True: + result = self._extend_selected_tree( + selected_space, + world, + tree, + target, + step_size, + ) + if result is None: + return None + if float(np.linalg.norm(result.config - target)) < self._goal_tolerance: + return result + def _validate_inputs( self, world: WorldSpec, @@ -226,8 +415,8 @@ def _extend_tree( new_config = nearest.config + step_size * (diff / dist) # Check validity of edge using context-free method - start_state = JointState(name=joint_names, position=nearest.config.tolist()) - end_state = JointState(name=joint_names, position=new_config.tolist()) + start_state = JointState({"name": joint_names, "position": nearest.config.tolist()}) + end_state = JointState({"name": joint_names, "position": new_config.tolist()}) if world.check_edge_collision_free( robot_id, start_state, end_state, self._collision_step_size ): @@ -277,7 +466,7 @@ def _extract_path( full_path_arrays = start_path + list(reversed(goal_path)) # Convert to list of JointState - return [JointState(name=joint_names, position=q.tolist()) for q in full_path_arrays] + return [JointState({"name": joint_names, "position": q.tolist()}) for q in full_path_arrays] def _simplify_path( self, diff --git a/dimos/manipulation/planning/planners/selected_joint_space.py b/dimos/manipulation/planning/planners/selected_joint_space.py new file mode 100644 index 0000000000..89d6efbadc --- /dev/null +++ b/dimos/manipulation/planning/planners/selected_joint_space.py @@ -0,0 +1,319 @@ +# 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. + +"""Selected planning-group joint-space projection helpers.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +from numpy.typing import NDArray + +from dimos.manipulation.planning.groups.identifiers import ( + is_global_joint_name, + local_joint_name_from_global, + make_global_joint_name, +) +from dimos.manipulation.planning.groups.models import PlanningGroupSelection +from dimos.manipulation.planning.spec.models import ( + JointPath, + LocalModelJointName, + RobotName, + WorldRobotID, +) +from dimos.manipulation.planning.spec.protocols import WorldSpec +from dimos.msgs.sensor_msgs.JointState import JointState + + +@dataclass(frozen=True) +class SelectedRobotProjection: + """Runtime state needed to project selected-group samples into one robot. + + This is not static model metadata like ``RobotModelConfig``. It captures the + current planning context: the world robot ID, base/current positions used for + non-selected joints, and joint-limit lookup tables for the selected-space + planner. + """ + + robot_id: WorldRobotID + robot_name: RobotName + local_joint_names: list[LocalModelJointName] + base_positions_by_local_name: dict[LocalModelJointName, float] + lower_limits_by_local_name: dict[LocalModelJointName, float] + upper_limits_by_local_name: dict[LocalModelJointName, float] + + +class SelectedJointSpace: + """Projection adapter between selected global joints and full robot states.""" + + def __init__( + self, + robot_projections: list[SelectedRobotProjection], + selected_joint_names: list[str], + ) -> None: + self.robot_projections = robot_projections + self.selected_joint_names = selected_joint_names + + @classmethod + def from_world( + cls, + world: WorldSpec, + selection: PlanningGroupSelection, + ) -> SelectedJointSpace: + return cls( + robot_projections=_build_robot_projections(world, selection), + selected_joint_names=list(selection.joint_names), + ) + + def joint_limits(self) -> tuple[NDArray[np.float64], NDArray[np.float64]]: + projections_by_robot_name = { + projection.robot_name: projection for projection in self.robot_projections + } + lower: list[float] = [] + upper: list[float] = [] + for global_name in self.selected_joint_names: + robot_name, local_name = _split_selected_global_joint_name( + global_name, projections_by_robot_name + ) + projection = projections_by_robot_name[robot_name] + lower.append(projection.lower_limits_by_local_name[local_name]) + upper.append(projection.upper_limits_by_local_name[local_name]) + return np.asarray(lower, dtype=np.float64), np.asarray(upper, dtype=np.float64) + + def config_collision_free( + self, + world: WorldSpec, + selected_positions: NDArray[np.float64], + ) -> bool: + with world.scratch_context() as ctx: + projected_states = self.project_config(selected_positions) + for projection in self.robot_projections: + world.set_joint_state( + ctx, projection.robot_id, projected_states[projection.robot_id] + ) + return all( + world.is_collision_free(ctx, projection.robot_id) + for projection in self.robot_projections + ) + + def edge_collision_free( + self, + world: WorldSpec, + start: NDArray[np.float64], + end: NDArray[np.float64], + step_size: float, + ) -> bool: + distance = float(np.linalg.norm(end - start)) + steps = max(1, int(np.ceil(distance / step_size))) + for step in range(steps + 1): + ratio = step / steps + candidate = start + ratio * (end - start) + if not self.config_collision_free(world, candidate): + return False + return True + + def project_config( + self, + selected_positions: NDArray[np.float64], + ) -> dict[WorldRobotID, JointState]: + selected_positions_by_global_name = dict( + zip(self.selected_joint_names, selected_positions.tolist(), strict=True) + ) + projected_states: dict[WorldRobotID, JointState] = {} + for projection in self.robot_projections: + positions: list[float] = [] + for local_name in projection.local_joint_names: + global_name = make_global_joint_name(projection.robot_name, local_name) + position = selected_positions_by_global_name.get( + global_name, + projection.base_positions_by_local_name[local_name], + ) + positions.append(float(position)) + projected_states[projection.robot_id] = JointState( + {"name": list(projection.local_joint_names), "position": positions} + ) + return projected_states + + def simplify_path( + self, + world: WorldSpec, + path: JointPath, + collision_step_size: float, + max_iterations: int = 100, + ) -> JointPath: + if len(path) <= 2: + return path + + simplified = list(path) + for _ in range(max_iterations): + if len(simplified) <= 2: + break + i = np.random.randint(0, len(simplified) - 2) + j = np.random.randint(i + 2, len(simplified)) + start = np.asarray(simplified[i].position, dtype=np.float64) + end = np.asarray(simplified[j].position, dtype=np.float64) + if self.edge_collision_free(world, start, end, collision_step_size): + simplified = simplified[: i + 1] + simplified[j:] + return simplified + + +def normalize_selection_target( + selection: PlanningGroupSelection, + target: JointState, + label: str, +) -> JointState: + """Normalize a selected-joint target to global selection order.""" + selected_global_names = list(selection.joint_names) + if not target.name: + if len(target.position) != len(selected_global_names): + raise ValueError( + f"{label} target has {len(target.position)} positions, " + f"expected {len(selected_global_names)}" + ) + return JointState({"name": selected_global_names, "position": list(target.position)}) + + if len(target.name) != len(target.position): + raise ValueError( + f"{label} target has {len(target.name)} names but {len(target.position)} positions" + ) + + names = list(target.name) + global_flags = [is_global_joint_name(name) for name in names] + if any(global_flags) and not all(global_flags): + raise ValueError(f"{label} target mixes global and local joint names: {names}") + + if all(global_flags): + expected_names = selected_global_names + else: + if len(selection.groups) != 1: + raise ValueError( + f"{label} target uses local joint names for a multi-group selection; " + "use global joint names" + ) + expected_names = list(selection.groups[0].local_joint_names) + + positions_by_name = dict(zip(names, target.position, strict=True)) + missing = [name for name in expected_names if name not in positions_by_name] + if missing: + raise ValueError(f"{label} target is missing joints: {missing}") + extra = sorted(set(names) - set(expected_names)) + if extra: + raise ValueError(f"{label} target has extra joints: {extra}") + + ordered_positions = [float(positions_by_name[name]) for name in expected_names] + return JointState({"name": selected_global_names, "position": ordered_positions}) + + +def _build_robot_projections( + world: WorldSpec, + selection: PlanningGroupSelection, +) -> list[SelectedRobotProjection]: + robot_ids_by_name = _robot_ids_by_name(world, selection.robot_names) + robot_projections: list[SelectedRobotProjection] = [] + with world.scratch_context() as ctx: + for robot_name in selection.robot_names: + robot_id = robot_ids_by_name[robot_name] + config = world.get_robot_config(robot_id) + local_joint_names = list(config.joint_names) + current_state = world.get_joint_state(ctx, robot_id) + base_positions_by_local_name = _positions_by_local_name( + current_state, + robot_name, + local_joint_names, + ) + lower, upper = world.get_joint_limits(robot_id) + if len(lower) != len(local_joint_names) or len(upper) != len(local_joint_names): + raise ValueError( + f"Robot '{robot_name}' joint limits do not match configured joints" + ) + robot_projections.append( + SelectedRobotProjection( + robot_id=robot_id, + robot_name=robot_name, + local_joint_names=local_joint_names, + base_positions_by_local_name=base_positions_by_local_name, + lower_limits_by_local_name=dict( + zip(local_joint_names, lower.tolist(), strict=True) + ), + upper_limits_by_local_name=dict( + zip(local_joint_names, upper.tolist(), strict=True) + ), + ) + ) + return robot_projections + + +def _positions_by_local_name( + joint_state: JointState, + robot_name: RobotName, + local_joint_names: list[LocalModelJointName], +) -> dict[LocalModelJointName, float]: + if not joint_state.name: + if len(joint_state.position) != len(local_joint_names): + raise ValueError( + f"Current state for robot '{robot_name}' has {len(joint_state.position)} positions, " + f"expected {len(local_joint_names)}" + ) + return dict(zip(local_joint_names, map(float, joint_state.position), strict=True)) + + positions_by_name = dict(zip(joint_state.name, joint_state.position, strict=True)) + positions_by_local_name: dict[LocalModelJointName, float] = {} + for local_name in local_joint_names: + global_name = make_global_joint_name(robot_name, local_name) + if local_name in positions_by_name: + positions_by_local_name[local_name] = float(positions_by_name[local_name]) + elif global_name in positions_by_name: + positions_by_local_name[local_name] = float(positions_by_name[global_name]) + else: + raise ValueError( + f"Current state for robot '{robot_name}' is missing joint '{local_name}'" + ) + return positions_by_local_name + + +def _robot_ids_by_name( + world: WorldSpec, + robot_names: tuple[RobotName, ...], +) -> dict[RobotName, WorldRobotID]: + robot_ids_by_name: dict[RobotName, WorldRobotID] = {} + for robot_name in robot_names: + matches = [ + robot_id + for robot_id in world.get_robot_ids() + if world.get_robot_config(robot_id).name == robot_name + ] + if not matches: + raise ValueError(f"Robot '{robot_name}' not found") + if len(matches) > 1: + raise ValueError(f"Robot name '{robot_name}' is not unique in planning world") + robot_ids_by_name[robot_name] = matches[0] + return robot_ids_by_name + + +def _split_selected_global_joint_name( + global_name: str, + projections_by_robot_name: dict[RobotName, SelectedRobotProjection], +) -> tuple[RobotName, LocalModelJointName]: + for robot_name, projection in projections_by_robot_name.items(): + try: + local_name = local_joint_name_from_global(robot_name, global_name) + except ValueError: + continue + if local_name not in projection.local_joint_names: + raise ValueError( + f"Selected joint '{global_name}' is not configured for robot '{robot_name}'" + ) + return robot_name, local_name + raise ValueError(f"Selected joint '{global_name}' does not belong to a selected robot") diff --git a/dimos/manipulation/planning/planners/test_rrt_planner_selection.py b/dimos/manipulation/planning/planners/test_rrt_planner_selection.py new file mode 100644 index 0000000000..4b330bc816 --- /dev/null +++ b/dimos/manipulation/planning/planners/test_rrt_planner_selection.py @@ -0,0 +1,199 @@ +# 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. + +"""Focused tests for selected-joint RRT planning.""" + +from __future__ import annotations + +from contextlib import nullcontext +from pathlib import Path + +import numpy as np +import pytest + +from dimos.manipulation.planning.groups.models import ( + PlanningGroup, + PlanningGroupDefinition, + PlanningGroupSelection, +) +from dimos.manipulation.planning.planners.rrt_planner import RRTConnectPlanner +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.enums import PlanningStatus +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.sensor_msgs.JointState import JointState + + +def _pose() -> PoseStamped: + return PoseStamped(position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0)) + + +def _group(name: str, joints: tuple[str, ...]) -> PlanningGroup: + return PlanningGroup( + id=f"arm/{name}", + robot_name="arm", + group_name=name, + joint_names=tuple(f"arm/{joint}" for joint in joints), + local_joint_names=joints, + base_link="base", + tip_link="tool", + ) + + +class _World: + is_finalized = True + + def __init__(self, current: list[float] | None = None) -> None: + self.current = current or [0.0, 0.0, 0.7] + self.projected_states: list[JointState] = [] + self.config = RobotModelConfig( + name="arm", + model_path=Path("robot.urdf"), + base_pose=_pose(), + joint_names=["joint_a", "joint_b", "gripper"], + base_link="base", + planning_groups=[ + PlanningGroupDefinition("arm", ("joint_a", "joint_b"), "base", "tool") + ], + ) + + def get_robot_ids(self) -> list[str]: + return ["robot"] + + def get_robot_config(self, robot_id: str) -> RobotModelConfig: + return self.config + + def scratch_context(self) -> nullcontext[None]: + return nullcontext(None) + + def get_joint_state(self, ctx: object, robot_id: str) -> JointState: + return JointState({"name": ["joint_a", "joint_b", "gripper"], "position": self.current}) + + def get_joint_limits(self, robot_id: str) -> tuple[np.ndarray, np.ndarray]: + return np.array([-1.0, -1.0, -1.0]), np.array([1.0, 1.0, 1.0]) + + def set_joint_state(self, ctx: object, robot_id: str, joint_state: JointState) -> None: + self.projected_states.append(joint_state) + + def is_collision_free(self, ctx: object, robot_id: str) -> bool: + return True + + +@pytest.mark.parametrize( + ("start", "goal", "expected_start", "expected_goal"), + [ + ( + JointState({"position": [0.1, 0.2]}), + JointState({"position": [0.3, 0.4]}), + [0.1, 0.2], + [0.3, 0.4], + ), + ( + JointState({"name": ["arm/joint_b", "arm/joint_a"], "position": [0.2, 0.1]}), + JointState({"name": ["arm/joint_b", "arm/joint_a"], "position": [0.4, 0.3]}), + [0.1, 0.2], + [0.3, 0.4], + ), + ( + JointState({"name": ["joint_b", "joint_a"], "position": [0.2, 0.1]}), + JointState({"name": ["joint_b", "joint_a"], "position": [0.4, 0.3]}), + [0.1, 0.2], + [0.3, 0.4], + ), + ], +) +def test_plan_selected_joint_path_normalizes_target_forms( + start: JointState, goal: JointState, expected_start: list[float], expected_goal: list[float] +) -> None: + group = _group("arm", ("joint_a", "joint_b")) + result = RRTConnectPlanner().plan_selected_joint_path( + _World(), PlanningGroupSelection.from_groups((group,)), start, goal + ) + + assert result.status == PlanningStatus.SUCCESS + assert result.path is not None + assert result.path[0].name == ["arm/joint_a", "arm/joint_b"] + assert result.path[0].position == expected_start + assert result.path[-1].position == expected_goal + + +@pytest.mark.parametrize( + ("start", "goal", "status", "message"), + [ + ( + JointState({"name": ["arm/joint_a"], "position": [0.0]}), + JointState({"position": [0.0, 0.0]}), + PlanningStatus.INVALID_START, + "missing", + ), + ( + JointState({"position": [0.0, 0.0]}), + JointState( + {"name": ["arm/joint_a", "arm/joint_b", "arm/extra"], "position": [0.0, 0.0, 0.0]} + ), + PlanningStatus.INVALID_GOAL, + "extra", + ), + ( + JointState({"name": ["arm/joint_a", "joint_b"], "position": [0.0, 0.0]}), + JointState({"position": [0.0, 0.0]}), + PlanningStatus.INVALID_START, + "mixes", + ), + ], +) +def test_plan_selected_joint_path_rejects_bad_targets( + start: JointState, goal: JointState, status: PlanningStatus, message: str +) -> None: + group = _group("arm", ("joint_a", "joint_b")) + result = RRTConnectPlanner().plan_selected_joint_path( + _World(), PlanningGroupSelection.from_groups((group,)), start, goal + ) + + assert result.status == status + assert message in result.message + + +def test_plan_selected_joint_path_rejects_local_names_for_multi_group_selection() -> None: + selection = PlanningGroupSelection.from_groups( + (_group("arm", ("joint_a",)), _group("gripper", ("gripper",))) + ) + + result = RRTConnectPlanner().plan_selected_joint_path( + _World(), + selection, + JointState({"name": ["joint_a", "gripper"], "position": [0.0, 0.0]}), + JointState({"position": [0.1, 0.2]}), + ) + + assert result.status == PlanningStatus.INVALID_START + assert "multi-group" in result.message + + +def test_plan_selected_joint_path_direct_edge_projects_full_state_with_unselected_joints() -> None: + world = _World(current=[0.0, 0.0, 0.77]) + group = _group("arm", ("joint_a", "joint_b")) + + result = RRTConnectPlanner().plan_selected_joint_path( + world, + PlanningGroupSelection.from_groups((group,)), + JointState({"position": [0.1, 0.2]}), + JointState({"position": [0.3, 0.4]}), + ) + + assert result.status == PlanningStatus.SUCCESS + assert world.projected_states + assert all(state.name == ["joint_a", "joint_b", "gripper"] for state in world.projected_states) + assert all(state.position[2] == 0.77 for state in world.projected_states) diff --git a/dimos/manipulation/planning/spec/config.py b/dimos/manipulation/planning/spec/config.py index 74dc3bd69b..1ef40d103b 100644 --- a/dimos/manipulation/planning/spec/config.py +++ b/dimos/manipulation/planning/spec/config.py @@ -21,6 +21,11 @@ from pydantic import Field from dimos.core.module import ModuleConfig +from dimos.manipulation.planning.groups.identifiers import ( + assert_local_joint_names, + assert_valid_robot_name, +) +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped @@ -30,10 +35,13 @@ class RobotModelConfig(ModuleConfig): Attributes: name: Human-readable robot name model_path: Path to robot model file (.urdf, .xacro, or .xml/MJCF) - base_pose: Pose of robot base in world frame (position + orientation) - joint_names: Ordered list of controlled joint names (in URDF namespace) - end_effector_link: Name of the end-effector link for FK/IK - base_link: Name of the base link (default: "base_link") + srdf_path: Optional path to SRDF file containing planning group definitions + base_pose: Placement transform. This is the canonical world placement for + robot instances. + joint_names: Ordered list of controllable joints in the local model + namespace. This is not a planning group. + base_link: Robot-scoped link that base_pose places in the world and + current backends use for weld/placement. package_paths: Dict mapping package names to filesystem Paths joint_limits_lower: Lower joint limits (radians) joint_limits_upper: Upper joint limits (radians) @@ -45,19 +53,20 @@ class RobotModelConfig(ModuleConfig): links may legitimately overlap (e.g., mimic joints). max_velocity: Maximum joint velocity for trajectory generation (rad/s) max_acceleration: Maximum joint acceleration for trajectory generation (rad/s^2) - joint_name_mapping: Maps coordinator joint names to URDF joint names. - Example: {"left/joint1": "joint1"} means coordinator's "left/joint1" - corresponds to URDF's "joint1". If empty, names are assumed to match. + joint_name_mapping: Maps coordinator joint names to local model joint names. + This is retained for current coordinator/monitor integrations while planning + APIs move toward globally scoped joint names. coordinator_task_name: Task name for executing trajectories via coordinator RPC. If set, trajectories can be executed via execute_trajectory() RPC. """ name: str model_path: Path - base_pose: PoseStamped + srdf_path: Path | None = None + base_pose: PoseStamped = Field(default_factory=PoseStamped) joint_names: list[str] - end_effector_link: str base_link: str = "base_link" + planning_groups: list[PlanningGroupDefinition] = Field(default_factory=list) package_paths: dict[str, Path] = Field(default_factory=dict) joint_limits_lower: list[float] | None = None joint_limits_upper: list[float] | None = None @@ -79,14 +88,44 @@ class RobotModelConfig(ModuleConfig): # Pre-grasp offset distance in meters (along approach direction) pre_grasp_offset: float = 0.10 + def model_post_init(self, __context: object) -> None: + """Validate delimiter-based naming constraints.""" + assert_valid_robot_name(self.name) + assert_local_joint_names(self.joint_names) + + @property + def end_effector_link(self) -> str: + """Compatibility pose target frame derived from planning groups. + + Current world, IK, and visualization layers still ask robot configs for + one end-effector link. The planning-group model stores that frame as a + group ``tip_link``; this shim keeps those layers working until they are + migrated to explicit planning-group IDs. + """ + pose_tip_links = [ + group.tip_link for group in self.planning_groups if group.tip_link is not None + ] + if not pose_tip_links: + raise ValueError( + f"RobotModelConfig '{self.name}' has no pose-target planning group; " + "define PlanningGroupDefinition.tip_link" + ) + unique_tip_links = list(dict.fromkeys(pose_tip_links)) + if len(unique_tip_links) > 1: + raise ValueError( + f"RobotModelConfig '{self.name}' has multiple pose-target planning groups; " + "use an explicit planning group ID" + ) + return unique_tip_links[0] + def get_urdf_joint_name(self, coordinator_name: str) -> str: - """Translate coordinator joint name to URDF joint name.""" + """Translate coordinator joint name to local model joint name.""" return self.joint_name_mapping.get(coordinator_name, coordinator_name) def get_coordinator_joint_name(self, urdf_name: str) -> str: - """Translate URDF joint name to coordinator joint name.""" - for coord_name, u_name in self.joint_name_mapping.items(): - if u_name == urdf_name: + """Translate local model joint name to coordinator joint name.""" + for coord_name, model_name in self.joint_name_mapping.items(): + if model_name == urdf_name: return coord_name return urdf_name @@ -94,4 +133,4 @@ def get_coordinator_joint_names(self) -> list[str]: """Get joint names in coordinator namespace.""" if not self.joint_name_mapping: return self.joint_names - return [self.get_coordinator_joint_name(j) for j in self.joint_names] + return [self.get_coordinator_joint_name(joint_name) for joint_name in self.joint_names] diff --git a/dimos/manipulation/planning/spec/enums.py b/dimos/manipulation/planning/spec/enums.py index 66a17ee199..3d5b668e4b 100644 --- a/dimos/manipulation/planning/spec/enums.py +++ b/dimos/manipulation/planning/spec/enums.py @@ -35,6 +35,7 @@ class IKStatus(Enum): JOINT_LIMITS = auto() COLLISION = auto() TIMEOUT = auto() + UNSUPPORTED = auto() class PlanningStatus(Enum): @@ -47,3 +48,4 @@ class PlanningStatus(Enum): INVALID_GOAL = auto() COLLISION_AT_START = auto() COLLISION_AT_GOAL = auto() + UNSUPPORTED = auto() diff --git a/dimos/manipulation/planning/spec/models.py b/dimos/manipulation/planning/spec/models.py index d412e9f766..b83dbbc2e6 100644 --- a/dimos/manipulation/planning/spec/models.py +++ b/dimos/manipulation/planning/spec/models.py @@ -30,9 +30,11 @@ import numpy as np from numpy.typing import NDArray + from dimos.manipulation.planning.groups.models import PlanningGroup from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.JointState import JointState + from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory RobotName: TypeAlias = str @@ -41,6 +43,15 @@ WorldRobotID: TypeAlias = str """Internal Drake world robot ID""" +PlanningGroupID: TypeAlias = str +"""Public planning group ID of the form {robot_name}/{group_name}.""" + +LocalModelJointName: TypeAlias = str +"""Joint name as it appears in URDF/SRDF before world binding.""" + +GlobalJointName: TypeAlias = str +"""Public joint name of the form {robot_name}/{local_joint_name}.""" + JointPath: TypeAlias = "list[JointState]" """List of joint states forming a path (each waypoint has names + positions)""" @@ -56,6 +67,25 @@ class PlanningSceneInfo: robots: Mapping[WorldRobotID, RobotModelConfig] """Robot model configurations keyed by world robot ID.""" + planning_groups: tuple[PlanningGroup, ...] = () + """Resolved immutable planning groups for the initialized scene.""" + + +@dataclass(frozen=True) +class VisualizationSession: + """One-shot immutable visualization initialization payload.""" + + scene: PlanningSceneInfo + operator: object | None = None + """Optional concrete ManipulationOperator; typed as object to avoid low-level cycles.""" + + +@dataclass(frozen=True) +class VisualizationStateFrame: + """Pushed current joint states for visualization backends.""" + + joint_states: Mapping[WorldRobotID, JointState] + Jacobian: TypeAlias = "NDArray[np.float64]" """6 x n Jacobian matrix (rows: [vx, vy, vz, wx, wy, wz])""" @@ -141,6 +171,24 @@ def is_success(self) -> bool: return self.status == PlanningStatus.SUCCESS +@dataclass +class GeneratedPlan: + """Canonical selected-planning-group plan exposed by ManipulationModule.""" + + group_ids: tuple[PlanningGroupID, ...] + trajectory: JointTrajectory + path: list[JointState] = field(default_factory=list) + status: PlanningStatus = PlanningStatus.NO_SOLUTION + planning_time: float = 0.0 + path_length: float = 0.0 + iterations: int = 0 + message: str = "" + + def is_success(self) -> bool: + """Check if the generated plan was successful.""" + return self.status == PlanningStatus.SUCCESS + + @dataclass class CollisionObjectMessage: """Message for adding/updating/removing obstacles. diff --git a/dimos/manipulation/planning/spec/protocols.py b/dimos/manipulation/planning/spec/protocols.py index c7ee95ee0a..f8762521d7 100644 --- a/dimos/manipulation/planning/spec/protocols.py +++ b/dimos/manipulation/planning/spec/protocols.py @@ -20,6 +20,7 @@ from __future__ import annotations +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable if TYPE_CHECKING: @@ -28,17 +29,20 @@ import numpy as np from numpy.typing import NDArray + from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupSelection from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.models import ( IKResult, - JointPath, Obstacle, + PlanningGroupID, PlanningResult, - PlanningSceneInfo, + VisualizationSession, + VisualizationStateFrame, WorldRobotID, ) from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.JointState import JointState + from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory @runtime_checkable @@ -170,6 +174,14 @@ def get_jacobian(self, ctx: Any, robot_id: WorldRobotID) -> NDArray[np.float64]: """Get end-effector Jacobian (6 x n_joints).""" ... + def get_group_ee_pose(self, ctx: Any, group_id: PlanningGroupID) -> PoseStamped: + """Get planning-group tip pose.""" + ... + + def get_group_jacobian(self, ctx: Any, group_id: PlanningGroupID) -> NDArray[np.float64]: + """Get planning-group Jacobian (6 x n_group_joints).""" + ... + @runtime_checkable class VisualizationSpec(Protocol): @@ -182,28 +194,26 @@ class VisualizationSpec(Protocol): visualization affordances. """ - def initialize_scene(self, scene: PlanningSceneInfo) -> None: - """Receive stable planning-scene metadata after world startup.""" + def initialize(self, session: VisualizationSession) -> None: + """Receive one-shot visualization session metadata after world startup.""" ... def get_visualization_url(self) -> str | None: """Get visualization URL if enabled.""" ... - def publish_visualization(self, ctx: Any | None = None) -> None: - """Publish current state to visualization.""" - ... - - def show_preview(self, robot_id: WorldRobotID) -> None: - """Show the preview representation for a robot.""" + def update_state(self, frame: VisualizationStateFrame) -> None: + """Receive current joint states keyed by initialized world robot ID.""" ... - def hide_preview(self, robot_id: WorldRobotID) -> None: - """Hide the preview representation for a robot.""" + def animate_trajectory( + self, trajectory: JointTrajectory, duration: float | None = None + ) -> None: + """Animate a raw globally named trajectory.""" ... - def animate_path(self, robot_id: WorldRobotID, path: JointPath, duration: float = 3.0) -> None: - """Animate a path in visualization.""" + def cancel_preview_animation(self, robot_ids: Sequence[WorldRobotID] | None = None) -> None: + """Cancel an active preview animation without waiting for its renderer to finish.""" ... def close(self) -> None: @@ -229,6 +239,20 @@ def solve( """Solve IK with optional collision checking.""" ... + def solve_pose_targets( + self, + world: WorldSpec, + pose_targets: Mapping[PlanningGroup, PoseStamped], + auxiliary_groups: Sequence[PlanningGroup] = (), + seed: JointState | None = None, + position_tolerance: float = 0.001, + orientation_tolerance: float = 0.01, + check_collision: bool = True, + max_attempts: int = 10, + ) -> IKResult: + """Solve planning-group-scoped pose targets.""" + ... + @runtime_checkable class PlannerSpec(Protocol): @@ -254,6 +278,18 @@ def plan_joint_path( """Plan a collision-free joint-space path.""" ... + def plan_selected_joint_path( + self, + world: WorldSpec, + selection: PlanningGroupSelection, + start: JointState, + goal: JointState, + timeout: float = 10.0, + max_iterations: int = 5000, + ) -> PlanningResult: + """Plan a collision-free path for an ordered planning-group selection.""" + ... + def get_name(self) -> str: """Get planner name.""" ... diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index ca426ba340..3474990e35 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -16,23 +16,26 @@ from __future__ import annotations +from collections.abc import Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass, field from pathlib import Path from threading import RLock, current_thread from typing import TYPE_CHECKING, Any +import xml.etree.ElementTree as ET import numpy as np +from dimos.manipulation.planning.groups.identifiers import ( + make_global_joint_names, + make_planning_group_id, +) +from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.manipulation.planning.groups.utils import joint_state_to_ordered_positions from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import ObstacleType -from dimos.manipulation.planning.spec.models import ( - JointPath, - Obstacle, - PlanningSceneInfo, - WorldRobotID, -) +from dimos.manipulation.planning.spec.models import Obstacle, PlanningGroupID, WorldRobotID from dimos.manipulation.planning.spec.protocols import VisualizationSpec, WorldSpec from dimos.manipulation.planning.utils.mesh_utils import prepare_urdf_for_drake from dimos.utils.logging_config import setup_logger @@ -45,6 +48,13 @@ from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.Transform import Transform from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory + +if TYPE_CHECKING: + from dimos.manipulation.planning.spec.models import ( + VisualizationSession, + VisualizationStateFrame, + ) try: from pydrake.geometry import ( @@ -194,6 +204,8 @@ def __init__(self, time_step: float = 0.0, enable_viz: bool = False) -> None: self._plant_context: Context | None = None self._scene_graph_context: Context | None = None self._finalized = False + self._preview_animation_generation = 0 + self._preview_animation_generations: dict[WorldRobotID, int] = {} # Obstacle source for dynamic obstacles self._obstacle_source_id: Any = None @@ -207,6 +219,10 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: raise RuntimeError("Cannot add robot after world is finalized") with self._lock: + if any(data.config.name == config.name for data in self._robots.values()): + raise ValueError(f"Robot name '{config.name}' is already registered") + self._validate_planning_group_config(config) + self._robot_counter += 1 robot_id = f"robot_{self._robot_counter}" @@ -215,9 +231,16 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: self._validate_joints(config, model_instance) - ee_frame = self._plant.GetBodyByName( - config.end_effector_link, model_instance - ).body_frame() + ee_link = config.base_link + try: + primary_group_id = self._primary_pose_group_id_for_config(config) + except ValueError: + primary_group_id = None + if primary_group_id is not None: + primary_group = self._planning_group_from_config(config, primary_group_id) + if primary_group.tip_link is not None: + ee_link = primary_group.tip_link + ee_frame = self._plant.GetBodyByName(ee_link, model_instance).body_frame() base_frame = self._plant.GetBodyByName(config.base_link, model_instance).body_frame() # Preview (yellow ghost) — always a separate instance per robot @@ -235,7 +258,6 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: base_frame=base_frame, preview_model_instance=preview_model_instance, ) - logger.info(f"Added robot '{robot_id}' ({config.name})") return robot_id @@ -256,7 +278,7 @@ def _load_model(self, config: RobotModelConfig) -> Any: xacro_args=config.xacro_args, convert_meshes=config.auto_convert_meshes, ) - prepared_path_obj = Path(prepared_path) + prepared_path_obj = self._strip_world_base_joint(Path(prepared_path), config) # Register package paths (not applicable to MJCF) if config.package_paths: @@ -271,9 +293,48 @@ def _load_model(self, config: RobotModelConfig) -> Any: model_instances = self._parser.AddModels(prepared_path_obj) if not model_instances: - raise ValueError(f"Failed to parse model: {prepared_path}") + raise ValueError(f"Failed to parse model: {prepared_path_obj}") return model_instances[0] + @staticmethod + def _strip_world_base_joint(model_path: Path, config: RobotModelConfig) -> Path: + if model_path.suffix != ".urdf": + return model_path + + tree = ET.parse(model_path) + root = tree.getroot() + joints = root.findall("joint") + joints_to_remove = [ + joint + for joint in joints + if joint.get("type") == "fixed" + and (parent := joint.find("parent")) is not None + and (child := joint.find("child")) is not None + and parent.get("link") == "world" + and child.get("link") == config.base_link + ] + + if not joints_to_remove: + return model_path + + for joint in joints_to_remove: + root.remove(joint) + + if not any( + element is not None and element.get("link") == "world" + for joint in root.findall("joint") + for element in (joint.find("parent"), joint.find("child")) + ): + for link in list(root.findall("link")): + if link.get("name") == "world": + root.remove(link) + + stripped_path = model_path.with_name( + f"{model_path.stem}_config_base_pose{model_path.suffix}" + ) + tree.write(stripped_path, encoding="utf-8", xml_declaration=True) + return stripped_path + def _weld_base_if_needed(self, config: RobotModelConfig, model_instance: Any) -> None: """Weld robot base to world if not already welded in URDF.""" base_body = self._plant.GetBodyByName(config.base_link, model_instance) @@ -317,6 +378,52 @@ def get_robot_config(self, robot_id: WorldRobotID) -> RobotModelConfig: raise KeyError(f"Robot '{robot_id}' not found") return self._robots[robot_id].config + @staticmethod + def _validate_planning_group_config(config: RobotModelConfig) -> None: + seen_group_names: set[str] = set() + for definition in config.planning_groups: + make_planning_group_id(config.name, definition.name) + if definition.name in seen_group_names: + raise ValueError(f"Planning group '{definition.name}' is already registered") + make_global_joint_names(config.name, definition.joint_names) + seen_group_names.add(definition.name) + + @staticmethod + def _planning_group_from_config( + config: RobotModelConfig, group_id: PlanningGroupID + ) -> PlanningGroup: + for definition in config.planning_groups: + if make_planning_group_id(config.name, definition.name) == group_id: + joint_names = tuple(make_global_joint_names(config.name, definition.joint_names)) + return PlanningGroup( + group_id, + config.name, + definition.name, + joint_names, + definition.joint_names, + definition.base_link, + definition.tip_link, + definition.source, + ) + raise KeyError(f"Unknown planning group ID: {group_id}") + + def _planning_group_from_id(self, group_id: PlanningGroupID) -> PlanningGroup: + for robot_data in self._robots.values(): + try: + return self._planning_group_from_config(robot_data.config, group_id) + except KeyError: + continue + raise KeyError(f"Unknown planning group ID: {group_id}") + + @staticmethod + def _primary_pose_group_id_for_config(config: RobotModelConfig) -> PlanningGroupID | None: + pose_groups = [group for group in config.planning_groups if group.has_pose_target] + if not pose_groups: + return None + if len(pose_groups) > 1: + raise ValueError(f"Robot '{config.name}' has multiple pose groups") + return make_planning_group_id(config.name, pose_groups[0].name) + def get_joint_limits( self, robot_id: WorldRobotID ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: @@ -683,10 +790,10 @@ def finalize(self) -> None: # Initial visualization publish (routed to Meshcat thread) if self._meshcat_visualizer is not None: - self.publish_visualization() + self._publish_visualization() # Hide all preview robots initially for robot_id in self._robots: - self.hide_preview(robot_id) + self._set_preview_visibility(robot_id, False) @property def is_finalized(self) -> bool: @@ -768,8 +875,7 @@ def sync_from_joint_state(self, robot_id: WorldRobotID, joint_state: JointState) if not self._finalized or self._plant_context is None: return # Silently ignore before finalization - # Extract positions as numpy array for internal use - positions = np.array(joint_state.position, dtype=np.float64) + positions = self._joint_state_to_q(robot_id, joint_state) with self._lock: self._set_positions_internal(self._plant_context, robot_id, positions) @@ -787,8 +893,7 @@ def set_joint_state( if not self._finalized: raise RuntimeError("World must be finalized first") - # Extract positions as numpy array for internal use - positions = np.array(joint_state.position, dtype=np.float64) + positions = self._joint_state_to_q(robot_id, joint_state) # Get plant context from diagram context plant_ctx = self._diagram.GetMutableSubsystemContext(self._plant, ctx) @@ -809,6 +914,28 @@ def _set_positions_internal( self._plant.SetPositions(plant_ctx, full_positions) + def _joint_state_to_q( + self, robot_id: WorldRobotID, joint_state: JointState + ) -> NDArray[np.float64]: + """Normalize unnamed, robot-local, mapped, or global JointState to robot joint order.""" + if robot_id not in self._robots: + raise KeyError(f"Robot '{robot_id}' not found") + robot_data = self._robots[robot_id] + return joint_state_to_ordered_positions( + joint_state, + joint_names=robot_data.config.joint_names, + joint_name_mapping=robot_data.config.joint_name_mapping, + ) + + def _robot_id_for_group(self, group_id: PlanningGroupID) -> WorldRobotID: + group = self._planning_group_from_id(group_id) + matches = [ + rid for rid, data in self._robots.items() if data.config.name == group.robot_name + ] + if not matches: + raise KeyError(f"No robot registered for planning group '{group_id}'") + return matches[0] + def get_joint_state(self, ctx: Context, robot_id: WorldRobotID) -> JointState: """Get robot joint state from given context.""" if not self._finalized: @@ -905,16 +1032,28 @@ def check_edge_collision_free( def get_ee_pose(self, ctx: Context, robot_id: WorldRobotID) -> PoseStamped: """Get end-effector pose.""" - if not self._finalized: - raise RuntimeError("World must be finalized first") - if robot_id not in self._robots: raise KeyError(f"Robot '{robot_id}' not found") - robot_data = self._robots[robot_id] + group_id = self._primary_pose_group_id_for_config(robot_data.config) + if group_id is None: + raise ValueError( + f"Robot '{robot_data.config.name}' has no pose-targetable planning group" + ) + return self.get_group_ee_pose(ctx, group_id) + + def get_group_ee_pose(self, ctx: Context, group_id: PlanningGroupID) -> PoseStamped: + """Get planning-group tip pose.""" + if not self._finalized: + raise RuntimeError("World must be finalized first") + + group = self._planning_group_from_id(group_id) + if group.tip_link is None: + raise ValueError(f"Planning group '{group_id}' has no tip link") + robot_data = self._robots[self._robot_id_for_group(group_id)] plant_ctx = self._diagram.GetSubsystemContext(self._plant, ctx) - ee_body = robot_data.ee_frame.body() + ee_body = self._plant.GetBodyByName(group.tip_link, robot_data.model_instance) X_WE = self._plant.EvalBodyPoseInWorld(plant_ctx, ee_body) # Extract position and quaternion from Drake transform @@ -955,30 +1094,58 @@ def get_jacobian(self, ctx: Context, robot_id: WorldRobotID) -> NDArray[np.float Rows: [vx, vy, vz, wx, wy, wz] (linear, then angular) """ - if not self._finalized: - raise RuntimeError("World must be finalized first") - if robot_id not in self._robots: raise KeyError(f"Robot '{robot_id}' not found") - robot_data = self._robots[robot_id] + group_id = self._primary_pose_group_id_for_config(robot_data.config) + if group_id is None: + raise ValueError( + f"Robot '{robot_data.config.name}' has no pose-targetable planning group" + ) + return self.get_group_jacobian(ctx, group_id) + + def get_group_jacobian(self, ctx: Context, group_id: PlanningGroupID) -> NDArray[np.float64]: + """Get geometric Jacobian (6 x group joints) in group-local order.""" + if not self._finalized: + raise RuntimeError("World must be finalized first") + + group = self._planning_group_from_id(group_id) + if group.tip_link is None: + raise ValueError(f"Planning group '{group_id}' has no tip link") + robot_data = self._robots[self._robot_id_for_group(group_id)] plant_ctx = self._diagram.GetSubsystemContext(self._plant, ctx) + tip_frame = self._plant.GetBodyByName( + group.tip_link, robot_data.model_instance + ).body_frame() # Compute full Jacobian J_full = self._plant.CalcJacobianSpatialVelocity( plant_ctx, JacobianWrtVariable.kQDot, - robot_data.ee_frame, + tip_frame, np.array([0.0, 0.0, 0.0]), # type: ignore[arg-type] # Point on end-effector self._plant.world_frame(), self._plant.world_frame(), ) - # Extract columns for this robot's joints - n_joints = len(robot_data.joint_indices) + # Extract columns for configured controllable joints only. + joint_indices_by_name = dict( + zip(robot_data.config.joint_names, robot_data.joint_indices, strict=True) + ) + missing = [ + joint_name + for joint_name in group.local_joint_names + if joint_name not in joint_indices_by_name + ] + if missing: + raise ValueError( + f"Planning group '{group_id}' references non-controllable joints: {missing}" + ) + group_joint_indices = [joint_indices_by_name[name] for name in group.local_joint_names] + n_joints = len(group_joint_indices) J_robot = np.zeros((6, n_joints)) - for i, joint_idx in enumerate(robot_data.joint_indices): + for i, joint_idx in enumerate(group_joint_indices): J_robot[:, i] = J_full[:, joint_idx] # Reorder rows: Drake uses [angular, linear], we want [linear, angular] @@ -988,7 +1155,7 @@ def get_jacobian(self, ctx: Context, robot_id: WorldRobotID) -> NDArray[np.float # Visualization - def initialize_scene(self, scene: PlanningSceneInfo) -> None: + def initialize(self, session: VisualizationSession) -> None: """Embedded Meshcat observes the Drake world directly; no extra sync needed.""" return None @@ -998,7 +1165,7 @@ def get_visualization_url(self) -> str | None: return self._meshcat.web_url() return None - def publish_visualization(self, ctx: Context | None = None) -> None: + def _publish_visualization(self, ctx: Context | None = None) -> None: """Publish current state to visualization.""" if self._meshcat_visualizer is None or self._meshcat is None: return @@ -1008,6 +1175,10 @@ def publish_visualization(self, ctx: Context | None = None) -> None: viz_ctx = self._diagram.GetSubsystemContext(self._meshcat_visualizer, ctx) self._meshcat.forced_publish(self._meshcat_visualizer, viz_ctx) + def update_state(self, frame: VisualizationStateFrame) -> None: + """Receive pushed state frame; embedded Meshcat uses Drake live context.""" + self._publish_visualization() + def _set_preview_positions( self, plant_ctx: Context, robot_id: WorldRobotID, positions: NDArray[np.float64] ) -> None: @@ -1021,57 +1192,127 @@ def _set_preview_positions( full_positions[idx] = positions[i] self._plant.SetPositions(plant_ctx, full_positions) - def show_preview(self, robot_id: WorldRobotID) -> None: - """Show the preview (yellow ghost) robot in Meshcat.""" - if self._meshcat is None: - return - robot_data = self._robots.get(robot_id) - if robot_data is None or robot_data.preview_model_instance is None: - return - model_name = self._plant.GetModelInstanceName(robot_data.preview_model_instance) - self._meshcat.SetProperty(f"visualizer/{model_name}", "visible", True) - - def hide_preview(self, robot_id: WorldRobotID) -> None: - """Hide the preview (yellow ghost) robot in Meshcat.""" + def _set_preview_visibility(self, robot_id: WorldRobotID, visible: bool) -> None: + """Set one preview robot's Meshcat visibility.""" if self._meshcat is None: return robot_data = self._robots.get(robot_id) if robot_data is None or robot_data.preview_model_instance is None: return model_name = self._plant.GetModelInstanceName(robot_data.preview_model_instance) - self._meshcat.SetProperty(f"visualizer/{model_name}", "visible", False) + self._meshcat.SetProperty(f"visualizer/{model_name}", "visible", visible) - def animate_path( - self, - robot_id: WorldRobotID, - path: JointPath, - duration: float = 3.0, + def cancel_preview_animation(self, robot_ids: Sequence[WorldRobotID] | None = None) -> None: + """Invalidate active preview frames and hide preview ghosts immediately.""" + with self._lock: + self._preview_animation_generation += 1 + affected = set(robot_ids) if robot_ids is not None else set(self._robots) + for robot_id in affected: + self._preview_animation_generations[robot_id] = ( + self._preview_animation_generations.get(robot_id, 0) + 1 + ) + if robot_id not in self._robots: + continue + self._set_preview_visibility(robot_id, False) + + def _robot_trajectory_indices( + self, trajectory: JointTrajectory + ) -> dict[WorldRobotID, list[tuple[int, str]]]: + robot_ids_by_name = { + robot.config.name: robot_id for robot_id, robot in self._robots.items() + } + indices: dict[WorldRobotID, list[tuple[int, str]]] = {} + for index, global_name in enumerate(trajectory.joint_names): + if "/" not in global_name: + raise ValueError(f"trajectory joint '{global_name}' is not globally named") + robot_name, local_name = global_name.split("/", 1) + robot_id = robot_ids_by_name.get(robot_name) + if robot_id is None: + raise ValueError(f"trajectory references unknown robot '{robot_name}'") + if local_name not in self._robots[robot_id].config.joint_names: + raise ValueError(f"trajectory references unknown joint '{global_name}'") + indices.setdefault(robot_id, []).append((index, local_name)) + return indices + + def animate_trajectory( + self, trajectory: JointTrajectory, duration: float | None = None ) -> None: - """Animate a path using the preview (yellow ghost) robot. - - The preview stays visible after animation completes. - """ - if self._meshcat is None or len(path) < 2: - return - - robot_data = self._robots.get(robot_id) - if robot_data is None or robot_data.preview_model_instance is None: + """Render raw globally named trajectory on its stored shared clock.""" + if self._meshcat is None or len(trajectory.points) < 2: return import time - self.show_preview(robot_id) - dt = duration / (len(path) - 1) - for joint_state in path: - positions = np.array(joint_state.position, dtype=np.float64) + robot_indices = self._robot_trajectory_indices(trajectory) + robot_ids = list(robot_indices) + playback_scale = 1.0 + if duration is not None: + if duration <= 0.0 or trajectory.duration <= 0.0: + raise ValueError("preview duration must be positive") + playback_scale = duration / trajectory.duration + baselines: dict[WorldRobotID, NDArray[np.float64]] = {} + joint_positions_by_name: dict[WorldRobotID, dict[str, int]] = {} + with self._lock: + assert self._plant_context is not None + assert self._live_context is not None + self._preview_animation_generation += 1 + animation_generations: dict[WorldRobotID, int] = {} + for robot_id in robot_ids: + self._preview_animation_generations[robot_id] = ( + self._preview_animation_generations.get(robot_id, 0) + 1 + ) + animation_generations[robot_id] = self._preview_animation_generations[robot_id] + robot_data = self._robots[robot_id] + self._set_preview_visibility(robot_id, True) + baselines[robot_id] = np.array( + self.get_joint_state(self._live_context, robot_id).position, + dtype=np.float64, + ) + joint_positions_by_name[robot_id] = dict( + zip( + robot_data.config.joint_names, + range(len(robot_data.config.joint_names)), + strict=True, + ) + ) + + try: + previous_time = trajectory.points[0].time_from_start + for frame_index, point in enumerate(trajectory.points): + with self._lock: + active_robot_ids = [ + robot_id + for robot_id in robot_ids + if self._preview_animation_generations.get(robot_id) + == animation_generations[robot_id] + ] + if not active_robot_ids: + return + assert self._plant_context is not None + for robot_id in active_robot_ids: + indexed_names = robot_indices[robot_id] + positions = baselines[robot_id].copy() + local_index = joint_positions_by_name[robot_id] + for trajectory_index, local_name in indexed_names: + positions[local_index[local_name]] = point.positions[trajectory_index] + self._set_preview_positions(self._plant_context, robot_id, positions) + self._publish_visualization() + if frame_index < len(trajectory.points) - 1: + next_time = trajectory.points[frame_index + 1].time_from_start + time.sleep((next_time - previous_time) * playback_scale) + previous_time = next_time + finally: with self._lock: - assert self._plant_context is not None - self._set_preview_positions(self._plant_context, robot_id, positions) - self.publish_visualization() - time.sleep(dt) + for robot_id in robot_ids: + if ( + self._preview_animation_generations.get(robot_id) + == animation_generations[robot_id] + ): + self._set_preview_visibility(robot_id, False) def close(self) -> None: """Shut down the viz thread.""" + self.cancel_preview_animation() if self._meshcat is not None: self._meshcat.close() diff --git a/dimos/manipulation/planning/world/roboplan_world.py b/dimos/manipulation/planning/world/roboplan_world.py index d252d3fdc9..f5dd421589 100644 --- a/dimos/manipulation/planning/world/roboplan_world.py +++ b/dimos/manipulation/planning/world/roboplan_world.py @@ -41,9 +41,21 @@ "Install the manipulation extra before selecting the roboplan backend." ) from exc +from dimos.manipulation.planning.groups.identifiers import ( + make_global_joint_names, + make_planning_group_id, +) +from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupSelection +from dimos.manipulation.planning.groups.utils import joint_state_to_ordered_positions +from dimos.manipulation.planning.planners.selected_joint_space import normalize_selection_target from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import ObstacleType, PlanningStatus -from dimos.manipulation.planning.spec.models import Obstacle, PlanningResult, WorldRobotID +from dimos.manipulation.planning.spec.models import ( + Obstacle, + PlanningGroupID, + PlanningResult, + WorldRobotID, +) from dimos.manipulation.planning.utils.mesh_utils import prepare_urdf_for_drake from dimos.manipulation.planning.utils.path_utils import compute_path_length from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped @@ -104,6 +116,9 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: raise ValueError("RoboPlanWorld currently supports one robot per Scene") if not Path(config.model_path).exists(): raise FileNotFoundError(f"Robot model not found: {Path(config.model_path).resolve()}") + if any(data.config.name == config.name for data in self._robots.values()): + raise ValueError(f"Robot name '{config.name}' is already registered") + self._validate_planning_group_config(config) self._validate_robot_config(config) self._robot_counter += 1 @@ -272,7 +287,17 @@ def check_edge_collision_free( def get_ee_pose(self, ctx: RoboPlanContext, robot_id: WorldRobotID) -> PoseStamped: """Get end-effector pose if RoboPlan exposes FK.""" robot = self._get_robot(robot_id) - mat = self.get_link_pose(ctx, robot_id, robot.config.end_effector_link) + group_id = self._primary_pose_group_id_for_config(robot.config) + if group_id is None: + raise ValueError(f"Robot '{robot.config.name}' has no pose-targetable planning group") + return self.get_group_ee_pose(ctx, group_id) + + def get_group_ee_pose(self, ctx: RoboPlanContext, group_id: PlanningGroupID) -> PoseStamped: + """Get planning-group tip pose if RoboPlan exposes FK.""" + group = self._planning_group_from_id(group_id) + if group.tip_link is None: + raise ValueError(f"Planning group '{group_id}' has no tip link") + mat = self.get_link_pose(ctx, self._robot_id_for_group(group_id), group.tip_link) pose = matrix_to_pose(mat) return PoseStamped( frame_id="world", @@ -299,17 +324,46 @@ def get_link_pose( def get_jacobian(self, ctx: RoboPlanContext, robot_id: WorldRobotID) -> NDArray[np.float64]: """Get end-effector Jacobian if RoboPlan exposes a compatible API.""" robot = self._get_robot(robot_id) + group_id = self._primary_pose_group_id_for_config(robot.config) + if group_id is None: + raise ValueError(f"Robot '{robot.config.name}' has no pose-targetable planning group") + return self.get_group_jacobian(ctx, group_id) + + def get_group_jacobian( + self, ctx: RoboPlanContext, group_id: PlanningGroupID + ) -> NDArray[np.float64]: + """Get planning-group Jacobian projected to group-local joint order.""" + group = self._planning_group_from_id(group_id) + if group.tip_link is None: + raise ValueError(f"Planning group '{group_id}' has no tip link") + robot_id = self._robot_id_for_group(group_id) + robot = self._get_robot(robot_id) q = ctx.q_by_robot.get(robot_id) if q is None: raise KeyError(f"Robot '{robot_id}' not found in context") scene = self._require_scene() - result = scene.computeFrameJacobian( - self._to_scene_q(robot_id, q), robot.config.end_effector_link, True - ) + result = scene.computeFrameJacobian(self._to_scene_q(robot_id, q), group.tip_link, True) arr = np.asarray(result, dtype=np.float64) if arr.shape[0] != 6: raise ValueError(f"Unexpected RoboPlan Jacobian shape: {arr.shape}; expected 6 x n") - return arr + scene_joint_order = self._query_scene_joint_order(scene, robot.config) + if scene_joint_order is not None and arr.shape[1] == len(scene_joint_order): + missing = [name for name in group.local_joint_names if name not in scene_joint_order] + if missing: + raise ValueError(f"Unknown joints for planning group '{group_id}': {missing}") + indices = [scene_joint_order.index(name) for name in group.local_joint_names] + return arr[:, indices] + if arr.shape[1] == len(robot.config.joint_names): + missing = [ + name for name in group.local_joint_names if name not in robot.config.joint_names + ] + if missing: + raise ValueError(f"Unknown joints for planning group '{group_id}': {missing}") + indices = [robot.config.joint_names.index(name) for name in group.local_joint_names] + return arr[:, indices] + raise ValueError( + f"Unexpected RoboPlan Jacobian shape: {arr.shape}; cannot project group '{group_id}'" + ) # PlannerSpec for native RoboPlan planning @@ -362,6 +416,82 @@ def plan_joint_path( message="RoboPlan path found", ) + def plan_selected_joint_path( + self, + world: WorldSpec, + selection: PlanningGroupSelection, + start: JointState, + goal: JointState, + timeout: float = 10.0, + max_iterations: int = 5000, + ) -> PlanningResult: + """Plan a single planning group using RoboPlan's native RRT.""" + if world is not self: + return PlanningResult( + status=PlanningStatus.NO_SOLUTION, + message="RoboPlan-native planner requires its RoboPlanWorld instance", + ) + start_time = time.time() + if not selection.groups: + return PlanningResult( + status=PlanningStatus.INVALID_GOAL, + message="No planning groups selected", + ) + if len(selection.groups) != 1: + return PlanningResult( + status=PlanningStatus.UNSUPPORTED, + message="RoboPlan-native planning supports exactly one selected planning group", + ) + + group = selection.groups[0] + try: + normalized_start = normalize_selection_target(selection, start, "start") + except ValueError as exc: + return PlanningResult(status=PlanningStatus.INVALID_START, message=str(exc)) + try: + normalized_goal = normalize_selection_target(selection, goal, "goal") + except ValueError as exc: + return PlanningResult(status=PlanningStatus.INVALID_GOAL, message=str(exc)) + + robot_id = self._robot_id_for_group(group.id) + q_start = np.asarray(normalized_start.position, dtype=np.float64) + q_goal = np.asarray(normalized_goal.position, dtype=np.float64) + try: + path_arrays = self._run_native_rrt( + robot_id, + q_start, + q_goal, + timeout, + group_name=group.group_name, + joint_names=list(group.local_joint_names), + ) + except ValueError as exc: + return PlanningResult( + status=PlanningStatus.NO_SOLUTION, + planning_time=time.time() - start_time, + message=f"RoboPlan-native planning failed: {exc}", + ) + if not path_arrays: + return PlanningResult( + status=PlanningStatus.NO_SOLUTION, + planning_time=time.time() - start_time, + message="RoboPlan-native planning failed: returned an empty path", + ) + path = [ + JointState( + name=list(selection.joint_names), + position=np.asarray(q).astype(float).tolist(), + ) + for q in path_arrays + ] + return PlanningResult( + status=PlanningStatus.SUCCESS, + path=path, + planning_time=time.time() - start_time, + path_length=compute_path_length(path), + message="RoboPlan path found", + ) + def get_name(self) -> str: """Get planner name.""" return "RoboPlan" @@ -407,6 +537,11 @@ def _generate_srdf(self, config: RobotModelConfig, urdf_path: Path) -> str: for joint_name in config.joint_names: lines.append(f' ') lines.append(" ") + for group in config.planning_groups: + lines.append(f' ') + for joint_name in group.joint_names: + lines.append(f' ') + lines.append(" ") for link1, link2 in self._collision_exclusion_pairs(config, urdf_path): lines.append( f' None: + """Validate planning groups before mutating backend state.""" + seen_group_names: set[str] = set() + for definition in config.planning_groups: + group_id = make_planning_group_id(config.name, definition.name) + if definition.name in seen_group_names: + raise ValueError(f"Planning group '{group_id}' is already registered") + make_global_joint_names(config.name, definition.joint_names) + seen_group_names.add(definition.name) + + def _planning_group_from_config( + self, config: RobotModelConfig, group_id: PlanningGroupID + ) -> PlanningGroup: + for definition in config.planning_groups: + if make_planning_group_id(config.name, definition.name) == group_id: + return PlanningGroup( + id=group_id, + robot_name=config.name, + group_name=definition.name, + joint_names=tuple(make_global_joint_names(config.name, definition.joint_names)), + local_joint_names=definition.joint_names, + base_link=definition.base_link, + tip_link=definition.tip_link, + source=definition.source, + ) + raise KeyError(f"Unknown planning group ID: {group_id}") + + def _planning_group_from_id(self, group_id: PlanningGroupID) -> PlanningGroup: + for robot in self._robots.values(): + try: + return self._planning_group_from_config(robot.config, group_id) + except KeyError: + continue + raise KeyError(f"Unknown planning group ID: {group_id}") + + def _primary_pose_group_id_for_config(self, config: RobotModelConfig) -> PlanningGroupID | None: + pose_group_ids = [ + make_planning_group_id(config.name, group.name) + for group in config.planning_groups + if group.has_pose_target + ] + if not pose_group_ids: + return None + if len(pose_group_ids) > 1: + raise ValueError( + f"Robot '{config.name}' has {len(pose_group_ids)} pose-targetable " + "planning groups; use an explicit planning group ID" + ) + return pose_group_ids[0] + def _get_robot(self, robot_id: WorldRobotID) -> _RoboPlanRobotData: if robot_id not in self._robots: raise KeyError(f"Robot '{robot_id}' not found") return self._robots[robot_id] + def _robot_id_for_group(self, group_id: PlanningGroupID) -> WorldRobotID: + group = self._planning_group_from_id(group_id) + matches = [ + rid for rid, data in self._robots.items() if data.config.name == group.robot_name + ] + if not matches: + raise KeyError(f"No robot registered for planning group '{group_id}'") + return matches[0] + def _joint_state_to_q( self, robot_id: WorldRobotID, joint_state: JointState ) -> NDArray[np.float64]: robot = self._get_robot(robot_id) - if len(joint_state.position) != len(robot.config.joint_names): - raise ValueError("JointState position length must match configured joint count") - if not joint_state.name: - return np.asarray(joint_state.position, dtype=np.float64) - name_to_pos = { - robot.config.get_urdf_joint_name(name): position - for name, position in zip(joint_state.name, joint_state.position, strict=True) - } - missing = [name for name in robot.config.joint_names if name not in name_to_pos] - if missing: - raise ValueError(f"JointState missing joints for RoboPlanWorld: {missing}") - return np.asarray( - [name_to_pos[name] for name in robot.config.joint_names], dtype=np.float64 + return joint_state_to_ordered_positions( + joint_state, + joint_names=robot.config.joint_names, + joint_name_mapping=robot.config.joint_name_mapping, ) def _require_finalized(self) -> None: @@ -628,27 +813,33 @@ def _run_native_rrt( q_start: NDArray[np.float64], q_goal: NDArray[np.float64], timeout: float, + *, + group_name: str | None = None, + joint_names: list[str] | None = None, ) -> list[NDArray[np.float64]]: scene = self._require_scene() robot = self._get_robot(robot_id) options = roboplan_rrt.RRTOptions() - options.group_name = robot.config.name + native_group_name = robot.config.name if group_name is None else group_name + native_joint_names = robot.config.joint_names if joint_names is None else joint_names + options.group_name = native_group_name options.max_planning_time = timeout options.collision_check_use_bisection = False planner = roboplan_rrt.RRT(scene, options) - start_config = self._to_native_joint_configuration(robot_id, q_start) - goal_config = self._to_native_joint_configuration(robot_id, q_goal) + start_config = self._to_native_joint_configuration(robot_id, q_start, native_joint_names) + goal_config = self._to_native_joint_configuration(robot_id, q_goal, native_joint_names) result = planner.plan(start_config, goal_config) if result is None: raise ValueError("RoboPlan RRT returned no path") return self._extract_native_path(result) def _to_native_joint_configuration( - self, robot_id: WorldRobotID, q: NDArray[np.float64] + self, robot_id: WorldRobotID, q: NDArray[np.float64], joint_names: list[str] | None = None ) -> roboplan_core.JointConfiguration: robot = self._get_robot(robot_id) return roboplan_core.JointConfiguration( - robot.config.joint_names, np.asarray(q, dtype=np.float64) + robot.config.joint_names if joint_names is None else joint_names, + np.asarray(q, dtype=np.float64), ) def _extract_native_path(self, result: roboplan_core.JointPath) -> list[NDArray[np.float64]]: diff --git a/dimos/manipulation/planning/world/test_drake_world_planning_groups.py b/dimos/manipulation/planning/world/test_drake_world_planning_groups.py new file mode 100644 index 0000000000..7832626fa5 --- /dev/null +++ b/dimos/manipulation/planning/world/test_drake_world_planning_groups.py @@ -0,0 +1,449 @@ +# 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 + +from pathlib import Path + +import numpy as np +import pytest + +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.world.drake_world import DRAKE_AVAILABLE, DrakeWorld +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint + +requires_drake = pytest.mark.skipif( + not DRAKE_AVAILABLE, + reason="Drake planning-group tests require the manipulation extra", +) + + +def _trajectory(names: list[str], first: list[float], second: list[float]) -> JointTrajectory: + return JointTrajectory( + joint_names=names, + points=[ + TrajectoryPoint(time_from_start=0.0, positions=first, velocities=[0.0] * len(names)), + TrajectoryPoint(time_from_start=2.0, positions=second, velocities=[0.0] * len(names)), + ], + ) + + +def _write_urdf(path: Path) -> None: + path.write_text( + """ + + + + + + + + + + + + + + + +""" + ) + + +def _write_urdf_with_world_base_joint(path: Path) -> None: + path.write_text( + """ + + + + + + + + + + + + + + + + + + + + +""" + ) + + +def _config( + path: Path, groups: list[PlanningGroupDefinition], joints: list[str] | None = None +) -> RobotModelConfig: + return RobotModelConfig( + name="arm", + model_path=path, + base_pose=PoseStamped(position=[0, 0, 0], orientation=[0, 0, 0, 1]), + joint_names=joints or ["joint1", "joint2"], + base_link="base_link", + planning_groups=groups, + ) + + +def _arm_group( + *joint_names: str, tip_link: str | None = "tool0", name: str = "arm" +) -> PlanningGroupDefinition: + return PlanningGroupDefinition( + name=name, joint_names=joint_names, base_link="base_link", tip_link=tip_link + ) + + +def test_drake_config_group_helpers_resolve_groups_without_drake_runtime(tmp_path: Path) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + config = _config(urdf, [_arm_group("joint2", "joint1", name="wrist")]) + + group = DrakeWorld._planning_group_from_config(config, "arm/wrist") + + assert DrakeWorld._primary_pose_group_id_for_config(config) == "arm/wrist" + assert group.id == "arm/wrist" + assert group.joint_names == ("arm/joint2", "arm/joint1") + assert group.local_joint_names == ("joint2", "joint1") + assert group.tip_link == "tool0" + + +def test_drake_config_group_helpers_validate_duplicate_and_ambiguous_groups( + tmp_path: Path, +) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + duplicate = _config( + urdf, + [_arm_group("joint1", name="same"), _arm_group("joint2", name="same")], + ) + ambiguous = _config( + urdf, + [_arm_group("joint1", name="a"), _arm_group("joint2", name="b")], + ) + + with pytest.raises(ValueError, match="already registered"): + DrakeWorld._validate_planning_group_config(duplicate) + with pytest.raises(ValueError, match="multiple pose"): + DrakeWorld._primary_pose_group_id_for_config(ambiguous) + with pytest.raises(KeyError, match="Unknown planning group ID"): + DrakeWorld._planning_group_from_config(ambiguous, "arm/missing") + + +@requires_drake +def test_drake_group_fk_uses_tip_link_and_legacy_unique_pose_group(tmp_path: Path) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + world = DrakeWorld() + robot_id = world.add_robot(_config(urdf, [_arm_group("joint1", "joint2")])) + world.finalize() + ctx = world.get_live_context() + world.set_joint_state( + ctx, robot_id, JointState({"name": ["joint1", "joint2"], "position": [0.0, 0.0]}) + ) + + group_pose = world.get_group_ee_pose(ctx, "arm/arm") + legacy_pose = world.get_ee_pose(ctx, robot_id) + + assert group_pose.position.x == pytest.approx(2.0) + assert legacy_pose.position.x == pytest.approx(group_pose.position.x) + assert world.get_jacobian(ctx, robot_id).shape == (6, 2) + + +@requires_drake +def test_drake_applies_config_base_pose_when_urdf_has_world_base_joint( + tmp_path: Path, +) -> None: + urdf = tmp_path / "robot_with_world.urdf" + _write_urdf_with_world_base_joint(urdf) + world = DrakeWorld(enable_viz=False) + left_id = world.add_robot( + RobotModelConfig( + name="left_arm", + model_path=urdf, + base_pose=PoseStamped(position=[0, 0.5, 0], orientation=[0, 0, 0, 1]), + joint_names=["joint1", "joint2"], + base_link="base_link", + planning_groups=[_arm_group("joint1", "joint2")], + ) + ) + right_id = world.add_robot( + RobotModelConfig( + name="right_arm", + model_path=urdf, + base_pose=PoseStamped(position=[0, -0.5, 0], orientation=[0, 0, 0, 1]), + joint_names=["joint1", "joint2"], + base_link="base_link", + planning_groups=[_arm_group("joint1", "joint2")], + ) + ) + world.finalize() + ctx = world.get_live_context() + + left_base_pose = world.get_link_pose(ctx, left_id, "base_link") + right_base_pose = world.get_link_pose(ctx, right_id, "base_link") + + assert left_base_pose[1, 3] == pytest.approx(0.5) + assert right_base_pose[1, 3] == pytest.approx(-0.5) + assert left_base_pose[1, 3] != pytest.approx(right_base_pose[1, 3]) + + +@requires_drake +def test_drake_group_jacobian_shape_and_group_local_order(tmp_path: Path) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + world = DrakeWorld() + robot_id = world.add_robot( + _config( + urdf, + [ + _arm_group("joint1", "joint2", name="wrist_forward"), + _arm_group("joint2", "joint1", name="wrist_reverse"), + ], + ) + ) + world.finalize() + ctx = world.get_live_context() + world.set_joint_state( + ctx, robot_id, JointState({"name": ["joint1", "joint2"], "position": [0.0, 0.0]}) + ) + + forward_jacobian = world.get_group_jacobian(ctx, "arm/wrist_forward") + reverse_jacobian = world.get_group_jacobian(ctx, "arm/wrist_reverse") + + assert reverse_jacobian.shape == (6, 2) + np.testing.assert_allclose(reverse_jacobian[:, 0], forward_jacobian[:, 1]) + np.testing.assert_allclose(reverse_jacobian[:, 1], forward_jacobian[:, 0]) + + +@requires_drake +def test_drake_legacy_wrappers_fail_at_call_time_for_no_or_ambiguous_pose(tmp_path: Path) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + no_pose = DrakeWorld() + no_pose_id = no_pose.add_robot(_config(urdf, [_arm_group("joint1", tip_link=None)])) + no_pose.finalize() + with pytest.raises(ValueError, match="no pose-targetable"): + no_pose.get_ee_pose(no_pose.get_live_context(), no_pose_id) + + ambiguous = DrakeWorld() + ambiguous_id = ambiguous.add_robot( + _config( + urdf, + [ + _arm_group("joint1", tip_link="link1", name="a"), + _arm_group("joint2", tip_link="tool0", name="b"), + ], + ) + ) + ambiguous.finalize() + with pytest.raises(ValueError, match="multiple pose"): + ambiguous.get_jacobian(ambiguous.get_live_context(), ambiguous_id) + + +@requires_drake +def test_drake_group_jacobian_rejects_non_controllable_group_joints(tmp_path: Path) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + world = DrakeWorld() + world.add_robot(_config(urdf, [_arm_group("joint1", "joint2")], joints=["joint1"])) + world.finalize() + + with pytest.raises(ValueError, match="non-controllable"): + world.get_group_jacobian(world.get_live_context(), "arm/arm") + + +@requires_drake +def test_drake_animate_trajectory_projects_all_robots_on_shared_ticks( + tmp_path: Path, monkeypatch +) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + world = DrakeWorld() + left_config = _config(urdf, [_arm_group("joint1")]).model_copy(update={"name": "left"}) + right_config = _config(urdf, [_arm_group("joint2")]).model_copy(update={"name": "right"}) + left_id = world.add_robot(left_config) + right_id = world.add_robot(right_config) + world.finalize() + world._meshcat = object() # type: ignore[assignment] + ctx = world.get_live_context() + world.set_joint_state(ctx, left_id, JointState(name=["joint1", "joint2"], position=[0.1, 0.2])) + world.set_joint_state(ctx, right_id, JointState(name=["joint1", "joint2"], position=[0.3, 0.4])) + updates: list[tuple[str, list[float]]] = [] + shown: list[tuple[str, ...]] = [] + hidden: list[tuple[str, ...]] = [] + sleeps: list[float] = [] + monkeypatch.setattr( + world, + "_set_preview_positions", + lambda _ctx, robot_id, positions: updates.append((robot_id, positions.tolist())), + ) + monkeypatch.setattr( + world, + "_set_preview_visibility", + lambda robot_id, visible: (shown if visible else hidden).append((robot_id,)), + ) + monkeypatch.setattr(world, "_publish_visualization", lambda: None) + monkeypatch.setattr("time.sleep", sleeps.append) + plan = type("Plan", (), {})() + plan.trajectory = _trajectory(["left/joint1", "right/joint2"], [1.0, 2.0], [3.0, 4.0]) + + world.animate_trajectory(plan.trajectory, duration=2.0) + + assert shown == [(left_id,), (right_id,)] + assert hidden == [(left_id,), (right_id,)] + assert updates == [ + (left_id, [1.0, 0.2]), + (right_id, [0.3, 2.0]), + (left_id, [3.0, 0.2]), + (right_id, [0.3, 4.0]), + ] + assert sleeps == [2.0] + + +@requires_drake +def test_drake_animate_trajectory_validates_before_visibility_and_cleans_up( + tmp_path: Path, monkeypatch +) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + world = DrakeWorld() + robot_id = world.add_robot(_config(urdf, [_arm_group("joint1")])) + world.finalize() + world._meshcat = object() # type: ignore[assignment] + world.set_joint_state( + world.get_live_context(), + robot_id, + JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), + ) + shown: list[tuple[str, ...]] = [] + hidden: list[tuple[str, ...]] = [] + monkeypatch.setattr( + world, + "_set_preview_visibility", + lambda robot_id, visible: (shown if visible else hidden).append((robot_id,)), + ) + monkeypatch.setattr(world, "_publish_visualization", lambda: None) + malformed = _trajectory(["unknown/joint1"], [0.0], [1.0]) + with pytest.raises(ValueError, match="unknown robot"): + world.animate_trajectory(malformed) + assert shown == [] + + valid = _trajectory(["arm/joint1"], [0.0], [1.0]) + + def fail_preview_update(*_args: object) -> None: + raise RuntimeError("boom") + + monkeypatch.setattr(world, "_set_preview_positions", fail_preview_update) + with pytest.raises(RuntimeError, match="boom"): + world.animate_trajectory(valid) + assert shown == [(robot_id,)] + assert hidden == [(robot_id,)] + + +@requires_drake +def test_drake_cancel_preview_hides_ghosts_before_animation_resumes( + tmp_path: Path, monkeypatch +) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + world = DrakeWorld() + robot_id = world.add_robot(_config(urdf, [_arm_group("joint1")])) + world.finalize() + world._meshcat = object() # type: ignore[assignment] + world.set_joint_state( + world.get_live_context(), + robot_id, + JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), + ) + hidden: list[tuple[str, ...]] = [] + hidden_snapshots_during_sleep: list[list[tuple[str, ...]]] = [] + monkeypatch.setattr( + world, + "_set_preview_visibility", + lambda robot_id, visible: None if visible else hidden.append((robot_id,)), + ) + monkeypatch.setattr(world, "_publish_visualization", lambda: None) + + def cancel_during_sleep(_duration: float) -> None: + world.cancel_preview_animation() + hidden_snapshots_during_sleep.append(list(hidden)) + + monkeypatch.setattr("time.sleep", cancel_during_sleep) + + world.animate_trajectory(_trajectory(["arm/joint1"], [0.0], [1.0])) + + assert hidden_snapshots_during_sleep == [[(robot_id,)]] + assert hidden[0] == (robot_id,) + + +@requires_drake +def test_drake_animate_trajectory_rejects_unknown_robot_before_visibility( + tmp_path: Path, monkeypatch +) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + world = DrakeWorld() + world.add_robot(_config(urdf, [_arm_group("joint1")])) + world.finalize() + world._meshcat = object() # type: ignore[assignment] + shown: list[tuple[str, ...]] = [] + with pytest.raises(ValueError, match="unknown robot"): + world.animate_trajectory(_trajectory(["missing/joint1"], [0.0], [1.0])) + + assert shown == [] + + +@requires_drake +def test_drake_animate_trajectory_cancellation_stops_stale_frames_and_hides_preview( + tmp_path: Path, monkeypatch +) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + world = DrakeWorld() + robot_id = world.add_robot(_config(urdf, [_arm_group("joint1")])) + world.finalize() + world._meshcat = object() # type: ignore[assignment] + world.set_joint_state( + world.get_live_context(), + robot_id, + JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), + ) + updates: list[list[float]] = [] + hidden: list[tuple[str, ...]] = [] + monkeypatch.setattr( + world, + "_set_preview_positions", + lambda _ctx, _robot_id, positions: updates.append(positions.tolist()), + ) + monkeypatch.setattr( + world, + "_set_preview_visibility", + lambda robot_id, visible: hidden.append((robot_id,)) if not visible else None, + ) + monkeypatch.setattr(world, "_publish_visualization", lambda: None) + monkeypatch.setattr("time.sleep", lambda _duration: world.cancel_preview_animation()) + + world.animate_trajectory(_trajectory(["arm/joint1"], [1.0], [2.0])) + + assert updates == [[1.0, 0.0]] + assert hidden[0] == (robot_id,) diff --git a/dimos/manipulation/skill_errors.py b/dimos/manipulation/skill_errors.py index 9a17085ec5..0c4e4a4b63 100644 --- a/dimos/manipulation/skill_errors.py +++ b/dimos/manipulation/skill_errors.py @@ -37,6 +37,12 @@ "GRASP_ATTEMPTS_EXHAUSTED", "GRIPPER_FAILED", "WORLD_MONITOR_UNAVAILABLE", + "EXECUTION_CANCELLED", + "EXECUTION_REJECTED", + "EXECUTION_CANCEL_REJECTED", + "EXECUTION_CANCEL_UNRESOLVED", + "EXECUTION_DISPATCH_FAILED", + "RESET_FAILED", ] # Union of codes a manipulation skill may emit (common + manipulation-specific). diff --git a/dimos/manipulation/test_execution_auxiliary.py b/dimos/manipulation/test_execution_auxiliary.py new file mode 100644 index 0000000000..bc190af221 --- /dev/null +++ b/dimos/manipulation/test_execution_auxiliary.py @@ -0,0 +1,13 @@ +from dimos.manipulation.execution_auxiliary import AuxiliaryCallBook +from dimos.manipulation.execution_effects import AuxiliaryDone + + +def test_expire_retains_timeout_result_through_late_completion() -> None: + book = AuxiliaryCallBook() + ticket = book.register("aux", 1.0, setter=True) + + assert book.expire(2.0, "timeout") + assert book.has_inflight() + assert book.complete(AuxiliaryDone(ticket.action_id, "late value")) + assert not book.has_unsettled() + assert book.take_result(ticket.action_id) == (False, None, "timeout") diff --git a/dimos/manipulation/test_execution_clock.py b/dimos/manipulation/test_execution_clock.py new file mode 100644 index 0000000000..79a841cb69 --- /dev/null +++ b/dimos/manipulation/test_execution_clock.py @@ -0,0 +1,82 @@ +"""Tests for the serialized, validated execution clock.""" + +import threading + +import pytest + +from dimos.manipulation.execution_clock import InvalidMonotonicClock, ValidatedMonotonicClock + + +class GatedSource: + def __init__(self, initial: float) -> None: + self.first_entry = threading.Event() + self.second_entry = threading.Event() + self.first_release = threading.Event() + self.second_release = threading.Event() + self._lock = threading.Lock() + self.calls = 0 + self.value = initial + + def __call__(self) -> float: + with self._lock: + self.calls += 1 + call = self.calls + if call == 2: + self.first_entry.set() + if not self.first_release.wait(1): + raise AssertionError("first clock sample was not released") + elif call == 3: + self.second_entry.set() + if not self.second_release.wait(1): + raise AssertionError("second clock sample was not released") + return self.value + + +def test_now_serializes_source_sampling_before_validation() -> None: + source = GatedSource(0.0) + clock = ValidatedMonotonicClock(source) + first_done = threading.Event() + second_done = threading.Event() + + def first() -> None: + clock.now() + first_done.set() + + def second() -> None: + clock.now() + second_done.set() + + one = threading.Thread(target=first) + two = threading.Thread(target=second) + one.start() + assert source.first_entry.wait(1) + two.start() + assert not source.second_entry.wait(0.02) + source.first_release.set() + assert first_done.wait(1) + assert source.second_entry.wait(1) + source.second_release.set() + assert second_done.wait(1) + one.join(1) + two.join(1) + + +@pytest.mark.parametrize("sample", [float("nan"), float("inf"), float("-inf")]) +def test_clock_rejects_nonfinite_initial_sample(sample: float) -> None: + with pytest.raises(InvalidMonotonicClock): + ValidatedMonotonicClock(lambda: sample) + + +@pytest.mark.parametrize("sample", [float("nan"), float("inf"), float("-inf")]) +def test_clock_rejects_nonfinite_later_sample(sample: float) -> None: + values = iter((0.0, sample)) + clock = ValidatedMonotonicClock(lambda: next(values)) + with pytest.raises(InvalidMonotonicClock): + clock.now() + + +def test_clock_rejects_decreasing_sample() -> None: + values = iter((1.0, 0.5)) + clock = ValidatedMonotonicClock(lambda: next(values)) + with pytest.raises(InvalidMonotonicClock): + clock.now() diff --git a/dimos/manipulation/test_execution_effects.py b/dimos/manipulation/test_execution_effects.py new file mode 100644 index 0000000000..b6afd2c0b6 --- /dev/null +++ b/dimos/manipulation/test_execution_effects.py @@ -0,0 +1,91 @@ +"""Direct coverage for the daemonized execution-effect runner.""" + +from dataclasses import FrozenInstanceError +import subprocess +import sys +import threading + +import pytest + +from dimos.manipulation.execution_effects import ( + AuxiliaryDone, + EffectDone, + ExecutionEffectRunner, + StopDone, +) + + +def test_four_blocked_effects_use_four_daemon_execution_rpc_workers() -> None: + runner = ExecutionEffectRunner() + entered = [threading.Event() for _ in range(4)] + release = threading.Event() + try: + for event in entered: + runner.submit_action( + "blocked", + lambda event=event: (event.set(), release.wait())[1], + lambda _: None, + ) + + assert all(event.wait(1) for event in entered) + workers = runner._executor._threads + assert len(workers) == 4 + assert {thread.name for thread in workers} == { + f"execution-rpc_{index}" for index in range(4) + } + assert all(thread.daemon for thread in workers) + finally: + release.set() + runner.shutdown() + + +def test_completion_envelopes_are_immutable_and_correlated() -> None: + runner = ExecutionEffectRunner() + completions: list[object] = [] + complete = threading.Event() + + def enqueue(value: object) -> None: + completions.append(value) + if len(completions) == 4: + complete.set() + + try: + runner.submit_action("action-7", lambda: "done", enqueue) + runner.submit_auxiliary("aux-3", lambda: 42, enqueue) + runner.submit_stop(lambda: None, enqueue) + runner.submit_action("action-8", lambda: None, enqueue) + + assert complete.wait(1) + assert EffectDone("action-7", "done") in completions + assert AuxiliaryDone("aux-3", 42) in completions + assert StopDone(True) in completions + assert EffectDone("action-8", None) in completions + with pytest.raises(FrozenInstanceError): + completions[0].action_id = "changed" # type: ignore[attr-defined] + finally: + runner.shutdown() + + +def test_permanently_blocked_effect_does_not_hold_interpreter_exit() -> None: + code = """ +import threading +from dimos.manipulation.execution_effects import ExecutionEffectRunner + +runner = ExecutionEffectRunner() +runner.submit_action("never", threading.Event().wait, lambda _: None) +""" + process = subprocess.Popen( + [sys.executable, "-c", code], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + try: + _, stderr = process.communicate(timeout=2) + except subprocess.TimeoutExpired: + process.kill() + process.communicate() + process.wait() + pytest.fail("permanently blocked effect held interpreter exit") + process.wait() + assert process.returncode == 0, stderr diff --git a/dimos/manipulation/test_execution_gateway.py b/dimos/manipulation/test_execution_gateway.py new file mode 100644 index 0000000000..6d0e57873d --- /dev/null +++ b/dimos/manipulation/test_execution_gateway.py @@ -0,0 +1,59 @@ +"""Black-box tests for coordinator gateway normalization.""" + +from typing import Any + +from dimos.manipulation.execution_gateway import ControlCoordinatorGateway +from dimos.manipulation.execution_models import Outcome +from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState + + +class Client: + def __init__(self, value: Any) -> None: + self.value = value + self.calls: list[tuple[str, str, dict[str, Any]]] = [] + + def task_invoke(self, task: str, method: str, kwargs: dict[str, Any]) -> Any: + self.calls.append((task, method, kwargs)) + return self.value + + def set_gripper_position(self, hardware_id: str, position: float) -> Any: + self.calls.append((hardware_id, "set_gripper_position", {"position": position})) + return self.value + + def get_gripper_position(self, hardware_id: str) -> Any: + self.calls.append((hardware_id, "get_gripper_position", {})) + return self.value + + +def test_gateway_normalizes_execute_cancel_and_status() -> None: + client = Client(True) + gateway = ControlCoordinatorGateway(client) + assert gateway.execute("task", {"trajectory": "request"}) == Outcome.ACCEPTED + assert gateway.cancel("task") == Outcome.CANCELLED + + client.value = TrajectoryState.EXECUTING + assert gateway.status("task") == Outcome.RUNNING + assert client.calls[0] == ("task", "execute", {"trajectory": "request"}) + + client.value = False + assert gateway.cancel("task") == Outcome.INACTIVE + + +def test_gateway_normalizes_unknown_values_safely() -> None: + client = Client("not-a-state") + gateway = ControlCoordinatorGateway(client) + assert gateway.status("task") == Outcome.UNKNOWN + client.value = object() + assert gateway.execute("task", {}) == Outcome.UNKNOWN + + +def test_gateway_uses_top_level_gripper_rpc_methods() -> None: + client = Client(True) + gateway = ControlCoordinatorGateway(client) + assert gateway.set_gripper_position("gripper", 0.4) == Outcome.ACCEPTED + client.value = 0.6 + assert gateway.get_gripper_position("gripper") == 0.6 + assert client.calls == [ + ("gripper", "set_gripper_position", {"position": 0.4}), + ("gripper", "get_gripper_position", {}), + ] diff --git a/dimos/manipulation/test_execution_models.py b/dimos/manipulation/test_execution_models.py new file mode 100644 index 0000000000..3a834e7404 --- /dev/null +++ b/dimos/manipulation/test_execution_models.py @@ -0,0 +1,27 @@ +"""Public model and snapshot projection tests.""" + +from dataclasses import FrozenInstanceError + +import pytest + +from dimos.manipulation.execution_models import LifecycleState, RuntimeContext, RuntimeSnapshot + + +def test_runtime_snapshot_is_an_immutable_public_projection() -> None: + context = RuntimeContext(state=LifecycleState.READY, diagnostic="ready") + snapshot = RuntimeSnapshot( + context.state, + context.ready_plan, + context.ready_plan_id, + context.planning_token, + context.active, + context.fault, + context.diagnostic, + context.shutdown, + context.shutdown_result, + context.revision, + ) + assert snapshot.state == LifecycleState.READY + assert snapshot.diagnostic == "ready" + with pytest.raises(FrozenInstanceError): + snapshot.diagnostic = "mutated" # type: ignore[misc] diff --git a/dimos/manipulation/test_execution_policy.py b/dimos/manipulation/test_execution_policy.py new file mode 100644 index 0000000000..cc216fcad3 --- /dev/null +++ b/dimos/manipulation/test_execution_policy.py @@ -0,0 +1,116 @@ +import pytest + +from dimos.manipulation.execution_models import ActionMethod, Outcome, TaskActivity +from dimos.manipulation.execution_policy import ( + reconcile_cancel_completion, + reconcile_execute_completion, + reconcile_reset_completion, + reconcile_status_completion, +) + + +@pytest.mark.parametrize( + ("helper", "outcome", "activity"), + [ + (reconcile_execute_completion, Outcome.ACCEPTED, TaskActivity.ACTIVE), + (reconcile_cancel_completion, Outcome.CANCELLED, TaskActivity.CANCELLED), + (reconcile_status_completion, Outcome.COMPLETED, TaskActivity.COMPLETED), + ], +) +def test_completion_policy_normalizes_public_outcomes( + helper: object, outcome: Outcome, activity: TaskActivity +) -> None: + decision = helper(outcome) # type: ignore[operator] + assert decision.activity == activity + assert not decision.fault + + +@pytest.mark.parametrize( + ( + "method", + "outcome", + "activity", + "expected_activity", + "cancel_required", + "request_cancel", + "reset_success", + ), + [ + ( + ActionMethod.STATUS, + Outcome.INACTIVE, + TaskActivity.ACTIVE, + TaskActivity.INACTIVE, + False, + False, + None, + ), + ( + ActionMethod.STATUS, + Outcome.CANCELLED, + TaskActivity.ACTIVE, + TaskActivity.CANCELLED, + False, + False, + None, + ), + ( + ActionMethod.STATUS, + Outcome.COMPLETED, + TaskActivity.ACTIVE, + TaskActivity.COMPLETED, + False, + False, + None, + ), + ( + ActionMethod.STATUS, + Outcome.RUNNING, + TaskActivity.COMPLETED, + TaskActivity.ACTIVE, + True, + True, + None, + ), + ( + ActionMethod.STATUS, + Outcome.ACCEPTED, + TaskActivity.INACTIVE, + TaskActivity.ACTIVE, + True, + True, + None, + ), + ( + ActionMethod.STATUS, + Outcome.UNKNOWN, + TaskActivity.ACTIVE, + TaskActivity.UNKNOWN, + True, + True, + False, + ), + (ActionMethod.CANCEL, Outcome.UNKNOWN, TaskActivity.ACTIVE, None, False, False, False), + (ActionMethod.RESET, Outcome.FAILED, TaskActivity.ACTIVE, None, None, False, False), + ], +) +def test_reset_completion_policy_is_pure_and_classifies_safety( + method: ActionMethod, + outcome: Outcome, + activity: TaskActivity, + expected_activity: TaskActivity | None, + cancel_required: bool | None, + request_cancel: bool, + reset_success: bool | None, +) -> None: + decision = reconcile_reset_completion( + method=method, + outcome=outcome, + activity=activity, + gate_sealed=True, + clock_failed=True, + ) + assert decision.cancel_required == cancel_required + assert decision.request_cancel is request_cancel + assert decision.activity == expected_activity + assert decision.reset_success == reset_success diff --git a/dimos/manipulation/test_execution_runtime.py b/dimos/manipulation/test_execution_runtime.py new file mode 100644 index 0000000000..d1e27b8cc1 --- /dev/null +++ b/dimos/manipulation/test_execution_runtime.py @@ -0,0 +1,2010 @@ +"""Black-box owner, reconciliation, deadline, and shutdown coverage.""" + +from dataclasses import dataclass, replace +from pathlib import Path +import threading +from typing import Any + +import pytest + +from dimos.manipulation.execution_effects import EffectDone +from dimos.manipulation.execution_models import ActionRecord, Operation, TaskRecord +from dimos.manipulation.execution_runtime import ( + ActionMethod, + ControlCoordinatorGateway, + ExecutionPlan, + ExecutionRuntime, + ExecutionTopology, + LifecycleState, + OperationHandle, + Outcome, + PreparedPlan, + ShutdownState, + TaskActivity, + TaskEntry, + _PendingAction, + prepare_generated_plan, +) +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.models import GeneratedPlan +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.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint +from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState + + +@dataclass +class FakeRPC: + values: dict[tuple[str, str], Any] + calls: list[tuple[str, str, dict[str, Any]]] + entered: threading.Event | None = None + release: threading.Event | None = None + + def task_invoke(self, task: str, method: str, kwargs: dict[str, Any]) -> Any: + self.calls.append((task, method, kwargs)) + if method == "execute" and self.entered is not None: + self.entered.set() + assert self.release is not None and self.release.wait(2) + return self.values.get((task, method)) + + def stop_rpc_client(self) -> None: + self.calls.append(("", "stop", {})) + + def set_gripper_position(self, hardware_id: str, position: float) -> Any: + self.calls.append((hardware_id, "set_gripper_position", {"position": position})) + return self.values.get((hardware_id, "set_gripper_position"), True) + + def get_gripper_position(self, hardware_id: str) -> Any: + self.calls.append((hardware_id, "get_gripper_position", {})) + return self.values.get((hardware_id, "get_gripper_position"), 0.5) + + +class BlockingGateway: + """Deterministic per-task barriers for effect/correlation assertions.""" + + def __init__( + self, + executes: dict[str, Outcome], + cancels: dict[str, Outcome] | None = None, + *, + block_execute: bool = False, + ) -> None: + self.executes = executes + self.cancels = cancels or {task: Outcome.CANCELLED for task in executes} + self.block_execute = block_execute + self.calls: list[tuple[str, ActionMethod]] = [] + self.entered = {task: threading.Event() for task in executes} + self.execute_entered = {task: threading.Event() for task in executes} + self.cancel_entered = {task: threading.Event() for task in executes} + self.status_entered = threading.Event() + self.release = {task: threading.Event() for task in executes} + self._lock = threading.Lock() + self.stop_calls = 0 + self.calls_at_stop = 0 + self.stop_entered = threading.Event() + + def _call(self, task: str, method: ActionMethod, outcome: Outcome) -> Outcome: + with self._lock: + self.calls.append((task, method)) + if method == ActionMethod.CANCEL: + self.entered[task].set() + self.cancel_entered[task].set() + assert self.release[task].wait(2) + return outcome + + def execute(self, task_name: str, request: Any) -> Outcome: + self.execute_entered[task_name].set() + if self.block_execute: + assert self.release[task_name].wait(2) + return self._call(task_name, ActionMethod.EXECUTE, self.executes[task_name]) + + def cancel(self, task_name: str) -> Outcome: + return self._call(task_name, ActionMethod.CANCEL, self.cancels[task_name]) + + def status(self, task_name: str) -> Outcome: + with self._lock: + self.calls.append((task_name, ActionMethod.STATUS)) + self.status_entered.set() + return Outcome.UNKNOWN + + def reset(self, task_name: str) -> Outcome: + return self._call(task_name, ActionMethod.RESET, Outcome.INACTIVE) + + def set_gripper_position(self, hardware_id: str, position: float) -> Outcome: + return self._call(hardware_id, ActionMethod.GRIPPER_SET, Outcome.ACCEPTED) + + def get_gripper_position(self, hardware_id: str) -> float: + with self._lock: + self.calls.append((hardware_id, ActionMethod.GRIPPER_GET)) + return 0.25 + + def stop(self) -> None: + self.stop_entered.set() + with self._lock: + self.stop_calls += 1 + self.calls_at_stop = len(self.calls) + for event in self.release.values(): + event.set() + + def count(self, task: str, method: ActionMethod) -> int: + with self._lock: + return self.calls.count((task, method)) + + +class BlockingAuxGateway(BlockingGateway): + def __init__(self) -> None: + super().__init__({"gripper": Outcome.ACCEPTED}) + self.aux_entered = threading.Event() + self.aux_release = threading.Event() + + def set_gripper_position(self, hardware_id: str, position: float) -> Outcome: + with self._lock: + self.calls.append((hardware_id, ActionMethod.GRIPPER_SET)) + self.aux_entered.set() + assert self.aux_release.wait(2) + return Outcome.ACCEPTED + + +def plan(*names: str) -> ExecutionPlan: + return ExecutionPlan( + "generated", tuple(TaskEntry(name, name, {"trajectory": name}) for name in names) + ) + + +def task_is(runtime: ExecutionRuntime, index: int, activity: TaskActivity) -> bool: + operation = runtime.snapshot().operation + return operation is not None and operation.tasks[index].activity == activity + + +def test_gateway_normalizes_task_invoke_and_real_states() -> None: + rpc = FakeRPC( + { + ("a", "execute"): True, + ("a", "cancel"): False, + ("a", "get_state"): TrajectoryState.EXECUTING, + }, + [], + ) + gateway = ControlCoordinatorGateway(rpc) + assert gateway.execute("a", {}) == Outcome.ACCEPTED + assert gateway.cancel("a") == Outcome.INACTIVE + assert gateway.status("a") == Outcome.RUNNING + + +def test_aggregate_dispatch_pre_registers_and_waits_for_each_acceptance() -> None: + rpc = FakeRPC({("a", "execute"): True, ("b", "execute"): True}, []) + runtime = ExecutionRuntime(lambda: ControlCoordinatorGateway(rpc), poll_interval=10) + try: + result = runtime.execute_explicit(plan("a", "b")) + assert result.accepted and result.value is not None + dispatch = runtime.wait_for_dispatch(result.value, timeout=1) + assert dispatch.accepted and dispatch.value is not None + assert len([call for call in rpc.calls if call[1] == "execute"]) == 2 + operation = runtime.snapshot().operation + assert operation is not None + assert len(operation.tasks) == 2 + assert runtime.snapshot().state == LifecycleState.RUNNING + assert all(task.activity == TaskActivity.ACTIVE for task in operation.tasks) + finally: + runtime.shutdown() + + +def test_ready_replacement_failure_and_consumed_identity() -> None: + runtime = ExecutionRuntime(lambda: BlockingGateway({"a": Outcome.ACCEPTED}), poll_interval=10) + try: + first_token = runtime.start_planning() + assert first_token is not None + assert runtime.complete_planning(first_token, plan("a")).accepted + replacement_token = runtime.start_planning() + assert replacement_token is not None + assert runtime.snapshot().ready_plan is None + assert not runtime.complete_planning( + replacement_token, ExecutionPlan("generated", ()) + ).accepted + assert runtime.snapshot().state == LifecycleState.IDLE + assert runtime.snapshot().ready_plan is None + token = runtime.start_planning() + assert token is not None + assert runtime.complete_planning(token, plan("a")).accepted + executed = runtime.execute_ready() + assert executed.accepted and executed.value is not None + assert runtime.snapshot().ready_plan is None + consumed = runtime.execute_ready() + assert not consumed.accepted + assert consumed.value is None + finally: + runtime.cancel() + runtime.shutdown() + + +def test_cancel_during_execute_handles_late_acceptance() -> None: + entered, release = (threading.Event(), threading.Event()) + rpc = FakeRPC({("a", "execute"): True, ("a", "cancel"): True}, [], entered, release) + runtime = ExecutionRuntime(lambda: ControlCoordinatorGateway(rpc), poll_interval=10) + try: + started = runtime.execute_explicit(plan("a")) + assert started.accepted and started.value is not None + assert entered.wait(1) + assert runtime.cancel().accepted + release.set() + terminal = runtime.wait_for_terminal(started.value, timeout=1) + assert terminal.accepted and terminal.value is not None + assert any(call[1] == "cancel" for call in rpc.calls) + finally: + release.set() + runtime.shutdown() + + +def test_planning_tokens_and_explicit_execution_are_atomic() -> None: + rpc = FakeRPC({}, []) + runtime = ExecutionRuntime( + lambda: ControlCoordinatorGateway(rpc), + topology=ExecutionTopology((("a", "robot", "a"),)), + poll_interval=10, + ) + try: + token = runtime.start_planning() + assert token is not None + assert not runtime.complete_planning("stale", plan("a")).accepted + planned = runtime.complete_planning(token, plan("a")) + assert planned.accepted and planned.value is not None + assert runtime.snapshot().state == LifecycleState.READY + handle = runtime.execute_ready().value + assert handle is not None + assert handle.plan_id == planned.value + assert runtime.snapshot().ready_plan is None + finally: + runtime.shutdown() + + +def test_deadline_preserves_unresolved_action_and_stop_rpc_is_deferred() -> None: + clock = AdjustableClock() + entered, release = (threading.Event(), threading.Event()) + rpc = FakeRPC({("a", "execute"): True}, [], entered, release) + runtime = ExecutionRuntime( + lambda: ControlCoordinatorGateway(rpc), + monotonic_clock=clock, + action_timeout=1, + poll_interval=100, + ) + runtime.execute_explicit(plan("a")) + assert entered.wait(1) + clock.set(2) + assert not runtime.poll().accepted + assert runtime.snapshot().state == LifecycleState.FAULT + assert runtime.snapshot().operation is not None + release.set() + runtime.shutdown() + threading.Event().wait(0.1) + assert ("", "stop", {}) in rpc.calls + + +def test_waiters_use_only_the_requested_result_map() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED, "b": Outcome.REJECTED}) + clock = AdjustableClock() + runtime = ExecutionRuntime( + lambda: gateway, monotonic_clock=clock, action_timeout=100, poll_interval=10 + ) + try: + handle = runtime.execute_explicit(plan("a", "b")).value + assert handle is not None + assert gateway.cancel_entered["a"].wait(1) + dispatch = runtime.wait_for_dispatch(handle, timeout=0.1) + assert ( + dispatch.accepted + and dispatch.value is not None + and (dispatch.value.outcome == Outcome.REJECTED) + ) + terminal = assert_public_wait_times_out( + lambda: runtime.wait_for_terminal(handle, timeout=1) + ) + assert not terminal.accepted and terminal.diagnostic == "timeout" + gateway.release["a"].set() + terminal = runtime.wait_for_terminal(handle, timeout=1) + assert ( + terminal.accepted + and terminal.value is not None + and (terminal.value.outcome == Outcome.REJECTED) + ) + finally: + gateway.release["a"].set() + runtime.shutdown() + + +def test_idle_shutdown_is_idempotent_and_rejects_new_work() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + first = runtime.shutdown(timeout=1) + assert first.accepted and first.value is not None and first.value.success + assert runtime.snapshot().shutdown == ShutdownState.CLOSED + assert not getattr(runtime.start_planning(), "accepted", False) + assert not runtime.execute_explicit(plan("a")).accepted + second = runtime.shutdown(timeout=1) + assert second.accepted and second.value == first.value + assert gateway.stop_calls == 1 + + +def test_gateway_stop_failure_never_reports_closed_success() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + + def fail_stop() -> None: + gateway.stop_calls += 1 + raise RuntimeError("close exploded") + + gateway.stop = fail_stop # type: ignore[method-assign] + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + result = runtime.shutdown(timeout=1) + assert not result.accepted + assert result.value is not None and (not result.value.success) + assert "gateway close failed: close exploded" in result.value.diagnostic + assert result.snapshot is not None + assert result.snapshot.shutdown == ShutdownState.CLOSING + assert result.snapshot.shutdown_result == result.value + assert gateway.stop_calls == 1 + assert runtime.shutdown(timeout=1).accepted is False + assert gateway.stop_calls == 1 + + +def test_invalid_later_clock_sample_does_not_fault_or_kill_owner() -> None: + samples = iter((0.0, 0.0, 0.0, float("nan"))) + invalid_seen = threading.Event() + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + + def clock() -> float: + value = next(samples, 0.0) + if value != value: + invalid_seen.set() + return value + + runtime = ExecutionRuntime(lambda: gateway, monotonic_clock=clock) + try: + started = runtime.execute_explicit(plan("a")) + assert started.accepted and started.value is not None + assert runtime.snapshot().state != LifecycleState.FAULT + assert runtime.snapshot().state == LifecycleState.CANCELLING + assert runtime.snapshot().shutdown == ShutdownState.CLOSING + polled = runtime.poll() + assert not polled.accepted + assert invalid_seen.wait(1) + assert gateway.cancel_entered["a"].wait(1) + assert gateway.count("a", ActionMethod.CANCEL) == 1 + gateway.release["a"].set() + assert gateway.stop_entered.wait(1) + assert polled.snapshot is not None + assert polled.snapshot.state != LifecycleState.FAULT + finally: + gateway.release["a"].set() + runtime.shutdown() + + +def test_failed_shutdown_drain_rejects_public_commands_after_owner_stops() -> None: + gateway = BlockingGateway({}) + + def fail_stop() -> None: + gateway.stop_calls += 1 + raise RuntimeError("close exploded") + + gateway.stop = fail_stop # type: ignore[method-assign] + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + result = runtime.shutdown(timeout=1) + assert not result.accepted + assert result.value is not None and not result.value.success + assert runtime.snapshot().shutdown == ShutdownState.CLOSING + runtime._owner.join(1) + assert not runtime._owner.is_alive() + + assert not runtime.cancel().accepted + assert not runtime.cancel_if_current(OperationHandle("plan", "operation", "attempt")).accepted + reset = runtime.reset() + assert not reset.accepted + assert reset.diagnostic == "runtime is closing" + assert runtime.shutdown(timeout=1).value == result.value + + +def test_invalid_shutdown_clock_handoff_does_not_deadlock_caller() -> None: + samples = iter((0.0, float("nan"))) + gateway = BlockingGateway({}) + runtime = ExecutionRuntime(lambda: gateway, monotonic_clock=lambda: next(samples)) + result = runtime.shutdown(timeout=1) + assert not result.accepted + assert result.value is not None + assert result.value.diagnostic == "invalid monotonic clock" + assert runtime.snapshot().shutdown == ShutdownState.CLOSING + + +def test_shutdown_interrupts_early_action_deadline_thread() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, action_timeout=10.0, poll_interval=10) + try: + started = runtime.execute_explicit(plan("a")) + assert started.accepted and started.value is not None + assert runtime.wait_for_dispatch(started.value, timeout=1).accepted + gateway.release["a"].set() + assert runtime.shutdown(timeout=1).accepted + finally: + gateway.release["a"].set() + if runtime.snapshot().shutdown != ShutdownState.CLOSED: + runtime.shutdown(timeout=1) + + +def test_cancel_if_current_cannot_cancel_replacement_operation() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + gateway.status = lambda task_name: Outcome.COMPLETED # type: ignore[method-assign] + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + original = runtime.execute_explicit(plan("a")) + assert original.accepted and original.value is not None + assert runtime.wait_for_dispatch(original.value, timeout=1).accepted + runtime.poll() + assert runtime.wait_for_terminal(original.value, timeout=1).accepted + replacement = runtime.execute_explicit(plan("a")) + assert replacement.accepted and replacement.value is not None + assert runtime.wait_for_dispatch(replacement.value, timeout=1).accepted + stale_cancel = runtime.cancel_if_current(original.value) + assert not stale_cancel.accepted + assert stale_cancel.diagnostic == "operation is no longer current" + active = runtime.snapshot().operation + assert active is not None + assert active.handle == replacement.value + assert gateway.count("a", ActionMethod.CANCEL) == 0 + finally: + gateway.release["a"].set() + if runtime.snapshot().shutdown != ShutdownState.CLOSED: + runtime.shutdown(timeout=1) + + +def test_shutdown_late_execute_acceptance_is_compensated_before_close() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}, block_execute=True) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + runtime.execute_explicit(plan("a")) + assert gateway.execute_entered["a"].wait(1) + result_holder: list[Any] = [] + thread = threading.Thread(target=lambda: result_holder.append(runtime.shutdown(timeout=1))) + thread.start() + gateway.release["a"].set() + thread.join(2) + assert ( + len(result_holder) == 1 and result_holder[0].accepted and result_holder[0].value.success + ) + assert gateway.count("a", ActionMethod.CANCEL) == 1 + assert runtime.snapshot().shutdown == ShutdownState.CLOSED + finally: + gateway.release["a"].set() + if runtime.snapshot().shutdown != ShutdownState.CLOSED: + runtime.shutdown() + + +def _prep_robot( + name: str, + joints: tuple[str, ...], + group: str, + task: str | None, + mapping: dict[str, str] | None = None, +) -> RobotModelConfig: + return RobotModelConfig( + name=name, + model_path=Path(f"/{name}.urdf"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion([0.0, 0.0, 0.0, 1.0])), + joint_names=list(joints), + planning_groups=[PlanningGroupDefinition(name=group, joint_names=joints, base_link="base")], + coordinator_task_name=task, + joint_name_mapping=mapping or {}, + ) + + +def _generated(names: list[str], groups: tuple[str, ...], width: int) -> GeneratedPlan: + return GeneratedPlan( + group_ids=groups, + trajectory=JointTrajectory( + joint_names=names, + points=[ + TrajectoryPoint( + time_from_start=0.0, positions=[0.0] * width, velocities=[0.1] * width + ), + TrajectoryPoint( + time_from_start=1.5, positions=[1.0] * width, velocities=[0.2] * width + ), + ], + ), + ) + + +def test_prepare_generated_plan_uses_ordered_groups_and_ignores_extraneous_robot_columns() -> None: + left = _prep_robot("left", ("l0", "l1"), "arm", "left_task", {"left_l0": "l0"}) + right = _prep_robot("right", ("r0",), "arm", "right_task") + extra = _prep_robot("extra", ("x0",), "arm", "extra_task") + topology = ExecutionTopology.from_robot_configs((left, right, extra)) + generated = _generated( + ["extra/x0", "right/r0", "left/l1", "left/l0"], ("left/arm", "right/arm"), 4 + ) + prepared = prepare_generated_plan(generated, topology) + assert prepared.generated_plan is generated + assert [entry.robot_name for entry in prepared.entries] == ["left", "right"] + assert [entry.task_name for entry in prepared.entries] == ["left_task", "right_task"] + assert prepared.entries[0].request["trajectory"].joint_names == ["l1", "left_l0"] + assert prepared.entries[0].request["trajectory"].points[1].time_from_start == 1.5 + assert prepared.entries[1].request["trajectory"].joint_names == ["r0"] + + +def test_prepare_rejects_missing_routes_tasks_malformed_and_missing_or_duplicate_joints() -> None: + left = _prep_robot("left", ("l0", "l1"), "arm", "left_task") + right = _prep_robot("right", ("r0",), "arm", "right_task") + topology = ExecutionTopology.from_robot_configs((left, right)) + with pytest.raises(ValueError, match="route set"): + ExecutionTopology.from_robot_configs((left, right), {"left/arm": ("left", "left_task")}) + with pytest.raises(ValueError, match="malformed"): + prepare_generated_plan(_generated(["left/l0", "bad"], ("left/arm",), 2), topology) + with pytest.raises(ValueError, match="missing robot joints"): + prepare_generated_plan(_generated(["left/l0"], ("left/arm",), 1), topology) + with pytest.raises(ValueError, match="duplicate global"): + prepare_generated_plan(_generated(["left/l0", "left/l0"], ("left/arm",), 2), topology) + + +def test_prepared_plan_is_real_type_and_partial_generated_entries_are_not_executable() -> None: + robot = _prep_robot("arm", ("j0",), "manipulator", "arm_task", {"coord_j0": "j0"}) + topology = ExecutionTopology.from_robot_configs((robot,)) + generated = _generated(["arm/j0"], ("arm/manipulator",), 1) + prepared = prepare_generated_plan(generated, topology) + assert isinstance(prepared, PreparedPlan) + assert id(prepared) != id(generated) + partial = PreparedPlan(generated, (), topology) + runtime = ExecutionRuntime(lambda: BlockingGateway({"arm": Outcome.ACCEPTED}), poll_interval=10) + try: + assert not runtime.execute_explicit(partial).accepted + finally: + runtime.shutdown() + + +def test_real_generated_plan_cannot_use_legacy_synthetic_execution_plan() -> None: + robot = _prep_robot("arm", ("j0",), "manipulator", "arm_task") + topology = ExecutionTopology.from_robot_configs((robot,)) + generated = _generated(["arm/j0"], ("arm/manipulator",), 1) + prepared = prepare_generated_plan(generated, topology) + runtime = ExecutionRuntime(lambda: BlockingGateway({"arm": Outcome.ACCEPTED}), poll_interval=10) + try: + legacy = ExecutionPlan(generated, prepared.entries) + assert not runtime.execute_explicit(legacy).accepted + assert runtime.execute_explicit(prepared).accepted + finally: + runtime.shutdown() + + +def test_prepared_validation_selects_only_generated_groups_from_larger_topology() -> None: + left = _prep_robot("left", ("l0",), "arm", "left_task") + right = _prep_robot("right", ("r0",), "arm", "right_task") + extra = _prep_robot("extra", ("x0",), "arm", "extra_task") + topology = ExecutionTopology.from_robot_configs((left, right, extra)) + prepared = prepare_generated_plan( + _generated(["left/l0", "right/r0"], ("right/arm", "left/arm"), 2), topology + ) + runtime = ExecutionRuntime( + lambda: BlockingGateway({"left": Outcome.ACCEPTED, "right": Outcome.ACCEPTED}), + poll_interval=10, + ) + try: + result = runtime.execute_explicit(prepared) + assert result.accepted + assert [entry.planning_group for entry in prepared.entries] == ["right/arm", "left/arm"] + finally: + runtime.shutdown() + + +def test_prepare_rejects_additional_selected_robot_joint_column() -> None: + robot = _prep_robot("arm", ("j0", "j1"), "manipulator", "arm_task") + topology = ExecutionTopology.from_robot_configs((robot,)) + with pytest.raises(ValueError, match="additional joint column"): + prepare_generated_plan( + _generated(["arm/j0", "arm/j1", "arm/extra"], ("arm/manipulator",), 3), topology + ) + + +def test_topology_rejects_duplicate_routes_bad_mapping_targets_and_unknown_group_joints() -> None: + robot = _prep_robot("arm", ("j0",), "manipulator", "arm_task") + with pytest.raises(ValueError, match="duplicate group route"): + ExecutionTopology.from_robot_configs( + (robot,), + [("arm/manipulator", "arm", "arm_task"), ("arm/manipulator", "arm", "arm_task")], + ) + with pytest.raises(ValueError, match="mapped local joint"): + ExecutionTopology.from_robot_configs( + (_prep_robot("arm", ("j0",), "manipulator", "arm_task", {"coord": "missing"}),) + ) + bad_group = _prep_robot("arm", ("j0",), "manipulator", "arm_task") + bad_group.planning_groups[0] = PlanningGroupDefinition( + name="manipulator", joint_names=("missing",), base_link="base" + ) + with pytest.raises(ValueError, match="planning group joints"): + ExecutionTopology.from_robot_configs((bad_group,)) + + +def test_planning_only_config_has_no_route_and_generated_execution_rejects_it() -> None: + planning_only = _prep_robot("planner", ("j0",), "manipulator", None) + topology = ExecutionTopology.from_robot_configs((planning_only,)) + assert topology.routes == () + with pytest.raises(ValueError, match="no coordinator task"): + prepare_generated_plan(_generated(["planner/j0"], ("planner/manipulator",), 1), topology) + + +def test_clear_ready_plan_is_a_dedicated_ready_to_idle_transition() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + token = runtime.start_planning() + assert token is not None + assert runtime.complete_planning(token, plan("a")).accepted + assert runtime.snapshot().state == LifecycleState.READY + cleared = runtime.clear_ready_plan() + assert cleared.accepted + snapshot = runtime.snapshot() + assert snapshot.state == LifecycleState.IDLE + assert snapshot.ready_plan is None and snapshot.ready_plan_id is None + assert not runtime.clear_ready_plan().accepted + finally: + runtime.shutdown() + + +def test_runtime_diagnostic_persists_idle_planning_and_execution_errors_then_clears_on_planning() -> ( + None +): + runtime = ExecutionRuntime(lambda: BlockingGateway({"a": Outcome.ACCEPTED}), poll_interval=10) + try: + token = runtime.start_planning() + assert token is not None + invalid = runtime.complete_planning(token, ExecutionPlan("generated", ())) + assert not invalid.accepted + assert runtime.snapshot().state == LifecycleState.IDLE + assert runtime.snapshot().diagnostic == "invalid prepared plan" + assert runtime.start_planning() is not None + assert runtime.snapshot().diagnostic is None + finally: + runtime.shutdown() + + +def test_rejected_execute_retains_contextual_diagnostic_in_snapshot_and_result() -> None: + gateway = BlockingGateway({"a": Outcome.REJECTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + handle = runtime.execute_explicit(plan("a")).value + assert handle is not None + result = runtime.wait_for_terminal(handle, timeout=1) + assert result.accepted and result.value is not None + assert result.value.outcome == Outcome.REJECTED + assert "task=a" in result.value.diagnostic + assert "action=" in result.value.diagnostic + assert "outcome=rejected" in result.value.diagnostic + snapshot = runtime.snapshot() + assert snapshot.state == LifecycleState.IDLE + assert snapshot.diagnostic == result.value.diagnostic + finally: + runtime.shutdown() + + +def test_status_terminal_failure_retains_contextual_diagnostic_in_terminal_result() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED, "b": Outcome.ACCEPTED}) + gateway.status = lambda task_name: Outcome.INACTIVE if task_name == "a" else Outcome.RUNNING # type: ignore[assignment,method-assign] + gateway.release["a"].set() + gateway.release["b"].set() + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + handle = runtime.execute_explicit(plan("a", "b")).value + assert handle is not None + assert runtime.wait_for_dispatch(handle, timeout=1).accepted + runtime.poll() + result = runtime.wait_for_terminal(handle, timeout=1) + assert result.accepted and result.value is not None + assert result.value.outcome == Outcome.FAILED + assert "task=a" in result.value.diagnostic + assert "outcome=inactive" in result.value.diagnostic + assert runtime.snapshot().state == LifecycleState.IDLE + assert runtime.snapshot().diagnostic == result.value.diagnostic + finally: + for event in gateway.release.values(): + event.set() + runtime.shutdown() + + +def test_auxiliary_gripper_calls_use_private_gateway_and_normalize_results() -> None: + gateway = BlockingGateway({"gripper": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + try: + set_result = runtime.set_gripper_position("gripper", 0.75) + get_result = runtime.get_gripper_position("gripper") + assert set_result.accepted and set_result.value == Outcome.ACCEPTED + assert get_result.accepted and get_result.value == 0.25 + assert gateway.count("gripper", ActionMethod.GRIPPER_SET) == 1 + assert gateway.count("gripper", ActionMethod.GRIPPER_GET) == 1 + assert not hasattr(runtime, "gateway") + finally: + runtime.shutdown() + + +def test_control_gateway_uses_top_level_gripper_rpc_methods() -> None: + rpc = FakeRPC( + {("gripper", "set_gripper_position"): True, ("gripper", "get_gripper_position"): 0.6}, [] + ) + gateway = ControlCoordinatorGateway(rpc) + assert gateway.set_gripper_position("gripper", 0.4) == Outcome.ACCEPTED + assert gateway.get_gripper_position("gripper") == 0.6 + assert rpc.calls == [ + ("gripper", "set_gripper_position", {"position": 0.4}), + ("gripper", "get_gripper_position", {}), + ] + + +def test_auxiliary_calls_reject_after_closing_and_inflight_auxiliary_drains_before_close() -> None: + gateway = BlockingAuxGateway() + runtime = ExecutionRuntime(lambda: gateway, poll_interval=10) + results: list[Any] = [] + caller = threading.Thread( + target=lambda: results.append(runtime.set_gripper_position("gripper", 0.5)) + ) + caller.start() + assert gateway.aux_entered.wait(1) + shutdown_results: list[Any] = [] + closer = threading.Thread(target=lambda: shutdown_results.append(runtime.shutdown(timeout=1))) + closer.start() + assert closer.is_alive() + gateway.aux_release.set() + caller.join(2) + closer.join(2) + assert results and results[0].accepted + assert shutdown_results and shutdown_results[0].accepted + assert gateway.calls_at_stop == len(gateway.calls) + assert gateway.stop_calls == 1 + assert not runtime.set_gripper_position("gripper", 0.1).accepted + assert not runtime.get_gripper_position("gripper").accepted + + +def test_auxiliary_deadline_waits_for_late_callback_before_shutdown_stop() -> None: + clock = AdjustableClock() + gateway = BlockingAuxGateway() + runtime = ExecutionRuntime( + lambda: gateway, + monotonic_clock=clock, + action_timeout=1, + poll_interval=100, + ) + gripper_results: list[Any] = [] + shutdown_results: list[Any] = [] + gripper: threading.Thread | None = None + closer: threading.Thread | None = None + shutdown_started = False + try: + gripper = threading.Thread( + target=lambda: gripper_results.append(runtime.set_gripper_position("gripper", 0.5)) + ) + gripper.start() + assert gateway.aux_entered.wait(1) + clock.set(2) + assert not runtime.poll().accepted + gripper.join(1) + assert gripper_results and not gripper_results[0].accepted + + closer = threading.Thread(target=lambda: shutdown_results.append(runtime.shutdown(1))) + shutdown_started = True + closer.start() + clock.set(3) + closer.join(1) + assert shutdown_results and not shutdown_results[0].accepted + assert not gateway.stop_entered.is_set() + + gateway.aux_release.set() + assert gateway.stop_entered.wait(1) + assert runtime.snapshot().shutdown == ShutdownState.CLOSING + shutdown_result = runtime.snapshot().shutdown_result + assert shutdown_result is not None and not shutdown_result.success + finally: + gateway.aux_release.set() + if gripper is not None: + gripper.join(2) + if closer is not None: + closer.join(2) + if not shutdown_started: + runtime.shutdown(timeout=1) + + +def test_gateway_factory_is_the_internal_runtime_injection_seam() -> None: + gateway = BlockingGateway({"gripper": Outcome.ACCEPTED}) + factories: list[bool] = [] + + def factory() -> BlockingGateway: + factories.append(True) + return gateway + + runtime = ExecutionRuntime(factory, poll_interval=10) + try: + assert factories == [True] + assert runtime.get_gripper_position("gripper").accepted + assert not hasattr(runtime, "client") + finally: + runtime.shutdown() + + +class AdjustableClock: + def __init__(self) -> None: + self.value = 0.0 + self.lock = threading.Lock() + self._barrier_thread: int | None = None + self._barrier_pending = False + self.sample_entered = threading.Event() + self.sample_release = threading.Event() + + def __call__(self) -> float: + with self.lock: + value = self.value + barrier = self._barrier_pending and self._barrier_thread == threading.get_ident() + self._barrier_pending = False if barrier else self._barrier_pending + if barrier: + self.sample_entered.set() + assert self.sample_release.wait(1) + return value + + def arm_current_thread_sample_barrier(self) -> None: + with self.lock: + self._barrier_thread = threading.get_ident() + self._barrier_pending = True + self.sample_entered.clear() + self.sample_release.clear() + + def set(self, value: float) -> None: + with self.lock: + self.value = value + + +def assert_public_wait_times_out(wait_call: Any) -> Any: + result: list[Any] = [] + ready = threading.Event() + + def observe() -> None: + ready.set() + result.append(wait_call()) + + waiter = threading.Thread(target=observe) + waiter.start() + assert ready.wait(1) + waiter.join(2) + assert not waiter.is_alive() and len(result) == 1 + return result[0] + + +class ScriptedGateway: + def __init__(self, statuses: dict[str, Outcome] | None = None) -> None: + self.statuses = statuses or {} + self.calls: list[tuple[str, ActionMethod]] = [] + self.lock = threading.Lock() + self.execute_entered = threading.Event() + self.execute_release = threading.Event() + self.block_execute = False + self.cancel_entered = threading.Event() + self.cancel_returned = threading.Event() + self.cancel_entered_by_task: dict[str, threading.Event] = {} + self.cancel_returned_by_task: dict[str, threading.Event] = {} + self.cancel_release = threading.Event() + self.block_cancel = False + self.status_entered = threading.Event() + self.status_release = threading.Event() + self.block_status = False + self.cancel_outcome = Outcome.CANCELLED + self.reset_outcome = Outcome.INACTIVE + self.reset_entered = threading.Event() + self.reset_release = threading.Event() + self.block_reset = False + self.stop_calls = 0 + + def _record(self, task: str, method: ActionMethod) -> None: + with self.lock: + self.calls.append((task, method)) + + def execute(self, task_name: str, request: Any) -> Outcome: + self._record(task_name, ActionMethod.EXECUTE) + self.execute_entered.set() + if self.block_execute: + assert self.execute_release.wait(1) + return Outcome.ACCEPTED + + def cancel(self, task_name: str) -> Outcome: + self._record(task_name, ActionMethod.CANCEL) + self.cancel_entered.set() + self.cancel_entered_by_task.setdefault(task_name, threading.Event()).set() + if self.block_cancel: + assert self.cancel_release.wait(1) + self.cancel_returned.set() + self.cancel_returned_by_task.setdefault(task_name, threading.Event()).set() + return self.cancel_outcome + + def status(self, task_name: str) -> Outcome: + self._record(task_name, ActionMethod.STATUS) + self.status_entered.set() + if self.block_status: + assert self.status_release.wait(1) + return self.statuses.get(task_name, Outcome.RUNNING) + + def reset(self, task_name: str) -> Outcome: + self._record(task_name, ActionMethod.RESET) + self.reset_entered.set() + if self.block_reset: + assert self.reset_release.wait(1) + return self.reset_outcome + + def set_gripper_position(self, hardware_id: str, position: float) -> Outcome: + return Outcome.ACCEPTED + + def get_gripper_position(self, hardware_id: str) -> float: + return 0.0 + + def stop(self) -> None: + self.stop_calls += 1 + self.execute_release.set() + self.cancel_release.set() + self.status_release.set() + + def count(self, task: str, method: ActionMethod) -> int: + with self.lock: + return self.calls.count((task, method)) + + +class SequencedStatusGateway(ScriptedGateway): + def __init__(self, statuses: list[Outcome]) -> None: + super().__init__({"a": Outcome.ACCEPTED}) + self.status_sequence = statuses + + def status(self, task_name: str) -> Outcome: + self._record(task_name, ActionMethod.STATUS) + self.status_entered.set() + return self.status_sequence.pop(0) + + +def scripted_plan(*names: str) -> ExecutionPlan: + return ExecutionPlan( + "generated", tuple(TaskEntry(name, name, {"trajectory": name}) for name in names) + ) + + +def start_scripted(runtime: ExecutionRuntime, gateway: ScriptedGateway, task: str = "a") -> Any: + started = runtime.execute_explicit(scripted_plan(task)) + assert started.accepted and started.value is not None + assert gateway.execute_entered.wait(1) + dispatched = runtime.wait_for_dispatch(started.value, timeout=1) + assert dispatched.accepted and dispatched.value is not None + assert runtime.snapshot().state == LifecycleState.RUNNING + return started.value + + +def close_scripted(runtime: ExecutionRuntime, gateway: ScriptedGateway) -> None: + gateway.stop() + if runtime.snapshot().shutdown != ShutdownState.CLOSED: + runtime.shutdown(timeout=1) + + +@pytest.mark.parametrize("active_status", [Outcome.RUNNING, Outcome.ACCEPTED]) +def test_reset_reconciles_active_status_then_inactive_and_succeeds( + active_status: Outcome, +) -> None: + clock = AdjustableClock() + gateway = SequencedStatusGateway([active_status, Outcome.INACTIVE]) + gateway.cancel_outcome = Outcome.UNKNOWN + runtime = ExecutionRuntime( + lambda: gateway, monotonic_clock=clock, physical_operation_timeout=60, poll_interval=100 + ) + try: + start_scripted(runtime, gateway) + clock.set(60) + runtime.poll() + assert gateway.cancel_returned.wait(1) + runtime.poll() + assert runtime.snapshot().state == LifecycleState.FAULT + + gateway.cancel_outcome = Outcome.CANCELLED + gateway.status_entered.clear() + gateway.cancel_returned.clear() + gateway.reset_entered.clear() + reset = runtime.reset() + assert reset.accepted and reset.value is not None + assert gateway.status_entered.wait(1) + assert gateway.count("a", ActionMethod.STATUS) == 1 + gateway.status_entered.clear() + assert gateway.cancel_returned.wait(1) + assert gateway.count("a", ActionMethod.CANCEL) == 2 + assert gateway.status_entered.wait(1) + assert gateway.count("a", ActionMethod.STATUS) == 2 + assert gateway.reset_entered.wait(1) + assert gateway.count("a", ActionMethod.RESET) == 1 + result = runtime.wait_for_reset(reset.value, timeout=1) + assert result.accepted and result.value is not None and result.value.success + assert result.snapshot is not None and result.snapshot.state == LifecycleState.IDLE + assert [method for _, method in gateway.calls] == [ + ActionMethod.EXECUTE, + ActionMethod.CANCEL, + ActionMethod.STATUS, + ActionMethod.CANCEL, + ActionMethod.STATUS, + ActionMethod.RESET, + ] + finally: + gateway.cancel_release.set() + gateway.reset_release.set() + close_scripted(runtime, gateway) + + +def test_queued_poll_uses_injected_clock_and_cancels_due_operation_once() -> None: + clock = AdjustableClock() + gateway = ScriptedGateway() + runtime = ExecutionRuntime( + lambda: gateway, + monotonic_clock=clock, + physical_operation_timeout=60, + action_timeout=1000, + poll_interval=100, + ) + try: + start_scripted(runtime, gateway) + clock.set(60) + poll_result: list[Any] = [] + queued = threading.Thread(target=lambda: poll_result.append(runtime.poll())) + queued.start() + queued.join(1) + assert not queued.is_alive() + assert len(poll_result) == 1 + assert poll_result[0].snapshot is not None + assert gateway.count("a", ActionMethod.STATUS) == 0 + assert gateway.count("a", ActionMethod.CANCEL) == 1 + finally: + runtime.shutdown(timeout=0) + + +def test_stale_a_deadline_at_sixty_cannot_cancel_b_started_at_thirty() -> None: + clock = AdjustableClock() + gateway = ScriptedGateway({"a": Outcome.COMPLETED, "b": Outcome.RUNNING}) + runtime = ExecutionRuntime( + lambda: gateway, monotonic_clock=clock, physical_operation_timeout=60, poll_interval=100 + ) + try: + a = start_scripted(runtime, gateway) + runtime.poll() + assert runtime.wait_for_terminal(a, timeout=1).accepted + clock.set(30) + b = start_scripted(runtime, gateway, "b") + gateway.status_entered.clear() + clock.set(61) + runtime.poll() + assert gateway.status_entered.wait(1) + snapshot = runtime.snapshot() + assert snapshot.operation is not None + assert snapshot.operation.handle == b + assert gateway.count("b", ActionMethod.STATUS) == 1 + assert gateway.count("b", ActionMethod.CANCEL) == 0 + finally: + close_scripted(runtime, gateway) + + +def test_cancel_clearance_survives_old_deadline_with_replacement_running() -> None: + clock = AdjustableClock() + gateway = ScriptedGateway() + runtime = ExecutionRuntime( + lambda: gateway, monotonic_clock=clock, physical_operation_timeout=60, poll_interval=100 + ) + try: + a = start_scripted(runtime, gateway) + assert runtime.cancel().accepted + assert gateway.cancel_entered.wait(1) + assert gateway.cancel_returned.wait(1) + runtime.poll() + assert runtime.wait_for_terminal(a, timeout=1).accepted + clock.set(30) + start_scripted(runtime, gateway, "b") + clock.set(61) + runtime.poll() + assert gateway.count("a", ActionMethod.CANCEL) == 1 + assert gateway.count("b", ActionMethod.CANCEL) == 0 + finally: + close_scripted(runtime, gateway) + + +def test_reset_rejects_while_fault_cleanup_rpc_is_unresolved() -> None: + clock = AdjustableClock() + gateway = ScriptedGateway({"a": Outcome.FAILED}) + gateway.block_cancel = True + runtime = ExecutionRuntime( + lambda: gateway, monotonic_clock=clock, physical_operation_timeout=60, poll_interval=100 + ) + try: + start_scripted(runtime, gateway) + runtime.poll() + assert gateway.cancel_entered.wait(1) + snapshot = runtime.snapshot() + assert snapshot.operation is not None + wait = threading.Event() + for _ in range(100): + if runtime.snapshot().state == LifecycleState.FAULT: + break + wait.wait(0.01) + assert runtime.snapshot().state == LifecycleState.FAULT + reset = runtime.reset() + assert not reset.accepted + assert reset.value is None + assert reset.diagnostic == "reset blocked by unresolved coordinator RPC" + assert runtime.snapshot().fault is not None + finally: + gateway.cancel_release.set() + runtime.shutdown(timeout=0) + + +def test_shutdown_clearance_blocks_cancel_and_returns_truthful_deadline_result() -> None: + clock = AdjustableClock() + gateway = ScriptedGateway() + gateway.block_cancel = True + runtime = ExecutionRuntime( + lambda: gateway, + monotonic_clock=clock, + physical_operation_timeout=60, + action_timeout=100, + poll_interval=100, + ) + result: list[Any] = [] + try: + start_scripted(runtime, gateway) + thread = threading.Thread(target=lambda: result.append(runtime.shutdown(timeout=100))) + thread.start() + assert gateway.cancel_entered.wait(1) + assert runtime.snapshot().shutdown == ShutdownState.CLOSING + clock.set(61) + gateway.cancel_release.set() + thread.join(1) + assert len(result) == 1 and result[0].value is not None + assert result[0].value.success + assert gateway.count("a", ActionMethod.CANCEL) == 1 + finally: + gateway.cancel_release.set() + close_scripted(runtime, gateway) + + +def test_reset_completes_before_terminal_waiter_observes_retained_diagnostic() -> None: + clock = AdjustableClock() + gateway = ScriptedGateway() + gateway.cancel_outcome = Outcome.UNKNOWN + gateway.statuses["a"] = Outcome.INACTIVE + runtime = ExecutionRuntime( + lambda: gateway, monotonic_clock=clock, physical_operation_timeout=60, poll_interval=100 + ) + try: + handle = start_scripted(runtime, gateway) + clock.set(60) + runtime.poll() + assert gateway.cancel_returned.wait(1) + runtime.poll() + assert runtime.snapshot().state == LifecycleState.FAULT + gateway.cancel_outcome = Outcome.CANCELLED + reset = runtime.reset() + assert reset.accepted and reset.value is not None + reset_result = runtime.wait_for_reset(reset.value, timeout=1) + assert reset_result.accepted and reset_result.value is not None + assert reset_result.value.success + waited = runtime.wait_for_terminal(handle, timeout=float("inf")) + assert not waited.accepted + assert waited.diagnostic == "cancellation is uncertain" + finally: + close_scripted(runtime, gateway) + + +def test_safe_physical_expiry_returns_failed_terminal_result_and_stable_diagnostic() -> None: + clock = AdjustableClock() + gateway = ScriptedGateway() + runtime = ExecutionRuntime( + lambda: gateway, monotonic_clock=clock, physical_operation_timeout=60, poll_interval=100 + ) + try: + handle = start_scripted(runtime, gateway) + clock.set(60) + runtime.poll() + assert gateway.cancel_returned.wait(1) + result = runtime.wait_for_terminal(handle, timeout=1) + assert result.accepted and result.value is not None + assert result.value.outcome == Outcome.FAILED + assert result.value.diagnostic == "physical operation deadline exceeded" + assert runtime.wait_for_terminal(handle, timeout=1).value == result.value + finally: + close_scripted(runtime, gateway) + + +def test_final_sequential_execute_acceptance_arms_physical_deadline() -> None: + clock = AdjustableClock() + gateway = BlockingGateway({"a": Outcome.ACCEPTED, "b": Outcome.ACCEPTED}, block_execute=True) + runtime = ExecutionRuntime( + lambda: gateway, + monotonic_clock=clock, + physical_operation_timeout=60, + action_timeout=1000, + poll_interval=100, + ) + try: + started = runtime.execute_explicit(plan("a", "b")) + assert started.accepted and started.value is not None + assert gateway.execute_entered["a"].wait(1) + clock.set(59) + runtime.poll() + assert gateway.count("a", ActionMethod.CANCEL) == 0 + gateway.release["a"].set() + assert gateway.execute_entered["b"].wait(1) + clock.set(60) + runtime.poll() + assert gateway.count("a", ActionMethod.CANCEL) == 0 + gateway.release["b"].set() + assert runtime.wait_for_dispatch(started.value, timeout=1).accepted + clock.set(120) + runtime.poll() + assert gateway.cancel_entered["a"].wait(1) + assert gateway.cancel_entered["b"].wait(1) + finally: + for release in gateway.release.values(): + release.set() + runtime.shutdown() + + +def test_blocked_sequential_dispatch_does_not_expire_before_final_acceptance() -> None: + clock = AdjustableClock() + gateway = BlockingGateway({"a": Outcome.ACCEPTED, "b": Outcome.ACCEPTED}, block_execute=True) + runtime = ExecutionRuntime( + lambda: gateway, + monotonic_clock=clock, + physical_operation_timeout=60, + action_timeout=1000, + poll_interval=100, + ) + try: + started = runtime.execute_explicit(plan("a", "b")) + assert started.accepted and gateway.execute_entered["a"].wait(1) + clock.set(600) + runtime.poll() + assert gateway.count("a", ActionMethod.CANCEL) == 0 + gateway.release["a"].set() + assert gateway.execute_entered["b"].wait(1) + assert gateway.count("a", ActionMethod.CANCEL) == 0 + gateway.release["b"].set() + assert runtime.wait_for_dispatch(started.value, timeout=1).accepted # type: ignore[arg-type] + assert gateway.count("a", ActionMethod.CANCEL) == 0 + assert gateway.count("b", ActionMethod.CANCEL) == 0 + finally: + for release in gateway.release.values(): + release.set() + runtime.shutdown() + + +def test_finite_terminal_wait_is_observer_only_before_physical_expiry() -> None: + clock = AdjustableClock() + gateway = ScriptedGateway() + runtime = ExecutionRuntime( + lambda: gateway, monotonic_clock=clock, physical_operation_timeout=60, poll_interval=100 + ) + try: + handle = start_scripted(runtime, gateway) + clock.set(59) + observed: list[Any] = [] + ready = threading.Event() + + def observe() -> None: + ready.set() + observed.append(runtime.wait_for_terminal(handle, timeout=1)) + + waiter = threading.Thread(target=observe) + waiter.start() + assert ready.wait(1) + clock.set(60) + waiter.join(2) + assert not waiter.is_alive() and len(observed) == 1 + waited = observed[0] + assert not waited.accepted and waited.diagnostic == "timeout" + assert gateway.count("a", ActionMethod.CANCEL) == 0 + clock.set(60) + runtime.poll() + assert gateway.cancel_returned.wait(1) + finally: + close_scripted(runtime, gateway) + + +def test_infinite_terminal_wait_returns_after_public_cancellation() -> None: + gateway = ScriptedGateway() + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) + observed: list[Any] = [] + ready = threading.Event() + try: + handle = start_scripted(runtime, gateway) + + def observe() -> None: + ready.set() + observed.append(runtime.wait_for_terminal(handle, timeout=float("inf"))) + + waiter = threading.Thread(target=observe) + waiter.start() + assert ready.wait(1) + assert runtime.cancel().accepted + waiter.join(1) + assert len(observed) == 1 + assert observed[0].accepted + assert observed[0].value is not None + assert observed[0].value.outcome == Outcome.CANCELLED + finally: + close_scripted(runtime, gateway) + + +def test_authoritative_completion_before_physical_expiry_wins() -> None: + clock = AdjustableClock() + gateway = ScriptedGateway({"a": Outcome.COMPLETED}) + runtime = ExecutionRuntime( + lambda: gateway, monotonic_clock=clock, physical_operation_timeout=60, poll_interval=100 + ) + try: + handle = start_scripted(runtime, gateway) + clock.set(59) + runtime.poll() + result = runtime.wait_for_terminal(handle, timeout=1) + assert result.accepted and result.value is not None + assert result.value.outcome == Outcome.COMPLETED + clock.set(60) + runtime.poll() + assert gateway.count("a", ActionMethod.CANCEL) == 0 + finally: + close_scripted(runtime, gateway) + + +def test_multitask_physical_expiry_cancels_each_task_and_reconciles_terminally() -> None: + clock = AdjustableClock() + gateway = ScriptedGateway() + runtime = ExecutionRuntime( + lambda: gateway, monotonic_clock=clock, physical_operation_timeout=60, poll_interval=100 + ) + try: + started = runtime.execute_explicit(scripted_plan("a", "b")) + assert started.accepted and started.value is not None + gateway.cancel_entered_by_task["a"] = threading.Event() + gateway.cancel_entered_by_task["b"] = threading.Event() + gateway.cancel_returned_by_task["a"] = threading.Event() + gateway.cancel_returned_by_task["b"] = threading.Event() + assert gateway.execute_entered.wait(1) + assert runtime.wait_for_dispatch(started.value, timeout=1).accepted + clock.set(60) + runtime.poll() + assert gateway.cancel_returned_by_task["a"].wait(1) + assert gateway.cancel_returned_by_task["b"].wait(1) + runtime.poll() + assert gateway.count("a", ActionMethod.CANCEL) == 1 + assert gateway.count("b", ActionMethod.CANCEL) == 1 + terminal = runtime.wait_for_terminal(started.value, timeout=1) + assert terminal.accepted and terminal.value is not None + assert terminal.value.outcome == Outcome.FAILED + finally: + close_scripted(runtime, gateway) + + +@pytest.mark.parametrize("cancel_outcome", [Outcome.UNKNOWN, Outcome.FAILED]) +def test_physical_expiry_uncertain_cancellation_is_correlated_to_fault( + cancel_outcome: Outcome, +) -> None: + clock = AdjustableClock() + gateway = ScriptedGateway() + gateway.cancel_outcome = cancel_outcome + runtime = ExecutionRuntime( + lambda: gateway, monotonic_clock=clock, physical_operation_timeout=60, poll_interval=100 + ) + try: + handle = start_scripted(runtime, gateway) + clock.set(60) + runtime.poll() + assert gateway.cancel_returned.wait(1) + observed = runtime.wait_for_terminal(handle, timeout=1) + assert not observed.accepted + assert observed.diagnostic in {"cancellation is uncertain", "cancellation failed"} + assert observed.snapshot is not None + assert observed.snapshot.state == LifecycleState.FAULT + finally: + close_scripted(runtime, gateway) + + +def test_shutdown_unknown_cancels_each_task_once_and_drains() -> None: + gateway = BlockingGateway( + {"a": Outcome.ACCEPTED, "b": Outcome.ACCEPTED}, + {"a": Outcome.UNKNOWN, "b": Outcome.UNKNOWN}, + ) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) + result: list[Any] = [] + try: + started = runtime.execute_explicit(plan("a", "b")) + assert started.accepted and started.value is not None + assert runtime.wait_for_dispatch(started.value, timeout=1).accepted + closer = threading.Thread(target=lambda: result.append(runtime.shutdown(timeout=1))) + closer.start() + assert gateway.cancel_entered["a"].wait(1) + assert gateway.cancel_entered["b"].wait(1) + gateway.release["a"].set() + gateway.release["b"].set() + closer.join(2) + assert result and not result[0].accepted + assert gateway.count("a", ActionMethod.CANCEL) == 1 + assert gateway.count("b", ActionMethod.CANCEL) == 1 + assert runtime.snapshot().shutdown == ShutdownState.CLOSING + assert runtime.snapshot().shutdown_result is not None + assert not runtime.snapshot().shutdown_result.success + finally: + for release in gateway.release.values(): + release.set() + + +def test_shutdown_latches_cancel_only_before_physical_expiry_and_admits_one_cancel() -> None: + clock = AdjustableClock() + gateway = ScriptedGateway() + gateway.block_cancel = True + runtime = ExecutionRuntime( + lambda: gateway, + monotonic_clock=clock, + physical_operation_timeout=60, + action_timeout=100, + poll_interval=100, + ) + result: list[Any] = [] + try: + start_scripted(runtime, gateway) + thread = threading.Thread(target=lambda: result.append(runtime.shutdown(timeout=100))) + thread.start() + assert gateway.cancel_entered.wait(1) + assert runtime.snapshot().shutdown == ShutdownState.CLOSING + assert not runtime.poll().accepted + clock.set(61) + gateway.cancel_release.set() + thread.join(1) + assert result and result[0].value is not None and result[0].value.success + assert gateway.count("a", ActionMethod.CANCEL) == 1 + finally: + gateway.cancel_release.set() + close_scripted(runtime, gateway) + + +def test_unresolved_shutdown_deadline_has_no_late_effects() -> None: + clock = AdjustableClock() + gateway = BlockingGateway({"a": Outcome.ACCEPTED}, block_execute=True) + runtime = ExecutionRuntime(lambda: gateway, monotonic_clock=clock, poll_interval=100) + shutdown_started = False + try: + started = runtime.execute_explicit(plan("a")) + assert started.accepted and gateway.execute_entered["a"].wait(1) + shutdown_started = True + result = runtime.shutdown(timeout=0) + assert not result.accepted + assert result.value is not None and not result.value.success + assert runtime.snapshot().shutdown == ShutdownState.CLOSING + cancel_count = gateway.count("a", ActionMethod.CANCEL) + gateway.release["a"].set() + assert gateway.stop_entered.wait(1) + assert gateway.count("a", ActionMethod.CANCEL) == cancel_count + 1 + finally: + gateway.release["a"].set() + if not shutdown_started: + runtime.shutdown(timeout=0) + + +@pytest.mark.parametrize( + ("execute_outcome", "expected_cancel"), + [(Outcome.ACCEPTED, 1), (Outcome.REJECTED, 0)], +) +def test_late_action_deadline_completion_is_correlated( + execute_outcome: Outcome, expected_cancel: int +) -> None: + clock = AdjustableClock() + gateway = BlockingGateway({"a": execute_outcome}, block_execute=True) + runtime = ExecutionRuntime( + lambda: gateway, monotonic_clock=clock, action_timeout=1, poll_interval=100 + ) + try: + started = runtime.execute_explicit(plan("a")) + assert started.accepted and started.value is not None + assert gateway.execute_entered["a"].wait(1) + clock.set(2) + runtime.poll() + gateway.release["a"].set() + dispatch = runtime.wait_for_dispatch(started.value, timeout=1) + assert dispatch.accepted and dispatch.value is not None + assert dispatch.value.outcome == Outcome.UNKNOWN + if expected_cancel: + assert gateway.cancel_entered["a"].wait(1) + gateway.release["a"].set() + assert gateway.count("a", ActionMethod.CANCEL) == expected_cancel + finally: + gateway.release["a"].set() + runtime.shutdown() + + +def test_unknown_partial_dispatch_compensates_only_admitted_effects() -> None: + gateway = BlockingGateway( + {"a": Outcome.ACCEPTED, "b": Outcome.UNKNOWN, "c": Outcome.ACCEPTED}, + {"a": Outcome.CANCELLED, "b": Outcome.CANCELLED}, + ) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) + try: + started = runtime.execute_explicit(plan("a", "b", "c")) + assert started.accepted and started.value is not None + assert gateway.cancel_entered["a"].wait(1) + assert gateway.cancel_entered["b"].wait(1) + assert not any( + task == "c" and method == ActionMethod.EXECUTE for task, method in gateway.calls + ) + for release in gateway.release.values(): + release.set() + dispatch = runtime.wait_for_dispatch(started.value, timeout=1) + assert dispatch.accepted and dispatch.value is not None + assert dispatch.value.outcome == Outcome.UNKNOWN + assert gateway.count("a", ActionMethod.CANCEL) == 1 + assert gateway.count("b", ActionMethod.CANCEL) == 1 + finally: + for release in gateway.release.values(): + release.set() + runtime.shutdown() + + +def _install_pending_action( + runtime: ExecutionRuntime, + *, + active_handle: OperationHandle, + pending_handle: OperationHandle, + task_id: str = "task", + active_task_id: str = "task", + current_action_id: str = "new-action", + pending_action_id: str = "old-action", +) -> tuple[ActionRecord, ActionRecord]: + old_action = ActionRecord(pending_action_id, ActionMethod.EXECUTE, 1.0, 1e9) + new_action = ActionRecord(current_action_id, ActionMethod.STATUS, 2.0, 1e9 + 1) + task = TaskRecord( + active_task_id, + "task", + plan("task").entries[0], + TaskActivity.ACTIVE, + new_action, + ) + operation = Operation(active_handle, plan("task"), (task,)) + runtime._commit(state=LifecycleState.RUNNING, active=operation) + runtime._pending[pending_action_id] = _PendingAction( + pending_handle, + "task", + replace(task, task_id=task_id, action=old_action), + ActionMethod.EXECUTE, + 1e9, + ) + runtime._pending[current_action_id] = _PendingAction( + active_handle, + "task", + task, + ActionMethod.STATUS, + 1e9 + 1, + ) + return old_action, new_action + + +def test_owner_completion_retires_old_pending_but_preserves_new_action() -> None: + runtime = ExecutionRuntime( + lambda: BlockingGateway({"task": Outcome.RUNNING}), poll_interval=100 + ) + try: + handle = OperationHandle("plan", "operation", "attempt") + _, newer = _install_pending_action(runtime, active_handle=handle, pending_handle=handle) + before = runtime.snapshot() + runtime._events.put(EffectDone("old-action", Outcome.ACCEPTED)) + runtime._submit(lambda: None) + after = runtime.snapshot() + assert "old-action" not in runtime._pending + assert "new-action" in runtime._pending + assert after == before + assert after.operation is not None + assert after.operation.tasks[0].action == newer + finally: + runtime._pending.clear() + runtime._commit(active=None, state=LifecycleState.IDLE) + runtime.shutdown() + + +@pytest.mark.parametrize("kind", ["unknown", "handle", "task"]) +def test_owner_uncorrelatable_completion_preserves_pending_and_snapshot(kind: str) -> None: + runtime = ExecutionRuntime( + lambda: BlockingGateway({"task": Outcome.RUNNING}), poll_interval=100 + ) + try: + handle = OperationHandle("plan", "operation", "attempt") + pending_handle = ( + OperationHandle("other", "operation", "attempt") if kind == "handle" else handle + ) + _install_pending_action( + runtime, + active_handle=handle, + pending_handle=pending_handle, + task_id="missing" if kind == "task" else "task", + ) + before = runtime.snapshot() + action_id = "unknown-action" if kind == "unknown" else "old-action" + runtime._events.put(EffectDone(action_id, Outcome.ACCEPTED)) + runtime._submit(lambda: None) + assert runtime.snapshot() == before + assert set(runtime._pending) == {"old-action", "new-action"} + finally: + runtime._pending.clear() + runtime._commit(active=None, state=LifecycleState.IDLE) + runtime.shutdown() + + +def test_concurrent_explicit_admission_accepts_one_and_correlates_rejection() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) + start = threading.Barrier(3) + results: list[Any] = [] + + def submit() -> None: + start.wait() + results.append(runtime.execute_explicit(plan("a"))) + + threads = [threading.Thread(target=submit) for _ in range(2)] + try: + for thread in threads: + thread.start() + start.wait() + for thread in threads: + thread.join(1) + assert len(results) == 2 + assert sum(result.accepted for result in results) == 1 + accepted = next(result for result in results if result.accepted) + assert accepted.value is not None + rejected = next(result for result in results if not result.accepted) + assert rejected.value is None + snapshot = runtime.snapshot() + assert snapshot.operation is not None + assert snapshot.operation.handle == accepted.value + gateway.release["a"].set() + assert gateway.count("a", ActionMethod.EXECUTE) == 1 + finally: + gateway.release["a"].set() + runtime.cancel() + runtime.shutdown() + + +def test_concurrent_ready_execution_consumes_ready_plan_once() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) + start = threading.Barrier(3) + results: list[Any] = [] + token = runtime.start_planning() + assert token is not None + completed = runtime.complete_planning(token, plan("a")) + assert completed.accepted and completed.value is not None + plan_id = completed.value + + def execute_ready() -> None: + start.wait() + results.append(runtime.execute_ready()) + + threads = [threading.Thread(target=execute_ready) for _ in range(2)] + try: + for thread in threads: + thread.start() + start.wait() + for thread in threads: + thread.join(1) + assert all(not thread.is_alive() for thread in threads) + assert len(results) == 2 + assert sum(result.accepted for result in results) == 1 + assert runtime.snapshot().ready_plan is None + accepted = next(result for result in results if result.accepted) + rejected = next(result for result in results if not result.accepted) + assert accepted.value is not None + assert accepted.value.plan_id == plan_id + assert not rejected.accepted + assert rejected.value is None + assert rejected.diagnostic == "no ready plan" + assert gateway.count("a", ActionMethod.EXECUTE) == 1 + finally: + gateway.release["a"].set() + runtime.cancel() + runtime.shutdown() + + +def test_active_operation_rejects_planning_and_explicit_execution_but_accepts_cancel() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) + try: + started = runtime.execute_explicit(plan("a")) + assert started.accepted and started.value is not None + assert gateway.execute_entered["a"].wait(1) + assert runtime.wait_for_dispatch(started.value, timeout=1).accepted + original = runtime.snapshot() + assert original.state == LifecycleState.RUNNING + assert not runtime.start_planning() + rejected = runtime.execute_explicit(plan("a")) + assert not rejected.accepted + assert rejected.value is None + assert runtime.snapshot().operation == original.operation + assert gateway.count("a", ActionMethod.EXECUTE) == 1 + assert runtime.snapshot().operation == original.operation + assert runtime.cancel().accepted + gateway.release["a"].set() + finally: + gateway.release["a"].set() + runtime.shutdown() + + +def test_sequential_rejection_skips_remaining_task_and_compensates_prior_acceptance() -> None: + gateway = BlockingGateway( + {"a": Outcome.ACCEPTED, "b": Outcome.REJECTED, "c": Outcome.ACCEPTED}, + {"a": Outcome.CANCELLED}, + ) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) + try: + started = runtime.execute_explicit(plan("a", "b", "c")) + assert started.accepted and started.value is not None + assert gateway.cancel_entered["a"].wait(1) + assert gateway.calls[:3] == [ + ("a", ActionMethod.EXECUTE), + ("b", ActionMethod.EXECUTE), + ("a", ActionMethod.CANCEL), + ] + assert not any( + task == "c" and method == ActionMethod.EXECUTE for task, method in gateway.calls + ) + gateway.release["a"].set() + dispatch = runtime.wait_for_dispatch(started.value, timeout=1) + terminal = runtime.wait_for_terminal(started.value, timeout=1) + assert dispatch.accepted and dispatch.value is not None + assert terminal.accepted and terminal.value is not None + assert dispatch.value.outcome == Outcome.REJECTED + assert terminal.value.outcome == Outcome.REJECTED + assert gateway.count("a", ActionMethod.EXECUTE) == 1 + assert gateway.count("a", ActionMethod.CANCEL) == 1 + assert gateway.count("b", ActionMethod.EXECUTE) == 1 + assert gateway.count("c", ActionMethod.EXECUTE) == 0 + assert gateway.count("c", ActionMethod.CANCEL) == 0 + finally: + for release in gateway.release.values(): + release.set() + runtime.shutdown() + + +def test_cancel_during_unresolved_dispatch_retains_cancelled_result_after_late_acceptance() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED, "b": Outcome.ACCEPTED}, block_execute=True) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) + try: + started = runtime.execute_explicit(plan("a", "b")) + assert started.accepted and started.value is not None + assert gateway.execute_entered["a"].wait(1) + assert runtime.cancel().accepted + gateway.release["a"].set() + dispatch = runtime.wait_for_dispatch(started.value, timeout=1) + terminal = runtime.wait_for_terminal(started.value, timeout=1) + assert dispatch.accepted and dispatch.value is not None + assert terminal.accepted and terminal.value is not None + assert dispatch.value.outcome == Outcome.CANCELLED + assert terminal.value.outcome == Outcome.CANCELLED + assert gateway.count("a", ActionMethod.EXECUTE) == 1 + assert gateway.calls[:2] == [ + ("a", ActionMethod.EXECUTE), + ("a", ActionMethod.CANCEL), + ] + assert gateway.count("a", ActionMethod.CANCEL) == 1 + assert gateway.count("b", ActionMethod.EXECUTE) == 0 + assert gateway.count("b", ActionMethod.CANCEL) == 0 + finally: + for release in gateway.release.values(): + release.set() + runtime.shutdown() + + +def test_unknown_execute_enters_fault_and_compensates_without_terminal_success() -> None: + gateway = BlockingGateway( + {"a": Outcome.UNKNOWN, "b": Outcome.REJECTED}, + {"a": Outcome.UNKNOWN, "b": Outcome.FAILED}, + ) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) + try: + started = runtime.execute_explicit(plan("a", "b")) + assert started.accepted and started.value is not None + assert gateway.cancel_entered["a"].wait(1) + snapshot = runtime.snapshot() + assert snapshot.state == LifecycleState.FAULT + assert snapshot.operation is not None and snapshot.operation.uncertain + assert gateway.count("a", ActionMethod.EXECUTE) == 1 + assert gateway.count("b", ActionMethod.EXECUTE) == 0 + assert gateway.count("a", ActionMethod.CANCEL) == 1 + assert gateway.count("b", ActionMethod.CANCEL) == 0 + gateway.release["a"].set() + dispatch = runtime.wait_for_dispatch(started.value, timeout=1) + assert dispatch.accepted and dispatch.value is not None + assert dispatch.value.outcome == Outcome.UNKNOWN + assert dispatch.value.diagnostic + assert "task=a" in dispatch.value.diagnostic + assert "outcome=unknown" in dispatch.value.diagnostic + terminal = runtime.wait_for_terminal(started.value, timeout=1) + assert not terminal.accepted + assert terminal.diagnostic == dispatch.value.diagnostic + finally: + gateway.release["a"].set() + runtime.shutdown() + + +def test_unknown_status_enters_fault_without_successful_terminal_result() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED, "b": Outcome.ACCEPTED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) + try: + started = runtime.execute_explicit(plan("a", "b")) + assert started.accepted and started.value is not None + handle = started.value + assert gateway.execute_entered["a"].wait(1) + assert runtime.wait_for_dispatch(started.value, timeout=1).accepted + runtime.poll() + assert gateway.status_entered.wait(1) + assert gateway.cancel_entered["a"].wait(1) + assert gateway.cancel_entered["b"].wait(1) + gateway.release["a"].set() + gateway.release["b"].set() + terminal = runtime.wait_for_terminal(handle, timeout=1) + assert not terminal.accepted + snapshot = runtime.snapshot() + assert snapshot.operation is not None + assert terminal.diagnostic == runtime.snapshot().diagnostic + assert terminal.diagnostic + assert "task=a" in terminal.diagnostic + assert "outcome=unknown" in terminal.diagnostic + assert "status is unsafe" in terminal.diagnostic + assert runtime.snapshot().state == LifecycleState.FAULT + assert gateway.count("a", ActionMethod.STATUS) == 1 + assert gateway.count("a", ActionMethod.CANCEL) == 1 + assert gateway.count("b", ActionMethod.CANCEL) == 1 + finally: + for release in gateway.release.values(): + release.set() + runtime.shutdown() + + +def test_failed_cancel_latches_fault_and_preserves_reset_required_operation() -> None: + gateway = BlockingGateway({"a": Outcome.ACCEPTED}, {"a": Outcome.FAILED}) + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) + try: + started = runtime.execute_explicit(plan("a")) + assert started.accepted and started.value is not None + assert gateway.execute_entered["a"].wait(1) + assert runtime.wait_for_dispatch(started.value, timeout=1).accepted + assert runtime.cancel().accepted + assert gateway.cancel_entered["a"].wait(1) + gateway.release["a"].set() + terminal = runtime.wait_for_terminal(started.value, timeout=1) + assert not terminal.accepted + assert terminal.diagnostic == "cancellation failed" + assert runtime.snapshot().state == LifecycleState.FAULT + snapshot = runtime.snapshot() + assert snapshot.operation is not None + assert snapshot.operation.tasks[0].reset_required + assert snapshot.diagnostic == "cancellation failed" + assert gateway.count("a", ActionMethod.CANCEL) == 1 + finally: + gateway.release["a"].set() + runtime.shutdown() + + +def test_reset_calls_share_one_public_handle_until_first_result_is_ready() -> None: + gateway = ScriptedGateway({"a": Outcome.FAILED}) + gateway.block_reset = True + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) + try: + start_scripted(runtime, gateway) + runtime.poll() + snapshot = runtime.snapshot() + assert snapshot.operation is not None + assert not runtime.wait_for_terminal(snapshot.operation.handle, timeout=1).accepted + gateway.statuses["a"] = Outcome.INACTIVE + assert gateway.cancel_returned.wait(1) + runtime.poll() + assert runtime.snapshot().state == LifecycleState.FAULT + first = runtime.reset() + assert first.accepted and first.value is not None + assert gateway.reset_entered.wait(1) + second = runtime.reset() + assert second.accepted and second.value == first.value + gateway.reset_release.set() + result = runtime.wait_for_reset(first.value, timeout=1) + assert result.accepted and result.value is not None and result.value.success + assert runtime.snapshot().state == LifecycleState.IDLE + assert gateway.count("a", ActionMethod.RESET) == 1 + finally: + gateway.reset_release.set() + close_scripted(runtime, gateway) + + +def test_unsafe_reset_remains_faulted_without_an_in_process_retry() -> None: + gateway = ScriptedGateway() + gateway.block_cancel = True + status_calls = 0 + + def status(task_name: str) -> Outcome: + nonlocal status_calls + status_calls += 1 + return Outcome.FAILED if status_calls == 1 else Outcome.UNKNOWN + + gateway.status = status # type: ignore[assignment,method-assign] + runtime = ExecutionRuntime(lambda: gateway, poll_interval=100) + try: + start_scripted(runtime, gateway) + runtime.poll() + snapshot = runtime.snapshot() + assert snapshot.operation is not None + assert not runtime.wait_for_terminal(snapshot.operation.handle, timeout=1).accepted + assert gateway.cancel_entered.wait(1) + reset = runtime.reset() + assert not reset.accepted and reset.value is None + assert reset.diagnostic == "reset blocked by unresolved coordinator RPC" + assert runtime.snapshot().state == LifecycleState.FAULT + assert runtime.snapshot().diagnostic + assert gateway.count("a", ActionMethod.RESET) == 0 + finally: + gateway.cancel_release.set() + close_scripted(runtime, gateway) + + +def test_public_results_retain_current_and_recent_handles_and_age_out_old_handles() -> None: + gateway = BlockingGateway({"a": Outcome.REJECTED}) + clock = AdjustableClock() + runtime = ExecutionRuntime(lambda: gateway, monotonic_clock=clock, poll_interval=100) + handles: list[Any] = [] + try: + for _ in range(20): + started = runtime.execute_explicit(plan("a")) + assert started.accepted and started.value is not None + handles.append(started.value) + dispatch = runtime.wait_for_dispatch(started.value, timeout=1) + terminal = runtime.wait_for_terminal(started.value, timeout=1) + assert dispatch.accepted and dispatch.value is not None + assert terminal.accepted and terminal.value is not None + assert dispatch.value.handle == started.value + assert terminal.value.handle == started.value + for handle in handles[-2:]: + dispatch = runtime.wait_for_dispatch(handle, timeout=1) + terminal = runtime.wait_for_terminal(handle, timeout=1) + assert dispatch.accepted and terminal.accepted + assert dispatch.value is not None and dispatch.value.outcome == Outcome.REJECTED + assert terminal.value is not None and terminal.value.outcome == Outcome.REJECTED + assert dispatch.value.handle == handle + assert terminal.value.handle == handle + aged_out = handles[0] + assert ( + assert_public_wait_times_out( + lambda: runtime.wait_for_dispatch(aged_out, timeout=1) + ).diagnostic + == "timeout" + ) + assert ( + assert_public_wait_times_out( + lambda: runtime.wait_for_terminal(aged_out, timeout=1) + ).diagnostic + == "timeout" + ) + wrong = replace(handles[-1], attempt_id="missing-attempt") + unknown = replace(handles[-1], operation_id="missing-operation") + assert ( + assert_public_wait_times_out( + lambda: runtime.wait_for_dispatch(wrong, timeout=1) + ).diagnostic + == "timeout" + ) + assert ( + assert_public_wait_times_out( + lambda: runtime.wait_for_terminal(wrong, timeout=1) + ).diagnostic + == "timeout" + ) + assert ( + assert_public_wait_times_out( + lambda: runtime.wait_for_dispatch(unknown, timeout=1) + ).diagnostic + == "timeout" + ) + assert ( + assert_public_wait_times_out( + lambda: runtime.wait_for_terminal(unknown, timeout=1) + ).diagnostic + == "timeout" + ) + finally: + runtime.shutdown() diff --git a/dimos/manipulation/test_execution_topology.py b/dimos/manipulation/test_execution_topology.py new file mode 100644 index 0000000000..c6478633f8 --- /dev/null +++ b/dimos/manipulation/test_execution_topology.py @@ -0,0 +1,199 @@ +"""Topology and generated-plan materialization tests.""" + +from pathlib import Path + +import pytest + +from dimos.manipulation.execution_models import Outcome +from dimos.manipulation.execution_runtime import ExecutionRuntime, ExecutionTopology +from dimos.manipulation.execution_topology import ( + ExecutionPlan, + PreparedPlan, + prepare_generated_plan, +) +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.models import GeneratedPlan +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.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint + + +class Gateway: + def __init__(self, outcomes: dict[str, Outcome]) -> None: + self.outcomes = outcomes + + def execute(self, task_name: str, request: object) -> Outcome: + return self.outcomes.get(task_name, Outcome.ACCEPTED) + + def cancel(self, task_name: str) -> Outcome: + return Outcome.CANCELLED + + def status(self, task_name: str) -> Outcome: + return Outcome.INACTIVE + + def reset(self, task_name: str) -> Outcome: + return Outcome.INACTIVE + + def set_gripper_position(self, hardware_id: str, position: float) -> Outcome: + return Outcome.ACCEPTED + + def get_gripper_position(self, hardware_id: str) -> float: + return 0.0 + + def stop(self) -> None: + return None + + +def _prep_robot( + name: str, + joints: tuple[str, ...], + group: str, + task: str | None, + mapping: dict[str, str] | None = None, +) -> RobotModelConfig: + return RobotModelConfig( + name=name, + model_path=Path(f"/{name}.urdf"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion([0.0, 0.0, 0.0, 1.0])), + joint_names=list(joints), + planning_groups=[PlanningGroupDefinition(name=group, joint_names=joints, base_link="base")], + coordinator_task_name=task, + joint_name_mapping=mapping or {}, + ) + + +def _generated(names: list[str], groups: tuple[str, ...], width: int) -> GeneratedPlan: + return GeneratedPlan( + group_ids=groups, + trajectory=JointTrajectory( + joint_names=names, + points=[ + TrajectoryPoint( + time_from_start=0.0, positions=[0.0] * width, velocities=[0.1] * width + ), + TrajectoryPoint( + time_from_start=1.5, positions=[1.0] * width, velocities=[0.2] * width + ), + ], + ), + ) + + +def test_prepare_generated_plan_uses_ordered_groups_and_ignores_extraneous_robot_columns() -> None: + left = _prep_robot("left", ("l0", "l1"), "arm", "left_task", {"left_l0": "l0"}) + right = _prep_robot("right", ("r0",), "arm", "right_task") + extra = _prep_robot("extra", ("x0",), "arm", "extra_task") + topology = ExecutionTopology.from_robot_configs((left, right, extra)) + generated = _generated( + ["extra/x0", "right/r0", "left/l1", "left/l0"], ("left/arm", "right/arm"), 4 + ) + + prepared = prepare_generated_plan(generated, topology) + assert prepared.generated_plan is generated + assert [entry.robot_name for entry in prepared.entries] == ["left", "right"] + assert [entry.task_name for entry in prepared.entries] == ["left_task", "right_task"] + assert prepared.entries[0].request["trajectory"].joint_names == ["l1", "left_l0"] + assert prepared.entries[0].request["trajectory"].points[1].time_from_start == 1.5 + assert prepared.entries[1].request["trajectory"].joint_names == ["r0"] + + +def test_prepare_rejects_missing_routes_tasks_malformed_and_missing_or_duplicate_joints() -> None: + left = _prep_robot("left", ("l0", "l1"), "arm", "left_task") + right = _prep_robot("right", ("r0",), "arm", "right_task") + topology = ExecutionTopology.from_robot_configs((left, right)) + with pytest.raises(ValueError, match="route set"): + ExecutionTopology.from_robot_configs((left, right), {"left/arm": ("left", "left_task")}) + with pytest.raises(ValueError, match="malformed"): + prepare_generated_plan(_generated(["left/l0", "bad"], ("left/arm",), 2), topology) + with pytest.raises(ValueError, match="missing robot joints"): + prepare_generated_plan(_generated(["left/l0"], ("left/arm",), 1), topology) + with pytest.raises(ValueError, match="duplicate global"): + prepare_generated_plan(_generated(["left/l0", "left/l0"], ("left/arm",), 2), topology) + + +def test_prepared_validation_selects_only_generated_groups_from_larger_topology() -> None: + left = _prep_robot("left", ("l0",), "arm", "left_task") + right = _prep_robot("right", ("r0",), "arm", "right_task") + extra = _prep_robot("extra", ("x0",), "arm", "extra_task") + topology = ExecutionTopology.from_robot_configs((left, right, extra)) + prepared = prepare_generated_plan( + _generated(["left/l0", "right/r0"], ("right/arm", "left/arm"), 2), topology + ) + runtime = ExecutionRuntime( + lambda: Gateway({"left": Outcome.ACCEPTED, "right": Outcome.ACCEPTED}), + poll_interval=10, + ) + try: + result = runtime.execute_explicit(prepared) + assert result.accepted + assert [entry.planning_group for entry in prepared.entries] == ["right/arm", "left/arm"] + finally: + runtime.shutdown() + + +def test_prepare_rejects_additional_selected_robot_joint_column() -> None: + robot = _prep_robot("arm", ("j0", "j1"), "manipulator", "arm_task") + topology = ExecutionTopology.from_robot_configs((robot,)) + with pytest.raises(ValueError, match="additional joint column"): + prepare_generated_plan( + _generated(["arm/j0", "arm/j1", "arm/extra"], ("arm/manipulator",), 3), topology + ) + + +def test_topology_rejects_duplicate_routes_bad_mapping_targets_and_unknown_group_joints() -> None: + robot = _prep_robot("arm", ("j0",), "manipulator", "arm_task") + with pytest.raises(ValueError, match="duplicate group route"): + ExecutionTopology.from_robot_configs( + (robot,), + [("arm/manipulator", "arm", "arm_task"), ("arm/manipulator", "arm", "arm_task")], + ) + with pytest.raises(ValueError, match="mapped local joint"): + ExecutionTopology.from_robot_configs( + (_prep_robot("arm", ("j0",), "manipulator", "arm_task", {"coord": "missing"}),) + ) + bad_group = _prep_robot("arm", ("j0",), "manipulator", "arm_task") + bad_group.planning_groups[0] = PlanningGroupDefinition( + name="manipulator", joint_names=("missing",), base_link="base" + ) + with pytest.raises(ValueError, match="planning group joints"): + ExecutionTopology.from_robot_configs((bad_group,)) + + +def test_planning_only_config_has_no_route_and_generated_execution_rejects_it() -> None: + planning_only = _prep_robot("planner", ("j0",), "manipulator", None) + topology = ExecutionTopology.from_robot_configs((planning_only,)) + assert topology.routes == () + with pytest.raises(ValueError, match="no coordinator task"): + prepare_generated_plan(_generated(["planner/j0"], ("planner/manipulator",), 1), topology) + + +def test_prepared_plan_is_real_type_and_partial_generated_entries_are_not_executable() -> None: + robot = _prep_robot("arm", ("j0",), "manipulator", "arm_task", {"coord_j0": "j0"}) + topology = ExecutionTopology.from_robot_configs((robot,)) + generated = _generated(["arm/j0"], ("arm/manipulator",), 1) + prepared = prepare_generated_plan(generated, topology) + assert isinstance(prepared, PreparedPlan) + assert prepared is not generated + partial = PreparedPlan(generated, (), topology) + runtime = ExecutionRuntime(lambda: Gateway({"arm": Outcome.ACCEPTED}), poll_interval=10) + try: + assert not runtime.execute_explicit(partial).accepted + finally: + runtime.shutdown() + + +def test_real_generated_plan_cannot_use_legacy_synthetic_execution_plan() -> None: + robot = _prep_robot("arm", ("j0",), "manipulator", "arm_task") + topology = ExecutionTopology.from_robot_configs((robot,)) + generated = _generated(["arm/j0"], ("arm/manipulator",), 1) + prepared = prepare_generated_plan(generated, topology) + runtime = ExecutionRuntime(lambda: Gateway({"arm": Outcome.ACCEPTED}), poll_interval=10) + try: + legacy = ExecutionPlan(generated, prepared.entries) + assert not runtime.execute_explicit(legacy).accepted + assert runtime.execute_explicit(prepared).accepted + finally: + runtime.shutdown() diff --git a/dimos/manipulation/test_generated_plan_materialization.py b/dimos/manipulation/test_generated_plan_materialization.py new file mode 100644 index 0000000000..dd341d8cc0 --- /dev/null +++ b/dimos/manipulation/test_generated_plan_materialization.py @@ -0,0 +1,202 @@ +# 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. + +"""Generated-plan materialization tests.""" + +from collections.abc import Iterator +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + +from dimos.manipulation._test_manipulation_helpers import ( + close_test_runtimes, + install_runtime, + make_module, +) +from dimos.manipulation.execution_runtime import PreparedPlan +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition +from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.enums import PlanningStatus +from dimos.manipulation.planning.spec.models import PlanningResult +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.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint + + +class RecordingGenerator: + calls: list[list[list[float]]] = [] + limits: tuple[list[float], list[float]] | None = None + fail = False + + def __init__( + self, num_joints: int, max_velocity: list[float], max_acceleration: list[float] + ) -> None: + self.num_joints = num_joints + RecordingGenerator.limits = (list(max_velocity), list(max_acceleration)) + + def generate(self, waypoints: list[list[float]]) -> JointTrajectory: + RecordingGenerator.calls.append(waypoints) + if RecordingGenerator.fail: + raise RuntimeError("boom") + return JointTrajectory( + points=[ + TrajectoryPoint( + time_from_start=float(index), + positions=list(point), + velocities=[0.0] * self.num_joints, + ) + for index, point in enumerate(waypoints) + ] + ) + + +def _robot(name: str, joints: list[str], velocity: float, acceleration: float) -> RobotModelConfig: + return RobotModelConfig( + name=name, + model_path=Path("/robot.urdf"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + joint_names=joints, + base_link="base", + planning_groups=[ + PlanningGroupDefinition( + name="group", joint_names=tuple(reversed(joints)), base_link="base", tip_link="tip" + ) + ], + max_velocity=velocity, + max_acceleration=acceleration, + coordinator_task_name=f"task_{name}", + ) + + +@pytest.fixture(autouse=True) +def _close_runtime() -> Iterator[None]: + yield + close_test_runtimes() + + +def _module(monkeypatch: pytest.MonkeyPatch): + RecordingGenerator.calls = [] + RecordingGenerator.limits = None + RecordingGenerator.fail = False + monkeypatch.setattr( + "dimos.manipulation.manipulation_module.JointTrajectoryGenerator", RecordingGenerator + ) + left = _robot("left", ["a", "b"], 1.0, 2.0) + right = _robot("right", ["c"], 3.0, 4.0) + module = make_module() + module._robots = { + "left": ("left_id", left, MagicMock()), + "right": ("right_id", right, MagicMock()), + } + module._world_monitor = MagicMock() + module._world_monitor.world = MagicMock() + module._world_monitor.planning_groups = PlanningGroupRegistry([left, right]) + module._planner = MagicMock() + install_runtime(module, [left, right]) + return module + + +def _path(names: list[str], first: list[float], second: list[float]) -> list[JointState]: + return [JointState(name=names, position=first), JointState(name=names, position=second)] + + +def test_materializes_once_with_reordered_groups_heterogeneous_limits_and_distinct_path( + monkeypatch, +): + module = _module(monkeypatch) + names = ["left/b", "left/a", "right/c"] + path = _path(names, [0.0, 0.0, 0.0], [0.2, 0.1, 0.3]) + module._planner.plan_selected_joint_path.return_value = PlanningResult( + status=PlanningStatus.SUCCESS, path=path + ) + + token = module._begin_group_planning() + assert token is not None + assert module._plan_selected_path(("left/group", "right/group"), path[0], path[-1], token) + assert RecordingGenerator.calls == [[[0.0, 0.0, 0.0], [0.2, 0.1, 0.3]]] + assert RecordingGenerator.limits == ([1.0, 1.0, 3.0], [2.0, 2.0, 4.0]) + ready = module._execution_runtime.snapshot().ready_plan + assert isinstance(ready, PreparedPlan) + assert ready.generated_plan.path is not ready.generated_plan.trajectory.points + assert ready.generated_plan.trajectory.joint_names == names + assert ready.generated_plan.trajectory.points[-1].time_from_start == 1.0 + + +@pytest.mark.parametrize( + ("path", "message"), + [ + (_path(["left/a", "left/b"], [0.0, 0.0], [1.0, 1.0]), "joint names"), + (_path(["left/b", "left/a"], [0.0], [1.0]), "dimension"), + (_path(["left/b", "left/a"], [0.0, float("nan")], [1.0, 1.0]), "non-finite"), + ], +) +def test_rejects_malformed_or_nonfinite_waypoints(monkeypatch, path, message): + module = _module(monkeypatch) + module._planner.plan_selected_joint_path.return_value = PlanningResult( + status=PlanningStatus.SUCCESS, path=path + ) + + token = module._begin_group_planning() + assert token is not None + assert not module._plan_selected_path(("left/group",), path[0], path[-1], token) + assert module._execution_runtime.snapshot().ready_plan is None + assert message in (module.get_error() or "") + + +def test_rejects_invalid_limits_and_generator_failure_without_caching(monkeypatch): + module = _module(monkeypatch) + module._robots["left"][1].max_velocity = 0.0 + names = ["left/b", "left/a"] + path = _path(names, [0.0, 0.0], [1.0, 1.0]) + module._planner.plan_selected_joint_path.return_value = PlanningResult( + status=PlanningStatus.SUCCESS, path=path + ) + + token = module._begin_group_planning() + assert token is not None + assert not module._plan_selected_path(("left/group",), path[0], path[-1], token) + assert module._execution_runtime.snapshot().ready_plan is None + assert RecordingGenerator.calls == [] + + module = _module(monkeypatch) + RecordingGenerator.fail = True + module._planner.plan_selected_joint_path.return_value = PlanningResult( + status=PlanningStatus.SUCCESS, path=path + ) + token = module._begin_group_planning() + assert token is not None + assert not module._plan_selected_path(("left/group",), path[0], path[-1], token) + assert module._execution_runtime.snapshot().ready_plan is None + assert len(RecordingGenerator.calls) == 1 + + +def test_zero_generation_after_caching_for_status_and_completion(monkeypatch): + module = _module(monkeypatch) + names = ["left/b", "left/a"] + path = _path(names, [0.0, 0.0], [1.0, 1.0]) + module._planner.plan_selected_joint_path.return_value = PlanningResult( + status=PlanningStatus.SUCCESS, path=path + ) + token = module._begin_group_planning() + assert token is not None + assert module._plan_selected_path(("left/group",), path[0], path[-1], token) + RecordingGenerator.calls = [] + + assert module.get_trajectory_status() is not None + assert RecordingGenerator.calls == [] diff --git a/dimos/manipulation/test_manipulation_module.py b/dimos/manipulation/test_manipulation_module.py index f8e914e3b3..bb1a24abd5 100644 --- a/dimos/manipulation/test_manipulation_module.py +++ b/dimos/manipulation/test_manipulation_module.py @@ -22,14 +22,22 @@ from __future__ import annotations import importlib.util -from unittest.mock import MagicMock +import threading +from types import SimpleNamespace +from typing import Any import pytest -from dimos.manipulation.manipulation_module import ( - ManipulationModule, - ManipulationState, +from dimos.manipulation import manipulation_module as module_impl +from dimos.manipulation._test_manipulation_helpers import FakeCoordinatorGateway +from dimos.manipulation.execution_runtime import ( + ExecutionRuntime, + LifecycleState, + Outcome, + PreparedPlan, ) +from dimos.manipulation.manipulation_module import ManipulationModule +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped @@ -62,8 +70,15 @@ def _get_xarm7_config() -> RobotModelConfig: model_path=desc_path / "urdf/xarm_device.urdf.xacro", base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), joint_names=["joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "joint7"], - end_effector_link="link7", base_link="link_base", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=("joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "joint7"), + base_link="link_base", + tip_link="link7", + ) + ], package_paths={"xarm_description": desc_path}, xacro_args={"dof": "7", "limited": "true"}, auto_convert_meshes=True, @@ -114,6 +129,9 @@ def module(xarm7_config): planning_timeout=10.0, visualization={"backend": "none"}, ) + gateway = FakeCoordinatorGateway() + mod._runtime_factory = lambda _factory, **kwargs: ExecutionRuntime(lambda: gateway, **kwargs) + mod._runtime_gateway = gateway mod.coordinator_joint_state = None mod.objects = None mod.start() @@ -121,6 +139,157 @@ def module(xarm7_config): mod.stop() +def test_execute_explicit_plan_cannot_dispatch_replaced_ready_plan(monkeypatch) -> None: + """A READY replacement between observation and submission cannot switch plans.""" + + plan_a = object() + plan_b = object() + prepared_a = object() + ready_b = PreparedPlan(generated_plan=plan_b, entries=()) # type: ignore[arg-type] + ready_a = PreparedPlan(generated_plan=plan_a, entries=()) # type: ignore[arg-type] + dispatched: list[object] = [] + ready_dispatches: list[object] = [] + + class FakeRuntime: + def snapshot(self): + # Model the replacement at precisely the old snapshot/command seam. + self.current_ready = ready_b + return SimpleNamespace(ready_plan=ready_a) + + def execute_ready(self): + ready_dispatches.append(self.current_ready.generated_plan) + return SimpleNamespace(accepted=True, value=object()) + + def execute_explicit(self, prepared): + dispatched.append(prepared) + return SimpleNamespace(accepted=True, value=object()) + + runtime = FakeRuntime() + module = object.__new__(ManipulationModule) + module._execution_runtime = runtime + module._execution_topology = object() + monkeypatch.setattr(module_impl, "prepare_generated_plan", lambda plan, topology: prepared_a) + + assert module._submit_execution(plan_a) is not None + assert dispatched == [prepared_a] + assert ready_dispatches == [] + + +def test_execute_waits_for_its_correlated_dispatch_result() -> None: + released = threading.Event() + entered = threading.Event() + expected_handle = object() + waits: list[object] = [] + + class FakeRuntime: + def execute_ready(self): + return SimpleNamespace(accepted=True, value=expected_handle) + + def wait_for_dispatch(self, handle, *, timeout): + waits.append(handle) + entered.set() + assert released.wait(1) + return SimpleNamespace( + accepted=True, + value=SimpleNamespace(outcome=Outcome.ACCEPTED), + ) + + module = object.__new__(ManipulationModule) + module._execution_runtime = FakeRuntime() + module._execution_topology = object() + result: list[bool] = [] + caller = threading.Thread(target=lambda: result.append(module.execute())) + caller.start() + assert entered.wait(1) + assert result == [] + released.set() + caller.join(1) + assert result == [True] + assert waits == [expected_handle] + + +def test_preview_waits_for_terminal_before_reporting_completion(monkeypatch) -> None: + released = threading.Event() + entered = threading.Event() + expected_handle = object() + preview_calls: list[tuple[object, object]] = [] + + class FakeRuntime: + def execute_explicit(self, prepared): + return SimpleNamespace(accepted=True, value=expected_handle) + + def wait_for_dispatch(self, handle, *, timeout): + return SimpleNamespace( + accepted=True, + value=SimpleNamespace(outcome=Outcome.ACCEPTED), + ) + + def wait_for_terminal(self, handle, *, timeout): + entered.set() + assert released.wait(1) + return SimpleNamespace( + accepted=True, + value=SimpleNamespace(outcome=Outcome.COMPLETED, diagnostic=""), + ) + + plan = object() + ready = PreparedPlan(generated_plan=plan, entries=()) # type: ignore[arg-type] + runtime = FakeRuntime() + module = object.__new__(ManipulationModule) + module._execution_runtime = runtime + module._execution_topology = object() + module.preview_path = lambda duration=None, robot_name=None, target_fps=30: ( + preview_calls.append((duration, robot_name)) or True + ) + monkeypatch.setattr(module_impl, "prepare_generated_plan", lambda value, topology: value) + + # The preview helper obtains the READY generated plan only to bind it through + # the runtime's explicit-plan path; terminal publication is the completion gate. + runtime.snapshot = lambda: SimpleNamespace(ready_plan=ready) + result: list[Any] = [] + caller = threading.Thread(target=lambda: result.append(module._preview_execute_wait())) + caller.start() + assert entered.wait(1) + assert result == [] + assert preview_calls == [(0.5, None)] + released.set() + caller.join(1) + assert len(result) == 1 and result[0].is_success() + + +def test_execution_projection_helpers_use_runtime_snapshot_only() -> None: + snapshot = SimpleNamespace( + state=LifecycleState.READY, + diagnostic="runtime diagnostic", + ready_plan_id="ready-1", + ready_plan=PreparedPlan(generated_plan=SimpleNamespace(path=[]), entries=()), + operation=None, + ) + + class FakeRuntime: + def __init__(self) -> None: + self.snapshot_calls = 0 + + def snapshot(self): + self.snapshot_calls += 1 + return snapshot + + runtime = FakeRuntime() + module = object.__new__(ManipulationModule) + module._execution_runtime = runtime + module._execution_topology = None + + assert module.get_state() == LifecycleState.READY.name + assert module.get_error() == "runtime diagnostic" + assert module.has_planned_path() is False + projection = module.get_execution_snapshot() + assert projection.ready_plan_id == "ready-1" + assert projection.has_ready_plan is True + assert runtime.snapshot_calls == 4 + assert not hasattr(module, "_execution_state") + assert not hasattr(module, "_execution_error") + + @pytest.mark.skipif(not _drake_available(), reason="Drake not installed") @pytest.mark.skipif(not _xarm_urdf_available(), reason="XArm URDF not available") class TestManipulationModuleIntegration: @@ -128,7 +297,7 @@ class TestManipulationModuleIntegration: def test_module_initialization(self, module): """Test module initializes with real Drake world.""" - assert module._state == ManipulationState.IDLE + assert module.get_state() == LifecycleState.IDLE.name assert module._world_monitor is not None assert module._planner is not None assert module._kinematics is not None @@ -158,13 +327,29 @@ def test_plan_to_joints(self, module, joint_state_zeros): success = module.plan_to_joints(target) assert success is True - assert module._state == ManipulationState.COMPLETED + assert module.get_state() == LifecycleState.READY.name assert module.has_planned_path() is True - assert "test_arm" in module._planned_trajectories - traj = module._planned_trajectories["test_arm"] - assert len(traj.points) > 1 - assert traj.duration > 0 + ready = module._execution_runtime.snapshot().ready_plan + assert isinstance(ready, PreparedPlan) + assert len(ready.generated_plan.trajectory.points) > 1 + assert ready.generated_plan.trajectory.duration > 0 + assert ready.generated_plan.group_ids == ("test_arm/manipulator",) + + def test_plan_to_explicit_joint_target(self, module, joint_state_zeros): + """Test planning to an explicit planning-group joint target.""" + module._on_joint_state(joint_state_zeros) + + success = module.plan_to_joint_targets( + {"test_arm/manipulator": JointState(position=[0.05] * 7)} + ) + + assert success is True + assert module.get_state() == LifecycleState.READY.name + ready = module._execution_runtime.snapshot().ready_plan + assert isinstance(ready, PreparedPlan) + assert ready.generated_plan.group_ids == ("test_arm/manipulator",) + assert module.has_planned_path() is True def test_add_and_remove_obstacle(self, module, joint_state_zeros): """Test adding and removing obstacles.""" @@ -192,6 +377,12 @@ def test_robot_info(self, module): assert info["end_effector_link"] == "link7" assert info["coordinator_task_name"] == "traj_arm" assert info["has_joint_name_mapping"] is True + groups = info["planning_groups"] + assert len(groups) == 1 + assert groups[0].id == "test_arm/manipulator" + + all_groups = module.list_planning_groups() + assert [group.id for group in all_groups] == ["test_arm/manipulator"] def test_ee_pose(self, module, joint_state_zeros): """Test getting end-effector pose.""" @@ -204,26 +395,24 @@ def test_ee_pose(self, module, joint_state_zeros): assert hasattr(pose, "y") assert hasattr(pose, "z") - def test_trajectory_name_translation(self, module, joint_state_zeros): - """Test that trajectory joint names are translated for coordinator.""" + def test_prepared_plan_contains_coordinator_joint_names(self, module, joint_state_zeros): + """Runtime materialization translates selected joints for the gateway.""" module._on_joint_state(joint_state_zeros) success = module.plan_to_joints(JointState(position=[0.05] * 7)) assert success is True - traj = module._planned_trajectories["test_arm"] - robot_config = module._robots["test_arm"][1] - - translated = module._translate_trajectory_to_coordinator(traj, robot_config) - - for name in translated.joint_names: - assert name.startswith("arm_") # Should have arm_ prefix + ready = module._execution_runtime.snapshot().ready_plan + assert isinstance(ready, PreparedPlan) + assert ready.entries[0].request["trajectory"].joint_names == list( + module._robots["test_arm"][1].joint_name_mapping.values() + ) @pytest.mark.skipif(not _drake_available(), reason="Drake not installed") @pytest.mark.skipif(not _xarm_urdf_available(), reason="XArm URDF not available") class TestCoordinatorIntegration: - """Test coordinator integration with mocked RPC client.""" + """Test coordinator integration through the runtime gateway.""" def test_execute_with_mock_coordinator(self, module, joint_state_zeros): """Test execute sends trajectory to coordinator.""" @@ -232,27 +421,18 @@ def test_execute_with_mock_coordinator(self, module, joint_state_zeros): success = module.plan_to_joints(JointState(position=[0.05] * 7)) assert success is True - # Mock the coordinator client - mock_client = MagicMock() - mock_client.task_invoke.return_value = True - module._coordinator_client = mock_client - result = module.execute() assert result is True - assert module._state == ManipulationState.COMPLETED + assert module.get_state() == LifecycleState.RUNNING.name - # Verify coordinator was called - mock_client.task_invoke.assert_called_once() - call_args = mock_client.task_invoke.call_args - task_name, method_name, kwargs = call_args[0] + task_name, request = module._runtime_gateway.execute_calls[0] assert task_name == "traj_arm" - assert method_name == "execute" - trajectory = kwargs["trajectory"] + trajectory = request["trajectory"] assert len(trajectory.points) > 1 - # Joint names should be translated - assert all(n.startswith("arm_") for n in trajectory.joint_names) + robot_config = module._robots["test_arm"][1] + assert trajectory.joint_names == list(robot_config.joint_name_mapping.values()) def test_execute_rejected_by_coordinator(self, module, joint_state_zeros): """Test handling of coordinator rejection.""" @@ -260,39 +440,31 @@ def test_execute_rejected_by_coordinator(self, module, joint_state_zeros): module.plan_to_joints(JointState(position=[0.05] * 7)) - # Mock coordinator to reject - mock_client = MagicMock() - mock_client.task_invoke.return_value = False - module._coordinator_client = mock_client + module._runtime_gateway.execute_outcome = Outcome.REJECTED result = module.execute() assert result is False - assert module._state == ManipulationState.FAULT - assert "rejected" in module._error_message.lower() + assert module.get_state() == LifecycleState.IDLE.name + assert "rejected" in module.get_error().lower() def test_state_transitions_during_execution(self, module, joint_state_zeros): """Test state transitions during plan and execute.""" - assert module._state == ManipulationState.IDLE + assert module.get_state() == LifecycleState.IDLE.name module._on_joint_state(joint_state_zeros) - # Plan - should go through PLANNING -> COMPLETED + # Plan commits a runtime-owned READY plan. module.plan_to_joints(JointState(position=[0.05] * 7)) - assert module._state == ManipulationState.COMPLETED + assert module.get_state() == LifecycleState.READY.name - # Reset works from COMPLETED + # Reset clears the READY plan. module.reset() - assert module._state == ManipulationState.IDLE + assert module.get_state() == LifecycleState.IDLE.name # Plan again module.plan_to_joints(JointState(position=[0.05] * 7)) - # Mock coordinator - mock_client = MagicMock() - mock_client.task_invoke.return_value = True - module._coordinator_client = mock_client - - # Execute - should go to EXECUTING then COMPLETED + # Execute is accepted by the runtime gateway. module.execute() - assert module._state == ManipulationState.COMPLETED + assert module.get_state() == LifecycleState.RUNNING.name diff --git a/dimos/manipulation/test_manipulation_monitor_preview.py b/dimos/manipulation/test_manipulation_monitor_preview.py new file mode 100644 index 0000000000..2e067f14ad --- /dev/null +++ b/dimos/manipulation/test_manipulation_monitor_preview.py @@ -0,0 +1,244 @@ +# 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. + +"""Monitor and preview unit tests for ManipulationModule.""" + +from __future__ import annotations + +from collections.abc import Iterator +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + +from dimos.manipulation._test_manipulation_helpers import ( + close_test_runtimes, + install_runtime, + make_module as _make_module, +) +from dimos.manipulation.execution_runtime import prepare_generated_plan +from dimos.manipulation.manipulation_module import ManipulationModule +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition +from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry +from dimos.manipulation.planning.monitor.world_monitor import WorldMonitor +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.enums import PlanningStatus +from dimos.manipulation.planning.spec.models import GeneratedPlan +from dimos.manipulation.planning.spec.protocols import VisualizationSpec +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.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint + + +@pytest.fixture(autouse=True) +def _close_test_runtimes() -> Iterator[None]: + yield + close_test_runtimes() + + +@pytest.fixture +def robot_config_with_mapping() -> RobotModelConfig: + """Create a robot config with joint name mapping.""" + return RobotModelConfig( + name="left_arm", + model_path=Path("/path/to/robot.urdf"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + joint_names=["joint1", "joint2", "joint3"], + base_link="link_base", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=("joint1", "joint2", "joint3"), + base_link="link_base", + tip_link="link_tcp", + ) + ], + joint_name_mapping={ + "left/joint1": "joint1", + "left/joint2": "joint2", + "left/joint3": "joint3", + }, + coordinator_task_name="traj_left", + ) + + +def _one_joint_config(name: str = "arm") -> RobotModelConfig: + return RobotModelConfig( + name=name, + model_path=Path("/path"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + joint_names=["j0"], + base_link="base_link", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", joint_names=("j0",), base_link="base_link", tip_link="ee" + ) + ], + coordinator_task_name=f"traj_{name}", + ) + + +def _install_generated_plan( + module: ManipulationModule, + config: RobotModelConfig, + traj_gen: MagicMock, + *points: list[float], +) -> None: + """Install a generated plan and enough monitor state to derive robot paths.""" + global_joint_names = [f"{config.name}/{joint}" for joint in config.joint_names] + module._robots = {config.name: ("robot_id", config, traj_gen)} + module._world_monitor = MagicMock() + module._world_monitor.planning_groups = PlanningGroupRegistry([config]) + module._world_monitor.get_current_joint_state.return_value = JointState( + name=config.joint_names, + position=[0.0 for _ in config.joint_names], + ) + generated_plan = GeneratedPlan( + trajectory=JointTrajectory( + joint_names=global_joint_names, + points=[ + TrajectoryPoint( + time_from_start=float(index), + positions=list(point), + velocities=[0.0 for _ in config.joint_names], + ) + for index, point in enumerate(points) + ], + ), + group_ids=(f"{config.name}/manipulator",), + status=PlanningStatus.SUCCESS, + path=[ + JointState( + name=global_joint_names, + position=list(point), + ) + for point in points + ], + ) + install_runtime(module, [config]) + token = module._execution_runtime.start_planning() + prepared = prepare_generated_plan(generated_plan, module._execution_topology) + assert token is not None + assert module._execution_runtime.complete_planning(token, prepared).accepted + + +def _make_module_with_monitor(*configs: RobotModelConfig) -> ManipulationModule: + """Create a ManipulationModule with a mocked world monitor and robots configured.""" + module = _make_module() + module._world_monitor = MagicMock() + module._init_joints = {} + for config in configs: + robot_id = f"robot_{config.name}" + module._robots[config.name] = (robot_id, config, MagicMock()) + return module + + +def _make_joint_state(positions: list[float], name: list[str] | None = None) -> JointState: + return JointState(name=name or [f"j{i}" for i in range(len(positions))], position=positions) + + +def _make_path(*points: list[float]) -> list[JointState]: + return [_make_joint_state(list(point)) for point in points] + + +def _make_trajectory(*points: tuple[float, list[float]]) -> JointTrajectory: + joint_names = [f"j{i}" for i in range(len(points[0][1]))] if points else [] + return JointTrajectory( + joint_names=joint_names, + points=[ + TrajectoryPoint(time_from_start=time_from_start, positions=positions) + for time_from_start, positions in points + ], + ) + + +def _make_world_monitor_with_viz(viz: VisualizationSpec | None) -> WorldMonitor: + world = MagicMock() + return WorldMonitor( + world=world, + visualization=viz, + ) + + +class FakeVisualization: + def __init__(self) -> None: + self.close_count = 0 + self.published = False + self.preview_shown: list[str] = [] + self.preview_hidden: list[str] = [] + self.animations: list[tuple[str, list[JointState], float]] = [] + self.preview_animation_cancellations = 0 + + def initialize(self, session) -> None: + pass + + def get_visualization_url(self) -> str | None: + return "123" + + def update_state(self, frame) -> None: + self.published = True + + def animate_trajectory( + self, trajectory: JointTrajectory, duration: float | None = None + ) -> None: + self.animations.append( + ( + tuple(trajectory.joint_names), + list(trajectory.points), + duration if duration is not None else 0.0, + ) + ) + + def cancel_preview_animation(self) -> None: + self.preview_animation_cancellations += 1 + + def close(self) -> None: + self.close_count += 1 + + +class TestOnJointState: + """Test _on_joint_state routing, splitting, and init capture.""" + + pass + + +class TestWorldMonitorVisualization: + pass + + +class TestManipulationPreview: + def test_dismiss_preview_noop_without_monitor(self): + module = _make_module() + + module._dismiss_preview(["arm/manipulator"]) + + def test_dismiss_preview_routes_to_monitor(self): + module = _make_module() + module._world_monitor = MagicMock() + + module._dismiss_preview(["arm/manipulator"]) + + module._world_monitor.cancel_preview_animation.assert_called_once_with() + + def test_preview_rejects_unaffected_compatibility_robot(self): + module = _make_module() + config = _one_joint_config() + traj_gen = MagicMock() + _install_generated_plan(module, config, traj_gen, [0.0], [1.0]) + + assert module.preview_plan(robot_name="other") is False + module._world_monitor.animate_trajectory.assert_not_called() diff --git a/dimos/manipulation/test_manipulation_unit.py b/dimos/manipulation/test_manipulation_unit.py index bc12bfe994..fc00666169 100644 --- a/dimos/manipulation/test_manipulation_unit.py +++ b/dimos/manipulation/test_manipulation_unit.py @@ -16,24 +16,32 @@ from __future__ import annotations +from collections.abc import Iterator from pathlib import Path -import threading from unittest.mock import MagicMock import pytest from pytest_mock import MockerFixture +from dimos.manipulation._test_manipulation_helpers import ( + close_test_runtimes, + install_runtime, + make_module as _make_module, +) +from dimos.manipulation.execution_runtime import prepare_generated_plan from dimos.manipulation.manipulation_module import ( ManipulationModule, ManipulationModuleConfig, - ManipulationState, ) +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition +from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry from dimos.manipulation.planning.kinematics.config import PinkKinematicsConfig from dimos.manipulation.planning.monitor.world_monitor import WorldMonitor from dimos.manipulation.planning.spec.config import RobotModelConfig -from dimos.manipulation.planning.spec.enums import IKStatus -from dimos.manipulation.planning.spec.models import IKResult, PlanningSceneInfo -from dimos.manipulation.planning.spec.protocols import VisualizationSpec +from dimos.manipulation.planning.spec.enums import IKStatus, PlanningStatus +from dimos.manipulation.planning.spec.models import ( + GeneratedPlan, +) from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.Quaternion import Quaternion @@ -51,14 +59,27 @@ def robot_config(): model_path=Path("/path/to/robot.urdf"), base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), joint_names=["joint1", "joint2", "joint3"], - end_effector_link="link_tcp", base_link="link_base", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=("joint1", "joint2", "joint3"), + base_link="link_base", + tip_link="link_tcp", + ) + ], max_velocity=1.0, max_acceleration=2.0, coordinator_task_name="traj_arm", ) +@pytest.fixture(autouse=True) +def _close_test_runtimes() -> Iterator[None]: + yield + close_test_runtimes() + + @pytest.fixture def robot_config_with_mapping(): """Create a robot config with joint name mapping (dual-arm scenario).""" @@ -67,8 +88,15 @@ def robot_config_with_mapping(): model_path=Path("/path/to/robot.urdf"), base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), joint_names=["joint1", "joint2", "joint3"], - end_effector_link="link_tcp", base_link="link_base", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=("joint1", "joint2", "joint3"), + base_link="link_base", + tip_link="link_tcp", + ) + ], joint_name_mapping={ "left/joint1": "joint1", "left/joint2": "joint2", @@ -94,90 +122,97 @@ def simple_trajectory(): ) -class _ManipulationModuleHarness(ManipulationModule): - def __init__(self) -> None: - self._state = ManipulationState.IDLE - self._lock = threading.Lock() - self._error_message = "" - self._planning_epoch = 0 - self._robots = {} - self._planned_paths = {} - self._planned_trajectories = {} - self._world_monitor = None - self._planner = None - self._kinematics = None - self._coordinator_client = None - self.config = MagicMock(planning_timeout=10.0) - - -def _make_module() -> ManipulationModule: - """Create a lightweight ManipulationModule harness for behavior tests.""" - return _ManipulationModuleHarness() - - -class TestStateMachine: - """Test state transitions.""" - - def test_cancel_interrupts_active_work(self): - """Cancel works for executing motion and in-progress planning.""" - module = _make_module() - - module._state = ManipulationState.IDLE - assert module.cancel() is False - - module._state = ManipulationState.PLANNING - assert module.cancel() is True - assert module._state == ManipulationState.IDLE - assert module._planning_epoch == 1 - - module._state = ManipulationState.EXECUTING - assert module.cancel() is True - assert module._state == ManipulationState.IDLE - - def test_reset_not_during_execution(self): - """Reset works in any state except EXECUTING.""" - module = _make_module() +def _one_joint_config(name: str = "arm") -> RobotModelConfig: + return RobotModelConfig( + name=name, + model_path=Path("/path"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + joint_names=["j0"], + base_link="base_link", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", joint_names=("j0",), base_link="base_link", tip_link="ee" + ) + ], + coordinator_task_name=f"traj_{name}", + ) - module._state = ManipulationState.FAULT - module._error_message = "Error" - result = module.reset() - assert result.is_success() - assert module._state == ManipulationState.IDLE - assert module._error_message == "" - module._state = ManipulationState.EXECUTING - result = module.reset() - assert not result.is_success() - assert result.error_code == "INVALID_STATE" +def _install_generated_plan( + module: ManipulationModule, + config: RobotModelConfig, + traj_gen: MagicMock, + *points: list[float], +) -> None: + """Install a generated plan and enough monitor state to derive robot paths.""" + global_joint_names = [f"{config.name}/{joint}" for joint in config.joint_names] + module._robots = {config.name: ("robot_id", config, traj_gen)} + module._world_monitor = MagicMock() + module._world_monitor.planning_groups = PlanningGroupRegistry([config]) + module._world_monitor.get_current_joint_state.return_value = JointState( + name=config.joint_names, + position=[0.0 for _ in config.joint_names], + ) + module._world_monitor.current_global_joint_state.return_value = JointState( + name=global_joint_names, + position=[0.0 for _ in config.joint_names], + ) + generated_plan = GeneratedPlan( + trajectory=JointTrajectory( + joint_names=global_joint_names, + points=[ + TrajectoryPoint( + time_from_start=float(index), + positions=list(point), + velocities=[0.0 for _ in config.joint_names], + ) + for index, point in enumerate(points) + ], + ), + group_ids=(f"{config.name}/manipulator",), + status=PlanningStatus.SUCCESS, + path=[ + JointState( + name=global_joint_names, + position=list(point), + ) + for point in points + ], + ) + install_runtime(module, [config]) + token = module._execution_runtime.start_planning() + prepared = prepare_generated_plan(generated_plan, module._execution_topology) + assert token is not None + assert module._execution_runtime.complete_planning(token, prepared).accepted - def test_fail_sets_fault_state(self): - """_fail helper sets FAULT state and message.""" - module = _make_module() - module._state = ManipulationState.PLANNING - result = module._fail("Test error") - assert result is False - assert module._state == ManipulationState.FAULT - assert module._error_message == "Test error" +def _generated_plan_trajectory(joint_names: list[str], *points: list[float]) -> JointTrajectory: + return JointTrajectory( + joint_names=joint_names, + points=[ + TrajectoryPoint( + time_from_start=float(index), + positions=list(point), + velocities=[0.0 for _ in joint_names], + ) + for index, point in enumerate(points) + ], + ) - def test_begin_planning_state_checks(self, robot_config): - """_begin_planning only allowed from IDLE or COMPLETED.""" - module = _make_module() - module._world_monitor = MagicMock() - module._robots = {"test_arm": ("robot_id", robot_config, MagicMock())} - # From IDLE - OK - module._state = ManipulationState.IDLE - assert module._begin_planning() == ("test_arm", "robot_id") - assert module._state == ManipulationState.PLANNING +def _make_trajectory(*points: tuple[float, list[float]]) -> JointTrajectory: + joint_names = [f"j{i}" for i in range(len(points[0][1]))] if points else [] + return JointTrajectory( + joint_names=joint_names, + points=[ + TrajectoryPoint(time_from_start=time_from_start, positions=positions) + for time_from_start, positions in points + ], + ) - # From COMPLETED - OK - module._state = ManipulationState.COMPLETED - assert module._begin_planning() == ("test_arm", "robot_id") - # From EXECUTING - Fail - module._state = ManipulationState.EXECUTING - assert module._begin_planning() is None +class TestStateMachine: + """Test state transitions.""" class TestRobotSelection: @@ -265,6 +300,7 @@ def test_kinematics_config_is_passed_to_factory( kinematics_name=None, kinematics=kinematics, ) + module.stop() def test_legacy_kinematics_name_still_selects_backend( self, robot_config, planning_initialization: PlanningInitializationHarness @@ -285,6 +321,7 @@ def test_legacy_kinematics_name_still_selects_backend( kinematics_name="pink", kinematics=module.config.kinematics, ) + module.stop() def test_nested_kinematics_config_parses_cli_override_shape(self) -> None: """Pydantic parses the nested CLI config shape used by -o overrides.""" @@ -302,155 +339,102 @@ def test_nested_kinematics_config_parses_cli_override_shape(self) -> None: assert config.kinematics.dt == 0.02 assert config.kinematics.posture_cost == 0.0 - def test_solve_ik_rpc_calls_configured_backend(self, robot_config): - """solve_ik returns the backend IKResult without path planning.""" + +class TestPlanningGroupApis: + """Test explicit planning-group API behavior.""" + + def test_list_planning_groups_and_robot_info_include_groups(self, robot_config): module = _make_module() + registry = PlanningGroupRegistry([robot_config]) module._robots = {"test_arm": ("robot_id", robot_config, MagicMock())} module._world_monitor = MagicMock() - module._world_monitor.world = MagicMock() - current = JointState(name=robot_config.joint_names, position=[0.0, 0.0, 0.0]) - module._world_monitor.get_current_joint_state.return_value = current - expected = IKResult( - status=IKStatus.SUCCESS, - joint_state=JointState(name=robot_config.joint_names, position=[0.1, 0.2, 0.3]), - position_error=0.0001, - orientation_error=0.0002, - iterations=3, - message="ok", + module._world_monitor.planning_groups = registry + module._init_joints = {} + + groups = module.list_planning_groups() + info = module.get_robot_info() + + assert [group.id for group in groups] == ["test_arm/manipulator"] + assert info is not None + assert info["planning_groups"] == groups + assert info["end_effector_link"] == "link_tcp" + assert info["has_joint_name_mapping"] is False + + def test_pose_wrappers_fail_safely_without_unique_pose_group(self, robot_config): + no_pose_config = RobotModelConfig( + name="test_arm", + model_path=robot_config.model_path, + base_pose=robot_config.base_pose, + joint_names=robot_config.joint_names, + base_link=robot_config.base_link, + planning_groups=[ + PlanningGroupDefinition( + name="joint_only", + joint_names=("joint1", "joint2", "joint3"), + base_link="link_base", + ) + ], ) - module._kinematics = MagicMock() - module._kinematics.solve.return_value = expected - - pose = Pose(position=Vector3(x=0.45, y=0.0, z=0.25), orientation=Quaternion()) - result = module.solve_ik(pose) - - assert result is expected - assert module._state == ManipulationState.COMPLETED - assert module._planned_paths == {} - module._kinematics.solve.assert_called_once() - _, kwargs = module._kinematics.solve.call_args - assert kwargs["world"] is module._world_monitor.world - assert kwargs["robot_id"] == "robot_id" - assert kwargs["seed"] is current - assert kwargs["check_collision"] is True - assert kwargs["target_pose"].frame_id == "world" - assert kwargs["target_pose"].position.x == 0.45 - - def test_solve_ik_rpc_returns_failure_without_joint_state(self, robot_config): - """solve_ik reports a failed IKResult when no seed state is available.""" module = _make_module() - module._robots = {"test_arm": ("robot_id", robot_config, MagicMock())} + module._robots = {"test_arm": ("robot_id", no_pose_config, MagicMock())} module._world_monitor = MagicMock() - module._world_monitor.get_current_joint_state.return_value = None + module._world_monitor.planning_groups = PlanningGroupRegistry([no_pose_config]) + module._world_monitor.get_ee_pose.side_effect = ValueError("no pose group") module._kinematics = MagicMock() pose = Pose(position=Vector3(x=0.45, y=0.0, z=0.25), orientation=Quaternion()) - result = module.solve_ik(pose) + assert module.get_ee_pose() is None + assert module.plan_to_pose(pose) is False + result = module.inverse_kinematics_single(pose) assert result.status == IKStatus.NO_SOLUTION - assert result.message == "No joint state" - assert module._state == ManipulationState.IDLE - module._kinematics.solve.assert_not_called() - - def test_solve_ik_rpc_uses_explicit_seed(self, robot_config): - """solve_ik initializes the backend from an explicit seed when provided.""" + assert "no pose-targetable planning group" in result.message + + def test_pose_wrappers_fail_safely_with_multiple_pose_groups(self, robot_config): + multi_pose_config = RobotModelConfig( + name="test_arm", + model_path=robot_config.model_path, + base_pose=robot_config.base_pose, + joint_names=robot_config.joint_names, + base_link=robot_config.base_link, + planning_groups=[ + PlanningGroupDefinition( + name="wrist", + joint_names=("joint1", "joint2"), + base_link="link_base", + tip_link="link_wrist", + ), + PlanningGroupDefinition( + name="tool", + joint_names=("joint1", "joint2", "joint3"), + base_link="link_base", + tip_link="link_tcp", + ), + ], + ) module = _make_module() - module._robots = {"test_arm": ("robot_id", robot_config, MagicMock())} + module._robots = {"test_arm": ("robot_id", multi_pose_config, MagicMock())} module._world_monitor = MagicMock() - module._world_monitor.world = MagicMock() - module._world_monitor.get_current_joint_state.return_value = JointState( - name=robot_config.joint_names, position=[0.0, 0.0, 0.0] - ) - explicit_seed = JointState(name=robot_config.joint_names, position=[0.2, 0.1, 0.0]) - expected = IKResult(status=IKStatus.SUCCESS, joint_state=explicit_seed) + module._world_monitor.planning_groups = PlanningGroupRegistry([multi_pose_config]) + module._world_monitor.get_ee_pose.side_effect = ValueError("multiple pose groups") module._kinematics = MagicMock() - module._kinematics.solve.return_value = expected pose = Pose(position=Vector3(x=0.45, y=0.0, z=0.25), orientation=Quaternion()) - result = module.solve_ik(pose, seed=explicit_seed) - assert result is expected - _, kwargs = module._kinematics.solve.call_args - assert kwargs["seed"] is explicit_seed - module._world_monitor.get_current_joint_state.assert_not_called() + assert module.get_ee_pose() is None + assert module.plan_to_pose(pose) is False + result = module.inverse_kinematics_single(pose) + assert result.status == IKStatus.NO_SOLUTION + assert "2 pose-targetable planning groups" in result.message class TestJointNameTranslation: """Test trajectory joint name translation for coordinator.""" - def test_no_mapping_returns_original(self, robot_config, simple_trajectory): - """Without mapping, trajectory is returned unchanged.""" - module = _make_module() - - result = module._translate_trajectory_to_coordinator(simple_trajectory, robot_config) - assert result is simple_trajectory # Same object - - def test_mapping_translates_names(self, robot_config_with_mapping, simple_trajectory): - """With mapping, joint names are translated.""" - module = _make_module() - - result = module._translate_trajectory_to_coordinator( - simple_trajectory, robot_config_with_mapping - ) - assert result.joint_names == ["left/joint1", "left/joint2", "left/joint3"] - assert len(result.points) == 2 # Points preserved - class TestExecute: """Test coordinator execution.""" - def test_execute_requires_trajectory(self, robot_config): - """Execute fails without planned trajectory.""" - module = _make_module() - module._robots = {"test_arm": ("id", robot_config, MagicMock())} - module._planned_trajectories = {} - - assert module.execute() is False - - def test_execute_requires_task_name(self): - """Execute fails without coordinator_task_name.""" - module = _make_module() - config_no_task = RobotModelConfig( - name="arm", - model_path=Path("/path"), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), - joint_names=["j1"], - end_effector_link="ee", - ) - module._robots = {"arm": ("id", config_no_task, MagicMock())} - module._planned_trajectories = {"arm": MagicMock()} - - assert module.execute() is False - - def test_execute_success(self, robot_config, simple_trajectory): - """Successful execute calls coordinator via task_invoke.""" - module = _make_module() - module._robots = {"test_arm": ("id", robot_config, MagicMock())} - module._planned_trajectories = {"test_arm": simple_trajectory} - - mock_client = MagicMock() - mock_client.task_invoke.return_value = True - module._coordinator_client = mock_client - - assert module.execute() is True - assert module._state == ManipulationState.COMPLETED - mock_client.task_invoke.assert_called_once_with( - "traj_arm", "execute", {"trajectory": simple_trajectory} - ) - - def test_execute_rejected(self, robot_config, simple_trajectory): - """Rejected execution sets FAULT state.""" - module = _make_module() - module._robots = {"test_arm": ("id", robot_config, MagicMock())} - module._planned_trajectories = {"test_arm": simple_trajectory} - - mock_client = MagicMock() - mock_client.task_invoke.return_value = False - module._coordinator_client = mock_client - - assert module.execute() is False - assert module._state == ManipulationState.FAULT - class TestRobotModelConfigMapping: """Test RobotModelConfig joint name mapping helpers.""" @@ -466,317 +450,3 @@ def test_bidirectional_mapping(self, robot_config_with_mapping): # URDF -> Coordinator assert config.get_coordinator_joint_name("joint1") == "left/joint1" assert config.get_coordinator_joint_name("unknown") == "unknown" - - -def _make_module_with_monitor(*configs: RobotModelConfig) -> ManipulationModule: - """Create a ManipulationModule with a mocked world monitor and robots configured.""" - module = _make_module() - module._world_monitor = MagicMock() - module._init_joints = {} - for config in configs: - robot_id = f"robot_{config.name}" - module._robots[config.name] = (robot_id, config, MagicMock()) - return module - - -def _make_joint_state(positions: list[float], name: list[str] | None = None) -> JointState: - return JointState(name=name or [f"j{i}" for i in range(len(positions))], position=positions) - - -def _make_path(*points: list[float]) -> list[JointState]: - return [_make_joint_state(list(point)) for point in points] - - -def _make_trajectory(*points: tuple[float, list[float]]) -> JointTrajectory: - joint_names = [f"j{i}" for i in range(len(points[0][1]))] if points else [] - return JointTrajectory( - joint_names=joint_names, - points=[ - TrajectoryPoint(time_from_start=time_from_start, positions=positions) - for time_from_start, positions in points - ], - ) - - -def _make_world_monitor_with_viz(viz: VisualizationSpec | None) -> WorldMonitor: - world = MagicMock() - return WorldMonitor( - world=world, - visualization=viz, - ) - - -class FakeVisualization: - def __init__(self) -> None: - self.close_count = 0 - self.published = False - self.preview_shown: list[str] = [] - self.preview_hidden: list[str] = [] - self.animations: list[tuple[str, list[JointState], float]] = [] - - def initialize_scene(self, scene: PlanningSceneInfo) -> None: - pass - - def get_visualization_url(self) -> str | None: - return "123" - - def publish_visualization(self, ctx: object | None = None) -> None: - self.published = True - - def show_preview(self, robot_id: str) -> None: - self.preview_shown.append(robot_id) - - def hide_preview(self, robot_id: str) -> None: - self.preview_hidden.append(robot_id) - - def animate_path(self, robot_id: str, path: list[JointState], duration: float = 3.0) -> None: - self.animations.append((robot_id, path, duration)) - - def close(self) -> None: - self.close_count += 1 - - -class TestOnJointState: - """Test _on_joint_state routing, splitting, and init capture.""" - - def test_routes_positions_to_monitor(self, robot_config_with_mapping): - """Joint positions from aggregated message are routed to the correct monitor.""" - module = _make_module_with_monitor(robot_config_with_mapping) - - msg = JointState( - name=["left/joint1", "left/joint2", "left/joint3"], - position=[0.1, 0.2, 0.3], - velocity=[1.0, 2.0, 3.0], - ) - module._on_joint_state(msg) - - # Verify world_monitor received the sub-message - module._world_monitor.on_joint_state.assert_called_once() - call_args = module._world_monitor.on_joint_state.call_args - sub_msg = call_args[0][0] - assert sub_msg.position == [0.1, 0.2, 0.3] - assert sub_msg.velocity == [1.0, 2.0, 3.0] - assert call_args[1]["robot_id"] == "robot_left_arm" - - def test_skips_robot_with_missing_joints(self, robot_config_with_mapping): - """Robots whose joints are absent from the message are skipped.""" - module = _make_module_with_monitor(robot_config_with_mapping) - - # Message has none of left_arm's joints - msg = JointState( - name=["right/joint1", "right/joint2"], - position=[0.5, 0.6], - ) - module._on_joint_state(msg) - - module._world_monitor.on_joint_state.assert_not_called() - - def test_captures_init_joints_on_first_call(self, robot_config_with_mapping): - """First joint state is stored as init joints; subsequent calls don't overwrite.""" - module = _make_module_with_monitor(robot_config_with_mapping) - - first_msg = JointState( - name=["left/joint1", "left/joint2", "left/joint3"], - position=[0.1, 0.2, 0.3], - ) - module._on_joint_state(first_msg) - assert "left_arm" in module._init_joints - assert module._init_joints["left_arm"].position == [0.1, 0.2, 0.3] - - # Second call should NOT overwrite - second_msg = JointState( - name=["left/joint1", "left/joint2", "left/joint3"], - position=[0.9, 0.8, 0.7], - ) - module._on_joint_state(second_msg) - assert module._init_joints["left_arm"].position == [0.1, 0.2, 0.3] - - def test_multi_robot_splits_correctly(self): - """With two robots, each gets only its own joints from the aggregated message.""" - left_config = RobotModelConfig( - name="left", - model_path=Path("/path/to/robot.urdf"), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), - joint_names=["j1", "j2"], - end_effector_link="ee", - base_link="base", - joint_name_mapping={"left/j1": "j1", "left/j2": "j2"}, - coordinator_task_name="traj_left", - ) - right_config = RobotModelConfig( - name="right", - model_path=Path("/path/to/robot.urdf"), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), - joint_names=["j1", "j2"], - end_effector_link="ee", - base_link="base", - joint_name_mapping={"right/j1": "j1", "right/j2": "j2"}, - coordinator_task_name="traj_right", - ) - module = _make_module_with_monitor(left_config, right_config) - - msg = JointState( - name=["left/j1", "left/j2", "right/j1", "right/j2"], - position=[1.0, 2.0, 3.0, 4.0], - velocity=[0.1, 0.2, 0.3, 0.4], - ) - module._on_joint_state(msg) - - assert module._world_monitor.on_joint_state.call_count == 2 - - # Collect calls by robot_id - calls = { - call[1]["robot_id"]: call[0][0] - for call in module._world_monitor.on_joint_state.call_args_list - } - assert calls["robot_left"].position == [1.0, 2.0] - assert calls["robot_right"].position == [3.0, 4.0] - assert calls["robot_left"].velocity == [0.1, 0.2] - assert calls["robot_right"].velocity == [0.3, 0.4] - - def test_no_monitor_returns_early(self, robot_config_with_mapping): - """When world_monitor is None, _on_joint_state returns without error.""" - module = _make_module() - module._robots = {"left_arm": ("id", robot_config_with_mapping, MagicMock())} - module._world_monitor = None - - # Should not raise - msg = JointState( - name=["left/joint1", "left/joint2", "left/joint3"], - position=[0.1, 0.2, 0.3], - ) - module._on_joint_state(msg) - - -class TestWorldMonitorVisualization: - def test_visualization_routing_and_stop_all_monitors(self): - viz = FakeVisualization() - monitor = _make_world_monitor_with_viz(viz) - state_monitor = MagicMock() - obstacle_monitor = MagicMock() - monitor._state_monitors = {"robot": state_monitor} - monitor._obstacle_monitor = obstacle_monitor - monitor._viz_thread = MagicMock() - monitor._viz_thread.is_alive.return_value = False - - assert monitor.get_visualization_url() == "123" - monitor.publish_visualization() - monitor.show_preview("robot") - monitor.hide_preview("robot") - path = _make_path([1.0], [2.0], [3.0]) - monitor.animate_path("robot", path, 4.5) - assert monitor.visualization is viz - assert viz.published is True - assert viz.preview_shown == ["robot"] - assert viz.preview_hidden == ["robot"] - assert viz.animations == [("robot", path, 4.5)] - - monitor.stop_all_monitors() - - assert viz.close_count == 1 - state_monitor.stop.assert_called_once() - obstacle_monitor.stop.assert_called_once() - - def test_visualization_none_is_noop(self): - monitor = _make_world_monitor_with_viz(None) - - assert monitor.get_visualization_url() is None - monitor.publish_visualization() - monitor.show_preview("robot") - monitor.hide_preview("robot") - monitor.animate_path("robot", [1], 1.0) - monitor.start_visualization_thread() - assert monitor._viz_thread is None - - -class TestManipulationPreview: - def test_dismiss_preview_noop_without_monitor(self): - module = _make_module() - - module._dismiss_preview("robot_id") - - def test_dismiss_preview_routes_to_monitor(self): - module = _make_module() - module._world_monitor = MagicMock() - - module._dismiss_preview("robot_id") - - module._world_monitor.hide_preview.assert_called_once_with("robot_id") - module._world_monitor.publish_visualization.assert_called_once_with() - - def test_preview_path_uses_trajectory_duration_and_interpolates(self): - module = _make_module() - module._world_monitor = MagicMock() - module._robots = {"arm": ("robot_id", MagicMock(), MagicMock())} - module._planned_paths = {"arm": _make_path([0.0], [2.0])} - module._planned_trajectories = {"arm": _make_trajectory((0.0, [0.0]), (2.0, [2.0]))} - - assert module.preview_path(robot_name="arm", target_fps=2.0) is True - - module._world_monitor.animate_path.assert_called_once() - robot_id, preview_path, duration = module._world_monitor.animate_path.call_args.args - assert robot_id == "robot_id" - assert duration == 2.0 - assert [state.position for state in preview_path] == [[0.0], [0.5], [1.0], [1.5], [2.0]] - - def test_preview_path_explicit_duration_overrides_and_fps_densifies(self): - module = _make_module() - module._world_monitor = MagicMock() - module._robots = {"arm": ("robot_id", MagicMock(), MagicMock())} - module._planned_paths = {"arm": _make_path([0.0], [9.0])} - module._planned_trajectories = {"arm": _make_trajectory((0.0, [0.0]), (9.0, [9.0]))} - - assert module.preview_path(duration=1.5, robot_name="arm", target_fps=2.0) is True - - module._world_monitor.animate_path.assert_called_once() - robot_id, preview_path, duration = module._world_monitor.animate_path.call_args.args - assert robot_id == "robot_id" - assert duration == 1.5 - assert [state.position for state in preview_path] == [[0.0], [3.0], [6.0], [9.0]] - - def test_preview_path_missing_trajectory_uses_default_duration(self): - module = _make_module() - module._world_monitor = MagicMock() - module._robots = {"arm": ("robot_id", MagicMock(), MagicMock())} - module._planned_paths = {"arm": _make_path([0.0], [1.0])} - module._planned_trajectories = {} - - assert module.preview_path(robot_name="arm", target_fps=10.0) is True - - module._world_monitor.animate_path.assert_called_once_with( - "robot_id", module._planned_paths["arm"], 3.0 - ) - - def test_preview_path_skips_interpolation_for_nonpositive_fps_or_duration(self): - module = _make_module() - module._world_monitor = MagicMock() - module._robots = {"arm": ("robot_id", MagicMock(), MagicMock())} - module._planned_paths = {"arm": _make_path([0.0], [1.0])} - module._planned_trajectories = {"arm": _make_trajectory((0.0, [0.0]), (2.0, [1.0]))} - - assert module.preview_path(robot_name="arm", target_fps=0.0) is True - assert module.preview_path(duration=0.0, robot_name="arm", target_fps=20.0) is True - - assert ( - module._world_monitor.animate_path.call_args_list[0].args[1] - == module._planned_paths["arm"] - ) - assert ( - module._world_monitor.animate_path.call_args_list[1].args[1] - == module._planned_paths["arm"] - ) - - def test_preview_path_returns_false_for_missing_inputs(self): - module = _make_module() - module._planned_paths = {"arm": _make_path([0.0], [1.0])} - module._robots = {"arm": ("robot_id", MagicMock(), MagicMock())} - - assert module.preview_path(robot_name="arm") is False - - module._world_monitor = MagicMock() - module._robots = {} - assert module.preview_path(robot_name="arm") is False - - module._robots = {"arm": ("robot_id", MagicMock(), MagicMock())} - module._planned_paths = {"arm": []} - assert module.preview_path(robot_name="arm") is False diff --git a/dimos/manipulation/test_plan_execution_reservation.py b/dimos/manipulation/test_plan_execution_reservation.py new file mode 100644 index 0000000000..3bf29c883b --- /dev/null +++ b/dimos/manipulation/test_plan_execution_reservation.py @@ -0,0 +1,306 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); + +"""Module-side tests for runtime-owned plan submission and lifecycle handling.""" + +from collections.abc import Iterator +from inspect import signature +from pathlib import Path +from unittest.mock import MagicMock + +from pydantic import ValidationError +import pytest + +from dimos.manipulation._test_manipulation_helpers import ( + FakeCoordinatorGateway, + close_test_runtimes, + install_ready_plan, + install_runtime, + make_module, +) +from dimos.manipulation.execution_runtime import ( + ExecutionRuntime, + LifecycleState, + Outcome, + prepare_generated_plan, +) +from dimos.manipulation.manipulation_module import ( + ManipulationExecutionSnapshot, + ManipulationModule, + ManipulationModuleConfig, +) +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition +from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.models import GeneratedPlan +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.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint + + +def test_physical_operation_timeout_config_defaults_to_sixty_and_rejects_invalid_values() -> None: + assert ManipulationModuleConfig().physical_operation_timeout == 60.0 + for value in (0.0, -1.0, float("inf"), float("nan")): + with pytest.raises(ValidationError): + ManipulationModuleConfig(physical_operation_timeout=value) + + +def test_module_runtime_factory_receives_distinct_action_and_physical_timeouts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + module = ManipulationModule( + robots=[_config()], + planning_timeout=3.0, + physical_operation_timeout=17.0, + visualization={"backend": "none"}, + ) + module.coordinator_joint_state = None + captured: dict[str, float] = {} + runtimes: list[ExecutionRuntime] = [] + + class WorldMonitor: + visualization = None + + def add_robot(self, config: RobotModelConfig) -> str: + del config + return "robot-id" + + def finalize(self) -> None: ... + + def start_state_monitor(self, robot_id: str) -> None: + del robot_id + + def set_visualization(self, visualization: object) -> None: + del visualization + + def initialize_visualization(self, operator: object) -> None: + del operator + + def stop_all_monitors(self) -> None: ... + + world_monitor = WorldMonitor() + specs = type( + "PlanningSpecs", + (), + {"world_monitor": world_monitor, "planner": object(), "kinematics": object()}, + )() + + def runtime_factory(_gateway_factory, **kwargs): + captured.update(kwargs) + runtime = ExecutionRuntime(lambda: FakeCoordinatorGateway(), **kwargs) + runtimes.append(runtime) + return runtime + + monkeypatch.setattr("dimos.manipulation.manipulation_module.create_world", lambda **_: object()) + monkeypatch.setattr( + "dimos.manipulation.manipulation_module.create_planning_specs", + lambda **_: specs, + ) + monkeypatch.setattr( + "dimos.manipulation.manipulation_module.create_manipulation_visualization", + lambda *_, **__: None, + ) + monkeypatch.setattr(ManipulationModule, "_runtime_factory", staticmethod(runtime_factory)) + + try: + module.start() + assert captured["action_timeout"] == 3.0 + assert captured["physical_operation_timeout"] == 17.0 + finally: + module.stop() + + +def test_public_move_to_joints_reports_timeout_fault_without_extra_cancel( + monkeypatch: pytest.MonkeyPatch, +) -> None: + gateway = FakeCoordinatorGateway( + status_outcome=Outcome.RUNNING, + status_outcomes=[Outcome.RUNNING], + cancel_outcome=Outcome.UNKNOWN, + ) + module, _ = _module(physical_operation_timeout=0.02, gateway=gateway) + plan = _plan() + + def plan_to_joints(_goal, _robot_name=None) -> bool: + install_ready_plan(module, plan) + return True + + monkeypatch.setattr(module, "plan_to_joints", plan_to_joints) + monkeypatch.setattr(module, "preview_path", lambda *args, **kwargs: True) + + result = module.move_to_joints("1.0") + + assert not result.is_success() + assert result.error_code == "EXECUTION_FAILED" + assert "uncertain" in result.message.lower() + assert module.get_state() == LifecycleState.FAULT.name + assert module.get_error() == result.message + snapshot = module.get_execution_snapshot() + assert snapshot.state == LifecycleState.FAULT.name + assert snapshot.diagnostic == result.message + assert gateway.cancel_calls == ["traj_arm"] + + module.get_state() + assert gateway.cancel_calls == ["traj_arm"] + + +def _config(task_name: str | None = "traj_arm") -> RobotModelConfig: + return RobotModelConfig( + name="arm", + model_path=Path("/path/to/robot.urdf"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + joint_names=["j0"], + base_link="base", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", joint_names=("j0",), base_link="base", tip_link="tool" + ) + ], + max_velocity=1.0, + max_acceleration=1.0, + coordinator_task_name=task_name, + ) + + +def _plan(group_id: str = "arm/manipulator") -> GeneratedPlan: + return GeneratedPlan( + group_ids=(group_id,), + path=[ + JointState(name=["arm/j0"], position=[0.0]), + JointState(name=["arm/j0"], position=[1.0]), + ], + trajectory=JointTrajectory( + joint_names=["arm/j0"], + points=[ + TrajectoryPoint(time_from_start=0.0, positions=[0.0], velocities=[0.0]), + TrajectoryPoint(time_from_start=1.0, positions=[1.0], velocities=[0.0]), + ], + ), + ) + + +def _module( + task_name: str | None = "traj_arm", + physical_operation_timeout: float = 60.0, + gateway: FakeCoordinatorGateway | None = None, +) -> tuple[ManipulationModule, FakeCoordinatorGateway]: + module = make_module() + module.config.physical_operation_timeout = physical_operation_timeout + config = _config(task_name) + module._robots = {"arm": ("robot_id", config, object())} + module._world_monitor = MagicMock() + module._world_monitor.planning_groups = PlanningGroupRegistry([config]) + return module, install_runtime(module, [config], gateway) + + +@pytest.fixture(autouse=True) +def _close_runtime() -> Iterator[None]: + yield + close_test_runtimes() + + +def test_execution_snapshot_is_atomic_and_reports_ready_identity() -> None: + module, _gateway = _module() + runtime = module._execution_runtime + original_snapshot = runtime.snapshot + runtime.snapshot = MagicMock(wraps=original_snapshot) + + view = module.get_execution_snapshot() + assert isinstance(view, ManipulationExecutionSnapshot) + assert view.state == LifecycleState.IDLE.name + assert view.diagnostic == "" + assert view.ready_plan_id is None + assert not view.has_ready_plan + runtime.snapshot.assert_called_once_with() + + runtime.snapshot.reset_mock() + token = runtime.start_planning() + prepared = prepare_generated_plan(_plan(), module._execution_topology) + assert runtime.complete_planning(token, prepared).accepted + runtime.snapshot.reset_mock() + view = module.get_execution_snapshot() + assert view.state == LifecycleState.READY.name + assert view.ready_plan_id is not None + assert view.has_ready_plan + runtime.snapshot.assert_called_once_with() + + +def test_execution_snapshot_is_safe_before_runtime_initialization() -> None: + view = make_module().get_execution_snapshot() + assert view == ManipulationExecutionSnapshot("IDLE", "", None, False) + + +def test_execution_api_has_no_robot_name_and_uses_exact_ready_identity() -> None: + module, gateway = _module() + plan = _plan() + token = module._execution_runtime.start_planning() + prepared = prepare_generated_plan(plan, module._execution_topology) + assert module._execution_runtime.complete_planning(token, prepared).accepted + + assert "robot_name" not in signature(module.execute).parameters + assert "robot_name" not in signature(module.execute_plan).parameters + assert module.execute_plan(plan=plan) is True + assert len(gateway.execute_calls) == 1 + assert module.get_state() in (LifecycleState.RUNNING.name, LifecycleState.IDLE.name) + assert module.cancel() is True + + +def test_trajectory_status_uses_operation_handle_identity() -> None: + module, _gateway = _module() + token = module._execution_runtime.start_planning() + prepared = prepare_generated_plan(_plan(), module._execution_topology) + assert module._execution_runtime.complete_planning(token, prepared).accepted + + result = module._execution_runtime.execute_ready() + assert result.accepted and result.value is not None + operation = module._execution_runtime.snapshot().operation + assert operation is not None + status = module.get_trajectory_status() + assert status is not None + assert status["operation_id"] == operation.handle.operation_id + assert module.cancel() is True + + +def test_execute_plan_explicit_plan_prepares_without_using_other_ready_plan() -> None: + module, gateway = _module() + ready = _plan() + explicit = _plan() + token = module._execution_runtime.start_planning() + prepared = prepare_generated_plan(ready, module._execution_topology) + assert module._execution_runtime.complete_planning(token, prepared).accepted + + assert module.execute_plan(plan=explicit) is True + assert len(gateway.execute_calls) == 1 + assert module.cancel() is True + + +def test_plan_without_coordinator_task_fails_during_preparation_without_dispatch() -> None: + module, gateway = _module(task_name=None) + assert module.execute_plan(plan=_plan()) is False + assert gateway.execute_calls == [] + assert module.get_state() == LifecycleState.IDLE.name + + +def test_cancel_ready_clears_plan_and_status_is_snapshot_only() -> None: + module, _gateway = _module() + plan = _plan() + token = module._execution_runtime.start_planning() + prepared = prepare_generated_plan(plan, module._execution_topology) + assert module._execution_runtime.complete_planning(token, prepared).accepted + assert module.get_trajectory_status() is not None + assert module.cancel() is True + assert module._execution_runtime.snapshot().ready_plan is None + + +def test_reset_fault_and_runtime_shutdown_are_owned_by_runtime() -> None: + module, gateway = _module() + gateway.execute_outcome = Outcome.REJECTED + assert module.execute_plan(plan=_plan()) is False + assert module.get_state() == LifecycleState.IDLE.name + assert module.reset().is_success() + assert module.get_state() == LifecycleState.IDLE.name + module.stop() diff --git a/dimos/manipulation/test_planning_factory.py b/dimos/manipulation/test_planning_factory.py index e695d9579f..99771f9fea 100644 --- a/dimos/manipulation/test_planning_factory.py +++ b/dimos/manipulation/test_planning_factory.py @@ -32,6 +32,7 @@ create_world, validate_backend_combination, ) +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.kinematics.config import JacobianKinematicsConfig from dimos.manipulation.planning.kinematics.jacobian_ik import JacobianIK from dimos.manipulation.planning.planners.rrt_planner import RRTConnectPlanner @@ -64,7 +65,14 @@ def robot_config() -> RobotModelConfig: model_path=Path("/path/to/robot.urdf"), base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), # type: ignore[call-arg] joint_names=["joint1", "joint2"], - end_effector_link="tcp", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=("joint1", "joint2"), + base_link="base_link", + tip_link="tcp", + ) + ], coordinator_task_name="traj_arm", ) diff --git a/dimos/manipulation/test_roboplan_world.py b/dimos/manipulation/test_roboplan_world.py index 68ad241067..f24d010ecc 100644 --- a/dimos/manipulation/test_roboplan_world.py +++ b/dimos/manipulation/test_roboplan_world.py @@ -25,6 +25,10 @@ import numpy as np import pytest +from dimos.manipulation.planning.groups.models import ( + PlanningGroupDefinition, + PlanningGroupSelection, +) from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import ObstacleType, PlanningStatus from dimos.manipulation.planning.spec.models import Obstacle @@ -252,7 +256,14 @@ def robot_config(tmp_path: Path) -> RobotModelConfig: model_path=model_path, base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), # type: ignore[call-arg] joint_names=["joint1", "joint2"], - end_effector_link="tcp", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=("joint1", "joint2"), + base_link="base", + tip_link="tcp", + ) + ], joint_limits_lower=[-1.0, -2.0], joint_limits_upper=[1.0, 2.0], ) @@ -367,6 +378,37 @@ def test_joint_name_mapping_is_applied_to_input_states( assert live_round_trip.position == [0.2, 0.3] +def test_global_joint_names_are_mapped_without_regressing_coordinator_names( + fake_roboplan: None, robot_config: RobotModelConfig +) -> None: + robot_config.joint_name_mapping = {"arm/j1": "joint1", "arm/j2": "joint2"} + world, robot_id = _make_world(fake_roboplan, robot_config) + world.finalize() + + world.sync_from_joint_state( + robot_id, JointState(name=["arm/j1", "arm/j2"], position=[0.4, 0.5]) + ) + assert world.get_joint_state(world.get_live_context(), robot_id).position == [0.4, 0.5] + + world.sync_from_joint_state( + robot_id, JointState(name=["arm/joint1", "arm/joint2"], position=[0.2, 0.3]) + ) + assert world.get_joint_state(world.get_live_context(), robot_id).position == [0.2, 0.3] + + +def test_duplicate_resolved_joint_names_fail_clearly( + fake_roboplan: None, robot_config: RobotModelConfig +) -> None: + robot_config.joint_name_mapping = {"alias": "joint1"} + world, robot_id = _make_world(fake_roboplan, robot_config) + world.finalize() + + with pytest.raises(ValueError, match="duplicate joint 'joint1'"): + world.sync_from_joint_state( + robot_id, JointState(name=["joint1", "alias"], position=[0.1, 0.2]) + ) + + def test_obstacle_mutation_updates_scene_and_stored_pose( fake_roboplan: None, robot_config: RobotModelConfig ) -> None: @@ -453,6 +495,204 @@ def test_fk_jacobian_and_explicit_min_distance_unsupported( world.get_min_distance(ctx, robot_id) +def test_group_fk_and_jacobian_use_group_tip_and_local_joint_order( + fake_roboplan: None, robot_config: RobotModelConfig, monkeypatch: pytest.MonkeyPatch +) -> None: + config = robot_config.model_copy( + update={ + "joint_names": ["joint1", "joint2", "joint3"], + "planning_groups": [ + PlanningGroupDefinition( + name="wrist", + joint_names=("joint3", "joint1"), + base_link="base", + tip_link="wrist_tip", + ) + ], + "joint_limits_lower": [-1.0, -2.0, -3.0], + "joint_limits_upper": [1.0, 2.0, 3.0], + } + ) + monkeypatch.setattr(FakeScene, "joint_group_joint_names", ["joint2", "joint1", "joint3"]) + monkeypatch.setattr(FakeScene, "position_limits_lower", [-2.0, -1.0, -3.0]) + monkeypatch.setattr(FakeScene, "position_limits_upper", [2.0, 1.0, 3.0]) + fk_frames: list[str] = [] + + def fake_fk( + self: FakeScene, q: np.ndarray, frame_name: str, base_frame: str = "" + ) -> np.ndarray: + fk_frames.append(frame_name) + mat = np.eye(4) + mat[0, 3] = float(np.sum(q)) + return mat + + def fake_jacobian( + self: FakeScene, q: np.ndarray, frame_name: str, local: bool = True + ) -> np.ndarray: + assert frame_name == "wrist_tip" + assert local is True + return np.arange(18, dtype=np.float64).reshape(6, 3) + + monkeypatch.setattr(FakeScene, "forwardKinematics", fake_fk) + monkeypatch.setattr(FakeScene, "computeFrameJacobian", fake_jacobian) + world, robot_id = _make_world(fake_roboplan, config) + world.finalize() + ctx = world.get_live_context() + world.set_joint_state( + ctx, + robot_id, + JointState({"name": ["joint1", "joint2", "joint3"], "position": [1.0, 2.0, 3.0]}), + ) + + pose = world.get_group_ee_pose(ctx, "arm/wrist") + jacobian = world.get_group_jacobian(ctx, "arm/wrist") + + assert fk_frames == ["wrist_tip"] + assert pose.position.x == pytest.approx(6.0) + np.testing.assert_allclose(jacobian, np.arange(18, dtype=np.float64).reshape(6, 3)[:, [2, 1]]) + + +def test_group_kinematics_reject_missing_tip_or_missing_context( + fake_roboplan: None, robot_config: RobotModelConfig +) -> None: + no_tip_config = robot_config.model_copy( + update={ + "planning_groups": [ + PlanningGroupDefinition( + name="joint_only", joint_names=("joint1", "joint2"), base_link="base" + ) + ] + } + ) + world, robot_id = _make_world(fake_roboplan, no_tip_config) + world.finalize() + + with pytest.raises(ValueError, match="no tip link"): + world.get_group_ee_pose(world.get_live_context(), "arm/joint_only") + with pytest.raises(ValueError, match="no tip link"): + world.get_group_jacobian(world.get_live_context(), "arm/joint_only") + + ctx = world.get_live_context() + del ctx.q_by_robot[robot_id] + with pytest.raises(KeyError, match=robot_id): + world.get_link_pose(ctx, robot_id, "tcp") + + jacobian_world, jacobian_robot_id = _make_world(fake_roboplan, robot_config) + jacobian_world.finalize() + jacobian_ctx = jacobian_world.get_live_context() + del jacobian_ctx.q_by_robot[jacobian_robot_id] + with pytest.raises(KeyError, match=jacobian_robot_id): + jacobian_world.get_group_jacobian(jacobian_ctx, "arm/manipulator") + + +def test_group_jacobian_validates_projection_shape_and_joint_names( + fake_roboplan: None, + robot_config: RobotModelConfig, + monkeypatch: pytest.MonkeyPatch, +) -> None: + world, robot_id = _make_world(fake_roboplan, robot_config) + world.finalize() + ctx = world.get_live_context() + world.set_joint_state(ctx, robot_id, JointState(name=["joint1", "joint2"], position=[0.0, 0.0])) + + monkeypatch.setattr(FakeScene, "joint_group_joint_names", ["joint1", "other"]) + monkeypatch.setattr( + FakeScene, + "computeFrameJacobian", + lambda self, q, frame_name, local=True: np.ones((6, 2)), + ) + with pytest.raises(ValueError, match="Unknown joints"): + world.get_group_jacobian(ctx, "arm/manipulator") + + monkeypatch.setattr(FakeScene, "joint_group_joint_names", ["joint1", "joint2"]) + monkeypatch.setattr( + FakeScene, + "computeFrameJacobian", + lambda self, q, frame_name, local=True: np.ones((5, 2)), + ) + with pytest.raises(ValueError, match="Unexpected RoboPlan Jacobian shape"): + world.get_group_jacobian(ctx, "arm/manipulator") + + monkeypatch.setattr( + FakeScene, + "computeFrameJacobian", + lambda self, q, frame_name, local=True: np.ones((6, 3)), + ) + with pytest.raises(ValueError, match="cannot project"): + world.get_group_jacobian(ctx, "arm/manipulator") + + +def test_group_jacobian_falls_back_to_configured_joint_order_when_scene_order_is_unavailable( + fake_roboplan: None, + robot_config: RobotModelConfig, + monkeypatch: pytest.MonkeyPatch, +) -> None: + world, robot_id = _make_world(fake_roboplan, robot_config) + world.finalize() + ctx = world.get_live_context() + world.set_joint_state(ctx, robot_id, JointState(name=["joint1", "joint2"], position=[0.0, 0.0])) + + def missing_group_info(self: FakeScene, name: str) -> FakeJointGroupInfo: + raise AttributeError("no joint group info") + + monkeypatch.setattr(FakeScene, "getJointGroupInfo", missing_group_info) + monkeypatch.setattr( + FakeScene, + "computeFrameJacobian", + lambda self, q, frame_name, local=True: np.arange(12, dtype=np.float64).reshape(6, 2), + ) + + np.testing.assert_allclose( + world.get_group_jacobian(ctx, "arm/manipulator"), + np.arange(12, dtype=np.float64).reshape(6, 2), + ) + + +def test_legacy_kinematics_wrappers_require_unique_pose_group( + fake_roboplan: None, robot_config: RobotModelConfig +) -> None: + no_pose_config = robot_config.model_copy( + update={ + "planning_groups": [ + PlanningGroupDefinition(name="base", joint_names=("joint1",), base_link="base") + ] + } + ) + no_pose_world, no_pose_id = _make_world(fake_roboplan, no_pose_config) + no_pose_world.finalize() + with pytest.raises(ValueError, match="no pose-targetable"): + no_pose_world.get_ee_pose(no_pose_world.get_live_context(), no_pose_id) + with pytest.raises(ValueError, match="no pose-targetable"): + no_pose_world.get_jacobian(no_pose_world.get_live_context(), no_pose_id) + + ambiguous_config = robot_config.model_copy( + update={ + "planning_groups": [ + PlanningGroupDefinition( + name="a", joint_names=("joint1",), base_link="base", tip_link="a_tip" + ), + PlanningGroupDefinition( + name="b", joint_names=("joint2",), base_link="base", tip_link="b_tip" + ), + ] + } + ) + ambiguous_world, ambiguous_id = _make_world(fake_roboplan, ambiguous_config) + ambiguous_world.finalize() + with pytest.raises(ValueError, match="pose-targetable planning groups"): + ambiguous_world.get_jacobian(ambiguous_world.get_live_context(), ambiguous_id) + + +def test_group_lookup_rejects_unknown_group_id( + fake_roboplan: None, robot_config: RobotModelConfig +) -> None: + world, _ = _make_world(fake_roboplan, robot_config) + world.finalize() + + with pytest.raises(KeyError, match="Unknown planning group ID"): + world.get_group_ee_pose(world.get_live_context(), "other/missing") + + def test_native_planner_converts_path(fake_roboplan: None, robot_config: RobotModelConfig) -> None: world, robot_id = _make_world(fake_roboplan, robot_config) world.finalize() @@ -480,6 +720,76 @@ def test_native_planner_names_path_from_robot_config_when_start_is_unnamed( assert [state.name for state in result.path] == [["joint1", "joint2"]] * 3 +def test_native_selected_planner_returns_global_selected_joint_names( + fake_roboplan: None, robot_config: RobotModelConfig +) -> None: + world, _ = _make_world(fake_roboplan, robot_config) + world.finalize() + group = world._planning_group_from_id("arm/manipulator") + selection = PlanningGroupSelection.from_groups((group,)) + + result = world.plan_selected_joint_path( + world, + selection, + JointState(name=["arm/joint1", "arm/joint2"], position=[0.0, 0.0]), + JointState(name=["arm/joint1", "arm/joint2"], position=[0.4, 0.2]), + timeout=1.0, + ) + + assert result.status == PlanningStatus.SUCCESS + assert [state.name for state in result.path] == [["arm/joint1", "arm/joint2"]] * 3 + assert [state.position for state in result.path] == [[0.0, 0.0], [0.2, 0.1], [0.4, 0.2]] + + +def test_native_selected_planner_accepts_local_joint_names( + fake_roboplan: None, robot_config: RobotModelConfig +) -> None: + world, _ = _make_world(fake_roboplan, robot_config) + world.finalize() + selection = PlanningGroupSelection.from_groups( + (world._planning_group_from_id("arm/manipulator"),) + ) + + result = world.plan_selected_joint_path( + world, + selection, + JointState(name=["joint2", "joint1"], position=[0.2, 0.0]), + JointState(name=["joint2", "joint1"], position=[0.4, 0.2]), + ) + + assert result.status == PlanningStatus.SUCCESS + assert result.path[0].name == ["arm/joint1", "arm/joint2"] + assert result.path[0].position == [0.0, 0.2] + + +def test_native_selected_planner_rejects_multi_group_selection( + fake_roboplan: None, robot_config: RobotModelConfig +) -> None: + config = robot_config.model_copy( + update={ + "planning_groups": [ + PlanningGroupDefinition("left", ("joint1",), "base", "left_tip"), + PlanningGroupDefinition("right", ("joint2",), "base", "right_tip"), + ] + } + ) + world, _ = _make_world(fake_roboplan, config) + world.finalize() + selection = PlanningGroupSelection.from_groups( + (world._planning_group_from_id("arm/left"), world._planning_group_from_id("arm/right")) + ) + + result = world.plan_selected_joint_path( + world, + selection, + JointState(name=list(selection.joint_names), position=[0.0, 0.0]), + JointState(name=list(selection.joint_names), position=[0.1, 0.1]), + ) + + assert result.status == PlanningStatus.UNSUPPORTED + assert "exactly one" in result.message + + def test_native_planner_rejects_empty_path( fake_roboplan: None, robot_config: RobotModelConfig, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/dimos/manipulation/visualization/factory.py b/dimos/manipulation/visualization/factory.py index 3e4e5c25f0..3d9d9a230d 100644 --- a/dimos/manipulation/visualization/factory.py +++ b/dimos/manipulation/visualization/factory.py @@ -53,10 +53,6 @@ def create_manipulation_visualization( ViserManipulationVisualizer, ) - return ViserManipulationVisualizer( - world_monitor=world_monitor, - manipulation_module=manipulation_module, - config=config, - ) + return ViserManipulationVisualizer(config=config) raise AssertionError(f"Unhandled manipulation visualization config: {config!r}") diff --git a/dimos/manipulation/visualization/operator.py b/dimos/manipulation/visualization/operator.py new file mode 100644 index 0000000000..77f707c084 --- /dev/null +++ b/dimos/manipulation/visualization/operator.py @@ -0,0 +1,356 @@ +# 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. + +"""UI-neutral facade for interactive manipulation.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +import math +from typing import TYPE_CHECKING, Any, cast + +from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.manipulation.planning.spec.models import GeneratedPlan, PlanningGroupID, RobotName +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.sensor_msgs.JointState import JointState + +if TYPE_CHECKING: + from dimos.manipulation.manipulation_module import ManipulationModule + from dimos.manipulation.planning.monitor.world_monitor import WorldMonitor + + +@dataclass(frozen=True) +class OperatorStatus: + """One coherent projection of the manipulation runtime snapshot.""" + + state: str + diagnostic: str = "" + ready_plan_status: str = "NONE" + ready_plan_id: str | None = None + # These two names are retained for callers that predate the snapshot API. + error: str | None = None + has_plan: bool | None = None + + def __post_init__(self) -> None: + if self.error is None: + object.__setattr__(self, "error", self.diagnostic) + if self.has_plan is None: + object.__setattr__(self, "has_plan", self.ready_plan_status == "READY") + + +@dataclass(frozen=True) +class JointTargetRequest: + """Canonical selected joint target request.""" + + group_ids: tuple[PlanningGroupID, ...] + target: JointState + + +@dataclass(frozen=True) +class PoseTargetRequest: + """Explicit world-frame pose target request.""" + + pose_targets: Mapping[PlanningGroupID, PoseStamped] + auxiliary_group_ids: tuple[PlanningGroupID, ...] = () + seed: JointState | None = None + + +@dataclass(frozen=True) +class TargetEvaluationResult: + """Advisory selected-domain target evaluation.""" + + success: bool + status: str + message: str + collision_free: bool = False + group_ids: tuple[PlanningGroupID, ...] = () + target_joints: JointState | None = None + group_diagnostics: Mapping[PlanningGroupID, str] = field(default_factory=dict) + group_poses: Mapping[PlanningGroupID, PoseStamped | None] = field(default_factory=dict) + + +class ManipulationOperator: + """Concrete synchronous facade over ManipulationModule and WorldMonitor.""" + + def __init__(self, module: ManipulationModule, world_monitor: WorldMonitor) -> None: + self._module = module + self._world_monitor = world_monitor + + def status(self) -> OperatorStatus: + """Return a coherent runtime snapshot without topology or telemetry.""" + snapshot: Any = self._module.get_execution_snapshot() + state_name = getattr(snapshot.state, "name", str(snapshot.state)) + return OperatorStatus( + state=state_name, + diagnostic=snapshot.diagnostic or "", + ready_plan_status="READY" if snapshot.has_ready_plan else "NONE", + ready_plan_id=(None if snapshot.ready_plan_id is None else str(snapshot.ready_plan_id)), + ) + + def get_init_joints(self, robot_name: RobotName) -> JointState | None: + """Return the operator-authoritative init joint state for a robot.""" + init = self._module.get_init_joints(robot_name) + return None if init is None else JointState(init) + + def evaluate_joint_target(self, request: JointTargetRequest) -> TargetEvaluationResult: + """Validate and evaluate a canonical global joint target.""" + groups, validation = self._validate_joint_request(request) + if validation is not None: + return validation + assert groups is not None + complete = self._complete_states(groups, request.target) + if complete is None: + return self._invalid(request.group_ids, "Incomplete robot target state") + return self._evaluate_global_target(groups, JointState(request.target), complete) + + def evaluate_pose_target(self, request: PoseTargetRequest) -> TargetEvaluationResult: + """Validate and evaluate explicit world-frame pose targets.""" + group_ids, validation = self._validate_pose_request(request) + if validation is not None: + return validation + ik = self._module.inverse_kinematics( + pose_targets=dict(request.pose_targets), + auxiliary_group_ids=request.auxiliary_group_ids, + seed=JointState(request.seed) if request.seed is not None else None, + check_collision=True, + ) + if not ik.is_success() or ik.joint_state is None: + return TargetEvaluationResult( + success=False, + status=ik.status.name, + message=ik.message, + collision_free=False, + group_ids=group_ids, + ) + groups = self._groups_for_ids(group_ids) + if groups is None: + return self._invalid(group_ids, "Unknown planning group") + return self._evaluate_global_target(groups, ik.joint_state) + + def plan_to_joints(self, request: JointTargetRequest) -> GeneratedPlan | None: + groups, validation = self._validate_joint_request(request) + if validation is not None: + return None + assert groups is not None + targets = { + group.id: JointState( + { + "name": list(group.joint_names), + "position": list( + request.target.position[offset : offset + len(group.joint_names)] + ), + } + ) + for group, offset in self._group_offsets(groups) + } + return self._module.generate_plan_to_joint_targets( + cast("Mapping[PlanningGroupID | PlanningGroup, JointState]", targets) + ) + + def plan_to_pose(self, request: PoseTargetRequest) -> GeneratedPlan | None: + group_ids, validation = self._validate_pose_request(request) + if validation is not None: + return None + poses = {group_id: stamped for group_id, stamped in request.pose_targets.items()} + return self._module.generate_plan_to_pose_targets( + cast("Mapping[PlanningGroupID | PlanningGroup, PoseStamped]", poses), + request.auxiliary_group_ids, + ) + + def preview(self, plan: GeneratedPlan, duration: float | None = None) -> bool: + return self._module.preview_plan(plan=plan, duration=duration) + + def execute(self, plan: GeneratedPlan) -> bool: + return self._module.execute_plan(plan=plan) + + def cancel(self) -> bool: + return self._module.cancel() + + def clear_plan(self) -> bool: + return self._module.clear_planned_path() + + def reset(self) -> bool: + result = self._module.reset() + return result.is_success() + + def _validate_joint_request( + self, request: JointTargetRequest + ) -> tuple[tuple[PlanningGroup, ...] | None, TargetEvaluationResult | None]: + groups = self._groups_for_ids(request.group_ids) + if groups is None: + return None, self._invalid( + request.group_ids, "Unknown, duplicate, or overlapping planning group" + ) + expected = tuple(name for group in groups for name in group.joint_names) + names = tuple(str(name) for name in request.target.name) + positions = tuple(float(value) for value in request.target.position) + if len(names) != len(positions): + return None, self._invalid(request.group_ids, "Joint target names and positions differ") + if len(set(names)) != len(names): + return None, self._invalid(request.group_ids, "Joint target contains duplicate joints") + if names != expected: + return None, self._invalid( + request.group_ids, "Joint target must use exact selected global joints in order" + ) + if any(not math.isfinite(value) for value in positions): + return None, self._invalid( + request.group_ids, "Joint target contains non-finite positions" + ) + return groups, None + + def _validate_pose_request( + self, request: PoseTargetRequest + ) -> tuple[tuple[PlanningGroupID, ...], TargetEvaluationResult | None]: + if not request.pose_targets: + return (), self._invalid((), "No pose target") + group_ids = tuple( + dict.fromkeys((*request.pose_targets.keys(), *request.auxiliary_group_ids)) + ) + groups = self._groups_for_ids(group_ids) + if groups is None: + return group_ids, self._invalid( + group_ids, "Unknown, duplicate, or overlapping planning group" + ) + pose_group_ids = set(request.pose_targets) + for group in groups: + if group.id in pose_group_ids and not group.has_pose_target: + return group_ids, self._invalid( + group_ids, f"Planning group '{group.id}' has no tip_link" + ) + for group_id, pose in request.pose_targets.items(): + if pose.frame_id != "world": + return group_ids, self._invalid( + group_ids, f"Unsupported pose frame for '{group_id}': {pose.frame_id}" + ) + if not self._pose_is_finite(pose): + return group_ids, self._invalid( + group_ids, f"Malformed pose target for '{group_id}'" + ) + if request.seed is not None: + seed_names = tuple(str(name) for name in request.seed.name) + if len(seed_names) != len(request.seed.position) or len(set(seed_names)) != len( + seed_names + ): + return group_ids, self._invalid(group_ids, "Malformed seed") + expected = tuple(name for group in groups for name in group.joint_names) + if seed_names != expected or any("/" not in name for name in seed_names): + return group_ids, self._invalid( + group_ids, "Seed must use exact selected global joints in order" + ) + if any(not math.isfinite(float(value)) for value in request.seed.position): + return group_ids, self._invalid(group_ids, "Seed contains non-finite positions") + return group_ids, None + + def _groups_for_ids( + self, group_ids: Sequence[PlanningGroupID] + ) -> tuple[PlanningGroup, ...] | None: + if not group_ids or len(set(group_ids)) != len(group_ids): + return None + try: + selection = self._world_monitor.planning_groups.select(tuple(group_ids)) + except (KeyError, ValueError): + return None + return selection.groups + + def _complete_states( + self, groups: Sequence[PlanningGroup], target: JointState + ) -> dict[RobotName, JointState] | None: + values = { + str(name): float(value) + for name, value in zip(target.name, target.position, strict=True) + } + complete: dict[RobotName, JointState] = {} + for robot_name in dict.fromkeys(group.robot_name for group in groups): + config = self._module.get_robot_config(robot_name) + robot_id = self._module.robot_id_for_name(robot_name) + baseline = ( + None if robot_id is None else self._world_monitor.get_current_joint_state(robot_id) + ) + if config is None or baseline is None or len(baseline.name) != len(baseline.position): + return None + baseline_values = { + str(name): float(value) + for name, value in zip(baseline.name, baseline.position, strict=True) + } + positions: list[float] = [] + for local_name in config.joint_names: + global_name = f"{robot_name}/{local_name}" + value = values.get(global_name, baseline_values.get(local_name)) + if value is None: + return None + positions.append(value) + complete[robot_name] = JointState( + {"name": list(config.joint_names), "position": positions} + ) + return complete + + def _evaluate_global_target( + self, + groups: Sequence[PlanningGroup], + target: JointState, + complete_states: Mapping[RobotName, JointState] | None = None, + ) -> TargetEvaluationResult: + complete = complete_states or self._complete_states(groups, target) + if complete is None: + return self._invalid( + tuple(group.id for group in groups), "Incomplete robot target state" + ) + diagnostics: dict[PlanningGroupID, str] = {} + poses: dict[PlanningGroupID, PoseStamped | None] = {} + valid = True + for group in groups: + robot_id = self._module.robot_id_for_name(group.robot_name) + state = complete[group.robot_name] + group_valid = bool( + robot_id is not None and self._world_monitor.is_state_valid(robot_id, state) + ) + valid = valid and group_valid + diagnostics[group.id] = ( + "Target is collision-free for this robot" + if group_valid + else "Target is in collision or violates limits" + ) + try: + poses[group.id] = self._world_monitor.get_group_ee_pose(group.id, state) + except ValueError: + poses[group.id] = None + return TargetEvaluationResult( + success=valid, + status="FEASIBLE" if valid else "COLLISION", + message="Target is collision-free for each robot" if valid else "Target is infeasible", + collision_free=valid, + group_ids=tuple(group.id for group in groups), + target_joints=JointState(target), + group_diagnostics=diagnostics, + group_poses=poses, + ) + + @staticmethod + def _invalid(group_ids: Sequence[PlanningGroupID], message: str) -> TargetEvaluationResult: + return TargetEvaluationResult(False, "INVALID", message, group_ids=tuple(group_ids)) + + @staticmethod + def _group_offsets(groups: Sequence[PlanningGroup]) -> tuple[tuple[PlanningGroup, int], ...]: + offsets: list[tuple[PlanningGroup, int]] = [] + offset = 0 + for group in groups: + offsets.append((group, offset)) + offset += len(group.joint_names) + return tuple(offsets) + + @staticmethod + def _pose_is_finite(pose: PoseStamped) -> bool: + values = [*pose.position, *pose.orientation] + return len(values) == 7 and all(math.isfinite(float(value)) for value in values) diff --git a/dimos/manipulation/visualization/test_factory.py b/dimos/manipulation/visualization/test_factory.py index 4d49f019cf..d4cf8db0d6 100644 --- a/dimos/manipulation/visualization/test_factory.py +++ b/dimos/manipulation/visualization/test_factory.py @@ -24,11 +24,12 @@ import pytest from dimos.manipulation.manipulation_module import ManipulationModuleConfig +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.models import ( - JointPath, Obstacle, - PlanningSceneInfo, + VisualizationSession, + VisualizationStateFrame, WorldRobotID, ) from dimos.manipulation.planning.spec.protocols import VisualizationSpec @@ -40,25 +41,25 @@ from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory class FakeVisualization: - def initialize_scene(self, scene: PlanningSceneInfo) -> None: + def initialize(self, session: VisualizationSession) -> None: return None def get_visualization_url(self) -> str | None: return None - def publish_visualization(self, ctx: object | None = None) -> None: + def update_state(self, frame: VisualizationStateFrame) -> None: return None - def show_preview(self, robot_id: WorldRobotID) -> None: + def animate_trajectory( + self, trajectory: JointTrajectory, duration: float | None = None + ) -> None: return None - def hide_preview(self, robot_id: WorldRobotID) -> None: - return None - - def animate_path(self, robot_id: WorldRobotID, path: JointPath, duration: float = 3.0) -> None: + def cancel_preview_animation(self) -> None: return None def close(self) -> None: @@ -78,7 +79,11 @@ def get_robot_config(self, robot_id: WorldRobotID) -> RobotModelConfig: model_path=Path("fake.urdf"), base_pose=PoseStamped(), joint_names=[], - end_effector_link="ee_link", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", joint_names=(), base_link="base_link", tip_link="ee_link" + ) + ], ) def get_joint_limits( @@ -152,6 +157,12 @@ def get_link_pose( def get_jacobian(self, ctx: object, robot_id: WorldRobotID) -> NDArray[np.float64]: return np.zeros((6, 0), dtype=np.float64) + def get_group_ee_pose(self, ctx: object, group_id: str) -> PoseStamped: + return PoseStamped() + + def get_group_jacobian(self, ctx: object, group_id: str) -> NDArray[np.float64]: + return np.zeros((6, 0), dtype=np.float64) + class FakeVisualizationWorld(FakeWorld, FakeVisualization): pass @@ -166,16 +177,18 @@ def test_config_defaults_to_no_visualization() -> None: def test_config_rejects_unknown_visualization_backend() -> None: with pytest.raises(ValidationError, match="visualization"): - ManipulationModuleConfig(visualization={"backend": "bad"}) + ManipulationModuleConfig.model_validate({"visualization": {"backend": "bad"}}) def test_config_validates_viser_visualization() -> None: - config = ManipulationModuleConfig( - visualization={ - "backend": "viser", - "visualization_host": "0.0.0.0", - "visualization_port": "8096", - "viser_panel_enabled": "false", + config = ManipulationModuleConfig.model_validate( + { + "visualization": { + "backend": "viser", + "visualization_host": "0.0.0.0", + "visualization_port": "8096", + "viser_panel_enabled": "false", + } }, ) @@ -186,7 +199,7 @@ def test_config_validates_viser_visualization() -> None: def test_config_meshcat_requires_world_visualization() -> None: - config = ManipulationModuleConfig(visualization={"backend": "meshcat"}) + config = ManipulationModuleConfig.model_validate({"visualization": {"backend": "meshcat"}}) assert isinstance(config.visualization, MeshcatVisualizationConfig) assert config.visualization.requires_world_visualization is True @@ -230,3 +243,18 @@ def test_create_visualization_meshcat_rejects_non_visualization_world() -> None: world_monitor=world_monitor, manipulation_module=MagicMock(), ) + + +def test_create_viser_visualization_has_group_preview_protocol_without_legacy_path_api() -> None: + pytest.importorskip("viser") + + visualization = create_manipulation_visualization( + ViserVisualizationConfig(), + world=FakeWorld(), + world_monitor=MagicMock(), + manipulation_module=MagicMock(), + ) + + assert isinstance(visualization, VisualizationSpec) + assert isinstance(FakeVisualization(), VisualizationSpec) + assert not hasattr(visualization, "animate_path") diff --git a/dimos/manipulation/visualization/test_operator.py b/dimos/manipulation/visualization/test_operator.py new file mode 100644 index 0000000000..a924a173c2 --- /dev/null +++ b/dimos/manipulation/visualization/test_operator.py @@ -0,0 +1,415 @@ +# 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. + +"""Focused tests for the manipulation visualization operator facade.""" + +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from dimos.agents.skill_result import SkillResult +from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupDefinition +from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.enums import IKStatus, PlanningStatus +from dimos.manipulation.planning.spec.models import ( + GeneratedPlan, + IKResult, + PlanningGroupID, + RobotName, +) +from dimos.manipulation.visualization.operator import ( + JointTargetRequest, + ManipulationOperator, + PoseTargetRequest, +) +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.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint + + +def _robot_config( + name: str = "arm", + joint_names: list[str] | None = None, + groups: tuple[PlanningGroup, ...] | None = None, +) -> RobotModelConfig: + joints = joint_names or ["j0", "j1"] + definitions = [ + PlanningGroupDefinition( + name="manipulator", + joint_names=tuple(joints), + base_link="base", + tip_link="tool", + ) + ] + if groups is not None: + definitions = [ + PlanningGroupDefinition( + name=group.group_name, + joint_names=group.local_joint_names, + base_link=group.base_link, + tip_link=group.tip_link, + ) + for group in groups + ] + return RobotModelConfig( + name=name, + model_path=Path("/robot.urdf"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), + joint_names=joints, + base_link="base", + planning_groups=definitions, + ) + + +class FakeModule: + def __init__(self) -> None: + self.state = "IDLE" + self.error = "" + self.has_plan = True + self.plan = GeneratedPlan( + group_ids=("arm/manipulator",), + trajectory=JointTrajectory( + joint_names=["arm/j0", "arm/j1"], + points=[TrajectoryPoint(0.0, [0.0, 0.0]), TrajectoryPoint(1.25, [0.4, 0.5])], + ), + path=[JointState({"name": ["arm/j0", "arm/j1"], "position": [0.0, 0.0]})], + status=PlanningStatus.SUCCESS, + ) + self.robot_configs: dict[RobotName, RobotModelConfig] = {"arm": _robot_config()} + self.robot_ids: dict[RobotName, str] = {"arm": "arm_id"} + self.plan_joint_targets: list[dict[PlanningGroupID, JointState]] = [] + self.plan_pose_targets: list[ + tuple[dict[PlanningGroupID, PoseStamped], tuple[PlanningGroupID, ...]] + ] = [] + self.ik_calls: list[ + tuple[ + dict[PlanningGroupID, PoseStamped], tuple[PlanningGroupID, ...], JointState | None + ] + ] = [] + self.plan_success = True + self.preview_success = True + self.execute_success = True + self.cancel_success = True + self.clear_success = True + self.reset_success = True + self.topology_calls = 0 + self.telemetry_calls = 0 + + def get_state(self) -> str: + return self.state + + def get_execution_snapshot(self) -> SimpleNamespace: + return SimpleNamespace( + state=self.state, + diagnostic=self.error, + has_ready_plan=self.has_plan, + ready_plan_id="fake-plan" if self.has_plan else None, + ) + + def get_error(self) -> str: + return self.error + + def has_planned_path(self) -> bool: + return self.has_plan + + def get_robot_config(self, robot_name: RobotName) -> RobotModelConfig | None: + self.topology_calls += 1 + return self.robot_configs.get(robot_name) + + def robot_id_for_name(self, robot_name: RobotName) -> str | None: + return self.robot_ids.get(robot_name) + + def get_current_joint_state(self, robot_name: RobotName) -> JointState | None: + self.telemetry_calls += 1 + return JointState(name=[f"{robot_name}/j0"], position=[0.0]) + + def inverse_kinematics( + self, + pose_targets: dict[PlanningGroupID, PoseStamped], + auxiliary_group_ids: tuple[PlanningGroupID, ...] = (), + seed: JointState | None = None, + check_collision: bool = True, + ) -> IKResult: + assert check_collision is True + self.ik_calls.append((pose_targets, auxiliary_group_ids, seed)) + return IKResult( + status=IKStatus.SUCCESS, + joint_state=JointState(name=["arm/j0", "arm/j1"], position=[0.4, 0.5]), + message="ok", + ) + + def plan_to_joint_targets(self, targets: dict[PlanningGroupID, JointState]) -> bool: + self.plan_joint_targets.append(targets) + return self.plan_success + + def generate_plan_to_joint_targets( + self, targets: dict[PlanningGroupID, JointState] + ) -> GeneratedPlan | None: + self.plan_joint_targets.append(targets) + return self.plan if self.plan_success else None + + def plan_to_pose_targets( + self, + targets: dict[PlanningGroupID, PoseStamped], + auxiliary_groups: tuple[PlanningGroupID, ...] = (), + ) -> bool: + self.plan_pose_targets.append((targets, auxiliary_groups)) + return self.plan_success + + def generate_plan_to_pose_targets( + self, + targets: dict[PlanningGroupID, PoseStamped], + auxiliary_groups: tuple[PlanningGroupID, ...] = (), + ) -> GeneratedPlan | None: + self.plan_pose_targets.append((targets, auxiliary_groups)) + return self.plan if self.plan_success else None + + def preview_plan( + self, plan: GeneratedPlan | None = None, duration: float | None = None + ) -> bool: + return self.preview_success + + def execute_plan(self, plan: GeneratedPlan | None = None) -> bool: + return self.execute_success + + def cancel(self) -> bool: + return self.cancel_success + + def clear_planned_path(self) -> bool: + return self.clear_success + + def reset(self) -> SkillResult[str]: + if self.reset_success: + return SkillResult.ok("reset") + return SkillResult.fail("ERR", "no reset") + + +class FakeWorldMonitor: + def __init__(self, registry: PlanningGroupRegistry) -> None: + self.planning_groups = registry + self.current_states: dict[str, JointState] = { + "arm_id": JointState(name=["j0", "j1"], position=[0.0, 0.0]) + } + self.valid = True + self.cancel_preview_calls = 0 + self.telemetry_calls = 0 + + def get_current_joint_state(self, robot_id: str) -> JointState | None: + self.telemetry_calls += 1 + return self.current_states.get(robot_id) + + def is_state_valid(self, robot_id: str, joint_state: JointState) -> bool: + return self.valid + + def get_group_ee_pose( + self, group_id: PlanningGroupID, joint_state: JointState | None = None + ) -> PoseStamped: + return PoseStamped( + frame_id="world", position=Vector3(1.0, 2.0, 3.0), orientation=Quaternion() + ) + + def cancel_preview_animation(self) -> None: + self.cancel_preview_calls += 1 + + +def _operator( + config: RobotModelConfig | None = None, +) -> tuple[ManipulationOperator, FakeModule, FakeWorldMonitor]: + robot_config = config or _robot_config() + module = FakeModule() + module.robot_configs = {robot_config.name: robot_config} + module.robot_ids = {robot_config.name: f"{robot_config.name}_id"} + monitor = FakeWorldMonitor(PlanningGroupRegistry([robot_config])) + monitor.current_states = { + f"{robot_config.name}_id": JointState( + name=robot_config.joint_names, position=[0.0] * len(robot_config.joint_names) + ) + } + return ManipulationOperator(module, monitor), module, monitor # type: ignore[arg-type] + + +def _joint_request( + names: list[str] | None = None, positions: list[float] | None = None +) -> JointTargetRequest: + return JointTargetRequest( + group_ids=("arm/manipulator",), + target=JointState(name=names or ["arm/j0", "arm/j1"], position=positions or [0.1, 0.2]), + ) + + +def _pose(frame_id: str = "world") -> PoseStamped: + return PoseStamped(frame_id=frame_id, position=Vector3(0.1, 0.2, 0.3), orientation=Quaternion()) + + +def test_status_is_compact_and_does_not_read_topology_or_telemetry() -> None: + operator, module, monitor = _operator() + + status = operator.status() + + assert status.state == "IDLE" + assert status.error == "" + assert status.has_plan is True + assert module.topology_calls == 0 + assert module.telemetry_calls == 0 + assert monitor.telemetry_calls == 0 + + +@pytest.mark.parametrize( + "state", + ("IDLE", "PLANNING", "READY", "DISPATCHING", "RUNNING", "CANCELLING", "FAULT"), +) +def test_status_projects_every_runtime_lifecycle_state_from_one_snapshot(state: str) -> None: + operator, module, _ = _operator() + module.state = state + module.has_plan = state == "READY" + module.error = "runtime fault" if state == "FAULT" else "" + + status = operator.status() + + assert status.state == state + assert status.diagnostic == ("runtime fault" if state == "FAULT" else "") + assert status.has_plan is (state == "READY") + assert status.ready_plan_status == ("READY" if state == "READY" else "NONE") + assert status.ready_plan_id == ("fake-plan" if state == "READY" else None) + + +def test_evaluate_joint_target_accepts_exact_global_selection_domain() -> None: + operator, _, _ = _operator() + + result = operator.evaluate_joint_target(_joint_request()) + + assert result.success is True + assert result.status == "FEASIBLE" + assert result.target_joints is not None + assert list(result.target_joints.name) == ["arm/j0", "arm/j1"] + assert list(result.target_joints.position) == [0.1, 0.2] + assert result.group_diagnostics["arm/manipulator"] == "Target is collision-free for this robot" + assert result.group_poses["arm/manipulator"] is not None + + +def test_joint_target_validation_rejects_bad_joint_requests() -> None: + cases = [ + _joint_request(["j0", "j1"], [0.1, 0.2]), + _joint_request(["arm/j0", "arm/j0"], [0.1, 0.2]), + _joint_request(["arm/j0"], [0.1]), + _joint_request(["arm/j0", "arm/j1", "arm/extra"], [0.1, 0.2, 0.3]), + _joint_request(["arm/j0", "arm/j1"], [0.1, float("nan")]), + JointTargetRequest( + ("missing/manipulator",), JointState(name=["missing/j0"], position=[0.1]) + ), + JointTargetRequest( + ("arm/manipulator", "arm/manipulator"), + JointState(name=["arm/j0", "arm/j1"], position=[0.1, 0.2]), + ), + ] + operator, _, _ = _operator() + + for request in cases: + result = operator.evaluate_joint_target(request) + assert result.success is False + assert result.status == "INVALID" + + +def test_joint_target_validation_rejects_overlapping_groups() -> None: + groups = ( + PlanningGroup("arm/first", "arm", "first", ("arm/j0",), ("j0",), "base"), + PlanningGroup("arm/second", "arm", "second", ("arm/j0",), ("j0",), "base"), + ) + operator, _, _ = _operator(_robot_config(groups=groups)) + request = JointTargetRequest( + ("arm/first", "arm/second"), JointState(name=["arm/j0", "arm/j0"], position=[0.1, 0.2]) + ) + + result = operator.evaluate_joint_target(request) + + assert result.success is False + assert result.status == "INVALID" + + +def test_pose_evaluation_accepts_world_frame_and_delegates_original_request() -> None: + operator, module, _ = _operator() + pose = _pose() + seed = JointState(name=["arm/j0", "arm/j1"], position=[0.0, 0.0]) + request = PoseTargetRequest({"arm/manipulator": pose}, seed=seed) + + result = operator.evaluate_pose_target(request) + + assert result.success is True + assert result.target_joints is not None + assert list(result.target_joints.name) == ["arm/j0", "arm/j1"] + assert module.ik_calls == [({"arm/manipulator": pose}, (), seed)] + + +def test_pose_validation_rejects_frame_capability_and_seed_errors() -> None: + no_pose_group = ( + PlanningGroup("arm/no_pose", "arm", "no_pose", ("arm/j0",), ("j0",), "base", None), + ) + no_pose_operator, _, _ = _operator(_robot_config(joint_names=["j0"], groups=no_pose_group)) + bad_seed_cases = [ + PoseTargetRequest({"arm/manipulator": _pose("camera")}), + PoseTargetRequest( + {"arm/manipulator": _pose()}, + seed=JointState(name=["j0", "j1"], position=[0.0, 0.0]), + ), + PoseTargetRequest( + {"arm/manipulator": _pose()}, + seed=JointState(name=["arm/j0", "arm/j0"], position=[0.0, 0.0]), + ), + ] + operator, _, _ = _operator() + + no_pose = no_pose_operator.evaluate_pose_target(PoseTargetRequest({"arm/no_pose": _pose()})) + assert no_pose.success is False + assert no_pose.status == "INVALID" + for request in bad_seed_cases: + result = operator.evaluate_pose_target(request) + assert result.success is False + assert result.status == "INVALID" + + +def test_planning_methods_return_exact_generated_plan() -> None: + operator, module, _ = _operator() + joint_request = _joint_request() + pose = _pose() + pose_request = PoseTargetRequest({"arm/manipulator": pose}) + + joint_result = operator.plan_to_joints(joint_request) + pose_result = operator.plan_to_pose(pose_request) + + assert joint_result is module.plan + assert list(module.plan_joint_targets[0]["arm/manipulator"].name) == ["arm/j0", "arm/j1"] + assert module.plan_pose_targets == [({"arm/manipulator": pose}, ())] + assert pose_result is module.plan + + +def test_actions_return_typed_results_and_cancel_fallback_ownership() -> None: + operator, module, monitor = _operator() + + assert operator.preview(module.plan, 0.5) is True + assert operator.execute(module.plan) is True + assert operator.clear_plan() is True + assert operator.reset() is True + cancel_result = operator.cancel() + assert cancel_result is True + assert monitor.cancel_preview_calls == 0 + + module.cancel_success = False + fallback = operator.cancel() + assert fallback is False + assert monitor.cancel_preview_calls == 0 diff --git a/dimos/manipulation/visualization/viser/adapter.py b/dimos/manipulation/visualization/viser/adapter.py deleted file mode 100644 index c3eb8e360b..0000000000 --- a/dimos/manipulation/visualization/viser/adapter.py +++ /dev/null @@ -1,174 +0,0 @@ -# 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 - -from collections.abc import Sequence -from typing import TYPE_CHECKING - -from dimos.manipulation.visualization.types import RobotInfo, TargetEvaluation -from dimos.msgs.sensor_msgs.JointState import JointState - -if TYPE_CHECKING: - from dimos.manipulation.manipulation_module import ManipulationModule - from dimos.manipulation.planning.monitor.world_monitor import WorldMonitor - from dimos.manipulation.planning.spec.config import RobotModelConfig - from dimos.manipulation.planning.spec.models import JointPath, RobotName, WorldRobotID - from dimos.msgs.geometry_msgs.Pose import Pose - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - - -def copy_joint_state(joint_state: JointState | None) -> JointState | None: - """Make a local copy of a JointState-like message for rendering.""" - return None if joint_state is None else JointState(joint_state) - - -class InProcessViserAdapter: - """Small in-process boundary between Viser callbacks and manipulation internals.""" - - def __init__( - self, - *, - world_monitor: WorldMonitor, - manipulation_module: ManipulationModule, - ) -> None: - self._world_monitor = world_monitor - self._module = manipulation_module - - def list_robots(self) -> list[RobotName]: - return list(self._module.list_robots()) - - def robot_items(self) -> list[tuple[RobotName, WorldRobotID, RobotModelConfig]]: - return self._module.robot_items() - - def robot_id_for_name(self, robot_name: RobotName) -> WorldRobotID | None: - return self._module.robot_id_for_name(robot_name) - - def robot_name_for_id(self, robot_id: WorldRobotID) -> RobotName | None: - return self._module.robot_name_for_id(robot_id) - - def get_robot_config(self, robot_name: RobotName) -> RobotModelConfig | None: - return self._module.get_robot_config(robot_name) - - def get_robot_info(self, robot_name: RobotName) -> RobotInfo | None: - info = self._module.get_robot_info(robot_name) - if info is None: - return None - return { - "name": str(info["name"]), - "world_robot_id": str(info["world_robot_id"]), - "joint_names": [str(name) for name in info["joint_names"]], - "end_effector_link": str(info["end_effector_link"]), - "base_link": str(info["base_link"]), - "max_velocity": float(info["max_velocity"]), - "max_acceleration": float(info["max_acceleration"]), - "has_joint_name_mapping": bool(info["has_joint_name_mapping"]), - "coordinator_task_name": None - if info["coordinator_task_name"] is None - else str(info["coordinator_task_name"]), - "home_joints": None - if info["home_joints"] is None - else [float(value) for value in info["home_joints"]], - "pre_grasp_offset": float(info["pre_grasp_offset"]), - "init_joints": None - if info["init_joints"] is None - else [float(value) for value in info["init_joints"]], - } - - def get_init_joints(self, robot_name: RobotName) -> JointState | None: - return copy_joint_state(self._module.get_init_joints(robot_name)) - - def get_current_joint_state(self, robot_name: RobotName) -> JointState | None: - robot_id = self.robot_id_for_name(robot_name) - if robot_id is None: - return None - return copy_joint_state(self._world_monitor.get_current_joint_state(robot_id)) - - def is_state_stale(self, robot_name: RobotName, max_age: float = 1.0) -> bool: - robot_id = self.robot_id_for_name(robot_name) - return True if robot_id is None else self._world_monitor.is_state_stale(robot_id, max_age) - - def get_ee_pose( - self, robot_name: RobotName, joint_state: JointState | None = None - ) -> PoseStamped | None: - robot_id = self.robot_id_for_name(robot_name) - if robot_id is None: - return None - return self._world_monitor.get_ee_pose(robot_id, copy_joint_state(joint_state)) - - def evaluate_joint_target(self, joints: JointState, robot_name: RobotName) -> TargetEvaluation: - """Evaluate a joint target through WorldMonitor helpers, not raw WorldSpec access.""" - result: TargetEvaluation = { - **self._module.evaluate_joint_target(copy_joint_state(joints), robot_name) - } - joint_state = result.get("joint_state") - result["joint_state"] = copy_joint_state( - joint_state if isinstance(joint_state, JointState) else None - ) - return result - - def evaluate_pose_target(self, pose: Pose, robot_name: RobotName) -> TargetEvaluation: - """Evaluate a Cartesian target through module/WorldMonitor helper boundaries.""" - result: TargetEvaluation = {**self._module.evaluate_pose_target(pose, robot_name)} - joint_state = result.get("joint_state") - result["joint_state"] = copy_joint_state( - joint_state if isinstance(joint_state, JointState) else None - ) - return result - - def get_planned_path(self, robot_name: RobotName) -> JointPath | None: - path = self._module.get_planned_path(robot_name) - if path is None: - return None - copied = [copy_joint_state(point) for point in path] - return [point for point in copied if point is not None] - - def get_planned_trajectory_duration(self, robot_name: RobotName) -> float | None: - return self._module.get_planned_trajectory_duration(robot_name) - - def get_module_state(self) -> str: - return str(self._module.get_state()) - - def get_error(self) -> str: - return self._module.get_error() - - def reset(self) -> bool: - return self._module.reset().is_success() - - def plan_to_pose(self, pose: Pose, robot_name: RobotName | None = None) -> bool: - return self._module.plan_to_pose(pose, robot_name) - - def plan_to_joints(self, joints: JointState, robot_name: RobotName | None = None) -> bool: - return self._module.plan_to_joints(joints, robot_name) - - def preview_path(self, robot_name: RobotName | None = None) -> bool: - return self._module.preview_path(robot_name=robot_name) - - def execute(self, robot_name: RobotName | None = None) -> bool: - return self._module.execute(robot_name) - - def cancel(self) -> bool: - return self._module.cancel() - - def clear_planned_path(self) -> bool: - return self._module.clear_planned_path() - - @staticmethod - def joints_from_values(joint_names: Sequence[str], values: Sequence[float]) -> JointState: - return JointState( - { - "name": list(joint_names), - "position": [float(value) for value in values], - } - ) diff --git a/dimos/manipulation/visualization/viser/animation.py b/dimos/manipulation/visualization/viser/animation.py index f5a574d82c..6754343e21 100644 --- a/dimos/manipulation/visualization/viser/animation.py +++ b/dimos/manipulation/visualization/viser/animation.py @@ -14,85 +14,49 @@ from __future__ import annotations -from collections.abc import Callable, Sequence -import time +from collections.abc import Sequence +from dataclasses import dataclass +from itertools import pairwise -from dimos.msgs.sensor_msgs.JointState import JointState +@dataclass(frozen=True) +class PreviewFrame: + """One timestamped local robot preview frame.""" -def interpolate_joint_path( - path: Sequence[JointState], duration: float, fps: float -) -> list[list[float]]: - """Interpolate a joint path into visualization frames.""" - waypoints = [list(waypoint.position) for waypoint in path if waypoint.position] - if not waypoints: - return [] - if len(waypoints) == 1 or duration <= 0.0: - return [waypoints[-1]] - frame_count = max(int(duration * max(fps, 1.0)) + 1, len(waypoints)) - segment_count = len(waypoints) - 1 - frames: list[list[float]] = [] - for frame_index in range(frame_count): - path_t = frame_index / max(frame_count - 1, 1) - scaled = path_t * segment_count - segment_index = min(int(scaled), segment_count - 1) - local_t = scaled - segment_index - start = waypoints[segment_index] - end = waypoints[segment_index + 1] - if len(start) != len(end): - continue - frames.append( - [ - start_value + (end_value - start_value) * local_t - for start_value, end_value in zip(start, end, strict=False) - ] - ) - if frames and frames[-1] != waypoints[-1]: - frames.append(waypoints[-1]) - return frames + time_from_start: float + positions: tuple[float, ...] -def sampled_joint_path_frames( - path: Sequence[JointState], duration: float, fps: float -) -> list[list[float]]: - """Return animation frames while preserving already sampled trajectories. +@dataclass(frozen=True) +class PreviewTrack: + """One fixed-baseline local robot track in a group-native preview.""" - ManipulationModule.preview_path() owns trajectory-aware interpolation because it has access - to JointTrajectory waypoint timing. If a path arrives already sampled near the target display - rate, Viser should play those samples directly instead of re-interpolating by waypoint index. - Sparse direct VisualizationSpec callers still get local interpolation as a fallback. - """ - waypoints = [list(waypoint.position) for waypoint in path if waypoint.position] - if not waypoints: - return [] - expected_frames = max(int(duration * max(fps, 1.0)) + 1, 1) if duration > 0.0 else 1 - if len(waypoints) >= expected_frames: - return waypoints - return interpolate_joint_path(path, duration, fps) + robot_id: str + joint_names: tuple[str, ...] + frames: tuple[PreviewFrame, ...] -class PreviewAnimator: - """Blocking preview-ghost path animator with Meshcat-compatible semantics. +@dataclass(frozen=True) +class GroupPreviewAnimation: + """Validated collection of robot tracks sharing one preview transaction.""" - This class is only for transient path playback. Persistent target ghosts are updated - directly by scene target methods and must not be routed through this animator. - """ + tracks: tuple[PreviewTrack, ...] - def __init__( - self, - set_joints: Callable[[Sequence[float]], None], - *, - sleep: Callable[[float], None] = time.sleep, - ) -> None: - self._set_joints = set_joints - self._sleep = sleep - def animate(self, path: Sequence[JointState], duration: float, fps: float) -> bool: - frames = sampled_joint_path_frames(path, duration, fps) - if not frames: - return False - step_delay = duration / max(len(frames) - 1, 1) if duration > 0.0 else 0.0 - for joints in frames: - self._set_joints(joints) - self._sleep(step_delay) - return True +def scaled_frame_delays(frames: Sequence[PreviewFrame], duration: float) -> tuple[float, ...]: + """Return stored inter-frame delays, optionally scaled to a requested duration.""" + if len(frames) < 2: + return () + original_duration = max(float(frames[-1].time_from_start), 0.0) + scale = duration / original_duration if duration > 0.0 and original_duration > 0.0 else 1.0 + return tuple( + max(float(next_frame.time_from_start) - float(frame.time_from_start), 0.0) * scale + for frame, next_frame in pairwise(frames) + ) + + +def preview_tick_times(preview: GroupPreviewAnimation) -> tuple[float, ...]: + """Union all stored track timestamps without synthesizing extra samples.""" + return tuple( + sorted({float(frame.time_from_start) for track in preview.tracks for frame in track.frames}) + ) diff --git a/dimos/manipulation/visualization/viser/gui.py b/dimos/manipulation/visualization/viser/gui.py index 2c0abac52d..00313c632b 100644 --- a/dimos/manipulation/visualization/viser/gui.py +++ b/dimos/manipulation/visualization/viser/gui.py @@ -14,10 +14,19 @@ from __future__ import annotations -from typing import TypeAlias - -from dimos.manipulation.visualization.types import RobotInfo, TargetEvaluation -from dimos.manipulation.visualization.viser.adapter import InProcessViserAdapter +from collections.abc import Mapping, MutableMapping, Sequence +from typing import TypeAlias, cast + +from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.models import PlanningGroupID, PlanningSceneInfo, RobotName +from dimos.manipulation.visualization.operator import ( + JointTargetRequest, + ManipulationOperator, + OperatorStatus, + PoseTargetRequest, + TargetEvaluationResult, +) from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig from dimos.manipulation.visualization.viser.runtime import VISER_INSTALL_HINT from dimos.manipulation.visualization.viser.scene import ViserManipulationScene @@ -35,6 +44,7 @@ TargetStatus, ) from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.JointState import JointState from dimos.utils.logging_config import setup_logger @@ -68,28 +78,52 @@ # Fallback joint-slider range (radians) when a robot config omits joint limits. DEFAULT_JOINT_LIMITS = (-3.14, 3.14) +PRIMARY_ACTION_COLOR = (0, 102, 179) +ACTIVE_GROUP_COLOR = PRIMARY_ACTION_COLOR +INACTIVE_GROUP_COLOR = (52, 52, 52) + + +def group_display_name(group: PlanningGroup) -> str: + return ( + str(group.robot_name) + if str(group.group_name) == "manipulator" + else f"{group.robot_name} {group.group_name}" + ) + + +def _copy_joint_state(state: JointState | None) -> JointState | None: + return None if state is None else JointState(state) class ViserPanelGui: - """Optional operator panel with parity for the original cc/viser-vis panel.""" + """Viser operator panel for manipulation target editing and plan control.""" def __init__( self, server: ViserServer, - adapter: InProcessViserAdapter, + scene_info: PlanningSceneInfo, + operator: ManipulationOperator | object, + current_states: MutableMapping[str, JointState], config: ViserVisualizationConfig, scene: ViserManipulationScene | None = None, ) -> None: self.server = server - self.adapter = adapter + self.scene_info = scene_info + self.operator = cast("ManipulationOperator", operator) + self.current_states = current_states + self._robots_by_name = { + config.name: (robot_id, config) for robot_id, config in scene_info.robots.items() + } + self._scene_groups_by_id = {group.id: group for group in scene_info.planning_groups} self.config = config self.scene = scene self.state = PanelState(runtime=PanelRuntime.STARTING) self._closed = False self._operation_sequence_id = 0 self._suppress_target_callbacks = False + self._default_group_initialized = False self._handles: dict[str, PanelHandle] = {} - self._joint_sliders: dict[str, GuiSliderHandle[float]] = {} + self._joint_sliders: dict[tuple[PlanningGroupID, str], GuiSliderHandle[float]] = {} self._worker = TargetEvaluationWorker( self._handle_target_evaluation_request, self._apply_target_evaluation_result, @@ -117,6 +151,8 @@ def close(self) -> None: return self._closed = True self.state.runtime = PanelRuntime.STOPPING + if self.scene is not None: + self.scene.cancel_preview_animation() self._worker.stop() self._operation_worker.stop(timeout=2.0) self._clear_joint_sliders() @@ -124,20 +160,179 @@ def close(self) -> None: self._handles.clear() self.state.runtime = PanelRuntime.STOPPED + def list_robots(self) -> list[RobotName]: + return [config.name for config in self.scene_info.robots.values()] + + def list_planning_groups(self) -> list[PlanningGroup]: + return list(self.scene_info.planning_groups) + + def robot_items(self) -> list[tuple[RobotName, str, RobotModelConfig]]: + return [ + (config.name, str(robot_id), config) + for robot_id, config in self.scene_info.robots.items() + ] + + def robot_id_for_name(self, robot_name: RobotName) -> str | None: + item = self._robots_by_name.get(robot_name) + return None if item is None else str(item[0]) + + def get_robot_config(self, robot_name: RobotName) -> RobotModelConfig | None: + item = self._robots_by_name.get(robot_name) + return None if item is None else item[1] + + def get_init_joints(self, robot_name: RobotName) -> JointState | None: + init = self.operator.get_init_joints(robot_name) + if init is None: + return None + config = self.get_robot_config(robot_name) + if config is None: + return JointState(init) + values = self._local_values_for_robot(robot_name, init) + if any(name not in values for name in config.joint_names): + return JointState(init) + return JointState( + { + "name": list(config.joint_names), + "position": [values[name] for name in config.joint_names], + } + ) + + def get_current_joint_state(self, robot_name: RobotName) -> JointState | None: + robot_id = self.robot_id_for_name(robot_name) + return None if robot_id is None else _copy_joint_state(self.current_states.get(robot_id)) + + def get_group_ee_pose(self, group_id: PlanningGroupID) -> PoseStamped | None: + group = self._scene_groups_by_id.get(group_id) + if group is None: + return None + targets = self._current_target_for_group(group) + if group_id not in targets: + return None + return self.evaluate_joint_target_set((group_id,), targets).group_poses.get(group_id) + + def _current_target_for_group(self, group: PlanningGroup) -> dict[PlanningGroupID, JointState]: + current = self.get_current_joint_state(group.robot_name) + if current is None or len(current.name) != len(current.position): + return {} + values = self._local_values_for_robot(group.robot_name, current) + if any(name not in values for name in group.local_joint_names): + return {} + return { + group.id: JointState( + { + "name": list(group.joint_names), + "position": [values[name] for name in group.local_joint_names], + } + ) + } + + def is_state_stale(self, robot_name: RobotName, max_age: float = 1.0) -> bool: + return self.get_current_joint_state(robot_name) is None + + def get_module_state(self) -> str: + return self.operator.status().state + + def get_error(self) -> str: + return self.operator.status().error or "" + + def reset(self) -> bool: + result = self.operator.reset() + if result: + self.state.error = "" + self.state.action_status = ActionStatus.IDLE + self.state.runtime_failure = False + self.state.local_action_failure = False + return result + + def evaluate_joint_target_set( + self, group_ids: Sequence[PlanningGroupID], targets: Mapping[PlanningGroupID, JointState] + ) -> TargetEvaluationResult: + names: list[str] = [] + positions: list[float] = [] + for group_id in group_ids: + target = targets.get(group_id) + if target is None: + return TargetEvaluationResult( + False, "INVALID", "Incomplete joint target", group_ids=tuple(group_ids) + ) + names.extend(str(name) for name in target.name) + positions.extend(float(value) for value in target.position) + return self.operator.evaluate_joint_target( + JointTargetRequest(tuple(group_ids), JointState({"name": names, "position": positions})) + ) + + def evaluate_pose_target_set( + self, + pose_targets: Mapping[PlanningGroupID, Pose], + auxiliary_group_ids: Sequence[PlanningGroupID] = (), + seed: JointState | None = None, + ) -> TargetEvaluationResult: + stamped = { + group_id: PoseStamped( + frame_id="world", position=pose.position, orientation=pose.orientation + ) + for group_id, pose in pose_targets.items() + } + return self.operator.evaluate_pose_target( + PoseTargetRequest(stamped, tuple(auxiliary_group_ids), _copy_joint_state(seed)) + ) + + def cancel(self) -> bool: + return self.operator.cancel() + + def clear_planned_path(self) -> bool: + return self.operator.clear_plan() + + def plan_to_selected_joints( + self, group_ids: Sequence[PlanningGroupID], targets: Mapping[PlanningGroupID, JointState] + ) -> bool: + names: list[str] = [] + positions: list[float] = [] + for group_id in group_ids: + target = targets.get(group_id) + if target is None: + return False + names.extend(str(name) for name in target.name) + positions.extend(float(value) for value in target.position) + plan = self.operator.plan_to_joints( + JointTargetRequest(tuple(group_ids), JointState({"name": names, "position": positions})) + ) + self.state.plan_state.plan = plan + self.state.plan_state.plan_id = None + return plan is not None + + def preview_path(self) -> bool: + plan = self.state.plan_state.plan + return plan is not None and self.operator.preview(plan) + + def execute(self) -> bool: + plan = self.state.plan_state.plan + return plan is not None and self.operator.execute(plan) + def refresh(self) -> None: if self._closed: return - robots = self.adapter.list_robots() + status = self.operator.status() + robots = self.list_robots() + groups = self.list_planning_groups() self.state.backend_status = ( BackendConnectionStatus.READY if robots else BackendConnectionStatus.WAITING_FOR_ROBOT ) - if self.state.selected_robot is None and robots: - self.state.selected_robot = robots[0] + if not self.state.selected_group_ids and groups and not self._default_group_initialized: + first = next((group for group in groups if group.has_pose_target), groups[0]) + self.state.selected_group_ids = (first.id,) + self.state.selected_robot = str(first.robot_name) self.state.target_status = TargetStatus.EMPTY + self._default_group_initialized = True + initialized_groups = set(self.state.group_joint_targets) + self._initialize_selected_group_targets() + if set(self.state.group_joint_targets) != initialized_groups: self._build_joint_sliders() - self._sync_robot_dropdown(robots) - self._refresh_selected_robot_state() + self._sync_group_selector(groups) + self._refresh_selected_robot_state(status) + self._apply_operator_status(status) self._ensure_scene_controls() + self._sync_target_ghost_visibility() self._sync_preset_dropdown() self._update_status_text() self._update_control_state() @@ -150,26 +345,26 @@ def _build(self) -> None: self._build_panel_controls(gui) def _build_panel_controls(self, gui: GuiApi) -> None: - self._handles["status"] = gui.add_markdown("Starting manipulation panel...") - robots = self.adapter.list_robots() + self._handles["status"] = gui.add_markdown("### Status\n**State:** Ready") self._build_scene_controls(gui) - robot_dropdown = gui.add_dropdown( - "Robot", - options=robots or [""], - initial_value=robots[0] if robots else "", + self._handles["planning_groups_heading"] = gui.add_markdown( + "### Planning Groups\nActive planning groups for pose goals, planning, and joint edits." ) - robot_dropdown.on_update(lambda event: self._select_robot(event.target.value)) - self._handles["robot"] = robot_dropdown + self._sync_group_selector(self.list_planning_groups()) + self._handles["target_heading"] = gui.add_markdown("### Target") preset_dropdown = gui.add_dropdown( - "Target Preset", + "Preset", options=["Select preset...", "Current"], initial_value="Select preset...", ) preset_dropdown.on_update(lambda event: self._apply_preset(event.target.value)) self._handles["preset"] = preset_dropdown - plan_button = gui.add_button("Plan", disabled=True) + self._handles["target_summary"] = gui.add_markdown("Feasibility: `unknown`") + self._handles["actions_heading"] = gui.add_markdown("### Actions") + plan_button = gui.add_button("Plan", disabled=True, color=PRIMARY_ACTION_COLOR) plan_button.on_click(lambda _: self._submit_plan()) self._handles["plan"] = plan_button + self._handles["plan_controls_heading"] = gui.add_markdown("**Plan controls**") preview_button = gui.add_button("Preview", disabled=True) preview_button.on_click(lambda _: self._submit_preview()) self._handles["preview"] = preview_button @@ -182,7 +377,62 @@ def _build_panel_controls(self, gui: GuiApi) -> None: clear_button = gui.add_button("Clear plan") clear_button.on_click(lambda _: self._submit_clear()) self._handles["clear"] = clear_button + joint_controls = gui.add_folder("Joint Control", expand_by_default=False) + self._handles["joint_control_folder"] = joint_controls + self._build_joint_sliders() + + def _sync_group_selector(self, groups: list[PlanningGroup]) -> None: + """Render source-order group toggle buttons without a robot dropdown.""" + selected = set(self.state.selected_group_ids) + seen: set[str] = set() + for group in sorted( + groups, key=lambda item: (not bool(item.has_pose_target), str(item.id)) + ): + group_id = str(group.id) + key = f"group:{group_id}" + seen.add(key) + label = group_display_name(group) + handle = self._handles.get(key) + color = ACTIVE_GROUP_COLOR if group_id in selected else INACTIVE_GROUP_COLOR + if handle is None: + handle = self.server.gui.add_button( + label, + color=color, + hint="Click to toggle this planning group in the target set.", + ) + + def on_click(_event: object, selected_group_id: str = group_id) -> None: + self._toggle_group_selected(selected_group_id) + + handle.on_click(on_click) + self._handles[key] = handle + else: + self._set_optional_handle_attr(handle, "label", label) + self._set_optional_handle_attr(handle, "color", color) + for key in [key for key in self._handles if key.startswith("group:") and key not in seen]: + handle = self._handles.pop(key) + remove = getattr(handle, "remove", None) + if callable(remove): + remove() + + def _toggle_group_selected(self, group_id: str) -> None: + groups = {str(group.id): group for group in self.list_planning_groups()} + if group_id not in groups: + return + current = list(self.state.selected_group_ids) + if group_id in current: + current.remove(group_id) + else: + current.append(group_id) + self.state.selected_group_ids = tuple(current) + self.state.advance_selection_epoch() + self._clear_invalidated_preview() + first = groups.get(current[0]) if current else None + self.state.selected_robot = None if first is None else str(first.robot_name) + self._prune_inactive_group_state() + self._initialize_selected_group_targets() self._build_joint_sliders() + self.refresh() def _build_scene_controls(self, gui: GuiApi) -> None: if self.scene is None: @@ -200,74 +450,143 @@ def _set_scene_grid_visible(self, visible: bool) -> None: return self.scene.set_reference_grid_visible(bool(visible)) - def _refresh_selected_robot_state(self) -> None: + def _refresh_selected_robot_state(self, status: OperatorStatus) -> None: robot_name = self.state.selected_robot if robot_name is None: - self.state.robot_info = None self.state.current_joints = None - self.state.current_ee_pose = None - self.state.manipulation_state = self.adapter.get_module_state() + self.state.manipulation_state = status.state return - self.state.robot_info = self.adapter.get_robot_info(robot_name) - current = self.adapter.get_current_joint_state(robot_name) + current = self.get_current_joint_state(robot_name) self.state.current_joints = list(current.position) if current is not None else None - self.state.current_ee_pose = self.adapter.get_ee_pose(robot_name) - self.state.manipulation_state = self.adapter.get_module_state() - adapter_error = self.adapter.get_error() - if adapter_error: - self.state.error = adapter_error + self.state.manipulation_state = status.state + + def _apply_operator_status(self, status: OperatorStatus) -> None: + """Project one operator snapshot onto the panel state.""" + self.state.manipulation_state = status.state + self.state.ready_plan_id = getattr(status, "ready_plan_id", None) + snapshot_status = hasattr(status, "diagnostic") + diagnostic = getattr(status, "diagnostic", getattr(status, "error", "")) + if diagnostic: + self.state.error = diagnostic + elif snapshot_status and self.state.action_status != ActionStatus.FAILED: + self.state.error = "" + if status.state == "FAULT": + self.state.runtime_failure = True + self.state.action_status = ActionStatus.FAILED + elif status.state == "CANCELLING": + self.state.action_status = ActionStatus.CANCELLING + elif status.state in {"DISPATCHING", "RUNNING"}: + self.state.action_status = ActionStatus.EXECUTING + elif status.state == "PLANNING" and self.state.action_status == ActionStatus.IDLE: + self.state.action_status = ActionStatus.RUNNING + elif status.state in {"IDLE", "READY"}: + was_runtime_failure = self.state.runtime_failure + self.state.runtime_failure = False + if self.state.action_status in { + ActionStatus.RUNNING, + ActionStatus.EXECUTING, + ActionStatus.CANCELLING, + }: + self.state.action_status = ActionStatus.IDLE + elif was_runtime_failure and not self.state.local_action_failure: + self.state.action_status = ActionStatus.IDLE + if not diagnostic: + self.state.error = "" + if status.state != "READY" and self.state.plan_state.status == PlanStatus.FRESH: + if status.state in {"DISPATCHING", "RUNNING", "CANCELLING", "IDLE", "FAULT"}: + self.state.plan_state.status = PlanStatus.STALE + if ( + self.state.plan_state.status == PlanStatus.FRESH + and self.state.plan_state.plan_id != self.state.ready_plan_id + ): + self.state.plan_state.status = PlanStatus.STALE def _ensure_scene_controls(self) -> None: - if self.scene is None or self.state.selected_robot is None: - return - robot_id = self.adapter.robot_id_for_name(self.state.selected_robot) - if robot_id is None: + if self.scene is None: return - ee_control = self.scene.ensure_target_controls(str(robot_id), self._on_transform_update) - if ee_control is not None: - self._handles["ee_control"] = ee_control - if ( - self.state.target_status == TargetStatus.EMPTY - and self.state.current_ee_pose is not None - ): - self.state.cartesian_target = self.state.current_ee_pose - self._suppress_target_callbacks = True - try: - self.scene.set_target_pose(str(robot_id), self.state.current_ee_pose) - finally: - self._suppress_target_callbacks = False + groups = self._groups_by_id() + pose_group_ids = tuple( + group_id + for group_id in self.state.selected_group_ids + if (group := groups.get(group_id)) is not None and group.has_pose_target + ) + for key in [key for key in self._handles if key.startswith("ee_control:")]: + if key.removeprefix("ee_control:") not in pose_group_ids: + self.scene.remove_target_controls(key.removeprefix("ee_control:")) + self._handles.pop(key, None) + for group_id in pose_group_ids: + group = groups[group_id] + + def on_transform_update( + target: TransformControlsHandle, + selected_group_id: PlanningGroupID = group_id, + ) -> None: + self._on_transform_update(selected_group_id, target) + + control = self.scene.ensure_target_controls( + str(group_id), + on_transform_update, + ) + if control is not None: + self._handles[f"ee_control:{group_id}"] = control + pose = self.state.pose_targets.get(group_id) + if pose is not None: + self._suppress_target_callbacks = True + try: + self.scene.set_target_pose(str(group_id), pose) + finally: + self._suppress_target_callbacks = False def _build_joint_sliders(self) -> None: - if self.state.selected_robot is None: - return gui = self.server.gui - config = self.adapter.get_robot_config(self.state.selected_robot) - if config is None: - return - current = self.adapter.get_current_joint_state(self.state.selected_robot) - values = list(current.position) if current is not None else [0.0] * len(config.joint_names) self._clear_joint_sliders() - joint_limits_lower = config.joint_limits_lower - joint_limits_upper = config.joint_limits_upper - for index, joint_name in enumerate(config.joint_names): - lower, upper = DEFAULT_JOINT_LIMITS - if joint_limits_lower is not None and index < len(joint_limits_lower): - lower = joint_limits_lower[index] - if joint_limits_upper is not None and index < len(joint_limits_upper): - upper = joint_limits_upper[index] - handle = gui.add_slider( - joint_name, - min=float(lower), - max=float(upper), - step=0.001, - initial_value=float(values[index] if index < len(values) else 0.0), - ) + if not self.state.selected_group_ids: + return + joint_folder = self._handles.get("joint_control_folder") + if joint_folder is not None: + folder = cast("GuiFolderHandle", joint_folder) + with folder: + self._build_joint_slider_handles(gui) + return + self._build_joint_slider_handles(gui) + + def _build_joint_slider_handles(self, gui: GuiApi) -> None: + for group_id in self.state.selected_group_ids: + group = self._groups_by_id().get(group_id) + if group is None: + continue + config = self.get_robot_config(group.robot_name) + target = self.state.group_joint_targets.get(group_id) + if config is None or target is None: + continue + config_indexes = {str(name): index for index, name in enumerate(config.joint_names)} + for _global_name, local_name, value in zip( + group.joint_names, group.local_joint_names, target.position, strict=True + ): + index = config_indexes.get(str(local_name)) + lower, upper = DEFAULT_JOINT_LIMITS + if index is not None and config.joint_limits_lower is not None: + lower = config.joint_limits_lower[index] + if index is not None and config.joint_limits_upper is not None: + upper = config.joint_limits_upper[index] + key = (group_id, str(local_name)) + handle = gui.add_slider( + f"{group_id}/{local_name}", + min=float(lower), + max=float(upper), + step=0.001, + initial_value=float(value), + ) - def on_update(_event: object, name: str = joint_name) -> None: - self._on_joint_slider_update(name) + def on_slider_update( + _event: object, + selected_group_id: PlanningGroupID = group_id, + name: str = str(local_name), + ) -> None: + self._on_joint_slider_update(selected_group_id, name) - handle.on_update(on_update) - self._joint_sliders[joint_name] = handle + handle.on_update(on_slider_update) + self._joint_sliders[key] = handle def _clear_joint_sliders(self) -> None: for handle in self._joint_sliders.values(): @@ -277,6 +596,127 @@ def _clear_joint_sliders(self) -> None: pass self._joint_sliders.clear() + def _groups_by_id(self) -> dict[PlanningGroupID, PlanningGroup]: + return {group.id: group for group in self.list_planning_groups()} + + def _selected_robot_names(self) -> tuple[str, ...]: + groups = self._groups_by_id() + return tuple( + dict.fromkeys( + str(groups[group_id].robot_name) + for group_id in self.state.selected_group_ids + if group_id in groups + ) + ) + + def _stale_robot_names(self, group_ids: tuple[PlanningGroupID, ...]) -> tuple[str, ...]: + """Return every affected robot whose monitored joint state is stale.""" + groups = self._groups_by_id() + robot_names = tuple( + dict.fromkeys( + str(groups[group_id].robot_name) for group_id in group_ids if group_id in groups + ) + ) + return tuple(name for name in robot_names if self.is_state_stale(name)) + + def _state_values_by_local_name(self, state: JointState | None) -> dict[str, float]: + if state is None or len(state.name) != len(state.position): + return {} + return { + str(name): float(value) for name, value in zip(state.name, state.position, strict=True) + } + + def _local_values_for_robot( + self, robot_name: str, state: JointState | None + ) -> dict[str, float]: + config = self.get_robot_config(robot_name) + if config is None or state is None or len(state.name) != len(state.position): + return {} + raw = self._state_values_by_local_name(state) + values: dict[str, float] = {} + for local_name in config.joint_names: + global_name = f"{robot_name}/{local_name}" + if local_name in raw: + values[local_name] = raw[local_name] + elif global_name in raw: + values[local_name] = raw[global_name] + return values + + def _initialize_selected_group_targets(self) -> None: + for group_id in self.state.selected_group_ids: + if group_id in self.state.group_joint_targets: + continue + group = self._groups_by_id().get(group_id) + if group is None: + continue + if self.is_state_stale(group.robot_name): + continue + values = self._local_values_for_robot( + str(group.robot_name), self.get_current_joint_state(group.robot_name) + ) + if any(str(name) not in values for name in group.local_joint_names): + continue + self.state.group_joint_targets[group_id] = JointState( + { + "name": list(group.joint_names), + "position": [float(values[str(name)]) for name in group.local_joint_names], + } + ) + if group.has_pose_target and group_id not in self.state.pose_targets: + pose = self.get_group_ee_pose(group_id) + if pose is not None: + self.state.pose_targets[group_id] = pose + self.state.group_poses[group_id] = pose + if self.state.cartesian_target is None: + self.state.cartesian_target = pose + self._refresh_target_joints_from_groups() + + def _prune_inactive_group_state(self) -> None: + selected = set(self.state.selected_group_ids) + for values in ( + self.state.pose_targets, + self.state.group_joint_targets, + self.state.group_poses, + ): + for group_id in tuple(values): + if group_id not in selected: + values.pop(group_id) + self._refresh_target_joints_from_groups() + + def _refresh_target_joints_from_groups(self) -> None: + names: list[str] = [] + positions: list[float] = [] + for group_id in self.state.selected_group_ids: + target = self.state.group_joint_targets.get(group_id) + if target is not None: + names.extend(str(name) for name in target.name) + positions.extend(float(value) for value in target.position) + self.state.target_joints = ( + JointState({"name": names, "position": positions}) if names else None + ) + + def _active_pose_targets(self) -> dict[PlanningGroupID, Pose]: + return { + group_id: self.state.pose_targets[group_id] + for group_id in self.state.selected_group_ids + if group_id in self.state.pose_targets + } + + def _preset_values_by_local_name(self, preset: str, robot_name: str) -> dict[str, float]: + if preset == "Current": + state = self.get_current_joint_state(robot_name) + elif preset == "Init": + state = self.get_init_joints(robot_name) + else: + config = self.get_robot_config(robot_name) + if config is None: + return {} + return { + str(name): float(value) + for name, value in zip(config.joint_names, config.home_joints or [], strict=False) + } + return self._local_values_for_robot(robot_name, state) + def _remove_panel_handles(self) -> None: for key, handle in list(self._handles.items()): remove = getattr(handle, "remove", None) @@ -284,51 +724,20 @@ def _remove_panel_handles(self) -> None: remove() self._handles.pop(key, None) - def _select_robot(self, robot_name: str) -> None: - if self._closed: - return - if (robot_name or None) == self.state.selected_robot: - self.refresh() - return - self.state.selected_robot = robot_name or None - self.state.target_status = TargetStatus.EMPTY - self.state.feasibility.status = FeasibilityStatus.UNKNOWN - self.state.plan_state = PanelPlanState() - self._build_joint_sliders() - self._sync_preset_dropdown() - self.refresh() - - def _sync_robot_dropdown(self, robots: list[str]) -> None: - handle = self._handles.get("robot") - if handle is None: - return - options = robots or [""] - for attr in ("options", "values"): - if hasattr(handle, attr): - try: - self._set_optional_handle_attr(handle, attr, options) - except Exception: - logger.warning("Could not set robot dropdown %s", attr, exc_info=True) - if hasattr(handle, "value") and self.state.selected_robot in robots: - try: - self._set_optional_handle_attr(handle, "value", self.state.selected_robot) - except Exception: - logger.warning("Could not set robot dropdown value", exc_info=True) - def _sync_preset_dropdown(self) -> None: handle = self._handles.get("preset") - if handle is None or self.state.selected_robot is None: + if handle is None or not self.state.selected_group_ids: return - info: RobotInfo | None = self.adapter.get_robot_info(self.state.selected_robot) - config = self.adapter.get_robot_config(self.state.selected_robot) options = ["Select preset..."] - if (info is not None and info["init_joints"] is not None) or self.adapter.get_init_joints( - self.state.selected_robot - ) is not None: + selected_robots = self._selected_robot_names() + if any(self.get_init_joints(robot_name) is not None for robot_name in selected_robots): options.append("Init") options.append("Current") - home_joints = config.home_joints if config is not None else None - if (info is not None and info["home_joints"] is not None) or home_joints is not None: + if any( + (config := self.get_robot_config(robot_name)) is not None + and config.home_joints is not None + for robot_name in selected_robots + ): options.append("Home") for attr in ("options", "values"): if hasattr(handle, attr): @@ -340,255 +749,413 @@ def _sync_preset_dropdown(self) -> None: def _apply_preset(self, preset: str) -> None: if self._closed: return - robot_name = self.state.selected_robot - if robot_name is None: - return - config = self.adapter.get_robot_config(robot_name) - if config is None: - return - if preset == "Current": - current = self.adapter.get_current_joint_state(robot_name) - values = list(current.position) if current is not None else [] - elif preset == "Init": - init = self.adapter.get_init_joints(robot_name) - values = list(init.position) if init is not None else [] - elif preset == "Home": - values = list(config.home_joints or []) - else: + if preset not in {"Current", "Init", "Home"}: return - self._set_slider_values(config.joint_names, values) - self.state.joint_target = [float(value) for value in values] + targets: dict[PlanningGroupID, JointState] = {} + slider_values: list[tuple[PlanningGroupID, tuple[str, ...], list[float]]] = [] + for group_id in self.state.selected_group_ids: + group = self._groups_by_id().get(group_id) + if group is None: + self._set_recoverable_error(f"Unknown planning group: {group_id}") + return + if preset == "Current" and self.is_state_stale(group.robot_name): + self._set_recoverable_error( + f"Cannot apply Current preset without fresh telemetry for: {group.robot_name}" + ) + return + values = self._preset_values_by_local_name(preset, str(group.robot_name)) + missing = [str(name) for name in group.local_joint_names if str(name) not in values] + if missing: + self._set_recoverable_error( + f"Cannot apply {preset} preset: missing joints for {group_id}: {', '.join(missing)}" + ) + return + positions = [float(values[str(name)]) for name in group.local_joint_names] + targets[group_id] = JointState({"name": list(group.joint_names), "position": positions}) + slider_values.append((group_id, group.local_joint_names, positions)) + self.state.group_joint_targets.update(targets) + if any( + (group_id, str(local_name)) not in self._joint_sliders + for group_id, local_names, _positions in slider_values + for local_name in local_names + ): + self._build_joint_sliders() + for group_id, local_names, positions in slider_values: + self._set_group_slider_values(group_id, local_names, positions) + self._refresh_target_joints_from_groups() self._submit_joint_target_evaluation() self.refresh() - def _set_slider_values(self, joint_names: list[str], values: list[float]) -> None: + def _set_group_slider_values( + self, group_id: PlanningGroupID, local_names: tuple[str, ...], values: list[float] + ) -> None: self._suppress_target_callbacks = True try: - for joint_name, value in zip(joint_names, values, strict=False): - handle = self._joint_sliders.get(joint_name) + for local_name, value in zip(local_names, values, strict=True): + handle = self._joint_sliders.get((group_id, str(local_name))) if handle is not None: handle.value = float(value) finally: self._suppress_target_callbacks = False - def _target_from_sliders(self, robot_name: str) -> JointState | None: - config = self.adapter.get_robot_config(robot_name) - if config is None: - self._set_error("No robot config") - return None - values = [ - float(self._joint_sliders[name].value) - for name in config.joint_names - if name in self._joint_sliders - ] - return self.adapter.joints_from_values(config.joint_names, values) - - def _on_joint_slider_update(self, _joint_name: str) -> None: + def _target_set_from_sliders(self) -> dict[PlanningGroupID, JointState] | None: + targets: dict[PlanningGroupID, JointState] = {} + for group_id in self.state.selected_group_ids: + group = self._groups_by_id().get(group_id) + if group is None: + self._set_error(f"Unknown planning group: {group_id}") + return None + positions: list[float] = [] + for local_name in group.local_joint_names: + handle = self._joint_sliders.get((group_id, str(local_name))) + if handle is None: + self._set_error(f"Missing target slider for {group_id}/{local_name}") + return None + positions.append(float(handle.value)) + targets[group_id] = JointState({"name": list(group.joint_names), "position": positions}) + return targets + + def _on_joint_slider_update(self, _group_id: PlanningGroupID, _local_name: str) -> None: if self._closed: return if self._suppress_target_callbacks: return self._submit_joint_target_evaluation() - def _on_transform_update(self, target: TransformControlsHandle) -> None: + def _on_transform_update( + self, group_id: PlanningGroupID, target: TransformControlsHandle + ) -> None: if self._closed: return - if self._suppress_target_callbacks or self.state.selected_robot is None: + if self._suppress_target_callbacks or group_id not in self.state.selected_group_ids: return pose = self._pose_from_transform_target(target) if pose is None: return self.state.cartesian_target = pose + self.state.pose_targets[group_id] = pose sequence_id = self.state.next_sequence_id() self._worker.submit( TargetEvaluationRequest( sequence_id=sequence_id, source="cartesian", - robot_name=self.state.selected_robot, - pose=pose, + selection_epoch=self.state.selection_epoch, + group_ids=self.state.selected_group_ids, + auxiliary_group_ids=tuple( + selected_group_id + for selected_group_id in self.state.selected_group_ids + if selected_group_id not in self._active_pose_targets() + ), + joints=( + None + if self.state.target_joints is None + else JointState(self.state.target_joints) + ), + pose_targets=dict(self._active_pose_targets()), ) ) self.refresh() def _submit_joint_target_evaluation(self) -> None: - robot_name = self.state.selected_robot - if robot_name is None: + targets = self._target_set_from_sliders() + if targets is None: return - target = self._target_from_sliders(robot_name) - if target is None: - return - self.state.joint_target = list(target.position) - self._move_joint_target_visuals(robot_name, target) + self.state.group_joint_targets = targets + self._refresh_target_joints_from_groups() + self._move_joint_target_visuals(targets) sequence_id = self.state.next_sequence_id() self._worker.submit( TargetEvaluationRequest( sequence_id=sequence_id, source="joints", - robot_name=robot_name, - joints=target, + selection_epoch=self.state.selection_epoch, + group_ids=self.state.selected_group_ids, + joint_targets=dict(targets), ) ) self.refresh() - def _move_joint_target_visuals(self, robot_name: str, target: JointState) -> None: + def _move_joint_target_visuals(self, targets: Mapping[PlanningGroupID, JointState]) -> None: """Optimistically move target visuals before collision/feasibility returns.""" - config = self.adapter.get_robot_config(robot_name) - robot_id = self.adapter.robot_id_for_name(robot_name) - if self.scene is not None and config is not None and robot_id is not None: - self.scene.set_target_joints(str(robot_id), config.joint_names, list(target.position)) - pose = self.adapter.get_ee_pose(robot_name, target) - if pose is not None: - self._suppress_target_callbacks = True - try: - self.scene.set_target_pose(str(robot_id), pose) - finally: - self._suppress_target_callbacks = False + if self.scene is None: + return + for robot_name, state in self._target_ghost_states(targets).items(): + config = self.get_robot_config(robot_name) + robot_id = self.robot_id_for_name(robot_name) + if config is not None and robot_id is not None: + self.scene.set_target_joints(str(robot_id), config.joint_names, state.position) + + def _target_ghost_states( + self, targets: Mapping[PlanningGroupID, JointState] + ) -> dict[str, JointState]: + groups = self._groups_by_id() + merged: dict[str, dict[str, float]] = {} + configs: dict[str, tuple[str, ...]] = {} + for group_id in self.state.selected_group_ids: + group = groups.get(group_id) + target = targets.get(group_id) + if group is None or target is None: + continue + robot_name = str(group.robot_name) + config = self.get_robot_config(robot_name) + current = self.get_current_joint_state(robot_name) + if config is None or current is None: + continue + values = self._local_values_for_robot(robot_name, current) + target_raw = self._state_values_by_local_name(target) + for local_name, global_name in zip( + group.local_joint_names, group.joint_names, strict=True + ): + if str(global_name) in target_raw: + values[str(local_name)] = target_raw[str(global_name)] + elif str(local_name) in target_raw: + values[str(local_name)] = target_raw[str(local_name)] + if all(name in values for name in config.joint_names): + merged[robot_name] = values + configs[robot_name] = tuple(config.joint_names) + return { + robot_name: JointState( + {"name": list(joint_names), "position": [values[name] for name in joint_names]} + ) + for robot_name, values in merged.items() + for joint_names in (configs[robot_name],) + } + + def _sync_target_ghost_visibility(self) -> None: + if self.scene is None: + return + active_robot_ids = { + str(robot_id) + for group_id in self.state.selected_group_ids + if (group := self._groups_by_id().get(group_id)) is not None + and group.has_pose_target + and (robot_id := self.robot_id_for_name(group.robot_name)) is not None + } + for _robot_name, robot_id, _config in self.robot_items(): + self.scene.set_target_active(str(robot_id), str(robot_id) in active_robot_ids) def _handle_target_evaluation_request( self, request: TargetEvaluationRequest - ) -> TargetEvaluation: + ) -> TargetEvaluationResult: if request.source == "cartesian": - if request.pose is None: - return {"success": False, "status": "INVALID", "message": "No pose target"} - return self.adapter.evaluate_pose_target(request.pose, request.robot_name) - if request.joints is None: - return {"success": False, "status": "INVALID", "message": "No joint target"} - return self.adapter.evaluate_joint_target(request.joints, request.robot_name) + if not request.pose_targets: + return TargetEvaluationResult(False, "INVALID", "No pose target") + return self.evaluate_pose_target_set( + request.pose_targets, request.auxiliary_group_ids, request.joints + ) + if not request.joint_targets: + return TargetEvaluationResult(False, "INVALID", "No joint target") + return self.evaluate_joint_target_set(request.group_ids, request.joint_targets) def _apply_target_evaluation_result( - self, request: TargetEvaluationRequest, result: TargetEvaluation + self, request: TargetEvaluationRequest, result: TargetEvaluationResult ) -> None: if self._closed: return - if request.sequence_id != self.state.latest_sequence_id: + if ( + request.sequence_id != self.state.latest_sequence_id + or request.selection_epoch != self.state.selection_epoch + or request.group_ids != self.state.selected_group_ids + ): + return + if self.state.manipulation_state == "FAULT": return - collision_free = bool(result.get("collision_free", False)) - success = bool(result.get("success", False)) + collision_free = result.collision_free + success = result.success self.state.feasibility.status = self._feasibility_status(result, success, collision_free) - self.state.feasibility.message = str(result.get("message", "")) + self.state.feasibility.message = result.message self.state.target_status = ( TargetStatus.FEASIBLE if success and collision_free else TargetStatus.INFEASIBLE ) self.state.error = "" if success and collision_free else self.state.feasibility.message + if result.target_joints is not None: + self.state.target_joints = JointState(result.target_joints) + self._split_target_joints_by_group(result.target_joints) + self.state.group_poses = { + str(group_id): pose + for group_id, pose in result.group_poses.items() + if isinstance(pose, Pose) + } if request.source == "joints": - joint_state = result.get("joint_state") - if isinstance(joint_state, JointState): - self.state.joint_target = list(joint_state.position) - if request.source == "cartesian": - joint_state = result.get("joint_state") - if isinstance(joint_state, JointState): - self.state.joint_target = list(joint_state.position) - pose = result.get("ee_pose") - if isinstance(pose, Pose): - self.state.cartesian_target = pose + self._sync_pose_targets_from_group_poses() + else: self._sync_controls_from_targets() self._update_target_visual_state() self.refresh() def _sync_controls_from_targets(self) -> None: - robot_name = self.state.selected_robot - if robot_name is None: - return - config = self.adapter.get_robot_config(robot_name) - if config is not None and self.state.joint_target is not None: - self._set_slider_values(list(config.joint_names), list(self.state.joint_target)) - robot_id = self.adapter.robot_id_for_name(robot_name) - if self.scene is not None and robot_id is not None: - self.scene.set_target_joints( - str(robot_id), config.joint_names, self.state.joint_target + for group_id, target in self.state.group_joint_targets.items(): + group = self._groups_by_id().get(group_id) + if group is not None: + self._set_group_slider_values( + group_id, group.local_joint_names, list(target.position) ) - # Do not write the Cartesian target back into the active transform - # control here. The gizmo is the source of truth for Cartesian edits; - # programmatic pose writes from delayed IK results can fight fast user - # dragging and make the gizmo jump back. + self._move_joint_target_visuals(self.state.group_joint_targets) + + def _split_target_joints_by_group(self, target_joints: JointState) -> None: + if len(target_joints.name) != len(target_joints.position): + return + positions = { + str(name): float(value) + for name, value in zip(target_joints.name, target_joints.position, strict=True) + } + for group_id in self.state.selected_group_ids: + group = self._groups_by_id().get(group_id) + if group is None or any(str(name) not in positions for name in group.joint_names): + continue + self.state.group_joint_targets[group_id] = JointState( + { + "name": list(group.joint_names), + "position": [positions[str(name)] for name in group.joint_names], + } + ) + + def _sync_pose_targets_from_group_poses(self) -> None: + groups = self._groups_by_id() + active_group_ids: list[PlanningGroupID] = [] + for group_id, pose in self.state.group_poses.items(): + group = groups.get(group_id) + if group is None or not group.has_pose_target: + continue + if group_id not in self.state.selected_group_ids: + continue + self.state.pose_targets[group_id] = pose + active_group_ids.append(group_id) + if self.scene is None: + return + self._suppress_target_callbacks = True + try: + for group_id in active_group_ids: + self.scene.set_target_pose(str(group_id), self.state.pose_targets[group_id]) + finally: + self._suppress_target_callbacks = False def _update_status_text(self) -> None: current = self.state.current_joints + status_label = self.state.error or self.state.module_state status = [ - "### Manipulation Panel", - f"Robot: `{self.state.selected_robot or 'none'}`", - f"Module: `{self.state.module_state}`", - f"Backend: `{self.state.backend_status.value}`", - f"Target: `{self.state.target_status.value}`", - f"Feasibility: `{self.state.feasibility.status.value}`", - f"Plan: `{self.state.plan_state.status.value}`", - f"Action: `{self.state.action_status.value}`", + "### Status", + f"**State:** {status_label}", + f"Target: `{self.state.target_status.value}` · Plan: `{self.state.plan_state.status.value}`", ] - if self.state.selected_robot is not None: - status.append( - f"State stale: `{self.adapter.is_state_stale(self.state.selected_robot)}`" - ) + stale_robots = self._stale_robot_names(self.state.selected_group_ids) + if self.state.selected_group_ids: + stale_detail = "False" if not stale_robots else f"True ({', '.join(stale_robots)})" + status.append(f"State stale: `{stale_detail}`") if current is not None: status.append(f"Current joints: `{[round(v, 3) for v in current]}`") if self.state.last_result: status.append(f"Last result: `{self.state.last_result}`") - if self.state.error: - status.append(f"Error: `{self.state.error}`") self._set_handle_value("status", "\n\n".join(status)) + self._set_handle_value( + "target_summary", + f"Feasibility: `{self.state.feasibility.status.value}`", + ) def _update_control_state(self) -> None: self._set_disabled("plan", not self.state.can_plan()) self._set_disabled("preview", not self.state.can_preview()) self._set_disabled( "execute", - not self.state.can_execute(self.config.current_match_tolerance), + not self._can_execute(), ) - self._set_disabled("cancel", not self.state.can_cancel()) + can_cancel = self.state.can_cancel() + self._set_disabled("cancel", not can_cancel) + self._set_visible("cancel", can_cancel) self._update_target_visual_state() def _update_target_visual_state(self) -> None: - if self.scene is None or self.state.selected_robot is None: - return - robot_id = self.adapter.robot_id_for_name(self.state.selected_robot) - if robot_id is None: + if self.scene is None: return - self.scene.set_target_visual_state( - str(robot_id), self.state.feasibility.status == FeasibilityStatus.FEASIBLE + feasible = self.state.feasibility.status == FeasibilityStatus.FEASIBLE + groups = self._groups_by_id() + selected_groups = tuple( + (group_id, groups[group_id]) + for group_id in self.state.selected_group_ids + if group_id in groups + ) + for group_id, group in selected_groups: + if group.has_pose_target: + self.scene.set_target_control_visual_state(str(group_id), feasible) + robot_ids = tuple( + dict.fromkeys( + str(robot_id) + for _group_id, group in selected_groups + if (robot_id := self.robot_id_for_name(str(group.robot_name))) is not None + ) ) + for robot_id in robot_ids: + self.scene.set_target_robot_visual_state(robot_id, feasible) + + def _can_execute(self) -> bool: + return self.state.can_execute() def _submit_plan(self) -> None: if self._closed: return - robot_name = self.state.selected_robot - if robot_name is None: - return if not self.state.can_plan(): self._set_recoverable_error( "Cannot plan until target is feasible and manipulation is idle" ) return + group_ids = self.state.selected_group_ids + selection_epoch = self.state.selection_epoch + target_sequence_id = self.state.latest_sequence_id + targets = self._target_set_from_sliders() + if targets is None: + return operation_id = self._next_operation_id() def operation() -> None: - if not self._operation_is_current(operation_id): + if not self._operation_is_current(operation_id, selection_epoch, target_sequence_id): return self.state.action_status = ActionStatus.RUNNING self.state.plan_state.status = PlanStatus.PLANNING - if self.state.manipulation_state == "FAULT" and not self.adapter.reset(): - self.state.plan_state.status = PlanStatus.FAILED - self._finish_operation("reset=False", clear_error=False, operation_id=operation_id) + stale_robots = self._stale_robot_names(group_ids) + if stale_robots: + if not self._operation_is_current( + operation_id, selection_epoch, target_sequence_id + ): + self._finish_operation( + "plan=False", operation_id=operation_id, selection_epoch=selection_epoch + ) + return + self.state.plan_state.status = PlanStatus.STALE + self.state.error = "Cannot plan without fresh telemetry for: " + ", ".join( + stale_robots + ) + self._finish_operation( + "plan=False", + clear_error=False, + operation_id=operation_id, + selection_epoch=selection_epoch, + ) return - target = self._target_from_sliders(robot_name) - if target is None: - self.state.plan_state.status = PlanStatus.FAILED + if not self._operation_is_current(operation_id, selection_epoch, target_sequence_id): self._finish_operation( - "plan_to_joints=False", clear_error=False, operation_id=operation_id + "plan=False", operation_id=operation_id, selection_epoch=selection_epoch ) return - ok = self.adapter.plan_to_joints(target, robot_name) - if not self._operation_is_current(operation_id): + ok = self.plan_to_selected_joints(group_ids, targets) + if not self._operation_is_current(operation_id, selection_epoch, target_sequence_id): + self._finish_operation( + "plan=False", operation_id=operation_id, selection_epoch=selection_epoch + ) return if ok: - path = self.adapter.get_planned_path(robot_name) + status = self.operator.status() self.state.plan_state.status = PlanStatus.FRESH - self.state.plan_state.robot = robot_name - self.state.plan_state.target_joints = list(target.position) - self.state.plan_state.target_pose = self.state.cartesian_target - self.state.plan_state.start_joints_snapshot = list(self.state.current_joints or []) - self.state.plan_state.planned_path = path + self.state.plan_state.group_ids = group_ids + self.state.plan_state.target_sequence_id = target_sequence_id + self.state.plan_state.plan_id = getattr(status, "ready_plan_id", None) else: self.state.plan_state.status = PlanStatus.FAILED - self._finish_operation(f"plan_to_joints={ok}", operation_id=operation_id) + self.state.plan_state.plan = None + self._finish_operation( + f"plan_to_joints={ok}", + operation_id=operation_id, + selection_epoch=selection_epoch, + ) self._operation_worker.submit( operation, on_error=lambda message: self._set_operation_error(message, operation_id) @@ -597,20 +1164,21 @@ def operation() -> None: def _submit_preview(self) -> None: if self._closed: return - robot_name = self.state.selected_robot - if robot_name is None: - return if not self.state.can_preview(): self._set_recoverable_error("No fresh plan to preview") return + selection_epoch = self.state.selection_epoch operation_id = self._next_operation_id() + self.state.action_status = ActionStatus.PREVIEWING + self.refresh() def operation() -> None: - if not self._operation_is_current(operation_id): + if not self._operation_is_current(operation_id, selection_epoch): return - self.state.action_status = ActionStatus.PREVIEWING - ok = self.adapter.preview_path(robot_name) - self._finish_operation(f"preview={ok}", operation_id=operation_id) + ok = self.preview_path() + self._finish_operation( + f"preview={ok}", operation_id=operation_id, selection_epoch=selection_epoch + ) self._operation_worker.submit( operation, @@ -618,30 +1186,59 @@ def operation() -> None: on_error=lambda message: self._set_operation_error(message, operation_id), ) + def _clear_invalidated_preview(self) -> None: + if self.state.action_status == ActionStatus.PREVIEWING: + self._operation_sequence_id += 1 + self.state.action_status = ActionStatus.IDLE + self.state.last_result = "preview=False" + def _submit_execute(self) -> None: if self._closed: return - robot_name = self.state.selected_robot - if robot_name is None: - return - if not self.state.can_execute(self.config.current_match_tolerance): - self._set_recoverable_error( - "Cannot execute: require feasible fresh plan and matching current joints" - ) + if not self._can_execute(): + self._set_recoverable_error("Cannot execute: require feasible fresh plan") return + selection_epoch = self.state.selection_epoch + group_ids = self.state.selected_group_ids + target_sequence_id = self.state.latest_sequence_id operation_id = self._next_operation_id() def operation() -> None: - if not self._operation_is_current(operation_id): + if not self._operation_is_current(operation_id, selection_epoch, target_sequence_id): + return + if ( + self.state.plan_state.group_ids != group_ids + or self.state.plan_state.target_sequence_id != target_sequence_id + ): + self.state.plan_state.status = PlanStatus.STALE + self._finish_operation( + "execute=False", + clear_error=False, + operation_id=operation_id, + selection_epoch=selection_epoch, + ) + return + if not self._operation_is_current(operation_id, selection_epoch, target_sequence_id): + self._finish_operation( + "execute=False", operation_id=operation_id, selection_epoch=selection_epoch + ) return self.state.action_status = ActionStatus.EXECUTING - self.state.plan_state.status = PlanStatus.EXECUTING - ok = self.adapter.execute(robot_name) - if not self._operation_is_current(operation_id): + ok = self.execute() + if not self._operation_is_current(operation_id, selection_epoch, target_sequence_id): + self._finish_operation( + "execute=False", operation_id=operation_id, selection_epoch=selection_epoch + ) return if not ok: self.state.plan_state.status = PlanStatus.FAILED - self._finish_operation(f"execute={ok}", operation_id=operation_id) + self._finish_operation( + f"execute={ok}", operation_id=operation_id, selection_epoch=selection_epoch + ) + else: + # True means dispatch was accepted, not that motion completed. + self.state.last_result = "execute=True (dispatched)" + self.refresh() self._operation_worker.submit( operation, on_error=lambda message: self._set_operation_error(message, operation_id) @@ -658,7 +1255,7 @@ def _submit_cancel(self) -> None: self._mark_cancelled_plan_state(cancelled_action) self._restart_operation_worker() try: - ok = self.adapter.cancel() + ok = self.cancel() except Exception as e: self._set_operation_error(str(e), operation_id) return @@ -667,10 +1264,7 @@ def _submit_cancel(self) -> None: def _mark_cancelled_plan_state(self, cancelled_action: ActionStatus) -> None: if self.state.plan_state.status == PlanStatus.PLANNING: self.state.plan_state.status = PlanStatus.FAILED - elif ( - cancelled_action == ActionStatus.EXECUTING - or self.state.plan_state.status == PlanStatus.EXECUTING - ): + elif cancelled_action == ActionStatus.EXECUTING: self.state.plan_state.status = PlanStatus.STALE def _restart_operation_worker(self) -> None: @@ -687,7 +1281,7 @@ def operation() -> None: if not self._operation_is_current(operation_id): return self.state.action_status = ActionStatus.CLEARING_PLAN - ok = self.adapter.clear_planned_path() + ok = self.clear_planned_path() if not self._operation_is_current(operation_id): return self.state.plan_state = PanelPlanState() @@ -701,19 +1295,38 @@ def _next_operation_id(self) -> int: self._operation_sequence_id += 1 return self._operation_sequence_id - def _operation_is_current(self, operation_id: int) -> bool: - return not self._closed and operation_id == self._operation_sequence_id + def _operation_is_current( + self, + operation_id: int, + selection_epoch: int | None = None, + target_sequence_id: int | None = None, + ) -> bool: + return ( + not self._closed + and operation_id == self._operation_sequence_id + and (selection_epoch is None or selection_epoch == self.state.selection_epoch) + and (target_sequence_id is None or target_sequence_id == self.state.latest_sequence_id) + ) def _finish_operation( - self, result: str, *, clear_error: bool = True, operation_id: int | None = None + self, + result: str, + *, + clear_error: bool = True, + operation_id: int | None = None, + selection_epoch: int | None = None, ) -> None: if self._closed or ( - operation_id is not None and not self._operation_is_current(operation_id) + operation_id is not None + and not self._operation_is_current(operation_id, selection_epoch) ): return + if self.state.manipulation_state == "FAULT": + return self.state.action_status = ActionStatus.IDLE if clear_error: self.state.error = "" + self.state.local_action_failure = False self.state.last_result = result self.refresh() @@ -725,13 +1338,18 @@ def _set_operation_error(self, message: str, operation_id: int) -> None: def _set_recoverable_error(self, message: str) -> None: if self._closed: return - self.state.error = message self.refresh() + if self.state.manipulation_state == "FAULT": + return + # This is a panel-local validation message, not a runtime diagnostic. + self.state.error = message + self._update_status_text() def _set_error(self, message: str) -> None: if self._closed: return self.state.action_status = ActionStatus.FAILED + self.state.local_action_failure = True self.state.error = message self.refresh() @@ -745,6 +1363,11 @@ def _set_disabled(self, key: str, disabled: bool) -> None: if isinstance(handle, GuiButtonHandle): self._set_optional_handle_attr(handle, "disabled", disabled) + def _set_visible(self, key: str, visible: bool) -> None: + handle = self._handles.get(key) + if handle is not None: + self._set_optional_handle_attr(handle, "visible", visible) + @staticmethod def _set_optional_handle_attr(handle: object, attr: str, value: object) -> None: setattr(handle, attr, value) @@ -755,9 +1378,9 @@ def _pose_from_transform_target(self, target: TransformControlsHandle) -> Pose | return Pose({"position": [px, py, pz], "orientation": [qx, qy, qz, qw]}) def _feasibility_status( - self, result: TargetEvaluation, success: bool, collision_free: bool + self, result: TargetEvaluationResult, success: bool, collision_free: bool ) -> FeasibilityStatus: - status = str(result.get("status", "")).upper() + status = result.status.upper() if success and collision_free: return FeasibilityStatus.FEASIBLE if status in {"COLLISION", "COLLISION_AT_START", "COLLISION_AT_GOAL"}: diff --git a/dimos/manipulation/visualization/viser/scene.py b/dimos/manipulation/visualization/viser/scene.py index cd50ab098e..296e18319f 100644 --- a/dimos/manipulation/visualization/viser/scene.py +++ b/dimos/manipulation/visualization/viser/scene.py @@ -15,22 +15,35 @@ from __future__ import annotations from collections.abc import Callable, Sequence +import hashlib +import os from pathlib import Path +import tempfile +from threading import RLock +import time from typing import Protocol, TypeAlias, cast +import xml.etree.ElementTree as ET from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.utils.mesh_utils import prepare_urdf_for_drake -from dimos.manipulation.visualization.viser.animation import PreviewAnimator +from dimos.manipulation.visualization.viser.animation import ( + GroupPreviewAnimation, + PreviewFrame, + preview_tick_times, + scaled_frame_delays, +) from dimos.manipulation.visualization.viser.runtime import ( VISER_INSTALL_HINT, VISER_URDF_INSTALL_HINT, ) from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.robot.model_parser import parse_model from dimos.utils.logging_config import setup_logger try: from viser import ( + FrameHandle, GridHandle, MeshHandle, TransformControlsEvent, @@ -55,6 +68,8 @@ logger = setup_logger() +_VISER_URDF_CACHE_DIR = Path(tempfile.gettempdir()) / "dimos_viser_urdf_cache" + GOAL_ROBOT_FEASIBLE_COLOR = (255, 122, 0) GOAL_ROBOT_INFEASIBLE_COLOR = (255, 30, 30) GOAL_ROBOT_FEASIBLE_OPACITY = 0.7 @@ -69,7 +84,7 @@ REFERENCE_GRID_CELL_COLOR = (44, 54, 58) REFERENCE_GRID_SECTION_COLOR = (90, 145, 165) -SceneHandle: TypeAlias = ViserUrdf | TransformControlsHandle | GridHandle | MeshHandle +SceneHandle: TypeAlias = ViserUrdf | TransformControlsHandle | GridHandle | MeshHandle | FrameHandle class _ColorHandle(Protocol): @@ -79,19 +94,21 @@ class _ColorHandle(Protocol): class ViserManipulationScene: """Viser scene graph helpers for current robot, ghost robot, and path rendering.""" - def __init__( - self, server: ViserServer, viser_urdf: type[ViserUrdf], *, preview_fps: float - ) -> None: + def __init__(self, server: ViserServer, viser_urdf: type[ViserUrdf]) -> None: self.server = server self.viser_urdf = viser_urdf - self.preview_fps = preview_fps self._configs_by_id: dict[str, RobotModelConfig] = {} self._urdfs: dict[str, ViserUrdf] = {} self._handles: dict[str, TransformControlsHandle] = {} + self._root_frames: dict[str, FrameHandle] = {} self._grid_handle: GridHandle | None = None self._grid_visible = True self._preview_visible: dict[str, bool] = {} + self._target_active: dict[str, bool] = {} self._target_tracks_current: dict[str, bool] = {} + self._scene_lock = RLock() + self._animation_generation = 0 + self._animation_generations: dict[str, int] = {} self._ensure_reference_grid() def has_reference_grid(self) -> bool: @@ -106,9 +123,18 @@ def set_reference_grid_visible(self, visible: bool) -> None: def register_robot(self, robot_id: str, config: RobotModelConfig) -> None: self._configs_by_id[robot_id] = config self._preview_visible.setdefault(robot_id, False) + self._animation_generations.setdefault(robot_id, 0) + self._target_active.setdefault(robot_id, False) self._target_tracks_current.setdefault(robot_id, True) self._ensure_robot_urdfs(robot_id, config) + def set_target_active(self, robot_id: str, active: bool) -> None: + """Show the target ghost only while a pose-target group is selected.""" + self._target_active[robot_id] = active + if not active: + self._target_tracks_current[robot_id] = True + self._set_target_visibility(robot_id, active) + def _ensure_reference_grid(self) -> None: try: scene = self.server.scene @@ -155,41 +181,104 @@ def dispatch(event: TransformControlsEvent) -> None: self._handles[handle_key] = handle return handle + def remove_target_controls(self, control_id: str) -> None: + self._remove_handle(f"{control_id}:ee_control") + def update_current_robot(self, robot_id: str, joint_state: JointState | None) -> None: - config = self._configs_by_id.get(robot_id) - if config is None or joint_state is None: - return - self._ensure_robot_urdfs(robot_id, config) - current = self._urdfs.get(f"{robot_id}:current") - self.set_urdf_joints(current, config.joint_names, joint_state.position) - if self._target_tracks_current.get(robot_id, True): - self._set_target_joints(robot_id, config.joint_names, joint_state.position) - self._set_target_visibility(robot_id, True) + with self._scene_lock: + config = self._configs_by_id.get(robot_id) + if config is None or joint_state is None: + return + self._ensure_robot_urdfs(robot_id, config) + current = self._urdfs.get(f"{robot_id}:current") + self.set_urdf_joints(current, config.joint_names, joint_state.position) + if self._target_tracks_current.get(robot_id, True): + self._set_target_joints(robot_id, config.joint_names, joint_state.position) + self._set_target_visibility(robot_id, self._target_active.get(robot_id, False)) + + def cancel_preview_animation(self, robot_ids: Sequence[str] | None = None) -> None: + """Prevent an old blocking animation from touching replacement handles.""" + with self._scene_lock: + self._animation_generation += 1 + affected = set(robot_ids) if robot_ids is not None else set(self._preview_visible) + for robot_id in affected: + self._animation_generations[robot_id] = ( + self._animation_generations.get(robot_id, 0) + 1 + ) + if robot_id not in self._preview_visible: + continue + self._preview_visible[robot_id] = False + self._set_preview_visibility(robot_id, False) - def show_preview(self, robot_id: str) -> None: - """Show the transient preview-animation ghost. + def animate_preview(self, preview: GroupPreviewAnimation, duration: float) -> bool: + """Play every robot from one normalized tick clock. - Target editing uses the separate target ghost and must not call this path. + Inputs are fully validated before ghosts become visible; a generation + replacement, clear, or close stops mutation before the next tick. """ - self._preview_visible[robot_id] = True - self._set_preview_visibility(robot_id, True) - - def hide_preview(self, robot_id: str) -> None: - """Hide the transient preview-animation ghost.""" - self._preview_visible[robot_id] = False - self._set_preview_visibility(robot_id, False) - - def animate_path(self, robot_id: str, path: Sequence[JointState], duration: float) -> bool: - config = self._configs_by_id.get(robot_id) - if config is None: + frames = {track.robot_id: track.frames for track in preview.tracks} + names = {track.robot_id: track.joint_names for track in preview.tracks} + if ( + not frames + or len(frames) != len(preview.tracks) + or any( + not values or robot_id not in self._configs_by_id + for robot_id, values in frames.items() + ) + ): return False - self.show_preview(robot_id) + tick_times = preview_tick_times(preview) + if not tick_times: + return False + with self._scene_lock: + self._animation_generation += 1 + generations: dict[str, int] = {} + for robot_id in frames: + self._animation_generations[robot_id] = ( + self._animation_generations.get(robot_id, 0) + 1 + ) + generations[robot_id] = self._animation_generations[robot_id] + self._preview_visible[robot_id] = True + self._set_preview_visibility(robot_id, True) try: - return PreviewAnimator( - lambda joints: self._set_preview_ghost_joints(robot_id, config.joint_names, joints) - ).animate(path, duration, self.preview_fps) + delays = scaled_frame_delays( + tuple( + PreviewFrame(time_from_start=tick_time, positions=()) + for tick_time in tick_times + ), + duration, + ) + frame_indices = {robot_id: 0 for robot_id in frames} + for tick, tick_time in enumerate(tick_times): + with self._scene_lock: + active_robot_ids = [ + robot_id + for robot_id in frames + if self._animation_generations.get(robot_id) == generations[robot_id] + ] + if not active_robot_ids: + return False + for robot_id in active_robot_ids: + robot_frames = frames[robot_id] + while ( + frame_indices[robot_id] + 1 < len(robot_frames) + and robot_frames[frame_indices[robot_id] + 1].time_from_start + <= tick_time + ): + frame_indices[robot_id] += 1 + source = frame_indices[robot_id] + self._set_preview_ghost_joints( + robot_id, names[robot_id], robot_frames[source].positions + ) + if tick < len(delays): + time.sleep(delays[tick]) + return True finally: - self.hide_preview(robot_id) + with self._scene_lock: + for robot_id in frames: + if self._animation_generations.get(robot_id) == generations[robot_id]: + self._preview_visible[robot_id] = False + self._set_preview_visibility(robot_id, False) def set_target_joints( self, robot_id: str, joint_names: Sequence[str], joints: Sequence[float] @@ -235,16 +324,26 @@ def set_target_pose(self, robot_id: str, pose: Pose | None) -> None: ) def set_target_visual_state(self, robot_id: str, feasible: bool) -> None: + """Set the legacy matching robot/control target visual state.""" + self.set_target_control_visual_state(robot_id, feasible) + self.set_target_robot_visual_state(robot_id, feasible) + + def set_target_control_visual_state(self, control_id: str, feasible: bool) -> None: + """Set feasibility color for one planning-group keyed target control.""" color = TARGET_CONTROL_FEASIBLE_COLOR if feasible else TARGET_CONTROL_INFEASIBLE_COLOR - mesh_color = GOAL_ROBOT_FEASIBLE_COLOR if feasible else GOAL_ROBOT_INFEASIBLE_COLOR - mesh_opacity = GOAL_ROBOT_FEASIBLE_OPACITY if feasible else GOAL_ROBOT_INFEASIBLE_OPACITY - handle = self._handles.get(f"{robot_id}:ee_control") + handle = self._handles.get(f"{control_id}:ee_control") if handle is not None: cast("_ColorHandle", handle).color = color + + def set_target_robot_visual_state(self, robot_id: str, feasible: bool) -> None: + """Set feasibility material for one robot-ID keyed target ghost.""" + mesh_color = GOAL_ROBOT_FEASIBLE_COLOR if feasible else GOAL_ROBOT_INFEASIBLE_COLOR + mesh_opacity = GOAL_ROBOT_FEASIBLE_OPACITY if feasible else GOAL_ROBOT_INFEASIBLE_OPACITY target = self._urdfs.get(f"{robot_id}:target") self._set_urdf_mesh_material(target, mesh_color, mesh_opacity) def close(self) -> None: + self.cancel_preview_animation() for key in list(self._handles): self._remove_handle(key) if self._grid_handle is not None: @@ -252,9 +351,13 @@ def close(self) -> None: self._grid_handle = None for urdf in self._urdfs.values(): self._remove_scene_handle(urdf) + for frame in self._root_frames.values(): + self._remove_scene_handle(frame) self._urdfs.clear() + self._root_frames.clear() self._configs_by_id.clear() self._preview_visible.clear() + self._target_active.clear() self._target_tracks_current.clear() def _ensure_robot_urdfs(self, robot_id: str, config: RobotModelConfig) -> None: @@ -264,11 +367,7 @@ def _ensure_robot_urdfs(self, robot_id: str, config: RobotModelConfig) -> None: key = f"{robot_id}:{kind}" if key in self._urdfs: continue - root_node_name = { - "current": f"/robots/{robot_id}/current", - "target": f"/targets/{robot_id}/target", - "preview": f"/previews/{robot_id}/ghost", - }[kind] + root_node_name = self._urdf_root_node_name(robot_id, kind, config) mesh_color_override = { "current": None, "target": GOAL_ROBOT_MESH_COLOR, @@ -284,7 +383,9 @@ def _ensure_robot_urdfs(self, robot_id: str, config: RobotModelConfig) -> None: self._set_urdf_mesh_material( self._urdfs[key], GOAL_ROBOT_FEASIBLE_COLOR, GOAL_ROBOT_FEASIBLE_OPACITY ) - self._set_handle_visibility(self._urdfs[key], True) + self._set_handle_visibility( + self._urdfs[key], self._target_active.get(robot_id, False) + ) elif kind == "preview": self._set_urdf_mesh_material( self._urdfs[key], PREVIEW_ROBOT_COLOR, PREVIEW_ROBOT_OPACITY @@ -295,7 +396,7 @@ def _ensure_robot_urdfs(self, robot_id: str, config: RobotModelConfig) -> None: def prepared_urdf_path(self, config: RobotModelConfig) -> Path: package_paths = {package: Path(path) for package, path in config.package_paths.items()} - return Path( + prepared_path = Path( prepare_urdf_for_drake( Path(str(config.model_path)), package_paths=package_paths, @@ -303,6 +404,137 @@ def prepared_urdf_path(self, config: RobotModelConfig) -> Path: convert_meshes=bool(config.auto_convert_meshes), ) ) + prepared_path = self._strip_visualization_world_root_attachment(config, prepared_path) + self._assert_base_link_is_urdf_root(config, prepared_path) + return prepared_path + + @staticmethod + def _strip_visualization_world_root_attachment( + config: RobotModelConfig, prepared_path: Path + ) -> Path: + """Detach a model-owned world root only for Viser base-pose rendering.""" + urdf_content = prepared_path.read_text() + try: + root = ET.fromstring(urdf_content) + except ET.ParseError: + return prepared_path + + attachments = [ + joint + for joint in root.findall("joint") + if joint.attrib.get("type") == "fixed" + and (parent := joint.find("parent")) is not None + and parent.attrib.get("link") == "world" + and (child := joint.find("child")) is not None + and child.attrib.get("link") == config.base_link + ] + if len(attachments) != 1: + return prepared_path + + root.remove(attachments[0]) + referenced_links = { + link + for joint in root.findall("joint") + for element in (joint.find("parent"), joint.find("child")) + if (link := element.get("link") if element is not None else None) is not None + } + if "world" in referenced_links: + return prepared_path + world_links = [link for link in root.findall("link") if link.attrib.get("name") == "world"] + if len(world_links) != 1: + return prepared_path + root.remove(world_links[0]) + + stripped_content = ET.tostring(root, encoding="unicode") + digest = hashlib.sha256( + f"viser-world-root-v1\0{config.base_link}\0{urdf_content}".encode() + ).hexdigest() + cache_path = _VISER_URDF_CACHE_DIR / f"{digest}.urdf" + if cache_path.exists(): + return cache_path + _VISER_URDF_CACHE_DIR.mkdir(parents=True, exist_ok=True) + temporary_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=_VISER_URDF_CACHE_DIR, + prefix=f".{digest}.", + suffix=".tmp", + delete=False, + ) as temporary_file: + temporary_file.write(stripped_content) + temporary_path = Path(temporary_file.name) + os.replace(temporary_path, cache_path) + finally: + if temporary_path is not None: + temporary_path.unlink(missing_ok=True) + return cache_path + + @staticmethod + def _assert_base_link_is_urdf_root(config: RobotModelConfig, prepared_path: Path) -> None: + root_link = parse_model(prepared_path).root_link + if root_link == config.base_link: + return + raise ValueError( + f"Viser visualization requires base_link '{config.base_link}' to match " + f"the prepared URDF root '{root_link}' because base_pose is applied to the URDF root" + ) + + def _urdf_root_node_name(self, robot_id: str, kind: str, config: RobotModelConfig) -> str: + root_node_name = { + "current": f"/robots/{robot_id}/current", + "target": f"/targets/{robot_id}/target", + "preview": f"/previews/{robot_id}/ghost", + }[kind] + if not self._has_non_identity_base_pose(config): + return root_node_name + self._ensure_base_pose_frame(robot_id, kind, config) + return f"{root_node_name}/base_pose/urdf" + + def _ensure_base_pose_frame(self, robot_id: str, kind: str, config: RobotModelConfig) -> None: + key = f"{robot_id}:{kind}:base_pose" + if key in self._root_frames: + return + pose = config.base_pose + frame_name = { + "current": f"/robots/{robot_id}/current/base_pose", + "target": f"/targets/{robot_id}/target/base_pose", + "preview": f"/previews/{robot_id}/ghost/base_pose", + }[kind] + self._root_frames[key] = self.server.scene.add_frame( + frame_name, + show_axes=False, + position=( + float(pose.position.x), + float(pose.position.y), + float(pose.position.z), + ), + wxyz=( + float(pose.orientation.w), + float(pose.orientation.x), + float(pose.orientation.y), + float(pose.orientation.z), + ), + ) + + @staticmethod + def _has_non_identity_base_pose(config: RobotModelConfig) -> bool: + pose = getattr(config, "base_pose", None) + if pose is None: + return False + return any( + abs(value) > 1e-12 + for value in ( + float(pose.position.x), + float(pose.position.y), + float(pose.position.z), + float(pose.orientation.x), + float(pose.orientation.y), + float(pose.orientation.z), + float(pose.orientation.w) - 1.0, + ) + ) def set_urdf_joints( self, urdf: ViserUrdf | None, joint_names: Sequence[str], joints: Sequence[float] @@ -328,8 +560,7 @@ def viser_joint_configuration( return [] values_by_name: dict[str, float] = {} for name, value in zip(joint_names, joints, strict=False): - values_by_name[name] = float(value) - values_by_name[name.rsplit("/", 1)[-1]] = float(value) + values_by_name[str(name)] = float(value) return [values_by_name.get(name, 0.0) for name in allowed_names] def viser_actuated_joint_names(self, urdf: ViserUrdf) -> tuple[str, ...]: diff --git a/dimos/manipulation/visualization/viser/state.py b/dimos/manipulation/visualization/viser/state.py index b46097df95..dc2e506536 100644 --- a/dimos/manipulation/visualization/viser/state.py +++ b/dimos/manipulation/visualization/viser/state.py @@ -21,7 +21,8 @@ import threading from typing import Literal -from dimos.manipulation.visualization.types import RobotInfo, TargetEvaluation +from dimos.manipulation.planning.spec.models import GeneratedPlan, PlanningGroupID +from dimos.manipulation.visualization.operator import TargetEvaluationResult from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.sensor_msgs.JointState import JointState from dimos.utils.logging_config import setup_logger @@ -65,7 +66,6 @@ class PlanStatus(str, Enum): PLANNING = "planning" FRESH = "fresh" STALE = "stale" - EXECUTING = "executing" FAILED = "failed" @@ -93,25 +93,31 @@ class FeasibilityState: class PanelPlanState: status: PlanStatus = PlanStatus.NONE robot: str | None = None - target_pose: Pose | None = None - target_joints: list[float] | None = None - start_joints_snapshot: list[float] | None = None - planned_path: list[JointState] | None = None + group_ids: tuple[PlanningGroupID, ...] = () + target_sequence_id: int = 0 + plan: GeneratedPlan | None = None + plan_id: str | None = None @dataclass class PanelState: selected_robot: str | None = None + selected_group_ids: tuple[PlanningGroupID, ...] = () + selection_epoch: int = 0 + pose_targets: dict[PlanningGroupID, Pose] = field(default_factory=dict) + group_joint_targets: dict[PlanningGroupID, JointState] = field(default_factory=dict) + target_joints: JointState | None = None + group_poses: dict[PlanningGroupID, Pose] = field(default_factory=dict) runtime: PanelRuntime = PanelRuntime.STOPPED backend_status: BackendConnectionStatus = BackendConnectionStatus.DISCONNECTED target_status: TargetStatus = TargetStatus.EMPTY action_status: ActionStatus = ActionStatus.IDLE + runtime_failure: bool = False + local_action_failure: bool = False manipulation_state: str = "DISCONNECTED" - robot_info: RobotInfo | None = None + ready_plan_id: str | None = None current_joints: list[float] | None = None - current_ee_pose: Pose | None = None cartesian_target: Pose | None = None - joint_target: list[float] | None = None feasibility: FeasibilityState = field(default_factory=FeasibilityState) latest_sequence_id: int = 0 plan_state: PanelPlanState = field(default_factory=PanelPlanState) @@ -125,18 +131,25 @@ def next_sequence_id(self) -> int: self.mark_plan_stale() return self.latest_sequence_id + def advance_selection_epoch(self) -> int: + """Invalidate callbacks and plans that belong to an older group selection.""" + self.selection_epoch += 1 + self.next_sequence_id() + self.plan_state = PanelPlanState() + return self.selection_epoch + def mark_plan_stale(self) -> None: - if self.plan_state.status == PlanStatus.FRESH: + if self.plan_state.status in {PlanStatus.FRESH, PlanStatus.PLANNING}: self.plan_state.status = PlanStatus.STALE def can_plan(self) -> bool: return ( self.runtime == PanelRuntime.RUNNING and self.backend_status == BackendConnectionStatus.READY - and self.selected_robot is not None + and bool(self.selected_group_ids) and self.action_status == ActionStatus.IDLE and self.target_status == TargetStatus.FEASIBLE - and self.manipulation_state in {"IDLE", "COMPLETED", "FAULT"} + and self.manipulation_state in {"IDLE", "READY"} and self.plan_state.status != PlanStatus.PLANNING ) @@ -145,21 +158,22 @@ def can_preview(self) -> bool: self.runtime == PanelRuntime.RUNNING and self.backend_status == BackendConnectionStatus.READY and self.action_status == ActionStatus.IDLE + and self.manipulation_state == "READY" and self.plan_state.status == PlanStatus.FRESH + and self.plan_state.plan_id is not None + and self.ready_plan_id is not None + and self.plan_state.plan_id == self.ready_plan_id ) def can_cancel(self) -> bool: - return self.action_status in { - ActionStatus.RUNNING, - ActionStatus.PREVIEWING, - ActionStatus.EXECUTING, - } or (self.manipulation_state == "EXECUTING") + return self.manipulation_state in { + "PLANNING", + "DISPATCHING", + "RUNNING", + "CANCELLING", + } - def can_execute( - self, - current_tolerance: float, - action_status: ActionStatus | None = None, - ) -> bool: + def can_execute(self, action_status: ActionStatus | None = None) -> bool: plan = self.plan_state effective_action_status = action_status or self.action_status if not ( @@ -167,21 +181,17 @@ def can_execute( and self.backend_status == BackendConnectionStatus.READY and effective_action_status == ActionStatus.IDLE and self.target_status == TargetStatus.FEASIBLE - and self.manipulation_state in {"IDLE", "COMPLETED"} + and self.manipulation_state == "READY" and plan.status == PlanStatus.FRESH - and plan.robot == self.selected_robot - and plan.start_joints_snapshot is not None - and self.current_joints is not None + and plan.plan is not None + and plan.plan_id is not None + and self.ready_plan_id is not None + and plan.plan_id == self.ready_plan_id + and plan.group_ids == self.selected_group_ids + and plan.target_sequence_id == self.latest_sequence_id ): return False - if len(plan.start_joints_snapshot) != len(self.current_joints): - return False - return all( - abs(expected - current) <= current_tolerance - for expected, current in zip( - plan.start_joints_snapshot, self.current_joints, strict=False - ) - ) + return True @property def connected(self) -> bool: @@ -203,9 +213,14 @@ def module_state(self) -> str: class TargetEvaluationRequest: sequence_id: int source: PreviewSource - robot_name: str + robot_name: str | None = None + selection_epoch: int = 0 + group_ids: tuple[PlanningGroupID, ...] = () pose: Pose | None = None joints: JointState | None = None + auxiliary_group_ids: tuple[PlanningGroupID, ...] = () + pose_targets: dict[PlanningGroupID, Pose] = field(default_factory=dict) + joint_targets: dict[PlanningGroupID, JointState] = field(default_factory=dict) class TargetEvaluationWorker: @@ -218,8 +233,8 @@ class TargetEvaluationWorker: def __init__( self, - handler: Callable[[TargetEvaluationRequest], TargetEvaluation], - apply_result: Callable[[TargetEvaluationRequest, TargetEvaluation], None], + handler: Callable[[TargetEvaluationRequest], TargetEvaluationResult], + apply_result: Callable[[TargetEvaluationRequest, TargetEvaluationResult], None], ) -> None: self._handler = handler self._apply_result = apply_result diff --git a/dimos/manipulation/visualization/viser/test_gui.py b/dimos/manipulation/visualization/viser/test_gui.py new file mode 100644 index 0000000000..c0a85d5465 --- /dev/null +++ b/dimos/manipulation/visualization/viser/test_gui.py @@ -0,0 +1,562 @@ +# 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 + +from collections.abc import Callable +from dataclasses import dataclass + +import pytest + +pytest.importorskip("viser", reason="Viser optional dependency is not installed") + +from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.manipulation.planning.spec.models import GeneratedPlan, PlanningSceneInfo +from dimos.manipulation.visualization.operator import OperatorStatus, TargetEvaluationResult +from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig +from dimos.manipulation.visualization.viser.gui import ViserPanelGui +from dimos.manipulation.visualization.viser.state import ( + ActionStatus, + BackendConnectionStatus, + FeasibilityStatus, + OperationWorker, + PanelRuntime, + PlanStatus, + TargetEvaluationRequest, + TargetEvaluationWorker, + TargetStatus, +) +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory + + +class EmptyServer: + pass + + +class FakeOperatorBackend: + def __init__(self) -> None: + self.cancel_calls = 0 + + def cancel(self) -> bool: + self.cancel_calls += 1 + return True + + +class FakeOperator: + def __init__(self, module: FakeOperatorBackend | None = None) -> None: + self.module = module or FakeOperatorBackend() + + def status(self) -> OperatorStatus: + return OperatorStatus( + state="READY", diagnostic="", ready_plan_status="READY", ready_plan_id="fake-plan" + ) + + def get_init_joints(self, robot_name: str) -> None: + return None + + def cancel(self) -> bool: + return self.module.cancel() + + def preview(self, *_args: object, **_kwargs: object) -> bool: + return True + + +@dataclass +class FakeStopOperationWorker(OperationWorker): + stop_calls: list[float | None] + + def __init__(self, stop_calls: list[float | None]) -> None: + self.stop_calls = stop_calls + + def stop(self, timeout: float | None = 2.0) -> None: + self.stop_calls.append(timeout) + + +@dataclass +class FakeStopEvaluationWorker(TargetEvaluationWorker): + stop_calls: list[float | None] + + def __init__(self, stop_calls: list[float | None]) -> None: + self.stop_calls = stop_calls + + def stop(self, timeout: float | None = 2.0) -> None: + self.stop_calls.append(timeout) + + +class FakeTimeoutSubmitWorker(OperationWorker): + def __init__(self, submissions: list[dict[str, float]]) -> None: + self.submissions = submissions + + def submit( + self, + operation: Callable[[], None], + *, + timeout_seconds: float | None = None, + on_error: Callable[[str], None] | None = None, + ) -> None: + kwargs = {} + if timeout_seconds is not None: + kwargs["timeout_seconds"] = timeout_seconds + self.submissions.append(kwargs) + + +class FakeOperationSubmitWorker(OperationWorker): + def __init__(self, submissions: list[Callable[[], None]]) -> None: + self.submissions = submissions + + def submit( + self, + operation: Callable[[], None], + *, + timeout_seconds: float | None = None, + on_error: Callable[[str], None] | None = None, + ) -> None: + self.submissions.append(operation) + + +class FakeOperationErrorWorker(OperationWorker): + def __init__(self, errors: list[Callable[[str], None]]) -> None: + self.errors = errors + + def submit( + self, + operation: Callable[[], None], + *, + timeout_seconds: float | None = None, + on_error: Callable[[str], None] | None = None, + ) -> None: + if on_error is not None: + self.errors.append(on_error) + + +class FakeRestartableOperationWorker(FakeOperationSubmitWorker): + def __init__( + self, submissions: list[Callable[[], None]], stop_calls: list[float | None] + ) -> None: + super().__init__(submissions) + self.stop_calls = stop_calls + + def stop(self, timeout: float | None = 2.0) -> None: + self.stop_calls.append(timeout) + + +def planning_group(robot: str, name: str, joints: tuple[str, ...]) -> PlanningGroup: + return PlanningGroup( + f"{robot}/{name}", + robot, + name, + tuple(f"{robot}/{joint}" for joint in joints), + joints, + "base", + None, + ) + + +def make_gui(module: FakeOperatorBackend | None = None) -> ViserPanelGui: + module = module or FakeOperatorBackend() + return ViserPanelGui( + EmptyServer(), + PlanningSceneInfo(robots={}), + FakeOperator(module), + {}, + ViserVisualizationConfig(), + ) + + +@pytest.mark.parametrize( + ("result", "success", "collision_free", "expected"), + [ + ( + TargetEvaluationResult(True, "FEASIBLE", "", True), + True, + True, + FeasibilityStatus.FEASIBLE, + ), + (TargetEvaluationResult(False, "COLLISION", ""), False, False, FeasibilityStatus.COLLISION), + ( + TargetEvaluationResult(False, "COLLISION_AT_START", ""), + False, + False, + FeasibilityStatus.COLLISION, + ), + ( + TargetEvaluationResult(False, "COLLISION_AT_GOAL", ""), + False, + False, + FeasibilityStatus.COLLISION, + ), + ( + TargetEvaluationResult(False, "NO_SOLUTION", ""), + False, + False, + FeasibilityStatus.IK_FAILED, + ), + ( + TargetEvaluationResult(False, "SINGULARITY", ""), + False, + False, + FeasibilityStatus.IK_FAILED, + ), + ( + TargetEvaluationResult(False, "JOINT_LIMITS", ""), + False, + False, + FeasibilityStatus.IK_FAILED, + ), + (TargetEvaluationResult(False, "TIMEOUT", ""), False, False, FeasibilityStatus.IK_FAILED), + ( + TargetEvaluationResult(False, "IK_SUCCEEDED", ""), + False, + False, + FeasibilityStatus.INVALID, + ), + ], +) +def test_gui_feasibility_status_uses_exact_status_mapping( + result: TargetEvaluationResult, + success: bool, + collision_free: bool, + expected: FeasibilityStatus, +) -> None: + gui = make_gui() + + assert gui._feasibility_status(result, success, collision_free) == expected + + +def test_group_status_composes_shared_panel_state_without_robot_dropdown() -> None: + gui = make_gui() + values: dict[str, str] = {} + gui.state.selected_group_ids = ("left/manipulator", "right/gripper") + gui.state.error = "planner unavailable" + gui.state.target_status = gui.state.target_status.FEASIBLE + gui.state.plan_state.status = gui.state.plan_state.status.FRESH + gui._stale_robot_names = lambda _group_ids: ("right",) # type: ignore[method-assign] + gui._set_handle_value = values.__setitem__ # type: ignore[method-assign] + + gui._update_status_text() + + assert "robot" not in gui._handles + assert values == { + "status": "### Status\n\n**State:** planner unavailable\n\n" + "Target: `feasible` · Plan: `fresh`\n\nState stale: `True (right)`", + "target_summary": "Feasibility: `unknown`", + } + + +def test_gui_close_uses_bounded_operation_worker_stop(monkeypatch: pytest.MonkeyPatch) -> None: + stop_timeouts: list[float | None] = [] + gui = make_gui() + gui._operation_worker.stop() + gui._worker.stop() + monkeypatch.setattr(gui, "_operation_worker", FakeStopOperationWorker(stop_timeouts)) + monkeypatch.setattr(gui, "_worker", FakeStopEvaluationWorker([])) + + gui.close() + + assert stop_timeouts == [2.0] + + +def test_gui_only_preview_submits_timeout_override(monkeypatch: pytest.MonkeyPatch) -> None: + submissions: list[dict[str, float]] = [] + gui = make_gui() + gui.config = ViserVisualizationConfig(preview_request_timeout=0.25) + gui._operation_worker.stop() + monkeypatch.setattr(gui, "_operation_worker", FakeTimeoutSubmitWorker(submissions)) + gui.state.runtime = PanelRuntime.RUNNING + gui.state.backend_status = BackendConnectionStatus.READY + gui.state.manipulation_state = "READY" + gui.state.ready_plan_id = "fake-plan" + gui.state.plan_state.status = PlanStatus.FRESH + gui.state.plan_state.plan_id = "fake-plan" + gui._submit_preview() + + assert submissions == [{"timeout_seconds": 0.25}] + + +def test_gui_preview_enters_previewing_before_worker_runs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + submissions: list[Callable[[], None]] = [] + gui = make_gui() + gui._operation_worker.stop() + monkeypatch.setattr(gui, "_operation_worker", FakeOperationSubmitWorker(submissions)) + monkeypatch.setattr(gui, "refresh", lambda: None) + gui.state.runtime = PanelRuntime.RUNNING + gui.state.backend_status = BackendConnectionStatus.READY + gui.state.target_status = TargetStatus.FEASIBLE + gui.state.manipulation_state = "READY" + gui.state.selected_group_ids = ("arm/manipulator",) + gui.state.plan_state.status = PlanStatus.FRESH + gui.state.plan_state.group_ids = gui.state.selected_group_ids + gui.state.plan_state.target_sequence_id = gui.state.latest_sequence_id + gui.state.ready_plan_id = "fake-plan" + gui.state.plan_state.plan_id = "fake-plan" + gui.state.plan_state.plan = GeneratedPlan( + group_ids=gui.state.selected_group_ids, + trajectory=JointTrajectory(), + path=[JointState({"name": [], "position": []})], + ) + + assert gui.state.can_execute() is True + + gui._submit_preview() + + assert gui.state.action_status == ActionStatus.PREVIEWING + assert gui.state.can_execute() is False + assert len(submissions) == 1 + + +def test_gui_selection_change_clears_invalidated_preview( + monkeypatch: pytest.MonkeyPatch, +) -> None: + submissions: list[Callable[[], None]] = [] + gui = make_gui() + gui._operation_worker.stop() + monkeypatch.setattr(gui, "_operation_worker", FakeOperationSubmitWorker(submissions)) + monkeypatch.setattr(gui, "refresh", lambda: None) + groups = [ + planning_group("arm", "manipulator", ("j1",)), + planning_group("arm", "gripper", ("j2",)), + ] + monkeypatch.setattr(gui, "list_planning_groups", lambda: groups) + monkeypatch.setattr(gui, "_build_joint_sliders", lambda: None) + gui.state.runtime = PanelRuntime.RUNNING + gui.state.backend_status = BackendConnectionStatus.READY + gui.state.target_status = TargetStatus.FEASIBLE + gui.state.manipulation_state = "READY" + gui.state.selected_group_ids = (groups[0].id,) + gui.state.plan_state.status = PlanStatus.FRESH + gui.state.plan_state.group_ids = gui.state.selected_group_ids + gui.state.plan_state.target_sequence_id = gui.state.latest_sequence_id + gui.state.ready_plan_id = "fake-plan" + gui.state.plan_state.plan_id = "fake-plan" + gui.state.plan_state.plan = GeneratedPlan( + group_ids=gui.state.selected_group_ids, + trajectory=JointTrajectory(), + path=[JointState({"name": [], "position": []})], + ) + + gui._submit_preview() + gui._toggle_group_selected(groups[1].id) + submissions[0]() + + assert gui.state.action_status == ActionStatus.IDLE + assert gui.state.last_result == "preview=False" + + +def test_gui_selection_change_ignores_invalidated_preview_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + errors: list[Callable[[str], None]] = [] + gui = make_gui() + gui._operation_worker.stop() + monkeypatch.setattr(gui, "_operation_worker", FakeOperationErrorWorker(errors)) + monkeypatch.setattr(gui, "refresh", lambda: None) + groups = [ + planning_group("arm", "manipulator", ("j1",)), + planning_group("arm", "gripper", ("j2",)), + ] + monkeypatch.setattr(gui, "list_planning_groups", lambda: groups) + monkeypatch.setattr(gui, "_build_joint_sliders", lambda: None) + gui.state.runtime = PanelRuntime.RUNNING + gui.state.backend_status = BackendConnectionStatus.READY + gui.state.target_status = TargetStatus.FEASIBLE + gui.state.manipulation_state = "READY" + gui.state.selected_group_ids = (groups[0].id,) + gui.state.plan_state.status = PlanStatus.FRESH + gui.state.plan_state.group_ids = gui.state.selected_group_ids + gui.state.plan_state.target_sequence_id = gui.state.latest_sequence_id + gui.state.ready_plan_id = "fake-plan" + gui.state.plan_state.plan_id = "fake-plan" + gui.state.plan_state.plan = GeneratedPlan( + group_ids=gui.state.selected_group_ids, + trajectory=JointTrajectory(), + path=[JointState({"name": [], "position": []})], + ) + + gui._submit_preview() + gui._toggle_group_selected(groups[1].id) + errors[0]("preview timed out") + + assert gui.state.action_status == ActionStatus.IDLE + assert gui.state.error == "" + assert gui.state.last_result == "preview=False" + + +def test_gui_cancel_bypasses_operation_worker(monkeypatch: pytest.MonkeyPatch) -> None: + submissions: list[Callable[[], None]] = [] + stop_calls: list[float | None] = [] + module = FakeOperatorBackend() + gui = make_gui(module) + gui._operation_worker.stop() + monkeypatch.setattr( + gui, "_operation_worker", FakeRestartableOperationWorker(submissions, stop_calls) + ) + gui.state.action_status = ActionStatus.PREVIEWING + + gui._submit_cancel() + gui.close() + + assert submissions == [] + assert stop_calls == [0.0] + assert module.cancel_calls == 1 + assert gui.state.action_status == ActionStatus.IDLE + assert gui.state.last_result == "cancel=True" + + +def test_gui_cancelled_planning_clears_active_plan_state(monkeypatch: pytest.MonkeyPatch) -> None: + submissions: list[Callable[[], None]] = [] + stop_calls: list[float | None] = [] + module = FakeOperatorBackend() + gui = make_gui(module) + gui._operation_worker.stop() + monkeypatch.setattr( + gui, "_operation_worker", FakeRestartableOperationWorker(submissions, stop_calls) + ) + stale_operation_id = gui._next_operation_id() + gui.state.action_status = ActionStatus.RUNNING + gui.state.plan_state.status = PlanStatus.PLANNING + assert gui.state.plan_state.status == PlanStatus.PLANNING + + gui._submit_cancel() + gui._finish_operation("plan_to_joints=True", operation_id=stale_operation_id) + gui.close() + + assert submissions == [] + assert module.cancel_calls == 1 + assert stop_calls == [0.0] + assert gui.state.action_status == ActionStatus.IDLE + assert gui.state.plan_state.status == PlanStatus.FAILED + assert gui.state.last_result == "cancel=True" + + +@pytest.mark.parametrize( + ("submit", "expected_error"), + [ + ("_submit_plan", "Cannot plan until target is feasible and manipulation is idle"), + ("_submit_preview", "No fresh plan to preview"), + ( + "_submit_execute", + "Cannot execute: require feasible fresh plan", + ), + ], +) +def test_gui_guard_errors_keep_action_idle( + submit: str, expected_error: str, monkeypatch: pytest.MonkeyPatch +) -> None: + submissions: list[Callable[[], None]] = [] + gui = make_gui() + gui._operation_worker.stop() + monkeypatch.setattr(gui, "_operation_worker", FakeOperationSubmitWorker(submissions)) + gui.state.runtime = PanelRuntime.RUNNING + gui.state.backend_status = BackendConnectionStatus.READY + gui.state.selected_robot = "arm" + gui.state.action_status = ActionStatus.IDLE + + getattr(gui, submit)() + + assert gui.state.action_status == ActionStatus.IDLE + assert gui.state.error == expected_error + assert submissions == [] + + +def test_gui_ignores_stale_timed_out_operation_finish() -> None: + gui = make_gui() + old_operation_id = gui._next_operation_id() + gui._set_operation_error("Operation timed out after 5.0s", old_operation_id) + gui.state.action_status = ActionStatus.FAILED + + gui._finish_operation("preview=True", operation_id=old_operation_id) + + assert gui.state.action_status == ActionStatus.FAILED + assert gui.state.error == "Operation timed out after 5.0s" + + +def test_gui_projects_all_runtime_states_and_fault_reset_preserves_local_failure() -> None: + gui = make_gui() + gui.state.runtime = PanelRuntime.RUNNING + gui.state.backend_status = BackendConnectionStatus.READY + gui.state.selected_group_ids = ("arm/manipulator",) + gui.state.target_status = TargetStatus.FEASIBLE + gui.state.action_status = ActionStatus.IDLE + + expected_actions = { + "IDLE": ActionStatus.IDLE, + "PLANNING": ActionStatus.RUNNING, + "READY": ActionStatus.IDLE, + "DISPATCHING": ActionStatus.EXECUTING, + "RUNNING": ActionStatus.EXECUTING, + "CANCELLING": ActionStatus.CANCELLING, + "FAULT": ActionStatus.FAILED, + } + for state, expected_action in expected_actions.items(): + gui._apply_operator_status( + OperatorStatus( + state=state, + diagnostic="fault" if state == "FAULT" else "", + ready_plan_status="READY" if state == "READY" else "NONE", + ready_plan_id="fake-plan" if state == "READY" else None, + ) + ) + assert gui.state.manipulation_state == state + assert gui.state.action_status == expected_action + + gui._apply_operator_status(OperatorStatus("IDLE")) + assert gui.state.can_plan() is True + + gui.state.action_status = ActionStatus.FAILED + gui.state.local_action_failure = True + gui._apply_operator_status(OperatorStatus("FAULT", diagnostic="fault")) + gui._apply_operator_status( + OperatorStatus("READY", ready_plan_status="READY", ready_plan_id="fake-plan") + ) + + assert gui.state.action_status == ActionStatus.FAILED + assert gui.state.local_action_failure is True + + +def test_late_local_callbacks_cannot_clear_runtime_fault() -> None: + gui = make_gui() + gui.state.manipulation_state = "FAULT" + gui.state.action_status = ActionStatus.FAILED + gui.state.runtime_failure = True + gui.state.error = "runtime fault" + gui.state.selected_group_ids = ("arm/manipulator",) + + gui._finish_operation("preview=True") + gui._apply_target_evaluation_result( + TargetEvaluationRequest(0, "joints", group_ids=("arm/manipulator",)), + TargetEvaluationResult(True, "FEASIBLE", "ok", True), + ) + + assert gui.state.manipulation_state == "FAULT" + assert gui.state.action_status == ActionStatus.FAILED + assert gui.state.error == "runtime fault" + assert gui.state.can_cancel() is False + + +def test_stale_execute_validation_cannot_replace_runtime_fault_diagnostic( + monkeypatch: pytest.MonkeyPatch, +) -> None: + gui = make_gui() + gui._apply_operator_status(OperatorStatus("FAULT", diagnostic="hardware fault")) + monkeypatch.setattr(gui, "refresh", lambda: None) + + gui._submit_execute() + + assert gui.state.manipulation_state == "FAULT" + assert gui.state.action_status == ActionStatus.FAILED + assert gui.state.error == "hardware fault" + assert gui.state.can_execute() is False diff --git a/dimos/manipulation/visualization/viser/test_gui_status.py b/dimos/manipulation/visualization/viser/test_gui_status.py deleted file mode 100644 index f461565bdf..0000000000 --- a/dimos/manipulation/visualization/viser/test_gui_status.py +++ /dev/null @@ -1,63 +0,0 @@ -# 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 pytest - -pytest.importorskip("viser", reason="Viser optional dependency is not installed") - -from dimos.manipulation.visualization.types import TargetEvaluation -from dimos.manipulation.visualization.viser.adapter import InProcessViserAdapter -from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig -from dimos.manipulation.visualization.viser.gui import ViserPanelGui -from dimos.manipulation.visualization.viser.state import FeasibilityStatus - - -class StatusOnlyServer: - pass - - -class StatusOnlyAdapter(InProcessViserAdapter): - def __init__(self) -> None: - pass - - -@pytest.mark.parametrize( - ("result", "success", "collision_free", "expected"), - [ - ({"status": "FEASIBLE"}, True, True, FeasibilityStatus.FEASIBLE), - ({"status": "COLLISION"}, False, False, FeasibilityStatus.COLLISION), - ({"status": "COLLISION_AT_START"}, False, False, FeasibilityStatus.COLLISION), - ({"status": "COLLISION_AT_GOAL"}, False, False, FeasibilityStatus.COLLISION), - ({"status": "NO_SOLUTION"}, False, False, FeasibilityStatus.IK_FAILED), - ({"status": "SINGULARITY"}, False, False, FeasibilityStatus.IK_FAILED), - ({"status": "JOINT_LIMITS"}, False, False, FeasibilityStatus.IK_FAILED), - ({"status": "TIMEOUT"}, False, False, FeasibilityStatus.IK_FAILED), - ({"status": "IK_SUCCEEDED"}, False, False, FeasibilityStatus.INVALID), - ], -) -def test_gui_feasibility_status_uses_exact_status_mapping( - result: TargetEvaluation, - success: bool, - collision_free: bool, - expected: FeasibilityStatus, -) -> None: - gui = ViserPanelGui( - StatusOnlyServer(), - StatusOnlyAdapter(), - ViserVisualizationConfig(), - ) - - assert gui._feasibility_status(result, success, collision_free) == expected diff --git a/dimos/manipulation/visualization/viser/test_operation_worker.py b/dimos/manipulation/visualization/viser/test_operation_worker.py deleted file mode 100644 index 7b7724ce2d..0000000000 --- a/dimos/manipulation/visualization/viser/test_operation_worker.py +++ /dev/null @@ -1,322 +0,0 @@ -# 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 - -from collections.abc import Callable -from dataclasses import dataclass -import threading - -import pytest - -pytest.importorskip("viser", reason="Viser optional dependency is not installed") - -from dimos.manipulation.visualization.types import RobotInfo -from dimos.manipulation.visualization.viser.adapter import InProcessViserAdapter -from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig -from dimos.manipulation.visualization.viser.gui import ViserPanelGui -from dimos.manipulation.visualization.viser.state import ( - ActionStatus, - BackendConnectionStatus, - OperationWorker, - PanelRuntime, - PlanStatus, - TargetEvaluationWorker, - TargetStatus, -) -from dimos.msgs.sensor_msgs.JointState import JointState - - -class EmptyServer: - pass - - -@dataclass -class FakeStopOperationWorker(OperationWorker): - stop_calls: list[float | None] - - def __init__(self, stop_calls: list[float | None]) -> None: - self.stop_calls = stop_calls - - def stop(self, timeout: float | None = 2.0) -> None: - self.stop_calls.append(timeout) - - -@dataclass -class FakeStopEvaluationWorker(TargetEvaluationWorker): - stop_calls: list[float | None] - - def __init__(self, stop_calls: list[float | None]) -> None: - self.stop_calls = stop_calls - - def stop(self, timeout: float | None = 2.0) -> None: - self.stop_calls.append(timeout) - - -class FakeTimeoutSubmitWorker(OperationWorker): - def __init__(self, submissions: list[dict[str, float]]) -> None: - self.submissions = submissions - - def submit( - self, - operation: Callable[[], None], - *, - timeout_seconds: float | None = None, - on_error: Callable[[str], None] | None = None, - ) -> None: - kwargs = {} - if timeout_seconds is not None: - kwargs["timeout_seconds"] = timeout_seconds - self.submissions.append(kwargs) - - -class FakeOperationSubmitWorker(OperationWorker): - def __init__(self, submissions: list[Callable[[], None]]) -> None: - self.submissions = submissions - - def submit( - self, - operation: Callable[[], None], - *, - timeout_seconds: float | None = None, - on_error: Callable[[str], None] | None = None, - ) -> None: - self.submissions.append(operation) - - -class FakeRestartableOperationWorker(FakeOperationSubmitWorker): - def __init__( - self, submissions: list[Callable[[], None]], stop_calls: list[float | None] - ) -> None: - super().__init__(submissions) - self.stop_calls = stop_calls - - def stop(self, timeout: float | None = 2.0) -> None: - self.stop_calls.append(timeout) - - -class FakeOperationAdapter(InProcessViserAdapter): - def __init__(self) -> None: - self.cancel_calls = 0 - - def list_robots(self) -> list[str]: - return [] - - def get_module_state(self) -> str: - return "IDLE" - - def get_robot_info(self, robot_name: str) -> RobotInfo | None: - return None - - def get_current_joint_state(self, robot_name: str) -> None: - return None - - def get_ee_pose(self, robot_name: str, joint_state: JointState | None = None) -> None: - return None - - def get_error(self) -> str: - return "" - - def get_robot_config(self, robot_name: str) -> None: - return None - - def is_state_stale(self, robot_name: str, max_age: float = 1.0) -> bool: - return False - - def cancel(self) -> bool: - self.cancel_calls += 1 - return True - - def plan_to_joints(self, joints: JointState, robot_name: str | None = None) -> bool: - return True - - -def test_operation_worker_uses_per_operation_timeout() -> None: - errors: list[str] = [] - worker = OperationWorker(errors.append, timeout_seconds=1.0) - worker.submit(lambda: None, timeout_seconds=0.25) - - request = worker._requests.get_nowait() - - assert worker._operation_timeout(request) == 0.25 - - -def test_operation_worker_uses_operation_error_callback_on_timeout() -> None: - default_errors: list[str] = [] - operation_errors: list[str] = [] - release = threading.Event() - finished = threading.Event() - worker = OperationWorker(default_errors.append) - - def operation() -> None: - release.wait(timeout=1.0) - finished.set() - - worker.submit( - operation, - timeout_seconds=0.001, - on_error=operation_errors.append, - ) - - worker._run_operation(worker._requests.get_nowait()) - release.set() - assert finished.wait(timeout=1.0) - - assert default_errors == [] - assert operation_errors == ["Operation timed out after 0.0s"] - - -def test_gui_close_uses_bounded_operation_worker_stop(monkeypatch: pytest.MonkeyPatch) -> None: - stop_timeouts: list[float | None] = [] - gui = ViserPanelGui( - EmptyServer(), - FakeOperationAdapter(), - ViserVisualizationConfig(), - ) - gui._operation_worker.stop() - gui._worker.stop() - monkeypatch.setattr(gui, "_operation_worker", FakeStopOperationWorker(stop_timeouts)) - monkeypatch.setattr(gui, "_worker", FakeStopEvaluationWorker([])) - - gui.close() - - assert stop_timeouts == [2.0] - - -def test_gui_only_preview_submits_timeout_override(monkeypatch: pytest.MonkeyPatch) -> None: - submissions: list[dict[str, float]] = [] - gui = ViserPanelGui( - EmptyServer(), - FakeOperationAdapter(), - ViserVisualizationConfig(preview_request_timeout=0.25), - ) - gui._operation_worker.stop() - monkeypatch.setattr(gui, "_operation_worker", FakeTimeoutSubmitWorker(submissions)) - gui.state.runtime = PanelRuntime.RUNNING - gui.state.backend_status = BackendConnectionStatus.READY - gui.state.selected_robot = "arm" - gui.state.target_status = TargetStatus.FEASIBLE - gui.state.manipulation_state = "IDLE" - - gui._submit_plan() - gui.state.plan_state.status = PlanStatus.FRESH - gui._submit_preview() - - assert "timeout_seconds" not in submissions[0] - assert submissions[1]["timeout_seconds"] == 0.25 - - -def test_gui_cancel_bypasses_operation_worker(monkeypatch: pytest.MonkeyPatch) -> None: - submissions: list[Callable[[], None]] = [] - stop_calls: list[float | None] = [] - adapter = FakeOperationAdapter() - gui = ViserPanelGui( - EmptyServer(), - adapter, - ViserVisualizationConfig(), - ) - gui._operation_worker.stop() - monkeypatch.setattr( - gui, "_operation_worker", FakeRestartableOperationWorker(submissions, stop_calls) - ) - gui.state.action_status = ActionStatus.PREVIEWING - - gui._submit_cancel() - gui.close() - - assert submissions == [] - assert stop_calls == [0.0] - assert adapter.cancel_calls == 1 - assert gui.state.action_status == ActionStatus.IDLE - assert gui.state.last_result == "cancel=True" - - -def test_gui_cancelled_planning_clears_active_plan_state(monkeypatch: pytest.MonkeyPatch) -> None: - submissions: list[Callable[[], None]] = [] - stop_calls: list[float | None] = [] - adapter = FakeOperationAdapter() - gui = ViserPanelGui( - EmptyServer(), - adapter, - ViserVisualizationConfig(), - ) - gui._operation_worker.stop() - monkeypatch.setattr( - gui, "_operation_worker", FakeRestartableOperationWorker(submissions, stop_calls) - ) - stale_operation_id = gui._next_operation_id() - gui.state.action_status = ActionStatus.RUNNING - gui.state.plan_state.status = PlanStatus.PLANNING - assert gui.state.plan_state.status == PlanStatus.PLANNING - - gui._submit_cancel() - gui._finish_operation("plan_to_joints=True", operation_id=stale_operation_id) - gui.close() - - assert submissions == [] - assert adapter.cancel_calls == 1 - assert stop_calls == [0.0] - assert gui.state.action_status == ActionStatus.IDLE - assert gui.state.plan_state.status == PlanStatus.FAILED - assert gui.state.last_result == "cancel=True" - - -@pytest.mark.parametrize( - ("submit", "expected_error"), - [ - ("_submit_plan", "Cannot plan until target is feasible and manipulation is idle"), - ("_submit_preview", "No fresh plan to preview"), - ( - "_submit_execute", - "Cannot execute: require feasible fresh plan and matching current joints", - ), - ], -) -def test_gui_guard_errors_keep_action_idle( - submit: str, expected_error: str, monkeypatch: pytest.MonkeyPatch -) -> None: - submissions: list[Callable[[], None]] = [] - gui = ViserPanelGui( - EmptyServer(), - FakeOperationAdapter(), - ViserVisualizationConfig(), - ) - gui._operation_worker.stop() - monkeypatch.setattr(gui, "_operation_worker", FakeOperationSubmitWorker(submissions)) - gui.state.runtime = PanelRuntime.RUNNING - gui.state.backend_status = BackendConnectionStatus.READY - gui.state.selected_robot = "arm" - gui.state.action_status = ActionStatus.IDLE - - getattr(gui, submit)() - - assert gui.state.action_status == ActionStatus.IDLE - assert gui.state.error == expected_error - assert submissions == [] - - -def test_gui_ignores_stale_timed_out_operation_finish() -> None: - gui = ViserPanelGui( - EmptyServer(), - FakeOperationAdapter(), - ViserVisualizationConfig(), - ) - old_operation_id = gui._next_operation_id() - gui._set_operation_error("Operation timed out after 5.0s", old_operation_id) - gui.state.action_status = ActionStatus.FAILED - - gui._finish_operation("preview=True", operation_id=old_operation_id) - - assert gui.state.action_status == ActionStatus.FAILED - assert gui.state.error == "Operation timed out after 5.0s" diff --git a/dimos/manipulation/visualization/viser/test_state.py b/dimos/manipulation/visualization/viser/test_state.py index 412c227050..ee49c8332b 100644 --- a/dimos/manipulation/visualization/viser/test_state.py +++ b/dimos/manipulation/visualization/viser/test_state.py @@ -14,21 +14,117 @@ from __future__ import annotations +from collections.abc import Callable +import threading + +from dimos.manipulation.planning.spec.models import PlanningGroupID from dimos.manipulation.visualization.viser.state import ( + ActionStatus, BackendConnectionStatus, + OperationWorker, PanelRuntime, PanelState, + PlanStatus, + TargetEvaluationWorker, TargetStatus, ) -def test_panel_can_plan_from_fault_after_planning_failure() -> None: +def test_panel_cannot_plan_from_fault_without_explicit_reset() -> None: state = PanelState( selected_robot="arm", + selected_group_ids=(PlanningGroupID("arm/manipulator"),), runtime=PanelRuntime.RUNNING, backend_status=BackendConnectionStatus.READY, target_status=TargetStatus.FEASIBLE, manipulation_state="FAULT", ) - assert state.can_plan() is True + assert state.can_plan() is False + + +def test_panel_cannot_plan_without_a_selected_group() -> None: + state = PanelState( + runtime=PanelRuntime.RUNNING, + backend_status=BackendConnectionStatus.READY, + target_status=TargetStatus.FEASIBLE, + manipulation_state="IDLE", + ) + + assert state.can_plan() is False + + +def test_sequence_change_marks_a_fresh_plan_stale() -> None: + state = PanelState(plan_state=PanelState().plan_state) + state.plan_state.status = PlanStatus.FRESH + + state.next_sequence_id() + + assert state.plan_state.status == PlanStatus.STALE + assert state.target_status == TargetStatus.CHECKING + + +def test_selection_epoch_change_resets_plan_and_invalidates_sequence() -> None: + state = PanelState(selected_group_ids=(PlanningGroupID("arm/manipulator"),)) + state.plan_state.status = PlanStatus.FRESH + + assert state.advance_selection_epoch() == 1 + + assert state.latest_sequence_id == 1 + assert state.plan_state.status == PlanStatus.NONE + + +def test_cancel_uses_only_the_runtime_lifecycle_gate() -> None: + state = PanelState() + + for lifecycle in ("PLANNING", "DISPATCHING", "RUNNING", "CANCELLING"): + state.manipulation_state = lifecycle + for action in ActionStatus: + state.action_status = action + assert state.can_cancel() is True + + for lifecycle in ("IDLE", "READY", "FAULT"): + state.manipulation_state = lifecycle + for action in (ActionStatus.RUNNING, ActionStatus.PREVIEWING, ActionStatus.EXECUTING): + state.action_status = action + assert state.can_cancel() is False + + +def test_operation_worker_uses_per_operation_timeout() -> None: + errors: list[str] = [] + worker = OperationWorker(errors.append, timeout_seconds=1.0) + worker.submit(lambda: None, timeout_seconds=0.25) + + request = worker._requests.get_nowait() + + assert worker._operation_timeout(request) == 0.25 + + +def test_operation_worker_uses_operation_error_callback_on_timeout() -> None: + default_errors: list[str] = [] + operation_errors: list[str] = [] + release = threading.Event() + finished = threading.Event() + worker = OperationWorker(default_errors.append) + + def operation() -> None: + release.wait(timeout=1.0) + finished.set() + + worker.submit( + operation, + timeout_seconds=0.001, + on_error=operation_errors.append, + ) + + worker._run_operation(worker._requests.get_nowait()) + release.set() + assert finished.wait(timeout=1.0) + + assert default_errors == [] + assert operation_errors == ["Operation timed out after 0.0s"] + + +class FakeTargetEvaluationWorker(TargetEvaluationWorker): + def __init__(self, calls: list[Callable[[], None]]) -> None: + self.calls = calls diff --git a/dimos/manipulation/visualization/viser/test_viser_visualization.py b/dimos/manipulation/visualization/viser/test_viser_visualization.py index 2994158022..52d201d9f4 100644 --- a/dimos/manipulation/visualization/viser/test_viser_visualization.py +++ b/dimos/manipulation/visualization/viser/test_viser_visualization.py @@ -1,4 +1,4 @@ -# Copyright 2025-2026 Dimensional Inc. +# 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. @@ -12,1576 +12,1220 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Hermetic contract tests for the group-aware Viser manipulation panel.""" + from __future__ import annotations from collections.abc import Callable, Iterator, Sequence from dataclasses import dataclass -import sys +from pathlib import Path import threading -from types import ModuleType, SimpleNamespace, TracebackType +from types import SimpleNamespace -import numpy as np import pytest pytest.importorskip("viser", reason="Viser optional dependency is not installed") -from dimos.manipulation.visualization.types import RobotInfo, TargetEvaluation -from dimos.manipulation.visualization.viser.adapter import InProcessViserAdapter +from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupSelection +from dimos.manipulation.planning.spec.enums import PlanningStatus +from dimos.manipulation.planning.spec.models import GeneratedPlan, PlanningSceneInfo +from dimos.manipulation.visualization.operator import OperatorStatus, TargetEvaluationResult from dimos.manipulation.visualization.viser.animation import ( - PreviewAnimator, - interpolate_joint_path, - sampled_joint_path_frames, + GroupPreviewAnimation, + PreviewFrame, + PreviewTrack, + scaled_frame_delays, ) from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig -from dimos.manipulation.visualization.viser.gui import ViserPanelGui -from dimos.manipulation.visualization.viser.scene import ViserManipulationScene +from dimos.manipulation.visualization.viser.gui import ( + ACTIVE_GROUP_COLOR, + INACTIVE_GROUP_COLOR, + ViserPanelGui, + group_display_name, +) +from dimos.manipulation.visualization.viser.scene import ( + GOAL_ROBOT_FEASIBLE_COLOR, + GOAL_ROBOT_INFEASIBLE_COLOR, + TARGET_CONTROL_FEASIBLE_COLOR, + TARGET_CONTROL_INFEASIBLE_COLOR, + ViserManipulationScene, +) from dimos.manipulation.visualization.viser.state import ( - ActionStatus, - FeasibilityStatus, - OperationWorker, PanelPlanState, PlanStatus, TargetEvaluationRequest, - TargetEvaluationWorker, TargetStatus, ) -from dimos.manipulation.visualization.viser.theme import _dimos_logo_data_url, apply_dimos_theme +from dimos.manipulation.visualization.viser.theme import apply_dimos_theme +from dimos.manipulation.visualization.viser.visualizer import ViserManipulationVisualizer from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.sensor_msgs.JointState import JointState - -GuiCallback = Callable[[SimpleNamespace], None] -ThemeValue = str | bool | tuple[int, int, int] | dict[str, str | dict[str, str]] | None -RobotConfigOverride = str | list[str] | list[float] | None - - -@dataclass -class RobotConfigStub: - name: str = "arm" - joint_names: list[str] | None = None - end_effector_link: str = "ee_link" - base_link: str = "base_link" - home_joints: list[float] | None = None - joint_limits_lower: list[float] | None = None - joint_limits_upper: list[float] | None = None - - def __post_init__(self) -> None: - if self.joint_names is None: - self.joint_names = ["j1", "j2"] - - -@dataclass -class SceneRobotConfigStub: - name: str = "arm" - model_path: str = "/tmp/arm.urdf" - package_paths: dict[str, str] | None = None - xacro_args: dict[str, str] | None = None - auto_convert_meshes: bool = False - joint_names: list[str] | None = None - - def __post_init__(self) -> None: - if self.package_paths is None: - self.package_paths = {} - if self.xacro_args is None: - self.xacro_args = {} - if self.joint_names is None: - self.joint_names = ["joint1"] - - -@dataclass -class NamedState: - name: str - - -@dataclass -class GuiMarkdownHandle: - value: str - removed: bool = False - - def remove(self) -> None: - self.removed = True - - -@dataclass -class GuiDropdownHandle: - label: str - options: list[str] - value: str - update_callback: GuiCallback | None = None - removed: bool = False - - def on_update(self, callback: GuiCallback) -> None: - self.update_callback = callback - - def remove(self) -> None: - self.removed = True +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint @dataclass -class GuiButtonHandle: - label: str +class Handle: + label: str = "" + value: object = None + options: list[str] | None = None disabled: bool = False - click_callback: GuiCallback | None = None + color: tuple[int, int, int] | None = None + min: float = 0.0 + max: float = 0.0 + step: float = 0.0 + visible: bool = True + callback: Callable[[object], None] | None = None removed: bool = False - def on_click(self, callback: GuiCallback) -> None: - self.click_callback = callback + def on_update(self, callback: Callable[[object], None]) -> None: + self.callback = callback - def remove(self) -> None: - self.removed = True - - -@dataclass -class GuiCheckboxHandle: - label: str - value: bool - update_callback: GuiCallback | None = None - removed: bool = False - - def on_update(self, callback: GuiCallback) -> None: - self.update_callback = callback + def on_click(self, callback: Callable[[object], None]) -> None: + self.callback = callback def remove(self) -> None: self.removed = True -@dataclass -class GuiSliderHandle: - label: str - min: float - max: float - step: float - value: float - removed: bool = False - update_callback: GuiCallback | None = None +class Folder(Handle): + def __init__(self, label: str, **kwargs: bool) -> None: + super().__init__(label=label) + self.kwargs = kwargs - def on_update(self, callback: GuiCallback) -> None: - self.update_callback = callback + def __enter__(self) -> Folder: + return self - def remove(self) -> None: - self.removed = True + def __exit__(self, *_: object) -> bool: + return False -class FakeHandle: +class Gui: def __init__(self) -> None: - self.visible: object | None = None - self.removed = False - self.name = "" - self.kwargs: dict[str, float | bool] = {} - - def remove(self) -> None: - self.removed = True - + self.folders: list[Folder] = [] + self.buttons: list[Handle] = [] + self.dropdowns: list[Handle] = [] + self.sliders: list[Handle] = [] + self.markdown: list[Handle] = [] + self.theme_kwargs: dict[str, object] | None = None + + def add_folder(self, label: str, **kwargs: bool) -> Folder: + folder = Folder(label, **kwargs) + self.folders.append(folder) + return folder + + def add_markdown(self, value: str) -> Handle: + handle = Handle(value=value) + self.markdown.append(handle) + return handle -class FakeUrdf: - def __init__(self, names: tuple[str, ...]) -> None: - self._urdf = SimpleNamespace(actuated_joint_names=names) - self._meshes = [] - self.cfg = None - self.removed = False + def add_button(self, label: str, **kwargs: object) -> Handle: + color = kwargs.get("color") + handle = Handle( + label=label, + disabled=bool(kwargs.get("disabled", False)), + color=color if isinstance(color, tuple) else None, + ) + self.buttons.append(handle) + return handle - def update_cfg(self, cfg: Sequence[float]) -> None: - self.cfg = list(cfg) + def add_dropdown(self, label: str, *, options: Sequence[str], initial_value: str) -> Handle: + handle = Handle(label=label, options=list(options), value=initial_value) + self.dropdowns.append(handle) + return handle - def remove(self) -> None: - self.removed = True + def add_checkbox(self, label: str, *, initial_value: bool) -> Handle: + return Handle(label=label, value=initial_value) + def add_slider(self, label: str, **kwargs: float) -> Handle: + handle = Handle(label=label, value=kwargs["initial_value"]) + handle.min, handle.max, handle.step = kwargs["min"], kwargs["max"], kwargs["step"] + self.sliders.append(handle) + return handle -class FakeJointState(JointState): - def __init__( - self, - name: Sequence[str], - position: Sequence[float] | None = None, - velocity: Sequence[float] | None = None, - effort: Sequence[float] | None = None, - ) -> None: - self.ts = 0.0 - self.frame_id = "" - self.name = list(name) - self.position = list(position or []) - self.velocity = list(velocity or []) - self.effort = list(effort or []) + def configure_theme(self, **kwargs: object) -> None: + self.theme_kwargs = kwargs -class FakeServer: +class Server: def __init__(self) -> None: + self.gui = Gui() self.scene = SimpleNamespace() - self.scene.add_transform_controls = self.add_transform_controls - def add_transform_controls(self, path: str, *, scale: float) -> FakeTransformHandle: - handle = FakeTransformHandle() - handle.path = path - handle.scale = scale - return handle +@dataclass +class Config: + name: str + joint_names: list[str] + joint_limits_lower: list[float] + joint_limits_upper: list[float] + home_joints: list[float] | None + base_link: str = "base" + end_effector_link: str = "tool" + model_path: Path | str = "robot.urdf" + package_paths: dict[str, str] | None = None + xacro_args: dict[str, str] | None = None + auto_convert_meshes: bool = False + max_velocity: float = 1.0 + max_acceleration: float = 1.0 + joint_name_mapping: dict[str, str] | None = None + coordinator_task_name: str | None = None + pre_grasp_offset: float = 0.0 -class FakeGridServer(FakeServer): - def __init__(self) -> None: - super().__init__() - self.grids = [] - self.scene.add_grid = self.add_grid - - def add_grid(self, name: str, **kwargs: float | bool) -> FakeHandle: - handle = FakeHandle() - handle.name = name - handle.kwargs = kwargs - handle.visible = kwargs.get("visible") - self.grids.append(handle) - return handle + def __post_init__(self) -> None: + if isinstance(self.model_path, str): + self.model_path = Path(self.model_path) + + +def group(robot: str, name: str, joints: tuple[str, ...], *, pose: bool = False) -> PlanningGroup: + return PlanningGroup( + f"{robot}/{name}", + robot, + name, + tuple(f"{robot}/{joint}" for joint in joints), + joints, + "base", + "tool" if pose else None, + ) + + +class Module: + def __init__(self, groups: list[PlanningGroup], states: dict[str, JointState]) -> None: + self.groups = groups + self.states = states + robots = {item.robot_name for item in groups} + self.configs = { + robot_name: Config(robot_name, ["j1", "j2"], [-1.0, -2.0], [1.0, 2.0], [0.0, 0.0]) + for robot_name in robots + } + self.plans: list[tuple[tuple[str, ...], dict[str, JointState]]] = [] + self.executions = 0 + self.cancelled = 0 + self.cleared = 0 + self.last_plan: GeneratedPlan | None = None + self.state = "IDLE" + + def make_plan(self, group_ids: tuple[str, ...]) -> GeneratedPlan: + names = [ + name for group in self.groups for name in group.joint_names if group.id in group_ids + ] + if not names: + names = ["robot/j1", "robot/j2"] + plan = GeneratedPlan( + group_ids=group_ids, + trajectory=JointTrajectory( + joint_names=names, + points=[ + TrajectoryPoint(0.0, [0.0] * len(names)), + TrajectoryPoint(1.0, [1.0] * len(names)), + ], + ), + path=[JointState({"name": names, "position": [0.0] * len(names)})], + status=PlanningStatus.SUCCESS, + ) + self.last_plan = plan + self.state = "READY" + return plan + def list_robots(self) -> list[str]: + return list(self.configs) -class FakeTransformHandle(FakeHandle): - def __init__(self) -> None: - super().__init__() - self.position = (0.0, 0.0, 0.0) - self.wxyz = (1.0, 0.0, 0.0, 0.0) - self.color = None - self.material_color = None - self.update_callback = None - self.path = "" - self.scale = 0.0 + def list_planning_groups(self) -> list[PlanningGroup]: + return self.groups - def on_update(self, callback: GuiCallback) -> None: - self.update_callback = callback + def robot_items(self) -> list[tuple[str, str, Config]]: + return [(name, f"id-{name}", config) for name, config in self.configs.items()] + def robot_id_for_name(self, name: str) -> str: + return f"id-{name}" -class FakeTransformServer(FakeServer): - def __init__(self) -> None: - super().__init__() - self.transform_controls = [] - self.scene.add_transform_controls = self.add_transform_controls - - def add_transform_controls(self, path: str, *, scale: float) -> FakeTransformHandle: - handle = FakeTransformHandle() - handle.path = path - handle.scale = scale - self.transform_controls.append(handle) - return handle + def get_robot_config(self, name: str) -> Config: + return self.configs[name] + def get_init_joints(self, name: str) -> JointState: + return JointState({"name": self.configs[name].joint_names, "position": [-0.5, -1.0]}) -class FakeFolder: - def __init__(self, label: str, kwargs: dict[str, bool]) -> None: - self.label = label - self.kwargs = kwargs - self.entered = False - self.exited = False - self.removed = False + def get_state(self) -> str: + return self.state - def __enter__(self) -> FakeFolder: - self.entered = True - return self + def get_error(self) -> str: + return "" - def __exit__( - self, - exc_type: type[BaseException] | None, - exc: BaseException | None, - traceback: TracebackType | None, - ) -> bool: - self.exited = True - return False + def reset(self) -> SimpleNamespace: + return SimpleNamespace(is_success=lambda: True) - def remove(self) -> None: - self.removed = True + def plan_to_joint_targets(self, targets: dict[str, JointState]) -> bool: + self.plans.append((tuple(targets), targets)) + return True + def preview_plan(self) -> bool: + return True -class FakeGuiServer: - def __init__(self) -> None: - self.theme_kwargs: dict[str, ThemeValue] | None = None - self.folders = [] - self.gui = SimpleNamespace( - add_markdown=lambda value: GuiMarkdownHandle(value=value), - add_dropdown=self.add_dropdown, - add_button=self.add_button, - add_checkbox=self.add_checkbox, - add_slider=self.add_slider, - add_folder=self.add_folder, - configure_theme=self.configure_theme, - ) - self.buttons: dict[str, GuiButtonHandle] = {} - self.checkboxes: dict[str, GuiCheckboxHandle] = {} - self.sliders: list[GuiSliderHandle] = [] + def execute(self) -> bool: + self.executions += 1 + self.state = "DISPATCHING" + return True - def configure_theme(self, **kwargs: ThemeValue) -> None: - self.theme_kwargs = kwargs + def cancel(self) -> bool: + self.cancelled += 1 + return True - def add_folder(self, label: str, **kwargs: bool) -> FakeFolder: - handle = FakeFolder(label, kwargs) - self.folders.append(handle) - return handle + def clear_planned_path(self) -> bool: + self.cleared += 1 + return True - def add_dropdown( - self, label: str, *, options: Sequence[str], initial_value: str - ) -> GuiDropdownHandle: - handle = GuiDropdownHandle(label=label, options=list(options), value=initial_value) - return handle - def add_button(self, label: str, *, disabled: bool = False) -> GuiButtonHandle: - handle = GuiButtonHandle(label=label, disabled=disabled) - self.buttons[label] = handle - return handle +class Monitor: + def __init__(self, module: Module) -> None: + self.module = module + self.invalid: set[str] = set() + self.stale: set[str] = set() + self.poses: dict[str, Pose] = {} - def add_checkbox(self, label: str, *, initial_value: bool) -> GuiCheckboxHandle: - handle = GuiCheckboxHandle(label=label, value=initial_value) - self.checkboxes[label] = handle - return handle + def get_current_joint_state(self, robot_id: str) -> JointState: + return JointState(self.module.states[robot_id.removeprefix("id-")]) - def add_slider( - self, - label: str, - *, - min: float, - max: float, - step: float, - initial_value: float, - ) -> GuiSliderHandle: - handle = GuiSliderHandle(label=label, min=min, max=max, step=step, value=initial_value) - self.sliders.append(handle) - return handle + def is_state_stale(self, robot_id: str, max_age: float = 1.0) -> bool: + return robot_id in self.stale + def is_state_valid(self, robot_id: str, _state: JointState) -> bool: + return robot_id not in self.invalid -def make_robot_config(**overrides: RobotConfigOverride) -> RobotConfigStub: - """Build a faithful RobotModelConfig stand-in with the fields the panel reads.""" - config = RobotConfigStub() - for name, value in overrides.items(): - setattr(config, name, value) - return config + def get_group_ee_pose(self, group_id: str, _state: JointState | None = None) -> Pose: + return self.poses.get( + group_id, + Pose({"position": [0.1, 0.2, 0.3], "orientation": [0.0, 0.0, 0.0, 1.0]}), + ) -class FakeManipulationModule(SimpleNamespace): - """Public ManipulationModule surface used by the in-process Viser adapter tests.""" +class Operator: + def __init__(self, module: Module, monitor: Monitor) -> None: + self.module = module + self.monitor = monitor - def list_robots(self) -> list[str]: - return list(getattr(self, "_robots", {}).keys()) + def status(self) -> OperatorStatus: + return OperatorStatus( + state=self.module.get_state(), + diagnostic=self.module.get_error(), + ready_plan_status="READY", + ready_plan_id="fake-plan", + ) - def robot_items(self) -> list[tuple[str, str, RobotConfigStub | SimpleNamespace]]: - return [ - (name, robot_id, config) - for name, (robot_id, config, _) in getattr(self, "_robots", {}).items() - ] + def get_init_joints(self, robot_name: str) -> JointState | None: + return self.module.get_init_joints(robot_name) - def robot_id_for_name(self, robot_name: str) -> str | None: - entry = getattr(self, "_robots", {}).get(robot_name) - return entry[0] if entry is not None else None - - def robot_name_for_id(self, robot_id: str) -> str | None: - for robot_name, (candidate_id, _, _) in getattr(self, "_robots", {}).items(): - if candidate_id == robot_id: - return robot_name - return None - - def get_robot_config(self, robot_name: str) -> RobotConfigStub | SimpleNamespace | None: - entry = getattr(self, "_robots", {}).get(robot_name) - return entry[1] if entry is not None else None - - def get_robot_info(self, robot_name: str) -> RobotInfo | None: - config = self.get_robot_config(robot_name) - if config is None: - return None - init = self.get_init_joints(robot_name) - home_joints = config.home_joints if hasattr(config, "home_joints") else None - return { - "name": config.name, - "world_robot_id": self.robot_id_for_name(robot_name) or robot_name, - "joint_names": list(config.joint_names), - "end_effector_link": config.end_effector_link, - "base_link": config.base_link, - "max_velocity": 1.0, - "max_acceleration": 1.0, - "has_joint_name_mapping": False, - "coordinator_task_name": None, - "home_joints": list(home_joints) if home_joints is not None else None, - "pre_grasp_offset": 0.0, - "init_joints": list(init.position) if init is not None else None, + def evaluate_joint_target(self, request: object) -> TargetEvaluationResult: + target = request.target # type: ignore[attr-defined] + group_ids = request.group_ids # type: ignore[attr-defined] + diagnostics = { + group_id: "Target is collision-free for this robot" for group_id in group_ids } + poses = {group_id: self.monitor.get_group_ee_pose(group_id) for group_id in group_ids} + return TargetEvaluationResult( + True, + "FEASIBLE", + "Target is collision-free for each robot", + True, + tuple(group_ids), + target, + diagnostics, + poses, + ) - def get_init_joints(self, robot_name: str) -> JointState | None: - return getattr(self, "_init_joints", {}).get(robot_name) + def evaluate_pose_target(self, request: object) -> TargetEvaluationResult: + group_ids = tuple( + dict.fromkeys((*request.pose_targets.keys(), *request.auxiliary_group_ids)) + ) # type: ignore[attr-defined] + js = JointState( + { + "name": [ + name + for group in self.module.groups + for name in group.joint_names + if group.id in group_ids + ], + "position": [ + 0.7 + for group in self.module.groups + for _ in group.joint_names + if group.id in group_ids + ], + } + ) + return TargetEvaluationResult( + True, + "FEASIBLE", + "ok", + True, + group_ids, + js, + {}, + {group_id: self.monitor.get_group_ee_pose(group_id) for group_id in group_ids}, + ) - def get_planned_path(self, robot_name: str) -> list[JointState] | None: - return getattr(self, "_planned_paths", {}).get(robot_name) + def plan_to_joints(self, request: object) -> GeneratedPlan: + self.module.plan_to_joint_targets( + {group_id: JointState({"name": [], "position": []}) for group_id in request.group_ids} + ) # type: ignore[attr-defined] + return self.module.make_plan(tuple(request.group_ids)) # type: ignore[attr-defined] - def get_planned_trajectory_duration(self, robot_name: str) -> float | None: - trajectory = getattr(self, "_planned_trajectories", {}).get(robot_name) - return None if trajectory is None else float(trajectory.duration) + def plan_to_pose(self, request: object) -> GeneratedPlan: + return self.module.make_plan(tuple(request.pose_targets)) # type: ignore[attr-defined] - def get_state(self) -> str: - state = getattr(self, "_state", "IDLE") - return str(getattr(state, "name", state)) + def preview(self, plan: GeneratedPlan, duration: float | None = None) -> bool: + return self.module.preview_plan() - def get_error(self) -> str: - return str(getattr(self, "_error_message", "")) - - def evaluate_joint_target(self, joints: JointState | None, robot_name: str) -> TargetEvaluation: - robot_id = self.robot_id_for_name(robot_name) - if robot_id is None or joints is None: - return {"success": False, "status": "NO_ROBOT", "joint_state": None} - world_monitor = getattr(self, "_world_monitor", None) - if world_monitor is None: - return {"success": False, "status": "UNAVAILABLE", "joint_state": None} - collision_free = world_monitor.is_state_valid(robot_id, joints) - return { - "success": True, - "status": "FEASIBLE" if collision_free else "COLLISION", - "message": "Target is collision-free" if collision_free else "Target is in collision", - "collision_free": collision_free, - "ee_pose": world_monitor.get_ee_pose(robot_id, joints), - "joint_state": joints, - } + def execute(self, plan: GeneratedPlan) -> bool: + return self.module.execute() - def evaluate_pose_target(self, _pose: Pose, _robot_name: str) -> TargetEvaluation: - return { - "success": False, - "joint_state": None, - "status": "UNAVAILABLE", - "message": "No fake pose IK", - "collision_free": False, - } + def cancel(self) -> bool: + return self.module.cancel() + def clear_plan(self) -> bool: + return self.module.clear_planned_path() -def make_adapter_with_robot() -> InProcessViserAdapter: - current = FakeJointState(["j1", "j2"], position=[0.3, 0.4]) - config = make_robot_config( - name="arm", - joint_names=["j1", "j2"], - joint_limits_lower=[-1.0, -2.0], - joint_limits_upper=[1.0, 2.0], - home_joints=[0.0, 0.0], - ) - world_monitor = SimpleNamespace( - get_current_joint_state=lambda _robot_id: current, - is_state_stale=lambda _robot_id, max_age=1.0: False, - is_state_valid=lambda _robot_id, _joint_state: True, - get_ee_pose=lambda _robot_id, joint_state=None: None, - ) - module = FakeManipulationModule( - _robots={"arm": ("robot-1", config, None)}, - _init_joints={"arm": FakeJointState(["j1", "j2"], position=[0.1, 0.2])}, - _planned_paths={}, - _planned_trajectories={}, - _state=NamedState(name="IDLE"), - _error_message="", - _world_monitor=world_monitor, - ) - return InProcessViserAdapter( - world_monitor=world_monitor, - manipulation_module=module, + def reset(self) -> bool: + result = self.module.reset() + return result.is_success() + + +def session_inputs(module: Module) -> tuple[PlanningSceneInfo, Operator, dict[str, JointState]]: + monitor = Monitor(module) + robots = {f"id-{name}": config for name, config in module.configs.items()} + current = {f"id-{name}": JointState(state) for name, state in module.states.items()} + return ( + PlanningSceneInfo(robots=robots, planning_groups=tuple(module.groups)), + Operator(module, monitor), + current, ) +def scene_gui(module: Module, server: Server, scene: ViserManipulationScene) -> ViserPanelGui: + scene_info, operator, current = session_inputs(module) + return ViserPanelGui(server, scene_info, operator, current, ViserVisualizationConfig(), scene) + + @pytest.fixture -def make_panel() -> Iterator[Callable[..., ViserPanelGui]]: - """Build and start a ViserPanelGui, closing it (and its worker threads) on teardown.""" +def panel() -> Iterator[ + Callable[[list[PlanningGroup], dict[str, JointState]], tuple[ViserPanelGui, Module, Server]] +]: panels: list[ViserPanelGui] = [] - def _make( - server: FakeGuiServer | FakeServer, - adapter: InProcessViserAdapter, - config: ViserVisualizationConfig | None = None, - scene: ViserManipulationScene | None = None, - ) -> ViserPanelGui: + def make( + groups: list[PlanningGroup], states: dict[str, JointState] + ) -> tuple[ViserPanelGui, Module, Server]: + module = Module(groups, states) + server = Server() + scene_info, operator, current = session_inputs(module) gui = ViserPanelGui( - server, adapter, config or ViserVisualizationConfig(panel_enabled=True), scene + server, scene_info, operator, current, ViserVisualizationConfig(panel_enabled=True) ) gui.start() panels.append(gui) - return gui + return gui, module, server - yield _make + yield make for gui in panels: gui.close() -def test_viser_config_enables_panel_by_default() -> None: - assert ViserVisualizationConfig().panel_enabled is True +def states(*robots: str) -> dict[str, JointState]: + return {robot: JointState({"name": ["j1", "j2"], "position": [0.1, 0.2]}) for robot in robots} -def test_gui_builds_controls_in_manipulation_panel_folder( - make_panel: Callable[..., ViserPanelGui], -) -> None: - server = FakeGuiServer() - adapter = make_adapter_with_robot() - gui = make_panel(server, adapter, ViserVisualizationConfig()) - assert server.folders - assert server.folders[0].label == "Manipulation Panel" - assert server.folders[0].kwargs == {"expand_by_default": True} - assert "status" in gui._handles - assert "robot" in gui._handles - assert "plan" in gui._handles - assert gui._operation_worker._timeout_seconds is None - - -def test_gui_scene_grid_checkbox_toggles_reference_grid( - make_panel: Callable[..., ViserPanelGui], +def test_panel_contract_group_order_defaults_and_controls( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], ) -> None: - grid_server = FakeGridServer() - scene = ViserManipulationScene( - grid_server, lambda *args, **kwargs: FakeUrdf(("joint1",)), preview_fps=10.0 - ) - server = FakeGuiServer() - adapter = make_adapter_with_robot() - make_panel(server, adapter, ViserVisualizationConfig(), scene) - assert grid_server.grids - assert server.checkboxes["Scene grid"].value is True - server.checkboxes["Scene grid"].update_callback( - SimpleNamespace(target=SimpleNamespace(value=False)) - ) - assert grid_server.grids[0].visible is False - server.checkboxes["Scene grid"].update_callback( - SimpleNamespace(target=SimpleNamespace(value=True)) - ) - assert grid_server.grids[0].visible is True - + pose = group("arm", "manipulator", ("j1",), pose=True) + auxiliary = group("arm", "gripper", ("j2",)) + gui, _module, server = panel([auxiliary, pose], states("arm")) -def test_gui_close_removes_handles_and_late_callbacks_are_noops( - make_panel: Callable[..., ViserPanelGui], -) -> None: - server = FakeGuiServer() - grid_server = FakeGridServer() - scene = ViserManipulationScene( - grid_server, lambda *args, **kwargs: FakeUrdf(("joint1",)), preview_fps=10.0 + assert [(folder.label, folder.kwargs) for folder in server.gui.folders] == [ + ("Manipulation Panel", {"expand_by_default": True}), + ("Joint Control", {"expand_by_default": False}), + ] + assert [button.label for button in server.gui.buttons] == [ + "arm", + "arm gripper", + "Plan", + "Preview", + "Execute", + "Cancel", + "Clear plan", + ] + assert "robot" not in gui._handles + assert ( + server.gui.markdown[1].value + == "### Planning Groups\nActive planning groups for pose goals, planning, and joint edits." ) - adapter = make_adapter_with_robot() - gui = make_panel(server, adapter, ViserVisualizationConfig(), scene) - robot_dropdown = gui._handles["robot"] - plan_button = server.buttons["Plan"] - grid = grid_server.grids[0] - handles = list(gui._handles.values()) - - gui.close() - if isinstance(robot_dropdown, GuiDropdownHandle) and robot_dropdown.update_callback is not None: - robot_dropdown.update_callback(SimpleNamespace(target=SimpleNamespace(value="arm"))) - if plan_button.click_callback is not None: - plan_button.click_callback(SimpleNamespace()) - gui._set_scene_grid_visible(False) - - assert all(getattr(handle, "removed", False) for handle in handles) - assert gui._handles == {} - assert grid.visible is True + assert [button.color for button in server.gui.buttons[:2]] == [ + ACTIVE_GROUP_COLOR, + INACTIVE_GROUP_COLOR, + ] + assert gui.state.selected_group_ids == ("arm/manipulator",) + assert server.gui.dropdowns[0].options == ["Select preset...", "Init", "Current", "Home"] + assert [ + (slider.label, slider.min, slider.max, slider.value) for slider in server.gui.sliders + ] == [("arm/manipulator/j1", -1.0, 1.0, 0.1)] + server.gui.buttons[1].callback(SimpleNamespace()) + assert gui.state.selected_group_ids == ("arm/manipulator", "arm/gripper") + assert [slider.label for slider in server.gui.sliders if not slider.removed] == [ + "arm/manipulator/j1", + "arm/gripper/j2", + ] -def test_gui_ignores_target_evaluation_after_close( - make_panel: Callable[..., ViserPanelGui], +def test_gui_target_ghost_states_use_exact_group_names() -> None: + left, right = group("left", "manipulator", ("j1",)), group("right", "manipulator", ("j1",)) + module = Module([left, right], states("left", "right")) + scene_info, operator, current = session_inputs(module) + gui = ViserPanelGui(Server(), scene_info, operator, current, ViserVisualizationConfig()) + gui.state.selected_group_ids = (left.id, right.id) + targets = { + left.id: JointState({"name": ["left/j1"], "position": [0.7]}), + right.id: JointState({"name": ["right/j1"], "position": [0.8]}), + } + ghost_states = gui._target_ghost_states(targets) + assert ghost_states["left"].position == [0.7, 0.2] + assert ghost_states["right"].position == [0.8, 0.2] + assert gui.evaluate_joint_target_set((left.id, right.id), targets).status == "FEASIBLE" + + +def test_target_callbacks_require_current_sequence_and_selection_epoch( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], ) -> None: - adapter = make_adapter_with_robot() - gui = make_panel(FakeGuiServer(), adapter) - gui.state.selected_robot = "arm" - sequence_id = gui.state.next_sequence_id() - request = TargetEvaluationRequest( - sequence_id=sequence_id, - source="joints", - robot_name="arm", - joints=FakeJointState(["j1", "j2"], position=[0.1, 0.2]), + first, second = ( + group("arm", "manipulator", ("j1",), pose=True), + group("arm", "gripper", ("j2",)), ) - gui.close() - - gui._apply_target_evaluation_result( - request, - { - "success": True, - "collision_free": True, - "status": "FEASIBLE", - "joint_state": FakeJointState(["j1", "j2"], position=[0.8, 0.9]), - }, + gui, _module, _server = panel([first, second], states("arm")) + request = TargetEvaluationRequest( + 1, "joints", selection_epoch=gui.state.selection_epoch, group_ids=(first.id,) ) - + gui.state.latest_sequence_id = 2 + gui._apply_target_evaluation_result(request, TargetEvaluationResult(True, "FEASIBLE", "", True)) + assert gui.state.target_status == TargetStatus.EMPTY + gui.state.latest_sequence_id = 1 + gui.state.advance_selection_epoch() + gui._apply_target_evaluation_result(request, TargetEvaluationResult(True, "FEASIBLE", "", True)) assert gui.state.target_status == TargetStatus.CHECKING - assert gui.state.joint_target is None - - -def test_dimos_theme_configures_supported_viser_chrome() -> None: - server = FakeGuiServer() - - assert apply_dimos_theme(server) is True - assert server.theme_kwargs is not None - assert server.theme_kwargs["brand_color"] == (22, 130, 163) - assert server.theme_kwargs["dark_mode"] is True - assert server.theme_kwargs["show_logo"] is False - assert server.theme_kwargs["show_share_button"] is False - assert server.theme_kwargs["control_layout"] == "collapsible" - assert server.theme_kwargs["control_width"] == "medium" - - -def test_dimos_theme_configures_titlebar_when_supported(monkeypatch: pytest.MonkeyPatch) -> None: - fake_viser = ModuleType("viser") - fake_theme = ModuleType("viser.theme") - fake_theme.TitlebarImage = lambda **kwargs: kwargs - fake_theme.TitlebarButton = lambda **kwargs: kwargs - fake_theme.TitlebarConfig = lambda **kwargs: kwargs - monkeypatch.setitem(sys.modules, "viser", fake_viser) - monkeypatch.setitem(sys.modules, "viser.theme", fake_theme) - server = FakeGuiServer() - - assert apply_dimos_theme(server) is True - assert server.theme_kwargs is not None - titlebar_content = server.theme_kwargs["titlebar_content"] - assert isinstance(titlebar_content, dict) - image = titlebar_content["image"] - assert isinstance(image, dict) - assert image["image_alt"] == "Dimensional" - assert image["image_url_light"].startswith("data:image/svg+xml;base64,") -def test_dimos_logo_asset_loads_as_data_url() -> None: - logo_url = _dimos_logo_data_url() - - assert logo_url is not None - assert logo_url.startswith("data:image/svg+xml;base64,") +def test_plan_target_sequence_invalidation_and_unfiltered_all_robot_execute( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], monkeypatch: pytest.MonkeyPatch +) -> None: + left, right = ( + group("left", "manipulator", ("j1",), pose=True), + group("right", "manipulator", ("j1",), pose=True), + ) + gui, module, _server = panel([left, right], states("left", "right")) + gui._toggle_group_selected(right.id) + gui.state.target_status = TargetStatus.FEASIBLE + monkeypatch.setattr( + gui, + "_operation_worker", + SimpleNamespace(submit=lambda operation, **_: operation(), stop=lambda **_: None), + ) + gui._submit_plan() + assert module.plans[-1][0] == (left.id, right.id) + assert gui.state.plan_state.group_ids == (left.id, right.id) + gui.state.next_sequence_id() + assert gui.state.plan_state.status == PlanStatus.STALE + gui.state.plan_state = PanelPlanState( + status=PlanStatus.FRESH, + group_ids=(left.id, right.id), + target_sequence_id=gui.state.latest_sequence_id, + plan=module.last_plan, + plan_id="fake-plan", + ) + gui.state.target_status = TargetStatus.FEASIBLE + gui._submit_execute() + assert module.executions == 1 -def test_dimos_theme_is_non_blocking_when_theme_api_fails() -> None: - class BrokenGui: - @staticmethod - def configure_theme(**_kwargs: ThemeValue) -> None: - raise TypeError("theme API changed") +def test_initialization_waits_for_complete_fresh_telemetry( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], +) -> None: + selected = group("arm", "manipulator", ("j1", "j2"), pose=True) + gui, module, _server = panel( + [selected], {"arm": JointState({"name": ["j1"], "position": [0.4]})} + ) - server = SimpleNamespace(gui=BrokenGui()) + assert selected.id not in gui.state.group_joint_targets + module.states["arm"] = JointState({"name": ["j1", "j2"], "position": [0.4, 0.5]}) + gui.refresh() + assert selected.id not in gui.state.group_joint_targets - assert apply_dimos_theme(server) is False + gui.current_states["id-arm"] = module.states["arm"] + gui.refresh() + assert gui.state.group_joint_targets[selected.id].position == [0.4, 0.5] -def test_dimos_theme_retries_without_titlebar_when_titlebar_content_fails( - monkeypatch: pytest.MonkeyPatch, +def test_incomplete_preset_preserves_existing_group_targets( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], ) -> None: - fake_viser = ModuleType("viser") - fake_theme = ModuleType("viser.theme") - fake_theme.TitlebarImage = lambda **kwargs: kwargs - fake_theme.TitlebarButton = lambda **kwargs: kwargs - fake_theme.TitlebarConfig = lambda **kwargs: kwargs - monkeypatch.setitem(sys.modules, "viser", fake_viser) - monkeypatch.setitem(sys.modules, "viser.theme", fake_theme) - titlebar_values: list[ThemeValue] = [] - - class FallbackGui: - @staticmethod - def configure_theme(**kwargs: ThemeValue) -> None: - titlebar_values.append(kwargs["titlebar_content"]) - if kwargs["titlebar_content"] is not None: - raise TypeError("titlebar unsupported") - - server = SimpleNamespace(gui=FallbackGui()) + selected = group("arm", "manipulator", ("j1", "j2"), pose=True) + gui, module, _server = panel([selected], states("arm")) + before = JointState(gui.state.group_joint_targets[selected.id]) + module.get_init_joints = lambda name: JointState({"name": ["j1"], "position": [-0.5]}) - assert apply_dimos_theme(server) is True - assert titlebar_values[0] is not None - assert titlebar_values[1] is None + gui._apply_preset("Init") + assert gui.state.group_joint_targets[selected.id] == before + assert "missing joints" in gui.state.error -class FakeMesh: - def __init__(self) -> None: - self.visible = None - self.color = None - self.material_color = None - self.opacity = None +def test_valid_init_preset_builds_sliders_after_incomplete_initial_telemetry( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], +) -> None: + selected = group("arm", "manipulator", ("j1", "j2"), pose=True) + gui, module, server = panel( + [selected], {"arm": JointState({"name": ["j1"], "position": [0.4]})} + ) -class FakeViserUrdfWithMeshes: - def __init__(self, names: tuple[str, ...] = ("joint1", "joint2", "joint3")) -> None: - self._urdf = SimpleNamespace(actuated_joint_names=names) - self._meshes = [FakeMesh(), FakeMesh()] - self.cfg = None + assert gui.state.group_joint_targets == {} + assert server.gui.sliders == [] - def update_cfg(self, cfg: Sequence[float]) -> None: - self.cfg = list(cfg) + module.configs["arm"].home_joints = [-0.5, -1.0] + gui._apply_preset("Init") + assert gui.state.group_joint_targets[selected.id].position == [-0.5, -1.0] + assert [slider.label for slider in server.gui.sliders if not slider.removed] == [ + "arm/manipulator/j1", + "arm/manipulator/j2", + ] -def test_viser_joint_configuration_maps_names_to_urdf_order() -> None: - server = FakeServer() - urdf = FakeUrdf(("shoulder", "elbow", "wrist")) - scene = ViserManipulationScene(server, lambda *args, **kwargs: urdf, preview_fps=10.0) - scene.prepared_urdf_path = lambda config: "dummy.urdf" - cfg = SimpleNamespace( - name="arm", - model_path="/tmp/arm.urdf", - package_paths={}, - xacro_args={}, - auto_convert_meshes=False, - joint_names=["arm/shoulder", "elbow"], +def test_incomplete_multi_group_preset_does_not_change_any_targets( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], +) -> None: + left, right = ( + group("left", "manipulator", ("j1",), pose=True), + group("right", "manipulator", ("j1",), pose=True), ) - scene.register_robot("robot1", cfg) - scene.set_urdf_joints(urdf, cfg.joint_names, [1.5, 2.5]) - assert urdf.cfg == [1.5, 2.5, 0.0] - - -def test_scene_adds_reference_grid_when_supported() -> None: - server = FakeGridServer() - scene = ViserManipulationScene( - server, lambda *args, **kwargs: FakeUrdf(("j1",)), preview_fps=10.0 + gui, module, _server = panel([left, right], states("left", "right")) + gui._toggle_group_selected(right.id) + before = { + group_id: JointState(target) for group_id, target in gui.state.group_joint_targets.items() + } + module.get_init_joints = lambda name: JointState( + {"name": ["j1"] if name == "left" else [], "position": [-0.5] if name == "left" else []} ) - assert scene.has_reference_grid() is True - assert len(server.grids) == 1 - grid = server.grids[0] - assert grid.name == "/reference_grid" - assert grid.kwargs["plane"] == "xy" - assert grid.kwargs["infinite_grid"] is True - assert grid.kwargs["visible"] is True + gui._apply_preset("Init") - scene.set_reference_grid_visible(False) - assert grid.visible is False - scene.set_reference_grid_visible(True) - assert grid.visible is True - - -def test_preview_visibility_only_affects_preview_ghost_and_close_removes_handles() -> None: - server = FakeServer() - urdfs = [FakeViserUrdfWithMeshes(("joint1",)) for _ in range(3)] - scene = ViserManipulationScene(server, lambda *args, **kwargs: urdfs.pop(0), preview_fps=10.0) - scene.prepared_urdf_path = lambda config: "dummy.urdf" - config = SimpleNamespace( - name="arm", - model_path="/tmp/arm.urdf", - package_paths={}, - xacro_args={}, - auto_convert_meshes=False, - joint_names=["joint1"], - ) - scene.register_robot("robot1", config) - target = scene._urdfs["robot1:target"] - preview = scene._urdfs["robot1:preview"] - assert all(mesh.visible is True for mesh in target._meshes) - assert all(mesh.visible is False for mesh in preview._meshes) - scene.show_preview("robot1") - assert all(mesh.visible is True for mesh in preview._meshes) - assert all(mesh.visible is True for mesh in target._meshes) - scene.hide_preview("robot1") - assert all(mesh.visible is False for mesh in preview._meshes) - assert all(mesh.visible is True for mesh in target._meshes) - scene.close() - assert scene._handles == {} - assert all(mesh.visible is False for mesh in preview._meshes) - - -def test_target_ghost_is_visible_and_tracks_current_until_target_moves_it() -> None: - server = FakeServer() - urdfs = [FakeViserUrdfWithMeshes(("joint1",)) for _ in range(3)] - scene = ViserManipulationScene(server, lambda *args, **kwargs: urdfs.pop(0), preview_fps=10.0) - scene.prepared_urdf_path = lambda config: "dummy.urdf" - config = SimpleNamespace( - name="arm", - model_path="/tmp/arm.urdf", - package_paths={}, - xacro_args={}, - auto_convert_meshes=False, - joint_names=["joint1"], - ) - scene.register_robot("robot1", config) - current = scene._urdfs["robot1:current"] - target = scene._urdfs["robot1:target"] - preview = scene._urdfs["robot1:preview"] - - assert all(mesh.visible is True for mesh in target._meshes) - assert all(mesh.visible is False for mesh in preview._meshes) - scene.update_current_robot("robot1", FakeJointState(["joint1"], position=[0.25])) - assert current.cfg == [0.25] - assert target.cfg == [0.25] - assert preview.cfg is None - - scene.set_target_joints("robot1", ["joint1"], [0.8]) - scene.update_current_robot("robot1", FakeJointState(["joint1"], position=[0.1])) - assert current.cfg == [0.1] - assert target.cfg == [0.8] - assert preview.cfg is None - - -def test_preview_animation_uses_separate_colored_ghost_and_hides_after_playback() -> None: - server = FakeServer() - urdfs = [FakeViserUrdfWithMeshes(("joint1",)) for _ in range(3)] - scene = ViserManipulationScene(server, lambda *args, **kwargs: urdfs.pop(0), preview_fps=10.0) - scene.prepared_urdf_path = lambda config: "dummy.urdf" - config = SimpleNamespace( - name="arm", - model_path="/tmp/arm.urdf", - package_paths={}, - xacro_args={}, - auto_convert_meshes=False, - joint_names=["joint1"], - ) - scene.register_robot("robot1", config) - target = scene._urdfs["robot1:target"] - preview = scene._urdfs["robot1:preview"] - - assert all(mesh.color == (255, 122, 0) for mesh in target._meshes) - assert all(mesh.color == (80, 180, 255) for mesh in preview._meshes) - assert all(mesh.opacity == 0.55 for mesh in preview._meshes) - - ok = scene.animate_path( - "robot1", - [ - FakeJointState(["joint1"], position=[0.0]), - FakeJointState(["joint1"], position=[1.0]), - ], - duration=0.0, - ) + assert gui.state.group_joint_targets == before + assert "missing joints" in gui.state.error - assert ok is True - assert preview.cfg == [1.0] - assert all(mesh.visible is False for mesh in preview._meshes) - assert all(mesh.visible is True for mesh in target._meshes) - -def test_scene_target_helpers_handle_missing_robot_and_pose() -> None: - server = FakeTransformServer() - scene = ViserManipulationScene( - server, lambda *args, **kwargs: FakeUrdf(("joint1",)), preview_fps=10.0 +def test_cancel_clear_and_close_invalidate_operations_and_preview_generation( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], monkeypatch: pytest.MonkeyPatch +) -> None: + selected = group("arm", "manipulator", ("j1",), pose=True) + gui, module, _server = panel([selected], states("arm")) + submitted: list[Callable[[], None]] = [] + gui._operation_worker.stop() + monkeypatch.setattr( + gui, + "_operation_worker", + SimpleNamespace( + submit=lambda operation, **_: submitted.append(operation), + stop=lambda **_: None, + start=lambda: None, + ), ) - - assert scene.animate_path("missing", [], duration=0.0) is False - assert scene.set_target_joints("missing", ["joint1"], [1.0]) is False - scene.set_target_pose("missing", Pose()) - handle = scene.ensure_target_controls("robot1", lambda _target: None) - scene.set_target_pose("robot1", None) - - assert handle is not None - assert handle.position == (0.0, 0.0, 0.0) - - -def test_scene_close_removes_grid_transform_and_urdf_handles() -> None: - server = FakeGridServer() - current = FakeUrdf(("joint1",)) - target = FakeUrdf(("joint1",)) - scene = ViserManipulationScene( - server, lambda *args, **kwargs: FakeUrdf(("joint1",)), preview_fps=10.0 + gui.state.target_status = TargetStatus.FEASIBLE + gui._submit_plan() + gui._submit_clear() + submitted[0]() + assert module.plans == [] + submitted[1]() + assert module.cleared == 1 and gui.state.plan_state.status == PlanStatus.NONE + gui.close() + status_before_callback = gui.state.target_status + gui._apply_target_evaluation_result( + TargetEvaluationRequest(0, "joints"), TargetEvaluationResult(True, "FEASIBLE", "", True) ) - handle = scene.ensure_target_controls("robot1", lambda _target: None) - scene._urdfs["robot1:current"] = current - scene._urdfs["robot1:target"] = target - - scene.close() - - assert handle is not None and handle.removed is True - assert current.removed is True - assert target.removed is True - assert server.grids[0].removed is True - assert scene.has_reference_grid() is False - + assert gui.state.target_status is status_before_callback -def test_sampled_joint_path_frames_preserves_dense_trajectory_samples() -> None: - dense_path = [FakeJointState(["j1"], position=[float(index)]) for index in range(32)] - - frames = sampled_joint_path_frames(dense_path, duration=1.0, fps=30.0) - - assert frames == [[float(index)] for index in range(32)] +class Mesh: + def __init__(self) -> None: + self.visible = False + self.color: tuple[int, int, int] | None = None + self.opacity: float | None = None -def test_sampled_joint_path_frames_interpolates_sparse_paths() -> None: - sparse_path = [ - FakeJointState(["j1"], position=[0.0]), - FakeJointState(["j1"], position=[1.0]), - ] - frames = sampled_joint_path_frames(sparse_path, duration=1.0, fps=4.0) +class Urdf: + def __init__(self, *_: object, **__: object) -> None: + self._urdf = SimpleNamespace(actuated_joint_names=("j1", "j2")) + self._meshes = [Mesh()] + self.cfg: list[float] | None = None - assert frames == [[0.0], [0.25], [0.5], [0.75], [1.0]] + def update_cfg(self, cfg: Sequence[float]) -> None: + self.cfg = list(cfg) + def remove(self) -> None: + pass -def test_joint_path_frame_edge_cases_and_empty_animation() -> None: - empty_position = FakeJointState(["j1"], position=[]) - single = FakeJointState(["j1"], position=[0.7]) - start = FakeJointState(["j1"], position=[0.0]) - middle = FakeJointState(["j1"], position=[1.0]) - mismatched_final = FakeJointState(["j1", "j2"], position=[2.0, 3.0]) - set_calls: list[list[float]] = [] - sleep_calls: list[float] = [] - assert interpolate_joint_path([empty_position], duration=1.0, fps=10.0) == [] - assert interpolate_joint_path([single], duration=1.0, fps=10.0) == [[0.7]] - assert interpolate_joint_path([start, middle, mismatched_final], duration=1.0, fps=2.0) == [ - [0.0], - [2.0, 3.0], - ] - assert sampled_joint_path_frames([empty_position], duration=1.0, fps=10.0) == [] - assert ( - PreviewAnimator(set_calls.append, sleep=sleep_calls.append).animate( - [empty_position], duration=1.0, fps=10.0 +def test_scene_active_only_ghosts_group_gizmos_feasibility_and_shared_ticks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + updates: list[tuple[str, list[float]]] = [] + server = Server() + server.scene.add_grid = lambda *_args, **_kwargs: Handle() + server.scene.add_transform_controls = lambda *_args, **_kwargs: Handle() + scene = ViserManipulationScene(server, Urdf) + scene.prepared_urdf_path = lambda _config: "robot.urdf" # type: ignore[method-assign] + config = Config("arm", ["j1", "j2"], [-1.0, -2.0], [1.0, 2.0], [0.0, 0.0]) + scene.register_robot("id-arm", config) + scene.set_target_active("id-arm", False) + assert scene._urdfs["id-arm:target"]._meshes[0].visible is False + scene.set_target_joints("id-arm", ["j1", "j2"], [0.8, 0.2]) + assert scene._urdfs["id-arm:target"].cfg == [0.8, 0.2] + scene.set_target_visual_state("id-arm", False) + assert scene._urdfs["id-arm:target"]._meshes[0].color == (255, 30, 30) + monkeypatch.setattr( + "dimos.manipulation.visualization.viser.scene.time.sleep", lambda _delay: None + ) + original = scene._set_preview_ghost_joints + scene._set_preview_ghost_joints = lambda robot, names, values: ( + updates.append((robot, list(values))), + original(robot, names, values), + ) # type: ignore[method-assign] + preview = GroupPreviewAnimation( + ( + PreviewTrack( + "id-arm", + ("j1", "j2"), + ( + PreviewFrame(0.0, (0.0, 0.2)), + PreviewFrame(1.0, (1.0, 0.2)), + ), + ), ) - is False - ) - assert set_calls == [] - assert sleep_calls == [] - - -def test_adapter_copies_joint_state_and_delegates_to_module() -> None: - copied = FakeJointState(["j1"], position=[1.0], velocity=[2.0], effort=[3.0]) - module = FakeManipulationModule( - _robots={"arm": ("robot-1", SimpleNamespace(), None)}, - _planned_paths={"arm": [copied]}, - _planned_trajectories={}, - plan_to_pose=lambda pose, robot_name=None: (pose, robot_name), - plan_to_joints=lambda joints, robot_name=None: (joints, robot_name), - preview_path=lambda robot_name=None: robot_name, - execute=lambda robot_name=None: robot_name, - cancel=lambda: True, - clear_planned_path=lambda: True, - ) - world_monitor = SimpleNamespace( - get_current_joint_state=lambda robot_id: copied, - is_state_stale=lambda robot_id, max_age=1.0: False, - is_state_valid=lambda robot_id, joint_state: True, - get_ee_pose=lambda robot_id, joint_state=None: (robot_id, joint_state), ) - module._world_monitor = world_monitor - adapter = InProcessViserAdapter(world_monitor=world_monitor, manipulation_module=module) - - planned = adapter.get_planned_path("arm") - assert planned is not None - assert planned[0] is not copied - assert planned[0].name is not copied.name - assert planned[0].position is not copied.position - - current = adapter.get_current_joint_state("arm") - assert current is not copied - assert current.name is not copied.name - - assert adapter.plan_to_pose("pose", "arm") == ("pose", "arm") - assert adapter.preview_path("arm") == "arm" - assert adapter.evaluate_joint_target(planned[0], "arm")["status"] == "FEASIBLE" + assert scene.animate_preview(preview, 1.0) is True + assert updates[-1] == ("id-arm", [1.0, 0.2]) -def test_adapter_evaluate_joint_target_uses_world_monitor_and_copies_input() -> None: - original = FakeJointState(["arm/j1", "j2"], position=[1.0, 2.0]) - seen = {} +def test_theme_and_reference_scene_contract() -> None: + server = Server() + assert apply_dimos_theme(server) is True + assert server.gui.theme_kwargs is not None + assert server.gui.theme_kwargs["brand_color"] == (0, 153, 255) + assert server.gui.theme_kwargs["dark_mode"] is True + assert server.gui.theme_kwargs["control_layout"] == "fixed" + assert ViserVisualizationConfig().panel_enabled is True - def is_state_valid(robot_id, joint_state) -> bool: - seen["robot_id"] = robot_id - seen["joint_state"] = joint_state - return True - world_monitor = SimpleNamespace( - get_current_joint_state=lambda robot_id: None, - is_state_stale=lambda robot_id, max_age=1.0: False, - is_state_valid=is_state_valid, - get_ee_pose=lambda robot_id, joint_state=None: (robot_id, joint_state), - ) - module = FakeManipulationModule( - _robots={"arm": ("robot-1", SimpleNamespace(), None)}, - _world_monitor=world_monitor, - ) - adapter = InProcessViserAdapter(world_monitor=world_monitor, manipulation_module=module) +def test_preview_selection_rejects_malformed_before_visibility() -> None: + selection = PlanningGroupSelection.from_groups((group("arm", "manipulator", ("j1",)),)) + assert selection.group_ids == ("arm/manipulator",) + # The scene transaction itself rejects missing tracks before revealing ghosts. + server = Server() + server.scene.add_grid = lambda *_args, **_kwargs: Handle() + scene = ViserManipulationScene(server, Urdf) + assert scene.animate_preview(GroupPreviewAnimation(()), 1.0) is False - result = adapter.evaluate_joint_target(original, "arm") - assert result["success"] is True - assert result["status"] == "FEASIBLE" - assert seen["robot_id"] == "robot-1" - assert seen["joint_state"] is not original - assert seen["joint_state"].name == ["arm/j1", "j2"] - assert seen["joint_state"].position == [1.0, 2.0] +def test_group_controls_use_source_labels_and_active_colors( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], +) -> None: + pose = group("arm", "manipulator", ("j1",), pose=True) + auxiliary = group("arm", "gripper", ("j2",)) + gui, _module, server = panel([auxiliary, pose], states("arm")) + + assert group_display_name(pose) == "arm" + assert group_display_name(auxiliary) == "arm gripper" + assert [button.label for button in server.gui.buttons[:2]] == ["arm", "arm gripper"] + assert [button.color for button in server.gui.buttons[:2]] == [ + ACTIVE_GROUP_COLOR, + INACTIVE_GROUP_COLOR, + ] + assert server.gui.buttons[1].callback is not None + server.gui.buttons[1].callback(SimpleNamespace()) + assert gui._handles[f"group:{auxiliary.id}"].color == ACTIVE_GROUP_COLOR -def test_obstacle_collision_marks_joint_target_infeasible() -> None: - obstacle = SimpleNamespace(name="blocking_box", blocked_joint_min=0.5) +def test_panel_preset_defaults_and_joint_slider_limits( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], +) -> None: + selected = group("arm", "manipulator", ("j1", "j2"), pose=True) + _gui, _module, server = panel([selected], states("arm")) + + assert server.gui.dropdowns[0].options == ["Select preset...", "Init", "Current", "Home"] + assert [ + (slider.label, slider.min, slider.max, slider.step, slider.value) + for slider in server.gui.sliders + ] == [ + ("arm/manipulator/j1", -1.0, 1.0, 0.001, 0.1), + ("arm/manipulator/j2", -2.0, 2.0, 0.001, 0.2), + ] - def is_state_valid(robot_id, joint_state) -> bool: - return bool(joint_state.position[0] < obstacle.blocked_joint_min) - world_monitor = SimpleNamespace( - get_current_joint_state=lambda robot_id: FakeJointState(["j1"], position=[0.0]), - is_state_stale=lambda robot_id, max_age=1.0: False, - is_state_valid=is_state_valid, - get_ee_pose=lambda robot_id, joint_state=None: SimpleNamespace( - position=SimpleNamespace(x=0.0, y=0.0, z=0.0) - ), - ) - module = FakeManipulationModule( - _robots={"arm": ("robot-1", SimpleNamespace(joint_names=["j1"]), None)}, - _world_monitor=world_monitor, - ) - adapter = InProcessViserAdapter(world_monitor=world_monitor, manipulation_module=module) +def test_init_and_home_presets_use_operator_init_and_config_home( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], +) -> None: + selected = group("arm", "manipulator", ("j1", "j2"), pose=True) + gui, module, _server = panel([selected], states("arm")) + module.configs["arm"].home_joints = [0.9, 0.8] - free = adapter.evaluate_joint_target(FakeJointState(["j1"], position=[0.25]), "arm") - colliding = adapter.evaluate_joint_target(FakeJointState(["j1"], position=[0.75]), "arm") + gui._apply_preset("Init") + assert gui.state.group_joint_targets[selected.id].position == [-0.5, -1.0] - assert free["success"] is True - assert free["status"] == "FEASIBLE" - assert free["collision_free"] is True - assert colliding["success"] is True - assert colliding["status"] == "COLLISION" - assert colliding["collision_free"] is False + gui._apply_preset("Home") + assert gui.state.group_joint_targets[selected.id].position == [0.9, 0.8] -def test_scene_registers_goal_robot_coloring_and_updates_visibility() -> None: - server = FakeServer() - scene = ViserManipulationScene( - server, - lambda *args, **kwargs: FakeViserUrdfWithMeshes(("joint1", "joint2")), - preview_fps=10.0, +def test_initial_pose_targets_are_group_id_keyed_for_same_robot_groups() -> None: + first = group("arm", "wrist", ("j1",), pose=True) + second = group("arm", "tool", ("j2",), pose=True) + module = Module([first, second], states("arm")) + monitor = Monitor(module) + monitor.poses[first.id] = Pose( + {"position": [1.0, 0.0, 0.0], "orientation": [0.0, 0.0, 0.0, 1.0]} ) - scene.prepared_urdf_path = lambda config: "dummy.urdf" - config = SimpleNamespace( - name="arm", - model_path="/tmp/arm.urdf", - package_paths={}, - xacro_args={}, - auto_convert_meshes=False, - joint_names=["joint1", "joint2"], - ) - - scene.register_robot("robot1", config) - target = scene._urdfs["robot1:target"] - preview = scene._urdfs["robot1:preview"] - - assert all(mesh.color == (255, 122, 0) for mesh in target._meshes) - assert all(mesh.opacity == 0.7 for mesh in target._meshes) - assert all(mesh.color == (80, 180, 255) for mesh in preview._meshes) - assert all(mesh.opacity == 0.55 for mesh in preview._meshes) - - scene.show_preview("robot1") - assert all(mesh.visible is True for mesh in preview._meshes) - scene.hide_preview("robot1") - assert all(mesh.visible is False for mesh in preview._meshes) - assert all(mesh.visible is True for mesh in target._meshes) - - -def test_scene_transform_controls_update_pose_callback_and_visual_state() -> None: - server = FakeTransformServer() - scene = ViserManipulationScene( - server, - lambda *args, **kwargs: FakeViserUrdfWithMeshes(("joint1", "joint2")), - preview_fps=10.0, + monitor.poses[second.id] = Pose( + {"position": [2.0, 0.0, 0.0], "orientation": [0.0, 0.0, 0.0, 1.0]} ) - scene.prepared_urdf_path = lambda config: "dummy.urdf" - config = SimpleNamespace( - name="arm", - model_path="/tmp/arm.urdf", - package_paths={}, - xacro_args={}, - auto_convert_meshes=False, - joint_names=["joint1", "joint2"], + server = Server() + scene_info = PlanningSceneInfo( + robots={"id-arm": module.configs["arm"]}, planning_groups=tuple(module.groups) ) - scene.register_robot("robot1", config) - updates = [] - - control = scene.ensure_target_controls("robot1", updates.append) - assert control is not None - assert server.transform_controls[0].path == "/targets/robot1/ee_control" - assert control.update_callback is not None - moved = SimpleNamespace(position=(1.0, 2.0, 3.0), wxyz=(1.0, 0.0, 0.0, 0.0)) - control.update_callback(SimpleNamespace(target=moved)) - assert updates == [moved] - - pose = Pose({"position": [0.1, 0.2, 0.3], "orientation": [0.0, 0.0, 0.0, 1.0]}) - scene.set_target_pose("robot1", pose) - assert control.position == (0.1, 0.2, 0.3) - assert control.wxyz == (1.0, 0.0, 0.0, 0.0) - - scene.set_target_visual_state("robot1", feasible=False) - target = scene._urdfs["robot1:target"] - preview = scene._urdfs["robot1:preview"] - assert control.color == (255, 40, 40) - assert all(mesh.color == (255, 30, 30) for mesh in target._meshes) - assert all(mesh.opacity == 0.75 for mesh in target._meshes) - assert all(mesh.color == (80, 180, 255) for mesh in preview._meshes) - - -def test_scene_target_controls_update_target_ghost_pose_and_feasibility() -> None: - server = FakeTransformServer() - scene = ViserManipulationScene( + current = {"id-arm": JointState(module.states["arm"])} + gui = ViserPanelGui( server, - lambda *args, **kwargs: FakeViserUrdfWithMeshes(("joint1", "joint2")), - preview_fps=10.0, - ) - scene.prepared_urdf_path = lambda config: "dummy.urdf" - config = SimpleNamespace( - name="arm", - model_path="/tmp/arm.urdf", - package_paths={}, - xacro_args={}, - auto_convert_meshes=False, - joint_names=["joint1", "joint2"], + scene_info, + Operator(module, monitor), + current, + ViserVisualizationConfig(panel_enabled=True), ) - scene.register_robot("robot1", config) - scene.ensure_target_controls("robot1", lambda target: None) - - pose = Pose({"position": [0.1, 0.2, 0.3], "orientation": [0.0, 0.0, 0.0, 1.0]}) - assert scene.set_target_joints("robot1", ["joint1", "joint2"], [0.7, 0.9]) is True - assert scene.set_target_pose("robot1", pose) is None - assert scene.set_target_visual_state("robot1", feasible=False) is None + try: + gui.start() + gui._toggle_group_selected(second.id) - target = scene._urdfs["robot1:target"] - handle = scene._handles["robot1:ee_control"] - assert target.cfg == [0.7, 0.9] - assert handle.position == (0.1, 0.2, 0.3) - assert handle.color == (255, 40, 40) + assert list(gui.state.pose_targets[first.id].position) == [1.0, 0.0, 0.0] + assert list(gui.state.pose_targets[second.id].position) == [2.0, 0.0, 0.0] + finally: + gui.close() -def test_gui_initializes_pose_selector_to_current_ee_pose( - make_panel: Callable[..., ViserPanelGui], +def test_panel_action_controls_are_present_in_source_order( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], ) -> None: - current = FakeJointState(["j1"], position=[0.25]) - current_pose = SimpleNamespace( - position=SimpleNamespace(x=0.1, y=0.2, z=0.3), - orientation=SimpleNamespace(w=0.9, x=0.1, y=0.2, z=0.3), - ) - config = make_robot_config(joint_names=["j1"], home_joints=[0.0]) - module = FakeManipulationModule( - _robots={"arm": ("robot-1", config, None)}, _planned_paths={}, _planned_trajectories={} - ) - world_monitor = SimpleNamespace( - get_current_joint_state=lambda robot_id: current, - is_state_stale=lambda robot_id, max_age=1.0: False, - get_ee_pose=lambda robot_id, joint_state=None: current_pose, - ) - adapter = InProcessViserAdapter(world_monitor=world_monitor, manipulation_module=module) - scene = ViserManipulationScene( - FakeTransformServer(), lambda *args, **kwargs: FakeViserUrdfWithMeshes(), preview_fps=10.0 - ) - gui = make_panel(FakeGuiServer(), adapter, ViserVisualizationConfig(panel_enabled=True), scene) - control = scene._handles["robot-1:ee_control"] - assert control.position == (0.1, 0.2, 0.3) - assert control.wxyz == (0.9, 0.1, 0.2, 0.3) - assert gui.state.cartesian_target is current_pose + selected = group("arm", "manipulator", ("j1",), pose=True) + _gui, _module, server = panel([selected], states("arm")) + + assert [button.label for button in server.gui.buttons[1:]] == [ + "Plan", + "Preview", + "Execute", + "Cancel", + "Clear plan", + ] + assert [folder.label for folder in server.gui.folders] == [ + "Manipulation Panel", + "Joint Control", + ] -def test_gui_preset_dropdown_and_controls_include_init_home_current_and_callbacks( - make_panel: Callable[..., ViserPanelGui], +def test_target_callbacks_require_current_target_identity( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], ) -> None: - current = FakeJointState(["arm/j1", "arm/j2"], position=[0.25, 0.5]) - config = make_robot_config(joint_names=["j1", "j2"], home_joints=[1.0, 2.0]) - module = FakeManipulationModule( - _robots={"arm": ("robot-1", config, None)}, - _init_joints={"arm": FakeJointState(["j1", "j2"], position=[-1.0, -2.0])}, - _planned_paths={}, - _planned_trajectories={}, + first, second = ( + group("arm", "manipulator", ("j1",), pose=True), + group("arm", "gripper", ("j2",)), ) - world_monitor = SimpleNamespace( - get_current_joint_state=lambda robot_id: current, - is_state_stale=lambda robot_id, max_age=1.0: False, - is_state_valid=lambda robot_id, joint_state: True, - get_ee_pose=lambda robot_id, joint_state=None: None, + gui, _module, _server = panel([first, second], states("arm")) + request = TargetEvaluationRequest( + gui.state.next_sequence_id(), + "joints", + selection_epoch=gui.state.selection_epoch, + group_ids=(second.id,), ) - adapter = InProcessViserAdapter(world_monitor=world_monitor, manipulation_module=module) - gui = make_panel(FakeGuiServer(), adapter) - assert gui._handles["preset"].options == ["Select preset...", "Init", "Current", "Home"] - assert list(gui._joint_sliders) == ["j1", "j2"] - gui._apply_preset("Home") - assert [gui._joint_sliders[name].value for name in ("j1", "j2")] == [1.0, 2.0] - gui._apply_preset("Current") - assert [gui._joint_sliders[name].value for name in ("j1", "j2")] == [0.25, 0.5] - gui._submit_execute() - assert "Cannot execute" in gui.state.error + gui._apply_target_evaluation_result(request, TargetEvaluationResult(True, "FEASIBLE", "", True)) -def test_gui_rebuilding_joint_sliders_removes_stale_viser_handles( - make_panel: Callable[..., ViserPanelGui], -) -> None: - current = FakeJointState(["j1", "j2"], position=[0.0, 0.0]) - config = make_robot_config(joint_names=["j1", "j2"], home_joints=[1.0, 2.0]) - module = FakeManipulationModule( - _robots={"arm": ("robot-1", config, None)}, _planned_paths={}, _planned_trajectories={} - ) - world_monitor = SimpleNamespace( - get_current_joint_state=lambda robot_id: current, - is_state_stale=lambda robot_id, max_age=1.0: False, - is_state_valid=lambda robot_id, joint_state: True, - get_ee_pose=lambda robot_id, joint_state=None: None, - ) - adapter = InProcessViserAdapter(world_monitor=world_monitor, manipulation_module=module) - server = FakeGuiServer() - gui = make_panel(server, adapter) - stale_sliders = list(server.sliders) - assert [slider.value for slider in stale_sliders] == [0.0, 0.0] - - current.position = [-0.738, -0.2826151825863572] - gui._build_joint_sliders() - - assert all(slider.removed is True for slider in stale_sliders) - assert [gui._joint_sliders[name].value for name in ("j1", "j2")] == [ - -0.738, - -0.2826151825863572, - ] - + assert gui.state.target_status == TargetStatus.CHECKING -def test_gui_parses_numpy_transform_control_arrays() -> None: - gui = ViserPanelGui(FakeGuiServer(), make_adapter_with_robot(), ViserVisualizationConfig()) - pose = gui._pose_from_transform_target( - SimpleNamespace( - position=np.array([1.0, 2.0, 3.0]), - wxyz=np.array([0.5, 0.1, 0.2, 0.3]), - ) +def test_scene_target_ghost_tracks_current_only_until_explicit_target() -> None: + server = Server() + server.scene.add_grid = lambda *_args, **_kwargs: Handle() + scene = ViserManipulationScene(server, Urdf) + scene.prepared_urdf_path = lambda _config: "robot.urdf" # type: ignore[method-assign] + config = Config("arm", ["j1", "j2"], [-1.0, -2.0], [1.0, 2.0], [0.0, 0.0]) + scene.register_robot("id-arm", config) + + scene.update_current_robot("id-arm", JointState({"name": ["j1", "j2"], "position": [0.1, 0.2]})) + assert scene._urdfs["id-arm:target"].cfg == [0.1, 0.2] + scene.set_target_joints("id-arm", ["j1", "j2"], [0.8, 0.9]) + scene.update_current_robot("id-arm", JointState({"name": ["j1", "j2"], "position": [0.2, 0.3]})) + + assert scene._urdfs["id-arm:target"].cfg == [0.8, 0.9] + + +def test_scene_target_feasibility_colors_ghost_and_gizmo() -> None: + server = Server() + server.scene.add_grid = lambda *_args, **_kwargs: Handle() + server.scene.add_transform_controls = lambda *_args, **_kwargs: Handle() + scene = ViserManipulationScene(server, Urdf) + scene.prepared_urdf_path = lambda _config: "robot.urdf" # type: ignore[method-assign] + config = Config("arm", ["j1", "j2"], [-1.0, -2.0], [1.0, 2.0], [0.0, 0.0]) + scene.register_robot("id-arm", config) + scene.ensure_target_controls("id-arm", lambda _target: None) + + scene.set_target_visual_state("id-arm", False) + + assert scene._urdfs["id-arm:target"]._meshes[0].color == (255, 30, 30) + assert scene._handles["id-arm:ee_control"].color == (255, 40, 40) + + +def test_panel_feasibility_colors_group_controls_and_deduplicated_robot_ghosts() -> None: + arm_primary, arm_secondary, other = ( + group("arm", "primary", ("j1",), pose=True), + group("arm", "secondary", ("j2",), pose=True), + group("other", "manipulator", ("j1",), pose=True), + ) + module = Module([arm_primary, arm_secondary, other], states("arm", "other")) + server = Server() + server.scene.add_grid = lambda *_args, **_kwargs: Handle() + server.scene.add_transform_controls = lambda *_args, **_kwargs: Handle() + scene = ViserManipulationScene(server, Urdf) + scene.prepared_urdf_path = lambda _config: "robot.urdf" # type: ignore[method-assign] + scene.register_robot("id-arm", module.configs["arm"]) + scene.register_robot("id-other", module.configs["other"]) + gui = scene_gui(module, server, scene) + gui.start() + gui._worker.stop() + gui._operation_worker.stop() + gui._toggle_group_selected(arm_secondary.id) + gui._toggle_group_selected(other.id) + + control_calls: list[str] = [] + robot_calls: list[str] = [] + original_control = scene.set_target_control_visual_state + original_robot = scene.set_target_robot_visual_state + scene.set_target_control_visual_state = lambda group_id, feasible: ( + control_calls.append(group_id), + original_control(group_id, feasible), + ) # type: ignore[method-assign] + scene.set_target_robot_visual_state = lambda robot_id, feasible: ( + robot_calls.append(robot_id), + original_robot(robot_id, feasible), + ) # type: ignore[method-assign] + request = TargetEvaluationRequest( + gui.state.next_sequence_id(), + "joints", + selection_epoch=gui.state.selection_epoch, + group_ids=gui.state.selected_group_ids, ) - assert pose is not None - assert list(pose.position) == [1.0, 2.0, 3.0] - assert list(pose.orientation) == [0.1, 0.2, 0.3, 0.5] - + gui._apply_target_evaluation_result(request, TargetEvaluationResult(True, "FEASIBLE", "", True)) -def test_panel_execution_requires_fresh_plan_and_refresh_updates_robot_controls( - make_panel: Callable[..., ViserPanelGui], -) -> None: - current = FakeJointState(["j1"], position=[1.2]) - config = make_robot_config(joint_names=["j1"], home_joints=[0.5]) - module = FakeManipulationModule( - _robots={"arm": ("robot-1", config, None)}, - _planned_paths={}, - _planned_trajectories={}, - execute=lambda robot_name=None: False, + assert control_calls == [arm_primary.id, arm_secondary.id, other.id] * 2 + assert robot_calls == ["id-arm", "id-other"] * 2 + assert all( + scene._handles[f"{item.id}:ee_control"].color == TARGET_CONTROL_FEASIBLE_COLOR + for item in (arm_primary, arm_secondary, other) ) - world_monitor = SimpleNamespace( - get_current_joint_state=lambda robot_id: current, - is_state_stale=lambda robot_id, max_age=1.0: False, - is_state_valid=lambda robot_id, joint_state: True, - get_ee_pose=lambda robot_id, joint_state=None: None, - ) - adapter = InProcessViserAdapter( - world_monitor=world_monitor, - manipulation_module=module, - ) - gui = make_panel(FakeGuiServer(), adapter) - gui.refresh() - assert gui.state.selected_robot == "arm" - assert list(gui._joint_sliders) == ["j1"] - gui._apply_preset("Home") - assert gui._joint_sliders["j1"].value == 0.5 - - gui._submit_execute() - assert "Cannot execute" in gui.state.error + assert scene._urdfs["id-arm:target"]._meshes[0].color == GOAL_ROBOT_FEASIBLE_COLOR + assert scene._urdfs["id-other:target"]._meshes[0].color == GOAL_ROBOT_FEASIBLE_COLOR - -def test_gui_moves_joint_target_immediately_and_stores_evaluated_joint_solution( - make_panel: Callable[..., ViserPanelGui], -) -> None: - current = FakeJointState(["j1", "j2"], position=[0.0, 0.0]) - target_pose = SimpleNamespace(position=SimpleNamespace(x=0.2, y=0.3, z=0.4)) - config = make_robot_config(joint_names=["j1", "j2"], home_joints=[0.5, 0.6]) - module = FakeManipulationModule( - _robots={"arm": ("robot-1", config, None)}, _planned_paths={}, _planned_trajectories={} - ) - world_monitor = SimpleNamespace( - get_current_joint_state=lambda robot_id: current, - is_state_stale=lambda robot_id, max_age=1.0: False, - is_state_valid=lambda robot_id, joint_state: True, - get_ee_pose=lambda robot_id, joint_state=None: target_pose, - ) - adapter = InProcessViserAdapter(world_monitor=world_monitor, manipulation_module=module) - target_updates = [] - target_pose_updates = [] - scene = SimpleNamespace( - has_reference_grid=lambda: False, - ensure_target_controls=lambda *args: None, - set_target_joints=lambda *args: target_updates.append(args) or True, - set_target_pose=lambda *args: target_pose_updates.append(args), - set_target_visual_state=lambda *args: None, - ) - gui = make_panel(FakeGuiServer(), adapter, ViserVisualizationConfig(panel_enabled=True), scene) - requests = [] - gui._worker.stop() - gui._worker = SimpleNamespace( - submit=lambda request: requests.append(request), stop=lambda: None + control_calls.clear() + robot_calls.clear() + request = TargetEvaluationRequest( + gui.state.next_sequence_id(), + "joints", + selection_epoch=gui.state.selection_epoch, + group_ids=gui.state.selected_group_ids, ) - gui._joint_sliders["j1"].value = 0.25 - gui._joint_sliders["j2"].value = 0.75 - gui._submit_joint_target_evaluation() - assert target_updates[-1] == ("robot-1", ["j1", "j2"], [0.25, 0.75]) - assert target_pose_updates[-1] == ("robot-1", target_pose) - assert requests[-1].source == "joints" - - stale_request = TargetEvaluationRequest(sequence_id=1, source="joints", robot_name="arm") - fresh_request = TargetEvaluationRequest(sequence_id=2, source="joints", robot_name="arm") - gui.state.latest_sequence_id = 2 gui._apply_target_evaluation_result( - stale_request, - { - "success": True, - "collision_free": True, - "joint_state": adapter.joints_from_values(["j1", "j2"], [9.0, 9.0]), - }, + request, TargetEvaluationResult(True, "COLLISION", "", False) ) - assert gui.state.joint_target == [0.25, 0.75] - gui._apply_target_evaluation_result( - fresh_request, - { - "success": True, - "collision_free": True, - "joint_state": adapter.joints_from_values(["j1", "j2"], [1.0, 2.0]), - }, + assert control_calls == [arm_primary.id, arm_secondary.id, other.id] * 2 + assert robot_calls == ["id-arm", "id-other"] * 2 + assert all( + scene._handles[f"{item.id}:ee_control"].color == TARGET_CONTROL_INFEASIBLE_COLOR + for item in (arm_primary, arm_secondary, other) ) - assert gui.state.target_status == TargetStatus.FEASIBLE - assert gui.state.feasibility.status == FeasibilityStatus.FEASIBLE - assert gui.state.joint_target == [1.0, 2.0] - assert [gui._joint_sliders[name].value for name in ("j1", "j2")] == [0.25, 0.75] - assert target_updates[-1] == ("robot-1", ["j1", "j2"], [0.25, 0.75]) + assert scene._urdfs["id-arm:target"]._meshes[0].color == GOAL_ROBOT_INFEASIBLE_COLOR + assert scene._urdfs["id-other:target"]._meshes[0].color == GOAL_ROBOT_INFEASIBLE_COLOR + gui.close() -def test_gui_cartesian_ik_result_does_not_rewrite_active_gizmo( - make_panel: Callable[..., ViserPanelGui], +def test_scene_shared_clock_uses_stored_unequal_robot_frames( + monkeypatch: pytest.MonkeyPatch, ) -> None: - current = FakeJointState(["j1", "j2"], position=[0.0, 0.0]) - config = make_robot_config(joint_names=["j1", "j2"], home_joints=[0.5, 0.6]) - module = FakeManipulationModule( - _robots={"arm": ("robot-1", config, None)}, _planned_paths={}, _planned_trajectories={} - ) - world_monitor = SimpleNamespace( - get_current_joint_state=lambda robot_id: current, - is_state_stale=lambda robot_id, max_age=1.0: False, - is_state_valid=lambda robot_id, joint_state: True, - get_ee_pose=lambda robot_id, joint_state=None: None, - ) - adapter = InProcessViserAdapter(world_monitor=world_monitor, manipulation_module=module) - target_joint_updates = [] - target_pose_updates = [] - scene = SimpleNamespace( - has_reference_grid=lambda: False, - ensure_target_controls=lambda *args: None, - set_target_joints=lambda *args: target_joint_updates.append(args) or True, - set_target_pose=lambda *args: target_pose_updates.append(args), - set_target_visual_state=lambda *args: None, + server = Server() + server.scene.add_grid = lambda *_args, **_kwargs: Handle() + scene = ViserManipulationScene(server, Urdf) + scene.prepared_urdf_path = lambda _config: "robot.urdf" # type: ignore[method-assign] + config = Config("arm", ["j1", "j2"], [-1.0, -2.0], [1.0, 2.0], [0.0, 0.0]) + scene.register_robot("left", config) + scene.register_robot("right", config) + updates: list[tuple[str, list[float]]] = [] + monkeypatch.setattr( + scene, + "_set_preview_ghost_joints", + lambda robot, _names, values: updates.append((robot, list(values))), ) - gui = make_panel(FakeGuiServer(), adapter, ViserVisualizationConfig(panel_enabled=True), scene) - gui.state.cartesian_target = Pose( - {"position": [0.1, 0.2, 0.3], "orientation": [0.0, 0.0, 0.0, 1.0]} + monkeypatch.setattr( + "dimos.manipulation.visualization.viser.scene.time.sleep", lambda _delay: None + ) + + assert scene.animate_preview( + GroupPreviewAnimation( + ( + PreviewTrack("left", ("j1",), (PreviewFrame(0.0, (0.0,)),)), + PreviewTrack( + "right", + ("j1",), + ( + PreviewFrame(0.0, (10.0,)), + PreviewFrame(1.0, (11.0,)), + ), + ), + ) + ), + 1.0, ) - request = TargetEvaluationRequest(sequence_id=1, source="cartesian", robot_name="arm") - gui.state.latest_sequence_id = 1 + assert updates == [ + ("left", [0.0]), + ("right", [10.0]), + ("left", [0.0]), + ("right", [11.0]), + ] - gui._apply_target_evaluation_result( - request, - { - "success": True, - "collision_free": True, - "joint_state": adapter.joints_from_values(["j1", "j2"], [1.0, 2.0]), - }, + +def test_animation_frame_helpers_scale_stored_timestamps() -> None: + frames = ( + PreviewFrame(0.0, (0.0,)), + PreviewFrame(0.25, (1.0,)), + PreviewFrame(1.0, (2.0,)), ) - assert gui.state.target_status == TargetStatus.FEASIBLE - assert [gui._joint_sliders[name].value for name in ("j1", "j2")] == [1.0, 2.0] - assert target_joint_updates[-1] == ("robot-1", ["j1", "j2"], [1.0, 2.0]) - assert target_pose_updates == [] + assert scaled_frame_delays(frames, 2.0) == (0.5, 1.5) -def test_gui_collision_evaluation_marks_target_infeasible_and_colors_scene( - make_panel: Callable[..., ViserPanelGui], +def test_panel_disables_plan_preview_and_execute_until_a_feasible_target( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], ) -> None: - current = FakeJointState(["j1"], position=[0.0]) - config = make_robot_config(joint_names=["j1"], home_joints=[0.0]) - module = FakeManipulationModule( - _robots={"arm": ("robot-1", config, None)}, _planned_paths={}, _planned_trajectories={} - ) - world_monitor = SimpleNamespace( - get_current_joint_state=lambda robot_id: current, - is_state_stale=lambda robot_id, max_age=1.0: False, - is_state_valid=lambda robot_id, joint_state: False, - get_ee_pose=lambda robot_id, joint_state=None: SimpleNamespace( - position=SimpleNamespace(x=0.0, y=0.0, z=0.0) - ), - ) - module._world_monitor = world_monitor - adapter = InProcessViserAdapter(world_monitor=world_monitor, manipulation_module=module) - visual_states = [] - scene = SimpleNamespace( - has_reference_grid=lambda: False, - ensure_target_controls=lambda *args: None, - set_target_joints=lambda *args: True, - set_target_pose=lambda *args: None, - set_target_visual_state=lambda *args: visual_states.append(args), - ) - gui = make_panel(FakeGuiServer(), adapter, ViserVisualizationConfig(panel_enabled=True), scene) - request = TargetEvaluationRequest(sequence_id=1, source="joints", robot_name="arm") - gui.state.latest_sequence_id = 1 - result = adapter.evaluate_joint_target(FakeJointState(["j1"], position=[1.0]), "arm") + selected = group("arm", "manipulator", ("j1",), pose=True) + _gui, _module, server = panel([selected], states("arm")) - gui._apply_target_evaluation_result(request, result) + assert [button.disabled for button in server.gui.buttons[1:4]] == [True, True, True] - assert result["status"] == "COLLISION" - assert gui.state.target_status == TargetStatus.INFEASIBLE - assert gui.state.feasibility.status == FeasibilityStatus.COLLISION - assert gui.state.error == "Target is in collision" - assert visual_states[-1] == ("robot-1", False) - -def test_gui_safe_execute_requires_fresh_matching_plan_and_clear_resets_path( - make_panel: Callable[..., ViserPanelGui], monkeypatch: pytest.MonkeyPatch +def test_panel_status_reports_target_and_plan_defaults( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], ) -> None: - current = FakeJointState(["j1"], position=[1.0]) - planned = [FakeJointState(["j1"], position=[1.0]), FakeJointState(["j1"], position=[2.0])] - executed = [] - cleared = [] - module = FakeManipulationModule( - _robots={ - "arm": ("robot-1", make_robot_config(joint_names=["j1"], home_joints=[1.0]), None) - }, - _planned_paths={"arm": planned}, - _planned_trajectories={}, - _state=NamedState(name="IDLE"), - execute=lambda robot_name=None: executed.append(robot_name) or True, - clear_planned_path=lambda: cleared.append(True) or True, - ) - world_monitor = SimpleNamespace( - get_current_joint_state=lambda robot_id: current, - is_state_stale=lambda robot_id, max_age=1.0: False, - is_state_valid=lambda robot_id, joint_state: True, - get_ee_pose=lambda robot_id, joint_state=None: SimpleNamespace( - position=SimpleNamespace(x=0.0, y=0.0, z=0.0) - ), - ) - adapter = InProcessViserAdapter(world_monitor=world_monitor, manipulation_module=module) - gui = make_panel( - FakeGuiServer(), - adapter, - ViserVisualizationConfig(panel_enabled=True, current_match_tolerance=0.05), - ) - gui._operation_worker.stop() - monkeypatch.setattr( - gui, - "_operation_worker", - SimpleNamespace( - submit=lambda operation, **_kwargs: operation(), stop=lambda timeout=2.0: None - ), - ) - gui.state.target_status = TargetStatus.FEASIBLE - gui.state.plan_state = PanelPlanState( - status=PlanStatus.FRESH, - robot="arm", - start_joints_snapshot=[1.2], - planned_path=planned, - ) - gui._submit_execute() - assert executed == [] - assert "Cannot execute" in gui.state.error + selected = group("arm", "manipulator", ("j1",), pose=True) + gui, _module, server = panel([selected], states("arm")) - gui.state.action_status = ActionStatus.IDLE - gui.state.error = "" - gui.state.plan_state.start_joints_snapshot = [1.0] - gui._submit_execute() - assert executed == ["arm"] + assert server.gui.markdown[0].value == "### Status\n**State:** Ready" + assert gui.state.error == "" - gui._submit_clear() - assert cleared == [True] - assert gui.state.plan_state.status == PlanStatus.NONE +def test_scene_reference_grid_has_expected_defaults_and_toggle() -> None: + server = Server() + grids: list[Handle] = [] + server.scene.add_grid = lambda *_args, **_kwargs: grids.append(Handle()) or grids[-1] + scene = ViserManipulationScene(server, Urdf) -def test_gui_plan_target_failure_recovers_action_state( - make_panel: Callable[..., ViserPanelGui], - monkeypatch: pytest.MonkeyPatch, -) -> None: - adapter = make_adapter_with_robot() - gui = make_panel(FakeGuiServer(), adapter) - gui._operation_worker.stop() - monkeypatch.setattr( - gui, - "_operation_worker", - SimpleNamespace( - submit=lambda operation, **_kwargs: operation(), stop=lambda timeout=2.0: None - ), - ) - gui.state.selected_robot = "missing" - gui.state.target_status = TargetStatus.FEASIBLE - gui.state.manipulation_state = "IDLE" - - gui._submit_plan() + assert scene.has_reference_grid() is True + scene.set_reference_grid_visible(False) + assert grids[0].visible is False + scene.set_reference_grid_visible(True) + assert grids[0].visible is True - assert gui.state.action_status == ActionStatus.IDLE - assert gui.state.plan_state.status == PlanStatus.FAILED - assert gui.state.error == "No robot config" - assert gui.state.last_result == "plan_to_joints=False" +def test_scene_returns_false_for_missing_robot_target_updates() -> None: + server = Server() + server.scene.add_grid = lambda *_args, **_kwargs: Handle() + scene = ViserManipulationScene(server, Urdf) -def test_gui_resets_fault_before_replanning( - make_panel: Callable[..., ViserPanelGui], - monkeypatch: pytest.MonkeyPatch, -) -> None: - calls = [] - adapter = make_adapter_with_robot() - gui = make_panel(FakeGuiServer(), adapter) - gui._operation_worker.stop() - monkeypatch.setattr( - gui, - "_operation_worker", - SimpleNamespace( - submit=lambda operation, **_kwargs: operation(), stop=lambda timeout=2.0: None - ), + assert scene.set_target_joints("missing", ["j1"], [0.1]) is False + assert ( + scene.animate_preview( + GroupPreviewAnimation( + (PreviewTrack("missing", ("j1",), (PreviewFrame(0.0, (0.0,)),)),) + ), + duration=0.0, + ) + is False ) - def reset() -> bool: - calls.append("reset") - return True - - def plan_to_joints(_joints: JointState, _robot_name: str | None = None) -> bool: - calls.append("plan") - return True - monkeypatch.setattr(adapter, "reset", reset) - monkeypatch.setattr(adapter, "plan_to_joints", plan_to_joints) - gui.state.target_status = TargetStatus.FEASIBLE - gui.state.manipulation_state = "FAULT" +def test_scene_cancel_generation_hides_preview_and_rejects_old_animation() -> None: + server = Server() + server.scene.add_grid = lambda *_args, **_kwargs: Handle() + scene = ViserManipulationScene(server, Urdf) + scene.prepared_urdf_path = lambda _config: "robot.urdf" # type: ignore[method-assign] + config = Config("arm", ["j1", "j2"], [-1.0, -2.0], [1.0, 2.0], [0.0, 0.0]) + scene.register_robot("id-arm", config) + scene._preview_visible["id-arm"] = True + scene._set_preview_visibility("id-arm", True) - gui._submit_plan() + scene.cancel_preview_animation() - assert calls == ["reset", "plan"] - assert gui.state.plan_state.status == PlanStatus.FRESH - assert gui.state.last_result == "plan_to_joints=True" + assert scene._preview_visible == {"id-arm": False} + assert scene._animation_generation == 1 -def test_operation_worker_coalesces_pending_requests() -> None: - errors = [] - calls = [] - worker = OperationWorker(errors.append) - worker.submit(lambda: calls.append("old")) - worker.submit(lambda: calls.append("new")) +def test_scene_base_pose_requires_urdf_root_to_match(monkeypatch: pytest.MonkeyPatch) -> None: + server = Server() + server.scene.add_grid = lambda *_args, **_kwargs: Handle() + scene = ViserManipulationScene(server, Urdf) + monkeypatch.setattr( + "dimos.manipulation.visualization.viser.scene.parse_model", + lambda _path: SimpleNamespace(root_link="world"), + ) - operation = worker._requests.get_nowait() - operation.operation() + with pytest.raises(ValueError, match="base_link 'base'.*URDF root 'world'"): + scene._assert_base_link_is_urdf_root(SimpleNamespace(base_link="base"), "robot.urdf") - assert calls == ["new"] - assert errors == [] +def test_scene_detects_non_identity_base_pose() -> None: + identity = SimpleNamespace( + base_pose=Pose({"position": [0.0, 0.0, 0.0], "orientation": [0.0, 0.0, 0.0, 1.0]}) + ) + translated = SimpleNamespace( + base_pose=Pose({"position": [1.0, 0.0, 0.0], "orientation": [0.0, 0.0, 0.0, 1.0]}) + ) -def test_operation_worker_stop_can_wait_for_in_flight_operation() -> None: - errors = [] - worker = OperationWorker(errors.append) - started = threading.Event() - release = threading.Event() - finished = threading.Event() - stopped = threading.Event() + assert ViserManipulationScene._has_non_identity_base_pose(identity) is False + assert ViserManipulationScene._has_non_identity_base_pose(translated) is True - def operation() -> None: - started.set() - release.wait(timeout=1.0) - finished.set() - worker.start() - worker.submit(operation) - assert started.wait(timeout=1.0) +@pytest.mark.parametrize("interruption", ["cancel", "replacement", "close"]) +def test_scene_inflight_preview_never_updates_after_generation_replacement( + monkeypatch: pytest.MonkeyPatch, interruption: str +) -> None: + server = Server() + server.scene.add_grid = lambda *_args, **_kwargs: Handle() + scene = ViserManipulationScene(server, Urdf) + scene.prepared_urdf_path = lambda _config: "robot.urdf" # type: ignore[method-assign] + config = Config("arm", ["j1", "j2"], [-1.0, -2.0], [1.0, 2.0], [0.0, 0.0]) + scene.register_robot("id-arm", config) + first_tick, release = threading.Event(), threading.Event() + updates: list[float] = [] + original = scene._set_preview_ghost_joints + + def record(robot_id: str, names: Sequence[str], values: Sequence[float]) -> None: + updates.append(float(values[0])) + original(robot_id, names, values) + + scene._set_preview_ghost_joints = record # type: ignore[method-assign] + sleep_calls = 0 + + def block_after_first_tick(_delay: float) -> None: + nonlocal sleep_calls + sleep_calls += 1 + if sleep_calls == 1: + first_tick.set() + assert release.wait(timeout=2.0) - stopper = threading.Thread( - target=lambda: (worker.stop(timeout=None), stopped.set()), - name="StopViserOperationTest", + monkeypatch.setattr( + "dimos.manipulation.visualization.viser.scene.time.sleep", block_after_first_tick + ) + old = GroupPreviewAnimation( + ( + PreviewTrack( + "id-arm", + ("j1",), + ( + PreviewFrame(0.0, (0.0,)), + PreviewFrame(1.0, (1.0,)), + ), + ), + ) ) - stopper.start() - assert not stopped.wait(timeout=0.05) + worker = threading.Thread(target=lambda: scene.animate_preview(old, 1.0)) + worker.start() + assert first_tick.wait(timeout=2.0) + if interruption == "cancel": + visualizer = ViserManipulationVisualizer( + config=ViserVisualizationConfig(panel_enabled=False), + ) + visualizer._scene = scene + visualizer.cancel_preview_animation() + stable_updates = list(updates) + elif interruption == "replacement": + updates.clear() + assert scene.animate_preview( + GroupPreviewAnimation( + (PreviewTrack("id-arm", ("j1",), (PreviewFrame(0.0, (10.0,)),)),) + ), + 0.0, + ) + stable_updates = list(updates) + else: + scene.close() + stable_updates = list(updates) release.set() - assert stopped.wait(timeout=1.0) - stopper.join(timeout=1.0) - - assert finished.is_set() - assert worker._thread is None - assert errors == [] - + worker.join(timeout=2.0) -def test_target_evaluation_worker_coalesces_pending_requests() -> None: - worker = TargetEvaluationWorker(lambda request: {}, lambda request, result: None) - old_request = TargetEvaluationRequest(sequence_id=1, source="joints", robot_name="arm") - new_request = TargetEvaluationRequest(sequence_id=2, source="joints", robot_name="arm") + assert not worker.is_alive() + assert updates == stable_updates - worker.submit(old_request) - worker.submit(new_request) - - assert worker._requests.get_nowait() is new_request +def test_transform_control_callback_preserves_pose_through_gui_and_backend( + panel: Callable[..., tuple[ViserPanelGui, Module, Server]], +) -> None: + selected = group("arm", "manipulator", ("j1",), pose=True) + module = Module([selected], states("arm")) + server = Server() + server.scene.add_grid = lambda *_args, **_kwargs: Handle() + server.scene.add_transform_controls = lambda *_args, **_kwargs: Handle() + scene = ViserManipulationScene(server, Urdf) + gui = scene_gui(module, server, scene) + submitted: list[TargetEvaluationRequest] = [] + gui._worker.submit = submitted.append # type: ignore[method-assign] + gui.start() + control = scene._handles[f"{selected.id}:ee_control"] + control.position = (1.0, 2.0, 3.0) + control.wxyz = (0.4, 0.1, 0.2, 0.3) + assert control.callback is not None + control.callback(SimpleNamespace(target=control)) + + request = submitted[-1] + assert list(gui.state.pose_targets[selected.id].position) == [1.0, 2.0, 3.0] + assert list(gui.state.pose_targets[selected.id].orientation) == [0.1, 0.2, 0.3, 0.4] + assert control.position == (1.0, 2.0, 3.0) + assert control.wxyz == (0.4, 0.1, 0.2, 0.3) + assert request.pose_targets[selected.id] == gui.state.pose_targets[selected.id] + gui.close() -def test_operation_worker_reports_timeout() -> None: - errors = [] - release = threading.Event() - finished = threading.Event() - worker = OperationWorker(errors.append, timeout_seconds=0.01) - def operation() -> None: - release.wait(timeout=1.0) - finished.set() +def test_joint_evaluation_updates_active_gizmo_from_computed_group_pose() -> None: + selected = group("arm", "manipulator", ("j1",), pose=True) + module = Module([selected], states("arm")) + server = Server() + server.scene.add_grid = lambda *_args, **_kwargs: Handle() + server.scene.add_transform_controls = lambda *_args, **_kwargs: Handle() + scene = ViserManipulationScene(server, Urdf) + gui = scene_gui(module, server, scene) + gui.start() + control = scene._handles[f"{selected.id}:ee_control"] + request = TargetEvaluationRequest( + gui.state.next_sequence_id(), + "joints", + selection_epoch=gui.state.selection_epoch, + group_ids=gui.state.selected_group_ids, + ) + computed_pose = Pose({"position": [0.7, 0.8, 0.9], "orientation": [0.1, 0.2, 0.3, 0.4]}) - worker.submit(operation, timeout_seconds=0.01) - worker._run_operation(worker._requests.get_nowait()) - release.set() + gui._apply_target_evaluation_result( + request, + TargetEvaluationResult( + True, "FEASIBLE", "", True, group_poses={selected.id: computed_pose} + ), + ) - assert errors == ["Operation timed out after 0.0s"] - assert finished.wait(timeout=1.0) + assert control.position == (0.7, 0.8, 0.9) + assert control.wxyz == (0.4, 0.1, 0.2, 0.3) + gui.close() diff --git a/dimos/manipulation/visualization/viser/test_visualizer_lifecycle.py b/dimos/manipulation/visualization/viser/test_visualizer_lifecycle.py index 19f1f2da60..58fc0d9da1 100644 --- a/dimos/manipulation/visualization/viser/test_visualizer_lifecycle.py +++ b/dimos/manipulation/visualization/viser/test_visualizer_lifecycle.py @@ -21,18 +21,24 @@ pytest.importorskip("viser", reason="Viser optional dependency is not installed") +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig -from dimos.manipulation.planning.spec.models import PlanningSceneInfo +from dimos.manipulation.planning.spec.models import ( + PlanningSceneInfo, + VisualizationSession, + VisualizationStateFrame, +) from dimos.manipulation.visualization.viser import ( runtime as runtime_module, visualizer as visualizer_module, ) -from dimos.manipulation.visualization.viser.adapter import InProcessViserAdapter from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig from dimos.manipulation.visualization.viser.runtime import ViserRuntime +from dimos.manipulation.visualization.viser.scene import ViserManipulationScene from dimos.manipulation.visualization.viser.visualizer import ViserManipulationVisualizer from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory class FakeDependency: @@ -63,7 +69,11 @@ def fake_robot_config(name: str) -> RobotModelConfig: model_path=Path(f"{name}.urdf"), base_pose=PoseStamped(), joint_names=[], - end_effector_link="ee_link", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", joint_names=(), base_link="base_link", tip_link="ee_link" + ) + ], ) @@ -74,8 +84,6 @@ def fail_runtime(_config: ViserVisualizationConfig) -> FakeServer: monkeypatch.setattr(visualizer_module, "ViserRuntime", fail_runtime) visualizer = ViserManipulationVisualizer( - world_monitor=FakeDependency(), - manipulation_module=FakeDependency(), config=ViserVisualizationConfig(panel_enabled=False), ) @@ -106,8 +114,6 @@ def __init__( self, server: FakeServer, viser_urdf: type[FakeViserUrdf], - *, - preview_fps: float, ) -> None: calls.append(("create", "scene")) @@ -121,7 +127,9 @@ class FakeGui: def __init__( self, server: FakeServer, - adapter: InProcessViserAdapter, + scene_info: PlanningSceneInfo, + operator: object, + current_states: dict[str, JointState], config: ViserVisualizationConfig, scene: FakeScene, ) -> None: @@ -141,8 +149,6 @@ def close(self) -> None: monkeypatch.setattr(visualizer_module, "ViserManipulationScene", FakeScene) monkeypatch.setattr(visualizer_module, "ViserPanelGui", FakeGui) visualizer = ViserManipulationVisualizer( - world_monitor=FakeDependency(), - manipulation_module=FakeDependency(), config=ViserVisualizationConfig(panel_enabled=True), ) scene = PlanningSceneInfo( @@ -152,7 +158,7 @@ def close(self) -> None: } ) - visualizer.initialize_scene(scene) + visualizer.initialize(VisualizationSession(scene, operator=FakeDependency())) assert calls == [ ("start", "runtime"), @@ -187,8 +193,6 @@ def __init__( self, server: FakeServer, viser_urdf: type[FakeViserUrdf], - *, - preview_fps: float, ) -> None: pass @@ -199,7 +203,9 @@ class FakeGui: def __init__( self, server: FakeServer, - adapter: InProcessViserAdapter, + scene_info: PlanningSceneInfo, + operator: object, + current_states: dict[str, JointState], config: ViserVisualizationConfig, scene: FakeScene, ) -> None: @@ -216,13 +222,13 @@ def close(self) -> None: monkeypatch.setattr(visualizer_module, "ViserManipulationScene", FakeScene) monkeypatch.setattr(visualizer_module, "ViserPanelGui", FakeGui) visualizer = ViserManipulationVisualizer( - world_monitor=FakeDependency(), - manipulation_module=FakeDependency(), config=ViserVisualizationConfig(panel_enabled=True), ) with pytest.raises(RuntimeError, match="gui failed"): - visualizer.initialize_scene(PlanningSceneInfo(robots={})) + visualizer.initialize( + VisualizationSession(PlanningSceneInfo(robots={}), operator=FakeDependency()) + ) assert closed == ["gui", "scene", "runtime"] assert visualizer.get_visualization_url() is None @@ -250,8 +256,6 @@ def __init__( self, server: FakeServer, viser_urdf: type[FakeViserUrdf], - *, - preview_fps: float, ) -> None: raise RuntimeError("scene failed") @@ -259,13 +263,13 @@ def __init__( monkeypatch.setattr(visualizer_module, "ViserUrdf", FakeViserUrdf) monkeypatch.setattr(visualizer_module, "ViserManipulationScene", FailingScene) visualizer = ViserManipulationVisualizer( - world_monitor=FakeDependency(), - manipulation_module=FakeDependency(), config=ViserVisualizationConfig(panel_enabled=False), ) with pytest.raises(RuntimeError, match="scene failed"): - visualizer.initialize_scene(PlanningSceneInfo(robots={})) + visualizer.initialize( + VisualizationSession(PlanningSceneInfo(robots={}), operator=FakeDependency()) + ) assert closed == ["runtime"] assert visualizer.get_visualization_url() is None @@ -293,8 +297,6 @@ def __init__( self, server: FakeServer, viser_urdf: type[FakeViserUrdf], - *, - preview_fps: float, ) -> None: pass @@ -305,7 +307,9 @@ class FailingGui: def __init__( self, server: FakeServer, - adapter: InProcessViserAdapter, + scene_info: PlanningSceneInfo, + operator: object, + current_states: dict[str, JointState], config: ViserVisualizationConfig, scene: FakeScene, ) -> None: @@ -326,11 +330,11 @@ def close(self) -> None: monkeypatch.setattr(visualizer_module, "ViserManipulationScene", FakeScene) monkeypatch.setattr(visualizer_module, "ViserPanelGui", FailingGui) visualizer = ViserManipulationVisualizer( - world_monitor=FakeDependency(), - manipulation_module=FakeDependency(), config=ViserVisualizationConfig(panel_enabled=True), ) - visualizer.initialize_scene(PlanningSceneInfo(robots={})) + visualizer.initialize( + VisualizationSession(PlanningSceneInfo(robots={}), operator=FakeDependency()) + ) with pytest.raises(RuntimeError, match="gui close failed"): visualizer.close() @@ -392,8 +396,6 @@ def __init__( self, server: FakeServer, viser_urdf: type[FakeViserUrdf], - *, - preview_fps: float, ) -> None: calls.append(("scene", "create")) @@ -401,48 +403,122 @@ def update_current_robot(self, robot_id: str, joint_state: JointState | None) -> assert joint_state == current calls.append(("update", robot_id)) - def show_preview(self, robot_id: str) -> None: - calls.append(("show", robot_id)) + def register_robot(self, robot_id: str, config: RobotModelConfig) -> None: + calls.append(("register", robot_id)) - def hide_preview(self, robot_id: str) -> None: - calls.append(("hide", robot_id)) + def cancel_preview_animation(self) -> None: + calls.append(("cancel", "preview")) - def animate_path(self, robot_id: str, path: list[JointState], duration: float) -> None: - assert path == [current] + def animate_preview(self, preview: object, duration: float) -> None: assert duration == 1.5 - calls.append(("animate", robot_id)) + calls.append(("animate", "groups")) def close(self) -> None: calls.append(("scene", "close")) - world_monitor = SimpleNamespace(get_current_joint_state=lambda _robot_id: current) - manipulation_module = SimpleNamespace( - robot_items=lambda: [("arm", "robot-1", fake_robot_config("arm"))], - robot_id_for_name=lambda robot_name: "robot-1" if robot_name == "arm" else None, - ) monkeypatch.setattr(visualizer_module, "ViserRuntime", FakeRuntime) monkeypatch.setattr(visualizer_module, "ViserUrdf", FakeViserUrdf) monkeypatch.setattr(visualizer_module, "ViserManipulationScene", FakeScene) visualizer = ViserManipulationVisualizer( - world_monitor=world_monitor, - manipulation_module=manipulation_module, config=ViserVisualizationConfig(panel_enabled=False), ) - visualizer.publish_visualization() - visualizer.show_preview("robot-1") - visualizer.hide_preview("robot-1") - visualizer.animate_path("robot-1", [current], duration=1.5) + assert hasattr(ViserManipulationVisualizer, "cancel_preview_animation") + visualizer.initialize( + VisualizationSession(PlanningSceneInfo({"robot-1": fake_robot_config("arm")})) + ) + visualizer.cancel_preview_animation() + visualizer.update_state(VisualizationStateFrame({"robot-1": current})) + visualizer.cancel_preview_animation() + visualizer.animate_trajectory(JointTrajectory(joint_names=["arm/joint1"]), duration=1.5) visualizer.close() - visualizer.publish_visualization() + visualizer.update_state(VisualizationStateFrame({"robot-1": current})) assert calls == [ ("runtime", "start"), ("scene", "create"), + ("register", "robot-1"), + ("cancel", "preview"), ("update", "robot-1"), - ("show", "robot-1"), - ("hide", "robot-1"), - ("animate", "robot-1"), + ("cancel", "preview"), + ("animate", "groups"), ("scene", "close"), ("runtime", "close"), ] + + +def test_scene_prepares_urdf_applies_base_pose_and_rejects_wrong_root( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + created: list[tuple[Path, str]] = [] + prepared: list[dict[str, object]] = [] + frames: list[dict[str, object]] = [] + fixed_world_root = tmp_path / "fixed-world.urdf" + fixed_world_root.write_text( + """ + + +""" + ) + non_fixed_world_root = tmp_path / "non-fixed-world.urdf" + non_fixed_world_root.write_text( + """ + + +""" + ) + + class Handle: + def remove(self) -> None: + return None + + class SceneApi: + def add_frame(self, name: str, **kwargs: object) -> Handle: + frames.append({"name": name, **kwargs}) + return Handle() + + class Server: + scene = SceneApi() + + class Urdf: + def __init__( + self, _server: Server, path: Path, *, root_node_name: str, **_kwargs: object + ) -> None: + created.append((path, root_node_name)) + self._meshes: list[object] = [] + + config = fake_robot_config("arm") + config.base_pose.position.x = 1.0 + + def prepare(path: Path, **kwargs: object) -> Path: + prepared.append(kwargs) + return fixed_world_root if path.name == "arm.urdf" else non_fixed_world_root + + monkeypatch.setattr( + "dimos.manipulation.visualization.viser.scene.prepare_urdf_for_drake", + prepare, + ) + + def parse_prepared_model(path: Path) -> SimpleNamespace: + content = path.read_text() + return SimpleNamespace(root_link="world" if 'name="world"' in content else "base_link") + + monkeypatch.setattr( + "dimos.manipulation.visualization.viser.scene.parse_model", parse_prepared_model + ) + scene = ViserManipulationScene(Server(), Urdf) + + scene.register_robot("robot-1", config) + + assert [root for _, root in created] == [ + "/robots/robot-1/current/base_pose/urdf", + "/targets/robot-1/target/base_pose/urdf", + "/previews/robot-1/ghost/base_pose/urdf", + ] + assert prepared == [{"package_paths": {}, "xacro_args": {}, "convert_meshes": False}] * 3 + assert all('name="world"' not in path.read_text() for path, _ in created) + assert all(frame["position"] == (1.0, 0.0, 0.0) for frame in frames) + wrong_root_config = fake_robot_config("wrong") + with pytest.raises(ValueError, match="prepared URDF root 'world'"): + scene.prepared_urdf_path(wrong_root_config) diff --git a/dimos/manipulation/visualization/viser/theme.py b/dimos/manipulation/visualization/viser/theme.py index 339be04492..d38767d1da 100644 --- a/dimos/manipulation/visualization/viser/theme.py +++ b/dimos/manipulation/visualization/viser/theme.py @@ -36,7 +36,7 @@ DIMOS_THEME_TITLE = "DimOS Manipulation" DIMOS_THEME_URL = "https://github.com/dimensionalOS/dimos" -DIMOS_BRAND_COLOR = (22, 130, 163) +DIMOS_BRAND_COLOR = (0, 153, 255) DIMOS_LOGO_PATH = Path(__file__).with_name("assets") / "dimensional-logo.svg" logger = setup_logger() @@ -86,8 +86,8 @@ def _configure_theme(server: ViserServer, titlebar_content: TitlebarConfig | Non try: server.gui.configure_theme( titlebar_content=titlebar_content, - control_layout="collapsible", - control_width="medium", + control_layout="fixed", + control_width="large", dark_mode=True, show_logo=False, show_share_button=False, diff --git a/dimos/manipulation/visualization/viser/visualizer.py b/dimos/manipulation/visualization/viser/visualizer.py index d12c98e53b..a9e6ee4c85 100644 --- a/dimos/manipulation/visualization/viser/visualizer.py +++ b/dimos/manipulation/visualization/viser/visualizer.py @@ -14,10 +14,15 @@ from __future__ import annotations +from collections.abc import Sequence from contextlib import suppress from typing import TYPE_CHECKING -from dimos.manipulation.visualization.viser.adapter import InProcessViserAdapter +from dimos.manipulation.visualization.viser.animation import ( + GroupPreviewAnimation, + PreviewFrame, + PreviewTrack, +) from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig from dimos.manipulation.visualization.viser.gui import ViserPanelGui from dimos.manipulation.visualization.viser.runtime import ( @@ -27,6 +32,8 @@ ) from dimos.manipulation.visualization.viser.scene import ViserManipulationScene from dimos.manipulation.visualization.viser.theme import apply_dimos_theme +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory from dimos.utils.logging_config import setup_logger try: @@ -41,12 +48,11 @@ raise ModuleNotFoundError(VISER_URDF_INSTALL_HINT) from e if TYPE_CHECKING: - from dimos.manipulation.manipulation_module import ManipulationModule - from dimos.manipulation.planning.monitor.world_monitor import WorldMonitor + from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.models import ( - JointPath, PlanningSceneInfo, - WorldRobotID, + VisualizationSession, + VisualizationStateFrame, ) logger = setup_logger() @@ -58,18 +64,19 @@ class ViserManipulationVisualizer: def __init__( self, *, - world_monitor: WorldMonitor, - manipulation_module: ManipulationModule, config: ViserVisualizationConfig | None = None, ) -> None: - self._world_monitor = world_monitor - self._manipulation_module = manipulation_module self.config = config or ViserVisualizationConfig() self._runtime: ViserRuntime | None = None self._server: ViserServer | None = None - self._adapter: InProcessViserAdapter | None = None self._scene: ViserManipulationScene | None = None self._gui: ViserPanelGui | None = None + self._session_scene: PlanningSceneInfo | None = None + self._operator: object | None = None + self._current_states: dict[str, JointState] = {} + self._robot_names_by_id: dict[str, str] = {} + self._robot_ids_by_name: dict[str, str] = {} + self._configs_by_name: dict[str, RobotModelConfig] = {} self._closed = False def _ensure_started(self) -> None: @@ -81,18 +88,19 @@ def _ensure_started(self) -> None: try: server = runtime.start() apply_dimos_theme(server) - adapter = InProcessViserAdapter( - world_monitor=self._world_monitor, - manipulation_module=self._manipulation_module, - ) - scene = ViserManipulationScene( - server, - ViserUrdf, - preview_fps=self.config.preview_fps, - ) + scene = ViserManipulationScene(server, ViserUrdf) gui = ( - ViserPanelGui(server, adapter, self.config, scene) + ViserPanelGui( + server, + self._session_scene, + self._operator, + self._current_states, + self.config, + scene, + ) if self.config.panel_enabled + and self._session_scene is not None + and self._operator is not None else None ) if gui is not None: @@ -108,20 +116,31 @@ def _ensure_started(self) -> None: runtime.close() self._runtime = None self._server = None - self._adapter = None self._scene = None self._gui = None self._closed = True raise self._runtime = runtime self._server = server - self._adapter = adapter self._scene = scene self._gui = gui self._closed = False logger.info(f"Viser manipulation visualization: {self.get_visualization_url()}") - def initialize_scene(self, scene: PlanningSceneInfo) -> None: + def initialize(self, session: VisualizationSession) -> None: + """Initialize Viser robot visuals from a one-shot visualization session.""" + self._operator = session.operator + self._session_scene = session.scene + self._robot_names_by_id = { + str(robot_id): config.name for robot_id, config in session.scene.robots.items() + } + self._robot_ids_by_name = { + config.name: str(robot_id) for robot_id, config in session.scene.robots.items() + } + self._configs_by_name = {config.name: config for config in session.scene.robots.values()} + self._initialize_scene(session.scene) + + def _initialize_scene(self, scene: PlanningSceneInfo) -> None: """Initialize Viser robot visuals from planning-scene metadata.""" if self._closed: return @@ -140,45 +159,94 @@ def initialize_scene(self, scene: PlanningSceneInfo) -> None: def get_visualization_url(self) -> str | None: return None if self._runtime is None else self._runtime.url - def publish_visualization(self, ctx: None = None) -> None: - """Update current robot render state. ctx is accepted for protocol compatibility.""" + def update_state(self, frame: VisualizationStateFrame) -> None: + """Update current robot render state from a pushed state frame.""" if self._closed: return self._ensure_started() - if self._adapter is None or self._scene is None: + if self._scene is None: return - for _robot_name, robot_id, _config in self._adapter.robot_items(): - current = self._adapter.get_current_joint_state(_robot_name) - self._scene.update_current_robot(str(robot_id), current) + for robot_id, current in frame.joint_states.items(): + robot_id_string = str(robot_id) + self._current_states[robot_id_string] = JointState(current) + self._scene.update_current_robot(robot_id_string, current) if self._gui is not None: self._gui.refresh() - def show_preview(self, robot_id: WorldRobotID) -> None: - if not self._closed: - self._ensure_started() - if self._scene is None: - return - self._scene.show_preview(str(robot_id)) - - def hide_preview(self, robot_id: WorldRobotID) -> None: - if not self._closed: - self._ensure_started() - if self._scene is None: - return - self._scene.hide_preview(str(robot_id)) - - def animate_path( - self, - robot_id: WorldRobotID, - path: JointPath, - duration: float = 3.0, + def animate_trajectory( + self, trajectory: JointTrajectory, duration: float | None = None ) -> None: if self._closed: return self._ensure_started() if self._scene is None: return - self._scene.animate_path(str(robot_id), list(path), duration) + preview = self._raw_preview_animation(trajectory) + if preview is not None: + self._scene.animate_preview( + preview, duration if duration is not None else max(float(trajectory.duration), 0.0) + ) + + def cancel_preview_animation(self, robot_ids: Sequence[str] | None = None) -> None: + """Cancel preview playback without starting a renderer or waiting for it. + + The world monitor deliberately invokes this outside its visualization + lock, so a renderer sleeping between frames can observe the scene + generation change immediately. Do not call ``_ensure_started()``: + cancelling before Viser has started must remain a no-op. Likewise, + retain the scene reference while ``close()`` is unwinding so a + concurrent cancellation can still invalidate an in-flight frame. + """ + scene = self._scene + if scene is not None: + if robot_ids is None: + scene.cancel_preview_animation() + else: + scene.cancel_preview_animation(robot_ids) + + def _raw_preview_animation(self, trajectory: JointTrajectory) -> GroupPreviewAnimation | None: + robot_indices: dict[str, list[tuple[int, str]]] = {} + for index, global_name in enumerate(trajectory.joint_names): + if "/" not in str(global_name): + return None + robot_name, local_name = str(global_name).split("/", 1) + if robot_name not in self._robot_ids_by_name: + return None + robot_indices.setdefault(robot_name, []).append((index, local_name)) + tracks: list[PreviewTrack] = [] + for robot_name, indexed_names in robot_indices.items(): + robot_id = self._robot_ids_by_name[robot_name] + config = self._configs_by_name[robot_name] + current = self._current_states.get(robot_id) + baseline = self._baseline_values(config, current) + if baseline is None: + return None + frames: list[PreviewFrame] = [] + for point in trajectory.points: + selected = { + local_name: float(point.positions[index]) for index, local_name in indexed_names + } + positions: list[float] = [] + for local_name in config.joint_names: + value = selected.get(local_name, baseline.get(local_name)) + if value is None: + return None + positions.append(float(value)) + frames.append(PreviewFrame(float(point.time_from_start), tuple(positions))) + tracks.append(PreviewTrack(robot_id, tuple(config.joint_names), tuple(frames))) + return GroupPreviewAnimation(tuple(tracks)) if tracks else None + + @staticmethod + def _baseline_values( + config: RobotModelConfig, current: JointState | None + ) -> dict[str, float] | None: + if current is None or len(current.name) != len(current.position): + return None + values = { + str(name): float(value) + for name, value in zip(current.name, current.position, strict=True) + } + return values if all(name in values for name in config.joint_names) else None def close(self) -> None: if self._closed: @@ -204,7 +272,6 @@ def close(self) -> None: errors.append(e) self._runtime = None self._server = None - self._adapter = None self._scene = None self._gui = None if errors: diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 47f56ff32d..ce4deace90 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -62,6 +62,8 @@ "drone-agentic": "dimos.robot.drone.blueprints.agentic.drone_agentic:drone_agentic", "drone-basic": "dimos.robot.drone.blueprints.basic.drone_basic:drone_basic", "dual-xarm6-planner": "dimos.robot.manipulators.xarm.blueprints.basic:dual_xarm6_planner", + "dual-xarm6-planner-coordinator": "dimos.robot.manipulators.xarm.blueprints.basic:dual_xarm6_planner_coordinator", + "dual-xarm6-planner-coordinator-mock-meshcat": "dimos.robot.manipulators.xarm.blueprints.basic:dual_xarm6_planner_coordinator_mock_meshcat", "keyboard-teleop-a1z": "dimos.robot.manipulators.a1z.blueprints.teleop:keyboard_teleop_a1z", "keyboard-teleop-a750": "dimos.robot.manipulators.a750.blueprints.teleop:keyboard_teleop_a750", "keyboard-teleop-openarm": "dimos.robot.manipulators.openarm.blueprints.teleop:keyboard_teleop_openarm", diff --git a/dimos/robot/manipulators/_modeling.py b/dimos/robot/manipulators/_modeling.py index 852a484b63..d9b528d0d5 100644 --- a/dimos/robot/manipulators/_modeling.py +++ b/dimos/robot/manipulators/_modeling.py @@ -16,6 +16,7 @@ from __future__ import annotations +import math from typing import TypeAlias from dimos.manipulation.planning.spec.models import RobotName @@ -31,10 +32,16 @@ JointNameMapping: TypeAlias = dict[CoordinatorJointName, UrdfJointName] -def base_pose(x: float = 0.0, y: float = 0.0, z: float = 0.0) -> PoseStamped: +def base_pose( + x: float = 0.0, + y: float = 0.0, + z: float = 0.0, + pitch: float = 0.0, +) -> PoseStamped: + half_pitch = pitch / 2.0 return PoseStamped( position=Vector3(x=x, y=y, z=z), - orientation=Quaternion(0.0, 0.0, 0.0, 1.0), + orientation=Quaternion([0.0, math.sin(half_pitch), 0.0, math.cos(half_pitch)]), ) diff --git a/dimos/robot/manipulators/a1z/config.py b/dimos/robot/manipulators/a1z/config.py index 355c8ee28f..1a8a032a08 100644 --- a/dimos/robot/manipulators/a1z/config.py +++ b/dimos/robot/manipulators/a1z/config.py @@ -19,6 +19,7 @@ from pathlib import Path from dimos.control.components import HardwareComponent, HardwareType, make_joints +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.robot.manipulators._modeling import ( base_pose, @@ -75,13 +76,21 @@ def make_a1z_model_config( coordinator_task_name: str | None = None, home_joints: list[float] | None = None, ) -> RobotModelConfig: + local_joint_names = joint_names(A1Z_DOF, prefix="arm_joint") return RobotModelConfig( name=name, model_path=A1Z_G1Z_MODEL_PATH if has_gripper else A1Z_FLANGE_MODEL_PATH, base_pose=base_pose(), - joint_names=joint_names(A1Z_DOF, prefix="arm_joint"), - end_effector_link=("gripper_eef_link" if has_gripper else "arm_link6"), + joint_names=local_joint_names, base_link="base_link", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=tuple(local_joint_names), + base_link="base_link", + tip_link=("gripper_eef_link" if has_gripper else "arm_link6"), + ) + ], package_paths=A1Z_PACKAGE_PATHS, auto_convert_meshes=True, collision_exclusion_pairs=A1Z_COLLISION_EXCLUSIONS, diff --git a/dimos/robot/manipulators/a750/config.py b/dimos/robot/manipulators/a750/config.py index 32bc57409a..4711c2f537 100644 --- a/dimos/robot/manipulators/a750/config.py +++ b/dimos/robot/manipulators/a750/config.py @@ -21,6 +21,7 @@ from dimos.control.components import HardwareComponent, HardwareType, make_joints from dimos.core.global_config import global_config +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.robot.manipulators._modeling import ( base_pose, @@ -102,13 +103,21 @@ def make_a750_model_config( coordinator_task_name: str | None = None, ) -> RobotModelConfig: dof = 6 + local_joint_names = joint_names(dof) return RobotModelConfig( name=name, model_path=A750_MODEL_PATH, base_pose=base_pose(), - joint_names=joint_names(dof), - end_effector_link="gripper_base", + joint_names=local_joint_names, base_link="base_link", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=tuple(local_joint_names), + base_link="base_link", + tip_link="gripper_base", + ) + ], package_paths=A750_PACKAGE_PATHS, auto_convert_meshes=True, collision_exclusion_pairs=A750_GRIPPER_COLLISION_EXCLUSIONS, diff --git a/dimos/robot/manipulators/openarm/config.py b/dimos/robot/manipulators/openarm/config.py index 5306d408bf..45b935b336 100644 --- a/dimos/robot/manipulators/openarm/config.py +++ b/dimos/robot/manipulators/openarm/config.py @@ -20,6 +20,7 @@ from typing import Any from dimos.control.components import HardwareComponent, HardwareType +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.robot.manipulators._modeling import base_pose from dimos.utils.data import LfsPath @@ -81,13 +82,21 @@ def openarm_hardware( def openarm_model_config(side: str, name: str | None = None) -> RobotModelConfig: validate_side(side) resolved_name = name or f"{side}_arm" + local_joint_names = openarm_joints(side) return RobotModelConfig( name=resolved_name, model_path=OPENARM_LEFT_MODEL if side == "left" else OPENARM_RIGHT_MODEL, base_pose=base_pose(), - joint_names=openarm_joints(side), - end_effector_link=f"openarm_{side}_link7", + joint_names=local_joint_names, base_link="openarm_body_link0", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=tuple(local_joint_names), + base_link="openarm_body_link0", + tip_link=f"openarm_{side}_link7", + ) + ], package_paths=OPENARM_PACKAGE_PATHS, collision_exclusion_pairs=OPENARM_COLLISION_EXCLUSIONS, auto_convert_meshes=True, @@ -112,13 +121,21 @@ def openarm_single_hardware( def openarm_single_model_config() -> RobotModelConfig: + local_joint_names = openarm_joints("left") return RobotModelConfig( name="arm", model_path=OPENARM_V10_FK_MODEL, base_pose=base_pose(), - joint_names=openarm_joints("left"), - end_effector_link="openarm_left_link7", + joint_names=local_joint_names, base_link="openarm_body_link0", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=tuple(local_joint_names), + base_link="openarm_body_link0", + tip_link="openarm_left_link7", + ) + ], package_paths=OPENARM_PACKAGE_PATHS, auto_convert_meshes=True, max_velocity=0.5, diff --git a/dimos/robot/manipulators/piper/config.py b/dimos/robot/manipulators/piper/config.py index f99e9589a5..68772166ad 100644 --- a/dimos/robot/manipulators/piper/config.py +++ b/dimos/robot/manipulators/piper/config.py @@ -20,6 +20,7 @@ from dimos.control.components import HardwareComponent, HardwareType, make_joints from dimos.core.global_config import global_config +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.robot.manipulators._modeling import ( base_pose, @@ -110,13 +111,21 @@ def make_piper_model_config( home_joints: list[float] | None = None, ) -> RobotModelConfig: dof = 6 + local_joint_names = joint_names(dof) return RobotModelConfig( name=name, model_path=PIPER_MODEL_PATH, base_pose=base_pose(), - joint_names=joint_names(dof), - end_effector_link="gripper_base", + joint_names=local_joint_names, base_link="base_link", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=tuple(local_joint_names), + base_link="base_link", + tip_link="gripper_base", + ) + ], package_paths=PIPER_PACKAGE_PATHS, auto_convert_meshes=True, collision_exclusion_pairs=PIPER_GRIPPER_COLLISION_EXCLUSIONS, diff --git a/dimos/robot/manipulators/test_blueprints.py b/dimos/robot/manipulators/test_blueprints.py index 3e7848beb1..59f9e0f794 100644 --- a/dimos/robot/manipulators/test_blueprints.py +++ b/dimos/robot/manipulators/test_blueprints.py @@ -20,7 +20,12 @@ from dimos.control.coordinator import ControlCoordinator, TaskConfig from dimos.core.coordination.blueprints import Blueprint from dimos.manipulation.manipulation_module import ManipulationModule, ManipulationModuleConfig -from dimos.manipulation.visualization.config import NoManipulationVisualizationConfig +from dimos.manipulation.visualization.config import ( + MeshcatVisualizationConfig, + NoManipulationVisualizationConfig, +) +from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig +from dimos.robot.get_all_blueprints import get_blueprint_by_name from dimos.robot.manipulators.a1z.blueprints.teleop import keyboard_teleop_a1z from dimos.robot.manipulators.a750.blueprints.teleop import keyboard_teleop_a750 from dimos.robot.manipulators.common.blueprints import eef_twist_task, planner @@ -32,6 +37,8 @@ from dimos.robot.manipulators.piper.blueprints.teleop import keyboard_teleop_piper from dimos.robot.manipulators.xarm.blueprints.basic import ( dual_xarm6_planner, + dual_xarm6_planner_coordinator, + dual_xarm6_planner_coordinator_mock_meshcat, xarm6_planner_only, xarm7_planner_coordinator, ) @@ -85,6 +92,34 @@ def test_xarm_planner_blueprints_default_to_no_visualization() -> None: assert isinstance(config.visualization, NoManipulationVisualizationConfig) +def test_dual_xarm6_planner_coordinator_blueprints_preserve_visualization_backends() -> None: + assert get_blueprint_by_name("dual-xarm6-planner-coordinator") is dual_xarm6_planner_coordinator + assert ( + get_blueprint_by_name("dual-xarm6-planner-coordinator-mock-meshcat") + is dual_xarm6_planner_coordinator_mock_meshcat + ) + + config = _manipulation_config(dual_xarm6_planner_coordinator) + legacy_config = _manipulation_config(dual_xarm6_planner_coordinator_mock_meshcat) + coordinator_kwargs = next( + atom.kwargs + for atom in dual_xarm6_planner_coordinator.blueprints + if atom.module is ControlCoordinator + ) + + assert isinstance(config.visualization, ViserVisualizationConfig) + assert isinstance(legacy_config.visualization, MeshcatVisualizationConfig) + assert [robot.name for robot in config.robots] == ["left_arm", "right_arm"] + assert [hardware.hardware_id for hardware in coordinator_kwargs["hardware"]] == [ + "left_arm", + "right_arm", + ] + assert [task.name for task in coordinator_kwargs["tasks"]] == [ + "traj_left_arm", + "traj_right_arm", + ] + + def test_eef_twist_task_helper_uses_hardware_joints_and_default_name() -> None: hardware = make_xarm_hardware("arm", 6, adapter_type="mock") diff --git a/dimos/robot/manipulators/xarm/blueprints/basic.py b/dimos/robot/manipulators/xarm/blueprints/basic.py index f66d50079b..c3f623356d 100644 --- a/dimos/robot/manipulators/xarm/blueprints/basic.py +++ b/dimos/robot/manipulators/xarm/blueprints/basic.py @@ -26,6 +26,7 @@ XARM7_SIM_PATH, make_xarm6_model_config, make_xarm7_model_config, + make_xarm_hardware, xarm6_hardware, xarm7_hardware, ) @@ -43,6 +44,43 @@ planning_timeout=10.0, ) +_mock_left_xarm6_hw = make_xarm_hardware("left_arm", 6) +_mock_right_xarm6_hw = make_xarm_hardware("right_arm", 6) + +dual_xarm6_planner_coordinator = autoconnect( + planner( + robots=[ + make_xarm6_model_config(name="left_arm", y_offset=0.5), + make_xarm6_model_config(name="right_arm", y_offset=-0.5), + ], + visualization={"backend": "viser"}, + ), + coordinator( + hardware=[_mock_left_xarm6_hw, _mock_right_xarm6_hw], + tasks=[ + trajectory_task(_mock_left_xarm6_hw), + trajectory_task(_mock_right_xarm6_hw), + ], + ), +) + +dual_xarm6_planner_coordinator_mock_meshcat = autoconnect( + planner( + robots=[ + make_xarm6_model_config(name="left_arm", y_offset=0.5), + make_xarm6_model_config(name="right_arm", y_offset=-0.5), + ], + visualization={"backend": "meshcat"}, + ), + coordinator( + hardware=[_mock_left_xarm6_hw, _mock_right_xarm6_hw], + tasks=[ + trajectory_task(_mock_left_xarm6_hw), + trajectory_task(_mock_right_xarm6_hw), + ], + ), +) + _xarm7_hw = xarm7_hardware("arm", gripper=True, mock_without_address=True) xarm7_planner_coordinator = autoconnect( diff --git a/dimos/robot/manipulators/xarm/config.py b/dimos/robot/manipulators/xarm/config.py index d918ee94a8..3f6beb7afc 100644 --- a/dimos/robot/manipulators/xarm/config.py +++ b/dimos/robot/manipulators/xarm/config.py @@ -21,6 +21,7 @@ from dimos.control.components import HardwareComponent, HardwareType, make_joints from dimos.core.global_config import global_config +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.robot.manipulators._modeling import ( base_pose, @@ -164,19 +165,28 @@ def make_xarm_model_config( xacro_args = { "dof": str(dof), "limited": "true", - "attach_xyz": f"{x_offset} {y_offset} {z_offset}", - "attach_rpy": f"0 {pitch} 0", + "attach_xyz": "0 0 0", + "attach_rpy": "0 0 0", } if add_gripper: xacro_args["add_gripper"] = "true" + local_joint_names = joint_names(dof) + tip_link = "link_tcp" if add_gripper else f"link{dof}" return RobotModelConfig( name=name, model_path=XARM_MODEL_PATH, - base_pose=base_pose(x_offset, y_offset, z_offset), - joint_names=joint_names(dof), - end_effector_link="link_tcp" if add_gripper else f"link{dof}", + base_pose=base_pose(x_offset, y_offset, z_offset, pitch), + joint_names=local_joint_names, base_link="link_base", + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=tuple(local_joint_names), + base_link="link_base", + tip_link=tip_link, + ) + ], package_paths=XARM_PACKAGE_PATHS, xacro_args=xacro_args, auto_convert_meshes=True, diff --git a/dimos/robot/test_all_blueprints.py b/dimos/robot/test_all_blueprints.py index adbc60ef00..af47f17f08 100644 --- a/dimos/robot/test_all_blueprints.py +++ b/dimos/robot/test_all_blueprints.py @@ -50,6 +50,8 @@ "coordinator-xarm6", "coordinator-xarm7", "dual-xarm6-planner", + "dual-xarm6-planner-coordinator", + "dual-xarm6-planner-coordinator-mock-meshcat", "learning-collect-quest-piper", "learning-collect-quest-xarm7", "teleop-hosted-go2", diff --git a/docs/capabilities/manipulation/adding_a_custom_arm.md b/docs/capabilities/manipulation/adding_a_custom_arm.md index df1299981d..54dde3f7d4 100644 --- a/docs/capabilities/manipulation/adding_a_custom_arm.md +++ b/docs/capabilities/manipulation/adding_a_custom_arm.md @@ -484,6 +484,7 @@ Place your URDF/xacro files under LFS data so they can be resolved via `LfsPath` from dimos.utils.data import LfsPath from dimos.manipulation.manipulation_module import manipulation_module from dimos.manipulation.planning.spec import RobotModelConfig +from dimos.manipulation.planning.spec.models import PlanningGroupDefinition from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.Quaternion import Quaternion from dimos.msgs.geometry_msgs.Vector3 import Vector3 @@ -507,7 +508,6 @@ def _make_base_pose(x=0.0, y=0.0, z=0.0) -> PoseStamped: def _make_yourarm_config( name: str = "arm", y_offset: float = 0.0, - joint_prefix: str = "", coordinator_task: str | None = None, ) -> RobotModelConfig: """Create YourArm robot config for planning. @@ -515,27 +515,31 @@ def _make_yourarm_config( Args: name: Robot name in the Drake planning world. y_offset: Y-axis offset for multi-arm setups. - joint_prefix: Prefix for joint name mapping to coordinator namespace. coordinator_task: Coordinator task name for trajectory execution via RPC. """ # These must match the joint names in your URDF joint_names = ["joint1", "joint2", "joint3", "joint4", "joint5", "joint6"] - joint_mapping = {f"{joint_prefix}{j}": j for j in joint_names} if joint_prefix else {} return RobotModelConfig( name=name, model_path=_YOURARM_URDF_PATH, - base_pose=_make_base_pose(y=y_offset), joint_names=joint_names, - end_effector_link="link6", # Last link in your URDF's kinematic chain - base_link="base_link", # Root link of your URDF + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", + joint_names=tuple(joint_names), + base_link="base_link", + tip_link="link6", + ) + ], + base_pose=_make_base_pose(y=y_offset), # world -> base_link placement + base_link="base_link", # Robot-scoped placement/weld/strip link package_paths={"yourarm_description": _YOURARM_PACKAGE_PATH}, xacro_args={}, # Xacro arguments if using .xacro files collision_exclusion_pairs=[], # Pairs of links that can touch (e.g., gripper fingers) auto_convert_meshes=True, # Convert DAE/STL meshes for Drake max_velocity=1.0, # Max velocity scaling factor max_acceleration=2.0, # Max acceleration scaling factor - joint_name_mapping=joint_mapping, coordinator_task_name=coordinator_task, ) ``` @@ -547,7 +551,7 @@ Add this to your `dimos/robot/yourarm/blueprints.py` alongside the coordinator b ```python skip yourarm_planner = manipulation_module( - robots=[_make_yourarm_config("arm", joint_prefix="arm_", coordinator_task="traj_arm")], + robots=[_make_yourarm_config("arm", coordinator_task="traj_arm")], planning_timeout=10.0, visualization={"backend": "meshcat"}, ) @@ -561,14 +565,23 @@ yourarm_planner = manipulation_module( | Field | Description | |-------|-------------| | `model_path` | Path to `.urdf` or `.xacro` file | -| `joint_names` | Ordered list of controlled joints (must match URDF) | -| `end_effector_link` | Link to use as the end-effector for IK | -| `base_link` | Root link of the robot model | +| `joint_names` | Ordered controllable local model joint set (must match URDF); not itself a planning group | +| `planning_groups` / `srdf_path` | Explicit planning groups or SRDF source; direct `RobotModelConfig(...)` helpers should pass explicit groups, while shared config helpers can discover groups from SRDF/fallback | +| `base_pose` / `base_link` | Optional robot placement: `base_pose` places `base_link` in the world for weld/strip behavior | | `package_paths` | Maps `package://` URIs to filesystem paths (for xacro) | -| `joint_name_mapping` | Maps coordinator names (e.g., `"arm_joint1"`) to URDF names (e.g., `"joint1"`) | | `coordinator_task_name` | Must match the `TaskConfig.name` in your coordinator blueprint | | `collision_exclusion_pairs` | List of `(link_a, link_b)` tuples for links that may legitimately touch (e.g., gripper fingers) | +Coordinator-facing joint states and trajectories use global joint names derived +mechanically as `{robot_name}/{local_joint_name}` (for example, `arm/joint1`). +Keep hardware-native name translation inside the hardware adapter; manipulation +planning config uses local model joint names. + +Planning-group `base_link`/`tip_link` values define kinematic chains and pose +target frames. `base_link` is only the robot-scoped link placed by +`base_pose`; do not use it as a substitute for planning-group chain metadata. +See [Planning Groups](/docs/capabilities/manipulation/planning_groups.md). + ## Step 5: Register Blueprints The blueprint registry in `dimos/robot/all_blueprints.py` is **auto-generated** by scanning the codebase for blueprint declarations. After adding your blueprints: diff --git a/docs/capabilities/manipulation/index.md b/docs/capabilities/manipulation/index.md index 9851c0e8fe..ece11778b8 100644 --- a/docs/capabilities/manipulation/index.md +++ b/docs/capabilities/manipulation/index.md @@ -149,8 +149,9 @@ Backend choices: - `meshcat`: embedded Drake/Meshcat visualizer. The planning world must be created with embedded visualization enabled, so this is selected through the visualization config. -- `viser`: in-process Viser visualizer. It renders current robot state, target controls, - transient preview ghosts, planned path previews, and optional panel controls. +- `viser`: in-process Viser visualizer. It renders pushed current robot state, + target controls, transient preview ghosts, synchronized trajectory previews, + and optional panel controls. - `none`: no manipulation planning visualization. CLI example: @@ -185,20 +186,27 @@ Viser support is included in the `manipulation` extra: uv sync --extra manipulation --inexact ``` -The Viser panel uses existing manipulation planning, preview, execute, cancel, and clear-plan -RPC methods through a small in-process adapter. GUI callbacks enqueue operations instead of -touching `WorldSpec`, IK, planner objects, or live Drake contexts directly. Rendering copies -mutable joint state/path containers at the read boundary, then updates the Viser scene after -manipulation/world accessors have returned. - -External manipulation visualizers are initialized from a backend-neutral planning-scene snapshot -after the planning world has added its robots. This snapshot maps world robot IDs to -`RobotModelConfig` metadata so Viser can prepare current, target, and transient preview robot -visuals without `WorldMonitor` depending on Viser-specific hooks. Embedded Meshcat visualization -does not need extra setup because it observes the Drake world directly. - -When the Viser panel is enabled, it can call the existing manipulation execution path after a -fresh feasible plan is available and the current robot joints still match the plan start. +The Viser panel talks to the concrete `ManipulationOperator` bound into its +`VisualizationSession`. GUI callbacks enqueue operations through that operator +for target evaluation, planning, preview, execution, cancellation, reset, and +clear-plan actions. The panel owns only target drafts, selection state, and +callback generations; it does not touch `WorldSpec`, IK, planner objects, +`ManipulationModule`, `WorldMonitor`, or live Drake contexts directly. + +External manipulation visualizers are initialized from a backend-neutral +`VisualizationSession` after the planning world has added its robots. The +session contains static `PlanningSceneInfo` metadata: world robot IDs, +`RobotModelConfig` values, and resolved planning groups. Runtime joint state is +then pushed through `VisualizationStateFrame` updates so renderers do not poll +world/module state or own freshness policy. Embedded Meshcat visualization does +not need extra setup because it observes the Drake world directly. + +Previews use the stored synchronized `JointTrajectory` from the generated plan. +Viser projects the globally named trajectory into robot-local preview ghosts and +plays the stored timestamped points directly; optional preview duration only +scales the stored delays. Execute freshness is enforced by the manipulation +module/operator immediately before dispatch, not by Viser-side telemetry +snapshots. ### Perception + Agent diff --git a/docs/capabilities/manipulation/planning_groups.md b/docs/capabilities/manipulation/planning_groups.md new file mode 100644 index 0000000000..7ce9edb122 --- /dev/null +++ b/docs/capabilities/manipulation/planning_groups.md @@ -0,0 +1,171 @@ +# Manipulation Planning Groups + +Planning groups are named, selectable kinematic chains used by manipulation +planning. They let APIs target a specific part of a robot, such as an arm or +torso, without confusing that group with the robot's hardware identity. + +## Concepts + +| Concept | Meaning | +|---------|---------| +| Robot name | The configured robot ID in `RobotModelConfig.name`. | +| Planning group | A named serial chain of controllable joints on one robot. | +| Planning group ID | Stable API ID in the form `{robot_name}/{group_name}`. | +| Local joint name | Joint name inside a robot model, such as `joint1`. | +| Global joint name | Boundary-level joint name in the form `{robot_name}/{local_joint_name}`. | +| Generated plan | Planning artifact containing selected group IDs, geometric waypoints, and one synchronized global-joint trajectory. | +| Auxiliary group | A selected group that contributes free DOFs to a pose plan without receiving its own pose target. | + +Local URDF/SRDF joint names stay inside robot-scoped configuration, model +parsing, and backend internals. Flat planning states and generated plan paths +use global joint names so multiple robots can safely share local names such as +`joint1`. + +## Discovering planning groups + +Robot configs can provide planning groups explicitly with +`RobotModelConfig.planning_groups`. Direct `RobotModelConfig(...)` construction +does not run discovery or synthesize groups in `model_post_init`; callers must +pass explicit `planning_groups` there. + +When code uses the discovery helper instead of explicit config, DimOS discovers +groups in this order: + +1. Explicit `srdf_path` provided to the helper. +2. Conservative SRDF auto-discovery near the model path, with a warning. +3. Fallback generation of one `{robot_name}/manipulator` group when the + configured controllable joints form exactly one unambiguous serial chain. +4. Error if no SRDF or fallback chain can provide a single valid group. + +Supported SRDF group forms: + +```xml + + + +``` + +```xml + + + + + +``` + +Unsupported SRDF forms are skipped with warnings: link groups, nested group +references, mixed group declarations, branching or non-serial groups, and SRDF +`` metadata. A chain group's `tip_link` is its pose target frame. +An ordered joint-list group can be pose-targeted only when DimOS can validate a +unique serial target frame. + +## Fallback behavior + +When discovery runs without an SRDF, fallback uses +`RobotModelConfig.joint_names` as the candidate controllable set. +This field is the robot's ordered local model joint set, not an implicit +planning group. + +Fallback succeeds only when those joints form one unambiguous serial chain. It +allows prismatic joints in the middle of the chain and strips only terminal tip +prismatic joints, which usually represent gripper fingers. The generated group +name is always `manipulator`. + +## Current APIs + +Use `list_planning_groups()` to discover group IDs and capabilities before +planning: + +```python skip +groups = manip.list_planning_groups() +pose_groups = [group for group in groups if group.has_pose_target] +group_id = pose_groups[0].id +``` + +Joint-space planning targets group IDs. Each target `JointState` may be +unnamed in the group's joint order, named with all local model joint names, or +named with all global joint names. Do not mix local and global names in one +target. + +```python skip +ok = manip.plan_to_joint_targets( + { + "left_arm/manipulator": JointState( + name=["joint1", "joint2", "joint3", "joint4", "joint5", "joint6"], + position=[0.0, -0.4, 0.2, 0.0, 0.3, 0.0], + ) + } +) +``` + +Pose planning targets pose-capable group IDs. Add auxiliary groups when another +chain should participate as free DOFs but does not have its own pose target. +Pose targets are `Pose` values keyed by planning group ID: + +```python skip +ok = manip.plan_to_pose_targets( + {"left_arm/manipulator": target_pose}, + auxiliary_groups=["torso/manipulator"], +) +``` + +After a successful planning call, preview and execution use the module's current +stored plan: + +```python skip +manip.preview_plan() +manip.execute_plan() +``` + +Callers that already hold a `GeneratedPlan` may pass it explicitly: + +```python skip +manip.preview_plan(plan) +manip.execute_plan(plan) +``` + +For robot-scoped compatibility APIs, unnamed joint vectors are interpreted in +the selected default planning group's joint order. If names are provided, they +may be all local model joint names or all global joint names. Missing joints, +extra joints, partial joint sets, and mixed local/global namespaces are rejected. + +## Generated plans and execution + +A `GeneratedPlan` stores: + +- selected planning group IDs; +- a geometric path of `JointState` waypoints keyed by global joint names; +- one materialized synchronized `JointTrajectory` over the same selected global + joint names; +- status, timing, path length, iteration count, and message metadata. + +Preview and execution consume the stored trajectory; they do not lazily +parameterize the geometric path. Preview forwards the raw globally named +trajectory through the visualization boundary, where renderers project it to +their robot-local visuals while preserving stored timestamps. Execution splits +the stored trajectory by affected trajectory task, translates selected joint +names at the coordinator boundary, and invokes each trajectory controller +without filling or commanding omitted joints. Controllers remain planning-group +agnostic, and trajectory tasks still claim their full configured joint set while +executing only the active planned subset. + +Multi-task dispatch is not atomic: if one trajectory task accepts and a later +task rejects, DimOS reports the rejection but does not roll back the accepted +task. + +## Robot placement config + +`RobotModelConfig.base_pose` and `RobotModelConfig.base_link` describe robot +placement: `base_pose` places `base_link` in the world and current backends +use that link for weld/placement and optional model-authored world-joint +stripping. This is robot placement metadata, not planning-chain metadata. + +Planning-group `base_link` and `tip_link` values are the only source for chain +bases and pose target frames. Robot-scoped end-effector config is no longer +supported; robot-level EE helper APIs are wrappers over a unique pose-targetable +planning group and should use explicit group APIs when multiple pose groups +exist. + +Robot placement can be encoded either in model assets or in `base_pose`, +depending on the blueprint. `joint_names` remains supported and should describe +the ordered controllable local model joint set. diff --git a/openspec/changes/extract-manipulation-execution-runtime/.openspec.yaml b/openspec/changes/extract-manipulation-execution-runtime/.openspec.yaml new file mode 100644 index 0000000000..29e56d8d3b --- /dev/null +++ b/openspec/changes/extract-manipulation-execution-runtime/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-07-20 diff --git a/openspec/changes/extract-manipulation-execution-runtime/design.md b/openspec/changes/extract-manipulation-execution-runtime/design.md new file mode 100644 index 0000000000..46d26cd201 --- /dev/null +++ b/openspec/changes/extract-manipulation-execution-runtime/design.md @@ -0,0 +1,115 @@ +# Design: Manipulation Execution Runtime + +## Context + +`ManipulationModule` currently owns planning, coordinator dispatch, cancellation, remote-task monitoring, and lifecycle bookkeeping. These responsibilities are interleaved with manually assigned state and separate execution flags, so concurrent requests can produce ambiguous dispatch, completion, or cancellation outcomes. Remote coordinator activity is also difficult to reconcile safely when a response is delayed, incomplete, or contradictory. + +## Baseline inventory (historical, task 1.1) + +The pre-runtime implementation used `ManipulationState.IDLE`, `PLANNING`, `EXECUTING`, `COMPLETED`, and `FAULT`, with module-owned `_state`, `_last_plan`, `_execution_generation`, `_planning_epoch`, `_execution_dispatch_lock`, and `_coordinator_client` fields. Planning entered through `plan_to_pose`, `plan_to_joints`, `plan_to_pose_targets`, and `plan_to_joint_targets`; execution entered through `execute`/`execute_plan` and split the stored plan into per-robot trajectories. + +Module-owned coordinator RPC sites included trajectory `execute`, `cancel`, and `get_state` calls, plus fallback completion polling and direct coordinator-client setup. Status consumers read `get_state`, `get_error`, and `get_trajectory_status`, while the operator/UI adapters consumed those module methods. + +The current replacement seams are `ExecutionRuntime` and its `LifecycleState`, `RuntimeSnapshot`, plan/operation/task records, and serialized `start_planning`/`complete_planning`/`execute_ready`/`execute_explicit`/`cancel`/`reset`/`poll` methods. The four planning APIs now submit planning epochs/results through those planning methods, and execution entry points submit prepared or ready plans through the runtime. `ControlCoordinatorGateway` now owns the runtime-facing `execute`/`cancel`/`status`/`reset` RPC boundary. `ManipulationModule.get_state`, `get_execution_snapshot`, `get_error`, `cancel`, `reset`, and execution entry points delegate to that snapshot/runtime seam; operator/UI status reads the atomic snapshot rather than legacy module flags. + +This change introduces an internal execution runtime owned by `ManipulationModule`. The runtime is the sole lifecycle owner for manipulation execution. The module remains responsible for running planning algorithms and submitting planning and execution events to the runtime; it does not independently assign lifecycle state or maintain a second execution lifecycle. + +The runtime provides the authoritative lifecycle and status snapshot for the manipulation layer while encapsulating coordinator trajectory-task interaction. The design is internal and implementation-agnostic: it specifies ownership, events, state semantics, and safety rules rather than a particular threading or queue library. + +## Goals / Non-Goals + +### Goals + +- Define one authoritative lifecycle with exactly these public states: `IDLE`, `PLANNING`, `READY`, `DISPATCHING`, `RUNNING`, `CANCELLING`, and `FAULT`. +- Serialize planning results, execute, cancel, reset, remote reconciliation, and status publication through a single-owner queue and thread (or equivalent single serialized owner). +- Reject execution requests received during an active attempt rather than queueing them for later dispatch. +- Use a formal reducer to derive lifecycle state and runtime data from events, rather than ad hoc state assignments. +- Make the runtime authoritative for the ready plan, operation IDs, attempts, task knowledge, faults, and the coordinator client. +- Make plans single-use and make replacement of a ready plan explicit and atomic. +- Dispatch every robot trajectory represented by a generated plan; caller-selected partial-robot execution is not supported. +- Reconcile coordinator task activity conservatively, latching `FAULT` whenever safe remote activity cannot be proven. +- Require reset to establish that relevant remote activity is inactive before leaving `FAULT`. +- Build and validate coordinator-name inverse mapping once, using canonical configuration references. + +### Non-Goals + +- This change does not redesign planning algorithms or trajectory representations. +- It does not add a replacement for manipulation-layer trajectory start-state freshness validation. +- It does not introduce partial-robot execution or retain a duplicate direct trajectory-execution path owned by `ManipulationModule`. Standalone diagnostic/operator CLI direct-RPC tooling in `dimos/manipulation/control/coordinator_client.py` is out of scope; moving or removing it requires a separate compatibility-scoped change. +- It does not define a new external runtime or coordinator dependency. +- It does not prescribe a UI layout, transport, or concrete queue/thread implementation. + +## Decisions + +### Runtime ownership and event processing + +The internal runtime is the sole lifecycle owner. `ManipulationModule` invokes planning algorithms and submits events such as planning started, plan produced, execute requested, cancel requested, reset requested, and coordinator observations. The module must not mutate lifecycle state, ready-plan ownership, operation bookkeeping, or fault state outside the runtime. + +All lifecycle-affecting work is serialized by a single-owner queue and thread, or an equivalent mechanism with the same single-writer guarantee. This includes execute, cancel, reset, coordinator polling, task reconciliation, and status snapshot publication. External callers may enqueue requests concurrently, but event order at the runtime owner is the order used for lifecycle decisions. + +Execution admission is state-based: while the runtime is `DISPATCHING`, `RUNNING`, or `CANCELLING`, a new execution request is completed as terminal `REJECTED` and is not retained for later dispatch. A coordinator refusal before a remote attempt is established is also a terminal `REJECTED` result and returns the runtime to `IDLE`; neither case is a safety fault. + +A formal reducer is the authority for applying events. It validates each event against the current state and runtime context, performs the associated side effects through the runtime-owned coordinator client where needed, and produces the next state plus an atomic status snapshot. Invalid or unsafe events do not create an alternate state path; they are rejected or become a fault according to the event's safety semantics. + +### Public lifecycle + +The public state vocabulary is limited to: + +| State | Meaning | +| --- | --- | +| `IDLE` | No ready plan or active manipulation operation is owned by the runtime. | +| `PLANNING` | A planning request is in progress; its eventual result is not yet executable. | +| `READY` | One complete plan is owned by the runtime and may be executed once. | +| `DISPATCHING` | The selected plan has been consumed and coordinator submission/reconciliation is in progress. | +| `RUNNING` | The current operation is known to be active remotely. | +| `CANCELLING` | Cancellation has been requested and the runtime is reconciling until inactivity or fault is established. | +| `FAULT` | The runtime cannot prove a safe, unambiguous lifecycle or remote-task condition; execution remains latched until reset succeeds. | + +The snapshot is atomic from consumers' perspective and includes the public state together with the applicable operation, attempt, plan, task, and fault information. Consumers must use the snapshot rather than infer lifecycle from independent flags. + +### Runtime-owned data and coordinator access + +The runtime owns the ready plan, operation IDs, attempt records, task knowledge, fault details, and coordinator client. An operation ID identifies a logical execution request; attempts identify coordinator submission/reconciliation efforts within that operation. This separates local request identity from remote task identity and prevents stale responses from being treated as current activity. + +The coordinator client used for `ManipulationModule` execution is accessed only by the runtime. A plan is dispatched through the runtime's single coordinator path; active module-owned duplicate direct trajectory execution paths are not retained. This does not require moving or removing `dimos/manipulation/control/coordinator_client.py`, which remains standalone diagnostic/operator CLI direct-RPC tooling; changing it is a separate compatibility-scoped change. The runtime reconciles submission responses and subsequent task observations against the current operation and attempt before changing state. + +### Plan replacement and single-use semantics + +Plans are single-use. A plan selected for dispatch is consumed before dispatch begins, so it cannot be executed again or reused after a failed or completed attempt. Planning while a ready plan exists may produce an explicit replacement event; applying that event atomically discards the prior ready plan and installs the new complete plan. There is never an externally observable interval in which both plans are executable. + +Only a plan containing every robot trajectory represented by its planning groups is accepted. Partial-robot selection and a `robot_name` execution-selection argument are not part of this lifecycle. A plan that omits a trajectory required by its own planning groups is rejected rather than converted into a partial execution. + +### Coordinator-name mapping + +The coordinator-name inverse mapping is constructed once during runtime initialization from canonical configuration references. Initialization validates that references are complete, names are unambiguous, and the mapping is internally consistent. Runtime operation uses this validated mapping and does not rebuild or infer it from mutable per-request data. + +### Freshness and uncertainty safety + +Manipulation-layer trajectory start-state freshness validation is removed in this change and has no replacement here. The runtime instead focuses on lifecycle and remote-task certainty. + +Remote uncertainty is fail-safe and latching. If coordinator task activity, ownership, submission outcome, cancellation outcome, or completion cannot be proven safe and attributable to the current operation, the reducer enters `FAULT`. A reset request may clear the fault only after the runtime proves that relevant remote activity is inactive; reset is not merely a local flag clear. Until then, the runtime remains in `FAULT` and does not dispatch a new plan. + +## Risks / Trade-offs + +- A single serialized owner simplifies race handling and makes reducer behavior deterministic, but polling, coordinator calls, or planning-result delivery can delay subsequent lifecycle events if not bounded or isolated appropriately. +- Conservative latching can expose `FAULT` for recoverable coordinator/network ambiguity and requires operator-visible reset, trading availability for avoidance of duplicate or uncontrolled motion. +- Consuming a plan before dispatch prevents accidental reuse but means a failed dispatch cannot silently retry the same plan; any retry must be represented as a new, explicit attempt under the runtime's rules. +- Atomic replacement makes plan ownership clear, but a newly completed plan can supersede a prior ready plan without execution of the prior one. +- Building the inverse mapping once improves consistency and detects configuration errors early, but configuration changes require runtime reinitialization rather than being picked up dynamically. +- Removing freshness validation leaves start-state correctness to lower layers or future work; this change must not imply that freshness is guaranteed elsewhere. + +## Migration Plan + +1. Add the internal runtime behind `ManipulationModule` and route all lifecycle-affecting requests, planning results, coordinator polling, cancellation, reset, and status publication through its event interface. +2. Move ready-plan storage, operation and attempt bookkeeping, task knowledge, fault handling, coordinator access, and state transitions into the runtime. Remove parallel module-level lifecycle flags and assignments. +3. Update manipulation state adapters and callers to consume the seven public states and atomic snapshots. Remove the `robot_name` execution-selection argument and any partial-robot execution paths. +4. Construct and validate the coordinator-name inverse mapping once from canonical configuration references during runtime initialization. +5. Remove manipulation-layer trajectory start-state freshness validation and active `ManipulationModule` duplicate direct trajectory execution paths, without adding replacement validation in this change. Leave `dimos/manipulation/control/coordinator_client.py` unchanged; moving or removing that standalone diagnostic/operator CLI direct-RPC tooling is a separate compatibility-scoped change. +6. Validate normal planning, atomic replacement, single-use dispatch, completion, cancellation, stale/ambiguous coordinator observations, and reset-after-inactivity behavior before removing the legacy lifecycle implementation. + +## Open Questions + +- What concrete polling interval and coordinator-call timeout best balance responsiveness with the runtime's single-owner serialization? +- Which coordinator observations constitute sufficient proof of inactivity for each reset path, and how long should that proof remain valid? +- How should faults and operation/attempt identifiers be exposed in existing UI and operator-facing diagnostics beyond the atomic snapshot? +- What future layer, if any, will own trajectory start-state freshness validation after it is removed from manipulation? diff --git a/openspec/changes/extract-manipulation-execution-runtime/proposal.md b/openspec/changes/extract-manipulation-execution-runtime/proposal.md new file mode 100644 index 0000000000..0eeedc65be --- /dev/null +++ b/openspec/changes/extract-manipulation-execution-runtime/proposal.md @@ -0,0 +1,29 @@ +## Why + +`ManipulationModule` currently combines planning, coordinator dispatch, cancellation, task-status monitoring, and lifecycle bookkeeping. Its manually assigned state, separate execution flags, and ambiguous dispatch/completion semantics make concurrent cancellation and remote-task uncertainty difficult to reason about safely. + +## What Changes + +- Extract coordinator trajectory dispatch and its full remote-task lifecycle into an internal execution runtime owned by `ManipulationModule`. +- Replace manual lifecycle assignments with a single-owner, event-driven state machine and atomic status snapshots. +- Replace the public lifecycle vocabulary with `IDLE`, `PLANNING`, `READY`, `DISPATCHING`, `RUNNING`, `CANCELLING`, and `FAULT`. +- Make plans single-use: replacement planning discards the prior ready plan, and dispatch consumes the selected plan. +- Serialize execute, cancel, reset, and status-polling work through the runtime; latch faults when coordinator task activity cannot be proven safe. +- Remove manipulation-layer trajectory start-state freshness validation without adding a replacement in this change. +- Remove partial-robot execution and the duplicate direct trajectory execution path. + +## Capabilities + +### New Capabilities +- `manipulation-lifecycle-runtime`: Defines the authoritative manipulation plan and execution lifecycle, state transitions, snapshots, and fault/reset behavior. +- `coordinator-trajectory-execution`: Defines safe plan-to-coordinator dispatch, task reconciliation, completion monitoring, and cancellation semantics. + +### Modified Capabilities + +None. + +## Impact + +- Affects `dimos/manipulation/manipulation_module.py`, manipulation operator/UI state adapters, and manipulation execution tests. +- Changes public state strings and removes the `robot_name` execution-selection argument; callers and UI state handling must be updated. +- Centralizes use of `ControlCoordinator` trajectory task RPCs in the internal runtime; no new external dependency is introduced. diff --git a/openspec/changes/extract-manipulation-execution-runtime/specs/coordinator-trajectory-execution/spec.md b/openspec/changes/extract-manipulation-execution-runtime/specs/coordinator-trajectory-execution/spec.md new file mode 100644 index 0000000000..d1c7ea1f5a --- /dev/null +++ b/openspec/changes/extract-manipulation-execution-runtime/specs/coordinator-trajectory-execution/spec.md @@ -0,0 +1,85 @@ +## ADDED Requirements + +### Requirement: Serialized coordinator task execution +The coordinator trajectory execution runtime SHALL serialize execute, cancel, reset, and status-polling operations. It SHALL allow at most one active plan and one active execution attempt at a time, and SHALL NOT issue planning or execution work while a remote execution is active. + +The runtime SHALL reject a second execution request while an execution attempt is active; execution requests SHALL NOT be queued for later dispatch. The rejection SHALL be terminal for the rejected request and SHALL not fault or interrupt the already active attempt. + +#### Scenario: Active execution rejects a second request +- **WHEN** an execution request arrives while the first attempt is active +- **THEN** the second request SHALL complete with a terminal `REJECTED` outcome without being queued +- **AND** the active attempt SHALL remain unchanged +- **AND** the coordinator SHALL receive no execute request for the rejected operation + +#### Scenario: Rejected execution does not fault the idle runtime +- **WHEN** a valid execution request reaches dispatch and the coordinator rejects it +- **THEN** the request SHALL complete with a terminal `REJECTED` outcome +- **AND** the manipulation lifecycle SHALL return to `IDLE` +- **AND** the runtime SHALL NOT transition to `FAULT` + +### Requirement: Global-to-local-to-coordinator mapping boundary +The runtime SHALL define an explicit mapping boundary from the global manipulation plan to the local coordinator trajectory/task request and back to the global manipulation result. The mapping SHALL preserve plan identity, attempt identity, trajectory/action data, and terminal outcome without exposing coordinator-specific task representation as the manipulation lifecycle contract. + +The runtime SHALL perform coordinator dispatch only from a fully mapped local request and SHALL map every coordinator result into the normalized global outcome before publishing it to manipulation callers. + +#### Scenario: Global plan is mapped for coordinator dispatch +- **WHEN** a global plan is selected for execution +- **THEN** the runtime SHALL create local coordinator requests containing every trajectory represented by the global plan and their correlation identities +- **AND** it SHALL dispatch that local request rather than passing an unbounded global plan object through the coordinator API + +#### Scenario: Coordinator completion is mapped back +- **WHEN** the coordinator reports a terminal result for the active task +- **THEN** the runtime SHALL map it to the corresponding global plan/attempt outcome +- **AND** callers SHALL receive the normalized manipulation outcome rather than raw coordinator task details + +### Requirement: Register the task before execute RPC +The runtime MUST register the active plan, attempt, and coordinator task correlation context before invoking the coordinator execute RPC. It MUST associate the RPC response with that pre-registered attempt and MUST retain the context when dispatch is accepted. + +#### Scenario: Accepted task is correlated without a race +- **WHEN** the runtime is about to invoke the coordinator execute RPC +- **THEN** its active-attempt/task registration SHALL already exist +- **AND** an immediate acceptance or completion response SHALL be attributed to that attempt + +### Requirement: Normalize cancellation and unknown outcomes +The runtime SHALL normalize coordinator cancellation responses into a cancellation outcome. It SHALL normalize unknown, malformed, unavailable, or contradictory coordinator task states into a safety-uncertain outcome rather than treating them as successful completion or ordinary cancellation. + +Cancellation SHALL be idempotent for an already terminal attempt, but cancellation of an active task SHALL not be reported complete until the coordinator confirms cancellation or confirms that the task is no longer active. If that confirmation cannot be obtained, the runtime SHALL latch remote safety uncertainty and SHALL block new execution until reset reconciliation. + +#### Scenario: Active cancellation waits for confirmation +- **WHEN** cancellation is requested for an active coordinator task +- **THEN** the runtime SHALL serialize the cancel request and enter cancellation handling +- **AND** it SHALL report cancellation only after the task is confirmed cancelled or inactive + +#### Scenario: Unknown task state is unsafe +- **WHEN** coordinator polling returns an unknown task state or cannot establish whether the task remains active +- **THEN** the runtime SHALL normalize the result as safety uncertainty +- **AND** it SHALL latch the runtime in `FAULT` until safe reset reconciliation + +### Requirement: Reset reconciles remote safety +Reset SHALL first reconcile coordinator task activity and SHALL clear local plan, attempt, and pending request state only after remote safety is proven. Reset SHALL leave the runtime faulted when remote activity is still active or cannot be determined, and SHALL be retryable without issuing a new trajectory execution. + +#### Scenario: Reset clears local state after safe reconciliation +- **WHEN** reset confirms that the coordinator has no active task +- **THEN** the runtime SHALL clear the active attempt, ready plan, and pending execution-request bookkeeping +- **AND** it SHALL return the manipulation lifecycle to `IDLE` + +#### Scenario: Reset does not hide active remote execution +- **WHEN** reset finds an active coordinator task +- **THEN** the runtime SHALL not clear the safety fault or start another execution +- **AND** it SHALL continue reconciliation or require cancellation before becoming idle + +### Requirement: Plan-complete execution only +The coordinator execution interface SHALL require every robot trajectory represented by the generated plan and SHALL not accept a partial-robot argument or robot-name selector. The runtime SHALL maintain a single trajectory execution path through the coordinator and SHALL remove duplicate direct execution paths owned by `ManipulationModule`. This requirement does not remove or relocate `dimos/manipulation/control/coordinator_client.py`, whose standalone diagnostic/operator CLI direct-RPC tooling is outside the `ManipulationModule` execution path; changing that tooling is a separate compatibility-scoped change. + +#### Scenario: Partial execution input has no dispatch effect +- **WHEN** an execution request omits a robot trajectory represented by its generated plan or supplies a robot selector +- **THEN** the runtime SHALL reject the request as invalid +- **AND** it SHALL issue no coordinator execute RPC + +### Requirement: No manipulation start-state freshness gate +The coordinator execution runtime SHALL not add or restore manipulation-layer validation that rejects a plan because its trajectory start state is not fresh. Start-state acceptance and physical safety remain the responsibility of the coordinator and its execution path. + +#### Scenario: Execution is not rejected by the removed freshness gate +- **WHEN** a complete mapped trajectory reaches the coordinator boundary with a non-fresh manipulation start-state observation +- **THEN** the runtime SHALL still submit the coordinator request +- **AND** it SHALL resolve the result according to coordinator acceptance, task status, cancellation, or safety uncertainty diff --git a/openspec/changes/extract-manipulation-execution-runtime/specs/manipulation-lifecycle-runtime/spec.md b/openspec/changes/extract-manipulation-execution-runtime/specs/manipulation-lifecycle-runtime/spec.md new file mode 100644 index 0000000000..89f1ff09b7 --- /dev/null +++ b/openspec/changes/extract-manipulation-execution-runtime/specs/manipulation-lifecycle-runtime/spec.md @@ -0,0 +1,96 @@ +## ADDED Requirements + +### Requirement: Authoritative lifecycle state +The manipulation execution runtime SHALL own the authoritative lifecycle state and SHALL expose only `IDLE`, `PLANNING`, `READY`, `DISPATCHING`, `RUNNING`, `CANCELLING`, and `FAULT` as lifecycle states. + +The runtime SHALL publish an atomic status snapshot containing the lifecycle state and the currently relevant plan, attempt, task, and fault information. Readers SHALL observe one internally consistent snapshot rather than independently read lifecycle flags. + +#### Scenario: Snapshot reports a coherent active execution +- **WHEN** a plan has been accepted for execution and its coordinator task is active +- **THEN** one status snapshot SHALL report `RUNNING` together with the associated plan/attempt and remote task identity +- **AND** the snapshot SHALL NOT combine fields from different plans or attempts + +#### Scenario: Invalid lifecycle transition is rejected +- **WHEN** an event requests a transition that is not legal from the current lifecycle state +- **THEN** the runtime SHALL leave the lifecycle state unchanged +- **AND** it SHALL report the rejected transition without silently changing an execution flag + +### Requirement: Legal lifecycle transitions +The runtime SHALL permit planning only from `IDLE` or `READY`, transitioning to `PLANNING`. A successful planning result SHALL transition to `READY`; a planning failure SHALL return to `IDLE` unless the runtime is already faulted. Dispatch of a selected ready plan SHALL transition through `DISPATCHING` and, after accepted dispatch, to `RUNNING`. A completed or cancelled attempt SHALL transition to `IDLE`. Cancellation of an active attempt SHALL transition to `CANCELLING` before reaching `IDLE` or `FAULT`. A safety-uncertain condition SHALL transition to and latch `FAULT`. + +The runtime SHALL NOT transition out of `FAULT` except through reset reconciliation that proves the remote task state is safe. + +#### Scenario: Replacement planning discards a ready plan +- **WHEN** the runtime is `READY` and a new planning request is accepted +- **THEN** the lifecycle SHALL become `PLANNING` +- **AND** the previous ready plan SHALL be discarded +- **AND** only the newly produced plan MAY become `READY` + +#### Scenario: Completed execution returns to idle +- **WHEN** the active attempt is physically reported complete +- **THEN** the runtime SHALL consume the attempt and transition to `IDLE` +- **AND** its next snapshot SHALL contain no executable ready plan + +### Requirement: Planning and single-use ready plans +The runtime SHALL serialize planning with dispatch, cancellation, reset, and status polling. It SHALL maintain at most one ready plan. Planning SHALL produce a plan with a unique plan identity and SHALL make it executable only after the planning operation succeeds. Dispatch SHALL consume the selected plan exactly once; a plan SHALL NOT be dispatched again after dispatch acceptance, completion, cancellation, or failure. + +The runtime SHALL reject planning and dispatch requests while an attempt is actively executing remotely (`DISPATCHING`, `RUNNING`, or `CANCELLING`), except that cancellation and reset operations SHALL remain available through their defined serialized paths. + +An execution request rejected because an attempt is active SHALL complete with a terminal `REJECTED` outcome, SHALL NOT be queued, and SHALL NOT transition the active runtime to `FAULT`. + +#### Scenario: Dispatch consumes a plan +- **WHEN** a ready plan is dispatched and the coordinator accepts the task +- **THEN** the plan SHALL no longer be available as a ready plan +- **AND** a repeated dispatch request using that plan identity SHALL be rejected + +#### Scenario: Planning is blocked during remote execution +- **WHEN** the runtime is `RUNNING` or `CANCELLING` +- **THEN** a planning request SHALL be rejected without creating or replacing a plan +- **AND** no trajectory execution request SHALL be issued + +#### Scenario: Execution is rejected during remote execution +- **WHEN** the runtime is `DISPATCHING`, `RUNNING`, or `CANCELLING` and a second execution request is received +- **THEN** the request SHALL complete with a terminal `REJECTED` outcome +- **AND** it SHALL remain unqueued and issue no coordinator execute request +- **AND** the active lifecycle state SHALL remain unchanged + +#### Scenario: Coordinator rejection returns to idle +- **WHEN** a valid execution request is rejected by the coordinator before an active remote attempt is established +- **THEN** the request SHALL complete with a terminal `REJECTED` outcome +- **AND** the lifecycle SHALL return to `IDLE` +- **AND** the runtime SHALL NOT enter `FAULT` + +### Requirement: Dispatch acceptance is distinct from physical completion +The runtime SHALL distinguish coordinator dispatch acceptance from physical trajectory completion. Dispatch acceptance SHALL establish an active attempt and `RUNNING` state, but SHALL NOT be reported as physical completion. The runtime SHALL transition to completion only after coordinator task status or an equivalent authoritative result reports that the robot physically completed the task. + +#### Scenario: Accepted dispatch remains running +- **WHEN** the coordinator accepts a trajectory task but has not reported physical completion +- **THEN** the runtime SHALL report `RUNNING` +- **AND** it SHALL continue task monitoring +- **AND** it SHALL NOT report the manipulation as complete + +### Requirement: Remote safety uncertainty is latched +If the runtime cannot prove that a coordinator task is safely inactive or reconciled, it SHALL latch a `FAULT` state and SHALL prevent new planning or execution. The fault SHALL remain latched across subsequent polls and SHALL be cleared only by reset reconciliation that establishes a safe remote state. + +#### Scenario: Unknown remote activity faults the runtime +- **WHEN** a status, cancellation, or reset operation returns an unknown/error result that does not prove the remote task is inactive +- **THEN** the runtime SHALL transition to `FAULT` +- **AND** subsequent plan and execute requests SHALL be rejected + +#### Scenario: Safe reset clears a latched fault +- **WHEN** reset reconciliation confirms that no coordinator task is active and the remote state is safe +- **THEN** the runtime SHALL clear the latched fault +- **AND** it SHALL transition to `IDLE` with no ready plan or active attempt + +### Requirement: No manipulation-layer start-state or partial-robot execution contract +The manipulation runtime SHALL NOT validate trajectory start-state freshness as a prerequisite for manipulation execution. It SHALL NOT accept a partial-robot execution selection or a `robot_name` execution argument; manipulation execution SHALL target every robot trajectory represented by the generated plan through the coordinator boundary. + +#### Scenario: Stale manipulation start state does not block dispatch +- **WHEN** a ready manipulation plan is otherwise valid but its trajectory start state is not fresh according to the removed manipulation-layer check +- **THEN** the runtime SHALL not reject the plan solely for that freshness condition +- **AND** execution SHALL proceed subject to coordinator acceptance and safety rules + +#### Scenario: Partial robot selection is rejected +- **WHEN** a caller attempts manipulation execution with a partial-robot selection or robot-specific execution argument +- **THEN** the request SHALL be rejected +- **AND** no coordinator trajectory request SHALL be issued diff --git a/openspec/changes/extract-manipulation-execution-runtime/tasks.md b/openspec/changes/extract-manipulation-execution-runtime/tasks.md new file mode 100644 index 0000000000..179c30cf00 --- /dev/null +++ b/openspec/changes/extract-manipulation-execution-runtime/tasks.md @@ -0,0 +1,71 @@ +## 1. Characterize existing behavior and update contracts + +- [x] 1.1 Inventory current manipulation lifecycle states, execution flags, coordinator RPC call sites, planning APIs, and status consumers; see the historical baseline below. +- [x] 1.2 Add characterization tests for current planning, dispatch acceptance, physical completion, cancellation, reset, and coordinator uncertainty behavior. +- [x] 1.3 Define the runtime event, normalized outcome, operation/attempt identity, and atomic snapshot APIs in the manipulation execution test contract. +- [x] 1.4 Update API tests and call sites to use `IDLE`, `PLANNING`, `READY`, `DISPATCHING`, `RUNNING`, `CANCELLING`, and `FAULT` only. +- [x] 1.5 Update API tests to reject `robot_name` and all partial-robot execution selections. + +## 2. Implement lifecycle state, reducer, and snapshots + +- [x] 2.1 Add the lifecycle state model with exactly the seven public states and explicit legal-transition validation. +- [x] 2.2 Add runtime event types for planning, plan replacement, execute, cancel, reset, coordinator responses, task observations, and polling. +- [x] 2.3 Implement the formal reducer so invalid transitions leave state unchanged and produce a rejected-transition result. +- [x] 2.4 Implement atomic status snapshots containing consistent state, plan, operation, attempt, task, and fault data. +- [x] 2.5 Add reducer tests for legal transitions, rejected transitions, coherent snapshots, and dispatch acceptance versus physical completion. + +## 3. Introduce the single-owner execution runtime + +- [x] 3.1 Add the internal runtime owned by `ManipulationModule` with a single-owner event queue and serialized processing loop. +- [x] 3.2 Move coordinator client ownership, ready-plan storage, operation IDs, attempts, task knowledge, and fault details into the runtime. +- [x] 3.3 Ensure execute, cancel, reset, planning results, coordinator polling, and snapshot publication all use the serialized runtime path. +- [x] 3.4 Add runtime tests proving that concurrent requests cannot create overlapping execution attempts or inconsistent snapshots. +- [x] 3.5 Add tests proving that planning and execution work are blocked while a remote attempt is active, while cancellation and reset remain available. + +## 4. Move planning results and ready-plan lifecycle + +- [x] 4.1 Route planning start requests from `ManipulationModule` to the runtime and transition accepted requests to `PLANNING`. +- [x] 4.2 Route planning success and failure results to runtime events with unique plan identities and normalized outcomes. +- [x] 4.3 Implement atomic replacement so a new successful planning request discards the prior ready plan before exposing the new plan as `READY`. +- [x] 4.4 Enforce validation that every robot trajectory represented by a plan is present, and reject incomplete plans without creating executable ready state. +- [x] 4.5 Implement single-use plan consumption before dispatch and reject repeated dispatch of a consumed plan identity. +- [x] 4.6 Add tests for planning from `IDLE` and `READY`, replacement, planning failure, plan identity, and single-use dispatch semantics. + +## 5. Implement mapping, dispatch, and task knowledge + +- [x] 5.1 Build and validate the coordinator-name inverse mapping once from canonical configuration references during runtime initialization. +- [x] 5.2 Add the global-plan to local-coordinator-request mapping and normalized coordinator-result mapping while preserving plan and attempt identities. +- [x] 5.3 Register active plan, operation, attempt, and task-correlation context before invoking the coordinator execute RPC. +- [x] 5.4 Implement the single complete-robot coordinator dispatch path through `DISPATCHING` to `RUNNING` on acceptance. +- [x] 5.5 Record coordinator task knowledge and reject stale or mismatched responses without attributing them to the active attempt. +- [x] 5.6 Add tests for mapping validation, complete trajectory requests, pre-registration, immediate responses, accepted-but-not-complete tasks, and normalized terminal results. + +## 6. Implement cancellation, polling, fault, and reset safety + +- [x] 6.1 Implement serialized status polling and task reconciliation for active attempts, including physical completion detection. +- [x] 6.2 Implement cancellation through `CANCELLING`, reporting cancellation only after coordinator cancellation or inactivity is confirmed. +- [x] 6.3 Normalize unknown, malformed, unavailable, and contradictory coordinator results as safety uncertainty and latch `FAULT`. +- [x] 6.4 Ensure latched `FAULT` rejects new planning and execution and remains latched across subsequent polls. +- [x] 6.5 Implement reset reconciliation that clears plan, attempt, queue, and fault state only after remote inactivity is proven, returning to `IDLE`. +- [x] 6.6 Ensure reset remains retryable and does not issue new trajectory execution while remote activity is active or uncertain. +- [x] 6.7 Add tests for rejection of execution during an active attempt (without queueing), terminal coordinator rejection returning to `IDLE`, active cancellation confirmation, unknown task states, fault latching, and safe/unsafe reset. + +## 7. Wire module, operator, and UI compatibility + +- [x] 7.1 Replace direct lifecycle assignments and execution flags in `ManipulationModule` with runtime event submission and snapshot consumption. +- [x] 7.2 Update manipulation operators and callers for the new execution API, normalized outcomes, and removal of robot-specific selection. +- [x] 7.3 Update UI/state adapters to consume atomic snapshots and map all seven public states without legacy state names. +- [x] 7.4 Add compatibility tests covering module status publication, operator actions, UI state mapping, fault visibility, and reset behavior. + +## 8. Remove legacy paths and validate the migration + +- [x] 8.1 Remove legacy module-level lifecycle bookkeeping, duplicate coordinator access, and active `ManipulationModule` direct trajectory execution paths. Do not move or remove `dimos/manipulation/control/coordinator_client.py`: it is standalone diagnostic/operator CLI direct-RPC tooling, not `ManipulationModule` execution; changing it is a separate compatibility-scoped change. +- [x] 8.2 Remove manipulation-layer trajectory start-state freshness validation without adding a replacement gate in this change. +- [x] 8.3 Remove partial-robot execution handling and the `robot_name` argument from implementation, adapters, and tests. +- [x] 8.4 Run targeted lifecycle runtime, coordinator execution, manipulation module, operator, and UI compatibility tests. +- [x] 8.5 Run formatting, lint, type checks, and the relevant OpenSpec or repository validation checks. Verified with `uv run --with mypy mypy dimos/manipulation`: 54 source files, no issues found. +- [x] 8.6 Review test coverage for every lifecycle state, transition, safety fault, reset path, plan replacement, and single-use invariant. + +## Verification note + +Self-hosted and Drake-dependent verification is environment-limited for this change and is not claimed unless run on the corresponding self-hosted/Drake environment. Local artifact review and environment-independent checks may be run separately; unavailable self-hosted/Drake checks must be reported as not run rather than inferred as passing. diff --git a/openspec/changes/simplify-manipulation-execution-runtime/.openspec.yaml b/openspec/changes/simplify-manipulation-execution-runtime/.openspec.yaml new file mode 100644 index 0000000000..c0a8162549 --- /dev/null +++ b/openspec/changes/simplify-manipulation-execution-runtime/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-07-21 diff --git a/openspec/changes/simplify-manipulation-execution-runtime/design.md b/openspec/changes/simplify-manipulation-execution-runtime/design.md new file mode 100644 index 0000000000..2e28874d8d --- /dev/null +++ b/openspec/changes/simplify-manipulation-execution-runtime/design.md @@ -0,0 +1,111 @@ +# Design: Simplified Manipulation Execution Runtime + +## Context + +The just-added manipulation runtime is approximately 2,084 LOC. It currently spreads coordination across a large reducer, poller and action-specific threads, deadline/reset/gateway-close threads, duplicated flags and histories, and module-side timeout reconciliation. That structure obscures the single-owner lifecycle contract even though the required behavior is already defined and must remain unchanged. + +The refactor is an internal, behavior-preserving rewrite with a compact initial-runtime scope. The public `ManipulationModule` API and all safety outcomes remain the compatibility boundary. Blocking coordinator RPCs must not block the owner loop, while lifecycle state must still have one writer and deterministic ordering. + +## Goals / Non-Goals + +### Goals + +- Preserve four explicit seams and the achieved policy/effects/auxiliary/passive module separations: + - plan materialization: 220–280 LOC; + - coordinator gateway: 80–120 LOC; + - models: 100–150 LOC in the final simplified architecture; + - runtime owner: no hard LOC threshold; keep responsibilities reviewable without splitting a required public façade from its private owner. +- Preserve all existing public `ManipulationModule` semantics, including serialized lifecycle and ready-plan behavior. +- Preserve minimal operation/task/method RPC correlation, sequential multi-task dispatch, compensation only for known or potentially-active tasks, physical completion polling, UNKNOWN fail-closed safety faults, and truthful shutdown. +- Keep lifecycle state and mutable ownership in one runtime owner loop. +- Make blocking RPC completion explicit through futures and result waits. +- Enforce a module-configured `physical_operation_timeout: float = 60.0` safety deadline for operations that remain `RUNNING` indefinitely. +- Accept an additive runtime `monotonic_clock` seam defaulting to `time.monotonic` for deterministic deadline behavior. +- Keep the initial runtime compact: retryable reset generations, rich action/reset generations, exact result-history depth, and eventual post-deadline shutdown success are deferred rather than implemented promises. +- Define success by reduced duplicated responsibility, tested policy/effect/auxiliary/passive boundaries, and no observable behavior regression; LOC is an observation, not a release gate. + +### Non-Goals + +- No redesign of planning algorithms, trajectory data, coordinator protocol, or public manipulation semantics. +- No freshness-based or robot-name-based execution selection; neither is a replacement safety gate. +- No new poller, per-action deadline, reset, gateway-close, or module-side timeout-reconciliation thread. +- No UI/layout work, unrelated source cleanup, or changes outside the requested change during implementation. + +## Decisions + +### Four narrow implementation seams + +Plan materialization converts planner output into a complete executable plan and owns no lifecycle state. Models contain immutable records/enums/results and the single mutable owner-state record. The gateway is a thin coordinator RPC adapter. The actor owns lifecycle handlers, minimal correlation, dispatch sequencing, polling, conditional compensation, and shutdown. + +The alternative of retaining the existing reducer and extracting helpers was rejected: it would preserve duplicate state representations and make ownership and transition behavior difficult to verify. The alternative of splitting into many event-specific services was rejected because it would recreate cross-thread ordering problems. + +### Fixed daemon blocking-RPC executor + +All potentially blocking coordinator RPCs run on exactly four daemon executor workers. Each submission registers operation, task, and method correlation before the effect and returns a future; the owner loop consumes completed results at defined handler boundaries. RPC deadlines are represented as future/result outcomes and safety faults, not as per-action timer threads. Unresolved work blocks replacement execution; no rich action/reset generation model is required. + +The executor is fixed for the runtime lifetime and is shut down truthfully: one shutdown initiator seals admission, shutdown stops accepting work, observer-only waits observe known futures, and incomplete/unknown remote work is reported rather than pretending it is complete. Concurrent shutdown callers are rejected. A shutdown deadline does not promise eventual success after the deadline; daemon workers prevent process teardown from being held hostage, but do not weaken fail-closed outcomes. + +Admission has one exception to strict owner-only writes: a caller-owned latch may seal gateway admission before the owner processes shutdown. That latch is an admission gate only; all `RuntimeContext` writes, including shutdown transitions, remain owner-loop writes. + +### Owner-loop timer scan + +One daemon owner loop serializes requests, completed future results, and a monotonic timer scan. The owner timer is authoritative for physical-completion polling, deadline observation, cancellation reconciliation, deferred reset handling, and shutdown observation. It does not perform RPCs itself; it schedules blocking work on the fixed executor and applies results only after correlation checks. + +This replaces poller, per-action deadline, reset, and gateway-close threads. A single timer scan is preferred over independent timers because it preserves ordering and avoids races between timeout, cancellation, reset, and shutdown. + +### Monotonic clock seam and public polling + +The runtime constructor keeps a compatible callable injection named `monotonic_clock`, defaulting to `time.monotonic`. It wraps that callable in a dedicated `ValidatedMonotonicClock` source wrapper. The wrapper serializes access, validates that every sample is finite and nondecreasing, and is the only clock exposed to runtime deadline code. An invalid sample violates the injected dependency contract and raises a dependency-contract error; it is not converted into lifecycle `FAULT` state or a remote safety outcome. + +Every deadline creation, deadline comparison, and relative wait calculation routes through `ValidatedMonotonicClock`. Wall-clock time, mixed clock domains, and ad hoc elapsed-time calculations are not allowed for runtime deadlines. The compatible constructor surface remains callable injection/default behavior; validation and serialization are supplied by the wrapper rather than by changing callers. + +The owner timer is the authoritative deadline processor. Public immediate polling is supported only while an operation is `RUNNING`; it does not provide a special narrow timer-precheck guarantee and does not promise `STATUS` in reset, shutdown, gateway, or other non-running states. + +### Focused handlers and one mutable state + +Phase 4 is decomposed into reviewable slices. Slice 1 adds the `_set_state` write +wrapper and extracts only planning, dispatch-admission, and lifecycle-deadline +branches. It intentionally retains the immutable context representation and all +Phase 3 scheduler/executor mechanics. Later slices will migrate state ownership, +remove redundant bookkeeping, and extract task/cancellation/reset/shutdown +handlers. + +Handlers are grouped by concern: planning/ready-plan materialization, execute admission, sequential task dispatch, task-result reconciliation, physical-operation deadline expiry, conditional cancellation/compensation, deferred reset handling, and shutdown. Each handler mutates only the owner-owned state record and publishes the existing snapshot/result contract. Correlation IDs are limited to operation/task/method RPC identity; stale or unresolved results fail closed and cannot replace active work. Result retention is small and current/recent, with no exact history-depth promise. + +The giant reducer, redundant flags, history/alias fields, and module-side timeout reconciliation are removed. Public methods remain adapters that submit owner-loop requests and wait for the corresponding future/result, rather than independently reconciling lifecycle state. This does not require a separate public façade/private owner split: the required boundary is behavioral ownership and tested effects, not a numeric actor size or an additional public layer. + +The physical-operation deadline starts when the operation enters `RUNNING` and is 60 seconds by default. It is not planning time, an action/RPC deadline, a caller observation timeout, or trajectory duration. On expiry the owner schedules cancellation and reconciliation; only authoritative reconciliation can yield completion, otherwise the owner publishes safe terminal `FAILED` or a correlated `FAULT`. The module never reconciles this deadline. + +### Sequential multi-task dispatch and compensation + +The actor dispatches the complete plan one coordinator task at a time. It registers operation, task, and method correlation before each RPC. If a later task fails after an earlier task is known or potentially active, compensation is invoked in order; if activity is UNKNOWN, the outcome remains fail-closed and replacement is blocked. No second operation is admitted while unresolved work remains. + +### Reset and shutdown scope + +Reset generations that can safely be retried are deferred to a later scope. An unsafe or timed-out reset remains `FAULT` and is not retried in-process. Shutdown has one initiator; concurrent callers are rejected, observer-only waits do not mutate lifecycle state, and expiry reports unresolved work without promising eventual success. Admission sealing is the sole caller-side latch exception before owner shutdown processing. + +## Risks / Trade-offs + +- **Blocking RPC exceeds its expected duration** → represent the result as an explicit deadline/unknown safety outcome, let the owner loop continue its timer scan, and preserve reset proof requirements. +- **A stale or unresolved future completes after reset or shutdown** → retain minimal operation/task/method correlation until classification; never apply it to new work, and fail closed when activity is unknown. +- **Daemon executor hides unfinished work at process exit** → truthful shutdown reports unresolved work and does not claim cancellation or completion without proof. +- **LOC targets encourage over-compression** → treat measurements as maintainability evidence only; prioritize reduced duplicated responsibility, tested boundaries, and the required safety matrix over an arbitrary actor-size gate. +- **Phase 1 faithful extraction exceeds the final models budget** → retain the maintainable 190-LOC models baseline during pure extraction; defer reduction to the later simplification phases and do not compress records unsafely. +- **Changing test seams masks behavior changes** → run characterization and black-box compatibility tests before and after each extraction; decouple white-box tests only after equivalent observable coverage exists. +- **A clock seam hides an unsafe clock** → put finite/nondecreasing validation and serialized access in `ValidatedMonotonicClock`; treat invalid samples as dependency-contract failures, not lifecycle faults, and test the wrapper independently. +- **Public polling races the timer owner** → make the owner timer authoritative and limit immediate public polling to `RUNNING` operations. + +## Migration Plan + +1. Freeze characterization coverage and extract pure modules with re-exports; no behavior change. +2. Introduce futures/result waits behind the existing gateway and verify correlation and outcomes. +3. Replace thread topology with the fixed executor and owner-loop timer scan. +4. Move logic into focused handlers and one mutable owner state, deleting redundant reducer data. +5. Clean module/API adapters while preserving public semantics and removing freshness/robot-name selection paths. +6. Keep only compact seam tests and the required safety matrix; defer rich reset-generation/history coverage. + +Rollback at each phase is a file-level revert before the next phase; no data migration or deployment migration is required. + +## Open Questions + +- None blocking implementation. Existing polling cadence, RPC deadline values, and reset proof criteria remain authoritative and must be carried over unchanged. diff --git a/openspec/changes/simplify-manipulation-execution-runtime/proposal.md b/openspec/changes/simplify-manipulation-execution-runtime/proposal.md new file mode 100644 index 0000000000..3e153556c3 --- /dev/null +++ b/openspec/changes/simplify-manipulation-execution-runtime/proposal.md @@ -0,0 +1,35 @@ +## Why + +The newly added manipulation execution runtime preserves the required safety and lifecycle behavior, but its 2,084 lines concentrate planning materialization, RPC adaptation, mutable lifecycle state, polling, deadlines, reset handling, and shutdown mechanics in one over-engineered implementation. Simplifying those seams now will make the behavior easier to audit and maintain without changing the public `ManipulationModule` contract or safety semantics. + +## What Changes + +- Retain and document the achieved separations between policy, effects, auxiliary, and passive modules, with explicit boundaries for plan materialization, the coordinator gateway, runtime models, and the runtime owner. +- Replace per-action/poller/deadline/reset/gateway-close threads with one fixed daemon blocking-RPC executor and one owner-loop timer scan. +- Replace the giant reducer and redundant flags/history/aliases with focused handlers over one owner-owned mutable runtime state. +- Use futures and observer-only waits for blocking coordinator operations, preserving minimal operation/task/method correlation and serialized lifecycle decisions. Unresolved work blocks replacement execution. +- Enforce a runtime-owned `physical_operation_timeout: float = 60.0` safety deadline for operations that remain `RUNNING`, with owner-scheduled cancellation/reconciliation and safe terminal failure or correlated fault. +- Enforce strict owner-only `RuntimeContext` writes; permit only a caller-owned admission latch to seal gateway admission before owner shutdown processing. +- Use a dedicated serialized `ValidatedMonotonicClock` wrapper behind compatible `monotonic_clock` callable injection/default behavior; reject invalid samples as dependency-contract failures rather than lifecycle faults. +- Preserve sequential multi-task dispatch, compensation only for known or potentially-active tasks, physical-completion polling, UNKNOWN fail-closed safety outcomes, the 60-second physical deadline, admission sealing, and truthful shutdown. +- Defer retryable reset generations; unsafe or timed-out reset remains `FAULT` with no in-process retry. Retain only a small current/recent result set rather than promising exact history depth. +- Permit one shutdown initiator only: concurrent shutdown callers are rejected, and no eventual success is promised after the shutdown deadline. +- Make the owner timer authoritative. Public immediate polling is only for `RUNNING` operations and has no special narrow timer-precheck guarantee. +- **BREAKING** Remove implementation-only freshness/robot-name execution selection paths; retain the existing public manipulation semantics without freshness- or robot-name-based execution selection. +- Keep all changes confined to the new internal runtime structure; do not add execution selection based on freshness or robot name. +- Do not impose a hard runtime-actor LOC threshold or require a public-façade/private-owner split; judge the refactor by reduced duplicated responsibility, tested boundaries, and unchanged behavior. + +## Capabilities + +### New Capabilities +- `manipulation-lifecycle-runtime`: Compact fail-closed lifecycle ownership, minimal correlation, safety, reset, and shutdown semantics for the simplified runtime. +- `coordinator-trajectory-execution`: Compact gateway, sequential dispatch, conditional compensation, polling, and blocking-RPC execution semantics. + +### Modified Capabilities + +## Impact + +- Affects only the implementation and tests of the recently added manipulation execution runtime and its `ManipulationModule` integration. +- Introduces no external dependency and no source changes outside the eventual runtime refactor. +- Preserves public `ManipulationModule` methods, lifecycle outcomes, serialized ready-plan behavior, and operator-visible safety results. +- Success is reduced duplicated responsibility across the existing policy/effects/auxiliary/passive separations, tested ownership and effect boundaries, and no behavior regression in the safety and lifecycle contract. diff --git a/openspec/changes/simplify-manipulation-execution-runtime/specs/coordinator-trajectory-execution/spec.md b/openspec/changes/simplify-manipulation-execution-runtime/specs/coordinator-trajectory-execution/spec.md new file mode 100644 index 0000000000..6aca794f0a --- /dev/null +++ b/openspec/changes/simplify-manipulation-execution-runtime/specs/coordinator-trajectory-execution/spec.md @@ -0,0 +1,51 @@ +## ADDED Requirements + +### Requirement: Use a fixed blocking-RPC executor and owner timer scan +The runtime SHALL execute blocking coordinator RPCs through one fixed executor with exactly four daemon workers returning futures. One owner-loop timer scan SHALL be authoritative for polling, deadline observation, cancellation reconciliation, deferred reset handling, and shutdown observation. The runtime SHALL NOT create poller, per-action deadline, reset, gateway-close, or module-side timeout-reconciliation threads. + +The runtime SHALL include the module-configured `physical_operation_timeout: float = 60.0` safety deadline for a coordinator state that remains `RUNNING` indefinitely. Its expiry SHALL be handled by the owner loop, which schedules cancellation/reconciliation and reports safe terminal `FAILED` or a correlated `FAULT`; it SHALL not be confused with planning, per-RPC action, caller wait, or trajectory-duration timeouts. + +The runtime constructor SHALL retain compatible callable injection through `monotonic_clock`, defaulting to `time.monotonic`, and SHALL route it through a dedicated serialized `ValidatedMonotonicClock` wrapper. The wrapper SHALL validate finite, nondecreasing samples and SHALL be the sole source used for deadline creation, deadline comparison, and relative wait calculations. An invalid sample SHALL fail the clock dependency contract, not create lifecycle `FAULT` state. + +#### Scenario: RPC does not block lifecycle ownership +- **WHEN** a coordinator RPC blocks +- **THEN** the RPC SHALL run on the fixed executor +- **AND** the owner loop SHALL continue processing its serialized timer/request path + +#### Scenario: Indefinite running state reaches the physical deadline +- **WHEN** a coordinator task remains `RUNNING` beyond `physical_operation_timeout` +- **THEN** the owner loop SHALL enqueue cancellation and reconciliation +- **AND** the eventual result SHALL be terminal `FAILED` or a correlated `FAULT` +- **AND** duration alone SHALL never be treated as physical completion + +### Requirement: Preserve sequential dispatch and compensation +The gateway/actor boundary SHALL preserve sequential multi-task dispatch for complete plans. It SHALL register minimal operation, task, and method correlation before each RPC effect. It SHALL invoke compensation only when earlier tasks are known or potentially active; UNKNOWN activity SHALL fail closed and block replacement execution. + +#### Scenario: Later task failure compensates earlier work +- **WHEN** a multi-task plan establishes an earlier task and a later task is rejected or becomes uncertain +- **THEN** the runtime SHALL preserve the existing compensation sequence +- **AND** the normalized operation result SHALL reflect compensation success or fail-closed safety uncertainty + +### Requirement: Preserve physical completion polling +Dispatch acceptance SHALL remain distinct from physical completion. The timer scan SHALL poll active coordinator tasks and transition to the existing terminal outcome only after authoritative completion or safe cancellation/inactivity proof. + +#### Scenario: Accepted task remains active +- **WHEN** a coordinator accepts a task without reporting physical completion +- **THEN** the runtime SHALL retain the active correlated task +- **AND** it SHALL report the existing running state rather than completion + +#### Scenario: Immediate public polling is running-only +- **WHEN** public `poll` is called +- **THEN** it SHALL issue immediate physical polling only for a `RUNNING` operation +- **AND** the owner timer SHALL process deadlines and other timer work +- **AND** `STATUS` SHALL not be guaranteed in reset, shutdown, gateway, or other non-running states + +### Requirement: Keep the gateway thin and synchronous at its boundary +The coordinator gateway SHALL contain only request/result adaptation and RPC invocation required by the actor. It SHALL not own lifecycle flags, timers, retries, or independent close state; futures and observer-only waits SHALL carry completion back to the owner loop. Unresolved work SHALL block replacement execution. + +The policy, effects, auxiliary, and passive module boundaries SHALL remain explicit and testable. The runtime SHALL NOT require a numeric actor-size limit or an additional public-façade/private-owner layer; behavioral ownership and effect isolation are the contract. + +#### Scenario: Gateway result is applied by owner +- **WHEN** the gateway future completes with acceptance, rejection, cancellation, completion, or uncertainty +- **THEN** the owner loop SHALL normalize and apply the result using minimal operation/task/method correlation +- **AND** the gateway SHALL not publish an independent lifecycle state diff --git a/openspec/changes/simplify-manipulation-execution-runtime/specs/manipulation-lifecycle-runtime/spec.md b/openspec/changes/simplify-manipulation-execution-runtime/specs/manipulation-lifecycle-runtime/spec.md new file mode 100644 index 0000000000..9447383be6 --- /dev/null +++ b/openspec/changes/simplify-manipulation-execution-runtime/specs/manipulation-lifecycle-runtime/spec.md @@ -0,0 +1,71 @@ +## ADDED Requirements + +### Requirement: Preserve serialized lifecycle ownership +The simplified runtime SHALL preserve the existing public `ManipulationModule` lifecycle semantics and SHALL have exactly one owner loop that mutates lifecycle, ready-plan, operation, task, fault, and shutdown state. Any attempt/action/reset bookkeeping is compact local correlation, not a rich generation model. + +The owner SHALL be the only writer of `RuntimeContext` and its mutable state. A caller-owned admission latch MAY be written outside the owner only to seal gateway admission before the owner processes shutdown; it SHALL not mutate lifecycle, operation, task, fault, or shutdown state. + +The runtime SHALL use one owner loop, exactly four daemon RPC workers, and observer-only caller waits. It SHALL retain only a small current/recent result set; no exact result-history depth is part of the contract. + +The policy, effects, auxiliary, and passive module separations SHALL remain explicit and testable. Success SHALL be assessed by reduced duplicated responsibility, tested ownership/effect boundaries, and no observable behavior regression, not by a hard actor LOC threshold or a required public-façade/private-owner split. + +The runtime constructor SHALL retain compatible callable injection through `monotonic_clock`, defaulting to `time.monotonic`, and SHALL wrap it in a dedicated `ValidatedMonotonicClock`. The wrapper SHALL serialize source access and SHALL accept only finite, nondecreasing samples. All runtime deadline creation, deadline comparison, and relative wait calculations SHALL route through the wrapper. An invalid sample SHALL be a dependency-contract failure, not a lifecycle `FAULT`, and SHALL not be converted into a remote safety outcome. + +#### Scenario: Concurrent lifecycle requests are serialized +- **WHEN** execute, cancel, reset, planning results, and timer events arrive concurrently +- **THEN** the owner loop SHALL apply them in one deterministic order +- **AND** no caller-visible snapshot SHALL combine state from different operations + +### Requirement: Preserve correlation and normalized outcomes +The runtime SHALL retain minimal operation, task, and method correlation through planning, sequential dispatch, polling, cancellation, compensation, deadlines, faults, and shutdown. Registration SHALL occur before each RPC effect. Stale, mismatched, or unresolved results SHALL NOT mutate the current operation, and unresolved work SHALL block replacement execution; rich action/reset generations are out of scope. + +#### Scenario: Late result cannot affect a new operation +- **WHEN** an RPC future from a prior operation completes after reset or a new operation begins +- **THEN** the owner loop SHALL classify it using its original correlation +- **AND** it SHALL not apply the result to the new operation + +### Requirement: Preserve safety and shutdown semantics +Unknown, malformed, unavailable, contradictory, or deadline-uncertain remote results SHALL fail closed and preserve the existing safety-fault behavior. Unsafe or timed-out reset SHALL remain `FAULT` with no in-process retry; retryable reset generations SHALL be deferred. Shutdown SHALL report unresolved work truthfully. + +#### Scenario: Deadline uncertainty faults safely +- **WHEN** a blocking coordinator result cannot be obtained before its established deadline +- **THEN** the runtime SHALL preserve the existing safety fault outcome +- **AND** it SHALL not report completion or safe reset without proof + +#### Scenario: Unsafe reset remains faulted +- **WHEN** reset is requested while remote inactivity is active, uncertain, or timed out +- **THEN** the runtime SHALL retain the fault/safety condition and issue no new execution +- **AND** it SHALL not retry the reset generation in-process + +#### Scenario: Shutdown is truthful +- **WHEN** shutdown begins with an unresolved RPC or remote attempt +- **THEN** the runtime SHALL not report that work as completed or cancelled without proof +- **AND** the public shutdown result SHALL reflect the unresolved condition + +#### Scenario: Shutdown has one initiator +- **WHEN** a shutdown is already being processed by one caller +- **THEN** a concurrent shutdown caller SHALL be rejected +- **AND** expiry of the shutdown deadline SHALL report unresolved work without promising eventual success + +### Requirement: Enforce a runtime-owned physical operation deadline +Each module-configured runtime SHALL expose `physical_operation_timeout: float = 60.0`. The owner SHALL start this safety deadline when a physical operation enters `RUNNING` and SHALL owner-schedule cancellation and reconciliation when it expires. This deadline is distinct from `planning_timeout`, per-RPC `action_timeout`, caller observation wait timeouts, and trajectory duration. + +#### Scenario: Physical operation exceeds its safety deadline +- **WHEN** a correlated operation remains `RUNNING` until `physical_operation_timeout` expires +- **THEN** the owner SHALL schedule cancellation and reconciliation +- **AND** it SHALL produce a safe terminal `FAILED` result or a correlated `FAULT` +- **AND** it SHALL not infer completion from trajectory duration or permit module-side reconciliation + +#### Scenario: Immediate polling is running-only +- **WHEN** the public `poll` method is called +- **THEN** it SHALL perform immediate physical polling only for a `RUNNING` operation +- **AND** the owner timer SHALL remain authoritative for deadlines, reset, shutdown, and other timer processing +- **AND** no special precheck SHALL guarantee `STATUS` during reset, shutdown, gateway, or other non-running states + +### Requirement: No execution selection by freshness or robot name +The simplified runtime SHALL preserve the contract that execution is not selected by freshness or robot name. It SHALL dispatch the complete plan represented by the planning groups and SHALL reject partial or robot-name-selected execution without issuing coordinator work. + +#### Scenario: Complete plan dispatch ignores freshness selection +- **WHEN** a complete ready plan reaches execution with a non-fresh manipulation-layer observation +- **THEN** the runtime SHALL not reject it solely for that observation +- **AND** it SHALL continue through normal coordinator acceptance and safety handling diff --git a/openspec/changes/simplify-manipulation-execution-runtime/tasks.md b/openspec/changes/simplify-manipulation-execution-runtime/tasks.md new file mode 100644 index 0000000000..7469cabc28 --- /dev/null +++ b/openspec/changes/simplify-manipulation-execution-runtime/tasks.md @@ -0,0 +1,71 @@ +## 1. Pure extraction and re-exports + +Baseline before Phase 1: `execution_runtime.py` was 2,084 LOC; topology/materialization occupied lines 82–337 and public models/results occupied lines 342–469. The public boundary was the `dimos.manipulation.execution_runtime` import path used by the manipulation module, tests, and helpers. Preserve the achieved policy/effects/auxiliary/passive module separations; LOC is recorded only as maintainability evidence, not as a hard actor gate or a requirement for a public-façade/private-owner split. + +- [x] 1.1 Record current public `ManipulationModule` behavior and LOC/boundary baselines without changing source behavior. +- [x] 1.2 Extract pure plan materialization into a 220–280 LOC internal module with no lifecycle ownership. +- [x] 1.3 Extract a faithful models and correlation/result baseline into an internal module (190 LOC in Phase 1; the 100–150 LOC target remains a later simplification target, not a reason for unsafe compression). +- [x] 1.4 Add compatibility re-exports so existing runtime imports and public module semantics remain unchanged. +- [x] 1.5 **Validation gate:** run plan/materialization, model, import, and existing manipulation characterization tests; no behavior changes are allowed. + +## 2. Futures and explicit result waits + +- [x] 2.1 Extract the coordinator gateway into an 80–120 LOC adapter containing only RPC request/result conversion. +- [x] 2.2 Route blocking coordinator calls through retained futures while registering minimal operation/task/method correlation before each RPC effect. The existing auxiliary-thread topology remains intentionally unchanged until Phase 3. +- [x] 2.3 Replace action-thread completion signaling at the owner boundary with explicit correlated completion results and observer-only waits with normalized outcomes. +- [x] 2.4 **Validation gate:** verify immediate acceptance, rejection, completion, cancellation, malformed, unavailable, and stale-result cases against characterization tests. + +## 3. Fixed executor and owner-loop timer scan + +- [x] 3.1 Add one fixed blocking-RPC executor with exactly four daemon workers, bounded lifetime, and truthful unresolved-work shutdown. +- [x] 3.2 Add one owner-loop timer scan for polling, deadlines, cancellation reconciliation, deferred reset handling, and shutdown observation. +- [x] 3.3 Remove poller, per-action deadline, reset, and gateway-close threads without changing polling cadence or deadline semantics. +- [x] 3.4 Keep module-side timeout reconciliation out of the new path. +- [x] 3.5 **Validation gate:** run concurrency, deadline/UNKNOWN fail-closed, physical-completion, deferred-reset, and truthful-shutdown tests; inspect four-worker ownership. +- [x] 3.6 Add the module-configured `physical_operation_timeout: float = 60.0` safety deadline for indefinitely `RUNNING` coordinator state; owner-schedule cancellation/reconciliation on expiry and produce safe terminal `FAILED` or correlated `FAULT`. Keep it distinct from planning, per-RPC action, caller observation, and trajectory-duration timeouts. +- [x] 3.7 Add the additive compatible `monotonic_clock` callable injection, defaulting to `time.monotonic`, and route it through a dedicated serialized `ValidatedMonotonicClock` wrapper; validate finite/nondecreasing samples, treat invalid samples as dependency-contract failures rather than lifecycle `FAULT`, and route every deadline creation/comparison and relative wait calculation through the wrapper. + +## 4. Focused handlers and mutable owner state + +- [x] 4.1a Add a behavior-preserving `_set_state(**changes)` owner-state wrapper. +- [x] 4.1b Extract planning lifecycle transitions into focused owner-only handlers. +- [x] 4.1c Enforce owner-only `RuntimeContext` writes in dispatch admission and lifecycle transitions; permit only a caller-owned latch that seals gateway admission before owner shutdown processing. +- [ ] 4.1d **Deferred/out of scope:** rich handler decomposition for reset generations and action-level transitions; retain only the compact task dispatch, result reconciliation, deadline, conditional compensation, and shutdown seam. +- [x] 4.2a Introduce the mutable owner-state record. +- [x] 4.2b Migrate lifecycle state writes while preserving immutable snapshots and strict owner-only `RuntimeContext` ownership. +- [x] 4.2c Migrate minimal operation/task/method correlation state writes under the owner-only rule. +- [ ] 4.3 **Deferred/replaced:** exact history/alias removal and rich action/reset bookkeeping are out of scope; retain only small current/recent results and the existing compact policy/effects/auxiliary/passive boundaries. +- [x] 4.4 Preserve sequential multi-task dispatch, compensation only for known or potentially-active tasks, physical completion polling, and safe rejection of overlapping execution. +- [x] 4.5a Verify planning and ready-plan replacement behavior. +- [ ] 4.5b **Deferred/replaced:** rich dispatch-admission and single-use-plan bookkeeping is out of scope; verify only unresolved-work admission blocking. +- [x] 4.5c Verify the required compact safety matrix: UNKNOWN fail-closed, sequential dispatch, conditional compensation, 60-second physical deadline, cancellation, unsafe/timed-out reset remaining `FAULT`, and truthful shutdown. +- [x] 4.5e Verify immediate public polling is `RUNNING`-only; no special narrow timer-precheck or `STATUS` guarantee applies in reset/shutdown/gateway states. +- [x] 4.5d **Validation gate:** run the required compact Phase-4 safety matrix after all slices. + +## 5. Module and API cleanup + +- [x] 5.1 Keep `ManipulationModule` as a public adapter over the runtime and remove module-owned lifecycle mutation and reconciliation. +- [x] 5.2 Remove freshness- and robot-name-based execution selection and partial-plan paths without adding a replacement selection gate. +- [x] 5.3 Preserve all existing public `ManipulationModule` semantics, status snapshots, normalized outcomes, and truthful shutdown behavior. +- [x] 5.4 **Validation gate:** run module, operator, UI/state-adapter, API, and compatibility tests; confirm no unrelated source paths changed. + +## 6. White-box test decoupling + +- [x] 6.1 Move white-box tests from reducer/thread internals to plan, model, gateway, actor-handler, and owner-loop seams. +- [x] 6.2 Add focused tests for minimal correlation isolation, future completion ordering, owner timer scans, conditional compensation, deferred reset handling, executor shutdown, and the monotonic clock seam. +- [x] 6.2a Add deterministic tests for stale and cleared physical deadlines, including nondecreasing clock samples and no cancellation when the deadline is absent or not due. +- [x] 6.2b Add deterministic tests that retained faults remain observable after physical-deadline handling and that expiry yields safe `FAILED` or correlated `FAULT`. +- [x] 6.2c Add public poll tests proving immediate polling is `RUNNING`-only and that `STATUS` is not guaranteed in reset/shutdown/gateway states. +- [x] 6.2d Add poll-versus-shutdown tests covering clock-driven physical expiry, admission sealing, one shutdown initiator, concurrent rejection, and truthful unresolved shutdown without eventual-success promise. +- [x] 6.2e Add isolated `ValidatedMonotonicClock` unit tests for default/injected callables, serialized access, finite-sample rejection, nondecreasing-sample rejection, and dependency-contract error classification without lifecycle mutation. +- [x] 6.2f Add deterministic lifecycle tests proving all deadline and relative-wait paths use the wrapper, including stale/cleared deadlines, retained faults, physical expiry, deferred reset behavior, and state-dependent poll/shutdown behavior. +- [x] 6.3 Retain black-box characterization coverage for every public lifecycle and safety contract. +- [x] 6.4 **Validation gate:** run targeted manipulation runtime and coordinator test suites with no skipped contract scenarios. + +## 7. Final verification + +- [x] 7.1 Record the compact seam measurement as maintainability evidence and verify reduced duplicated responsibility across the policy/effects/auxiliary/passive boundaries; there is no hard actor LOC gate or required public-façade/private-owner split. +- [x] 7.2 Verify strict owner-only `RuntimeContext` writes, the sole admission-latch exception, exactly four daemon RPC workers, absence of forbidden thread types, and absence of module-side timeout reconciliation. +- [x] 7.3 Verify physical-operation deadline expiry schedules owner cancellation/reconciliation and yields safe terminal `FAILED` or correlated `FAULT`, never duration-inferred completion. +- [x] 7.4 Run formatting, lint, type checks, targeted manipulation/coordinator tests, the full repository test, and OpenSpec strict validation; formatting/lint/type checks, the manipulation target suite (364 passed, 10 skipped, 16 deselected), OpenSpec validation, and diff review passed. The full repository suite was environment-limited by an unrelated MCP `localhost:9990` port conflict/connection resets (2910 passed, 30 skipped, 289 deselected, 2 failed, 24 errors). +- [x] 7.5 Review the final diff and task ledger so no source implementation task is marked complete merely because work was planned, and only supported completed work remains checked.