diff --git a/dimos/control/tasks/cartesian_ik_task/cartesian_ik_task.py b/dimos/control/tasks/cartesian_ik_task/cartesian_ik_task.py index 073e0fb38c..4b901c4b02 100644 --- a/dimos/control/tasks/cartesian_ik_task/cartesian_ik_task.py +++ b/dimos/control/tasks/cartesian_ik_task/cartesian_ik_task.py @@ -170,11 +170,9 @@ def __init__(self, name: str, config: CartesianIKTaskConfig) -> None: self._joint_names = frozenset(config.joint_names) self._joint_names_list = list(config.joint_names) self._num_joints = len(config.joint_names) - expected_joints = config.control_ik.robot_model.get_coordinator_joint_names() - if config.joint_names != expected_joints: - raise ValueError( - f"CartesianIKTask {name}: task joints must match RobotModelConfig coordinator joints" - ) + model_joint_names = list(config.control_ik.robot_model.joint_names) + if list(config.joint_names) != model_joint_names: + raise ValueError(f"CartesianIKTask {name}: task joints must match model joints exactly") # Create IK solver from model self._ik: PinkControlIK = create_pink_control_ik(config.control_ik) diff --git a/dimos/control/tasks/cartesian_ik_task/pink_control_ik.py b/dimos/control/tasks/cartesian_ik_task/pink_control_ik.py index a49e225b8a..b952577b87 100644 --- a/dimos/control/tasks/cartesian_ik_task/pink_control_ik.py +++ b/dimos/control/tasks/cartesian_ik_task/pink_control_ik.py @@ -207,7 +207,7 @@ def _end_effector_link(robot: RobotModelConfig) -> str: if len(tip_links) != 1: raise ValueError( f"Pink control requires exactly one pose-targetable planning group; " - f"robot '{robot.name}' has {len(tip_links)}" + f"model has {len(tip_links)}" ) return tip_links[0] @@ -216,7 +216,7 @@ def _build_mapping( model: pinocchio.Model, robot: RobotModelConfig, ) -> _CoordinateMapping: - joint_names = tuple(robot.get_coordinator_joint_names()) + joint_names = tuple(robot.joint_names) if not joint_names or len(set(joint_names)) != len(joint_names): raise ValueError("control task joints must be unique and non-empty") @@ -224,15 +224,15 @@ def _build_mapping( v_indices: list[int] = [] q_widths: list[int] = [] joint_ids: set[int] = set() - for urdf_name in (robot.get_urdf_joint_name(name) for name in joint_names): - if not model.existJointName(urdf_name): - raise ValueError(f"control joint mapping references unknown joint: {urdf_name}") - joint_id = int(model.getJointId(urdf_name)) + for joint_name in joint_names: + if not model.existJointName(joint_name): + raise ValueError(f"control model references unknown joint: {joint_name}") + joint_id = int(model.getJointId(joint_name)) if joint_id <= 0 or joint_id >= len(model.joints): - raise ValueError(f"invalid control joint index for {urdf_name}") + raise ValueError(f"invalid control joint index for {joint_name}") joint = model.joints[joint_id] if int(joint.nv) != 1 or int(joint.nq) not in (1, 2): - raise ValueError(f"control joint must be one-DoF: {urdf_name}") + raise ValueError(f"control joint must be one-DoF: {joint_name}") q_indices.append(int(joint.idx_q)) v_indices.append(int(joint.idx_v)) q_widths.append(int(joint.nq)) diff --git a/dimos/control/tasks/cartesian_ik_task/test_cartesian_ik_task.py b/dimos/control/tasks/cartesian_ik_task/test_cartesian_ik_task.py index 08672e1fb6..75c3889770 100644 --- a/dimos/control/tasks/cartesian_ik_task/test_cartesian_ik_task.py +++ b/dimos/control/tasks/cartesian_ik_task/test_cartesian_ik_task.py @@ -40,7 +40,6 @@ def _robot(path: Path) -> RobotModelConfig: return RobotModelConfig( - name="tiny", model_path=path, base_pose=PoseStamped(position=[0, 0, 0], orientation=[0, 0, 0, 1]), joint_names=["joint1"], @@ -135,6 +134,26 @@ def test_factory_rejects_invalid_default_pink_configuration() -> None: control_task_registry.create("cartesian_ik", config, hardware={}) +def test_cartesian_task_rejects_task_joints_that_do_not_match_model_order( + tmp_path: Path, mocker +) -> None: + backend = _FakeControlIK() + mocker.patch( + "dimos.control.tasks.cartesian_ik_task.cartesian_ik_task.create_pink_control_ik", + return_value=backend, + ) + model = _robot(tmp_path / "unused.urdf") + + with pytest.raises(ValueError, match="task joints must match model joints exactly"): + CartesianIKTask( + "cartesian", + CartesianIKTaskConfig( + joint_names=["different_joint"], + control_ik=PinkControlIKConfig(robot_model=model), + ), + ) + + @pytest.mark.parametrize( "module_name", [ diff --git a/dimos/control/tasks/cartesian_ik_task/test_pink_control_ik.py b/dimos/control/tasks/cartesian_ik_task/test_pink_control_ik.py index 32801fb1cb..f357e091b2 100644 --- a/dimos/control/tasks/cartesian_ik_task/test_pink_control_ik.py +++ b/dimos/control/tasks/cartesian_ik_task/test_pink_control_ik.py @@ -84,7 +84,6 @@ def _robot( joint_names = joints or ["joint1", "joint2"] joint_count = len(joint_names) return RobotModelConfig( - name="tiny", model_path=path, base_pose=PoseStamped(position=[0, 0, 0], orientation=[0, 0, 0, 1]), joint_names=joint_names, @@ -184,7 +183,7 @@ def prepare( } -def test_pink_validates_named_frame_and_exact_joint_mapping(tmp_path: Path) -> None: +def test_pink_validates_named_frame_and_exact_model_joints(tmp_path: Path) -> None: model_path = _write_urdf(tmp_path) with pytest.raises(ValueError, match="end-effector frame"): @@ -192,9 +191,7 @@ def test_pink_validates_named_frame_and_exact_joint_mapping(tmp_path: Path) -> N PinkControlIKConfig(robot_model=_robot(model_path, frame="missing")), ) - mismatched = _robot(model_path).model_copy( - update={"joint_name_mapping": {"joint1": "missing", "joint2": "joint2"}} - ) + mismatched = _robot(model_path).model_copy(update={"joint_names": ["missing", "joint2"]}) with pytest.raises(ValueError, match="unknown joint"): create_pink_control_ik( PinkControlIKConfig(robot_model=mismatched), diff --git a/dimos/control/tasks/eef_twist_task/test_eef_twist_task.py b/dimos/control/tasks/eef_twist_task/test_eef_twist_task.py index c07bb26ca6..b40481065a 100644 --- a/dimos/control/tasks/eef_twist_task/test_eef_twist_task.py +++ b/dimos/control/tasks/eef_twist_task/test_eef_twist_task.py @@ -81,7 +81,6 @@ def fake_ik(mocker) -> FakeIK: def _fake_robot_model() -> RobotModelConfig: local_joints = ["joint1", "joint2", "joint3"] return RobotModelConfig( - name="fake", model_path="fake.urdf", base_pose=PoseStamped(position=[0, 0, 0], orientation=[0, 0, 0, 1]), joint_names=local_joints, @@ -93,10 +92,6 @@ def _fake_robot_model() -> RobotModelConfig: tip_link="tool", ) ], - joint_name_mapping={ - f"arm/joint{index}": joint_name - for index, joint_name in enumerate(local_joints, start=1) - }, home_joints=[0.0, 0.0, 0.0], ) @@ -106,7 +101,7 @@ def task(fake_ik: FakeIK) -> EEFTwistTask: return EEFTwistTask( "eef", EEFTwistTaskConfig( - joint_names=["arm/joint1", "arm/joint2", "arm/joint3"], + joint_names=["joint1", "joint2", "joint3"], control_ik=PinkControlIKConfig(robot_model=_fake_robot_model()), timeout=0.3, max_joint_delta_deg=15.0, @@ -122,7 +117,7 @@ def _state( values = [0.0, 0.0, 0.0] if positions is None else positions return CoordinatorState( joints=JointStateSnapshot( - joint_positions={f"arm/joint{i + 1}": value for i, value in enumerate(values)}, + joint_positions={f"joint{i + 1}": value for i, value in enumerate(values)}, ), t_now=t_now, dt=dt, @@ -143,7 +138,7 @@ def test_first_nonzero_command_activates_seeds_from_fk_and_outputs_servo_positio assert output is not None assert output.mode == ControlMode.SERVO_POSITION - assert output.joint_names == ["arm/joint1", "arm/joint2", "arm/joint3"] + assert output.joint_names == ["joint1", "joint2", "joint3"] assert output.positions == [0.01, 0.02, 0.03] assert fake_ik.solve_calls[0].translation[0] > 0.0 @@ -264,7 +259,7 @@ def test_preemption_discards_last_commanded_solve_seed(task: EEFTwistTask, fake_ assert task.on_ee_twist_command(_twist(), t_now=1.0) assert task.compute(_state(1.01)) is not None - task.on_preempted("higher_priority", frozenset(["arm/joint1"])) + task.on_preempted("higher_priority", frozenset(["joint1"])) output = task.compute(_state(1.02, positions=[0.5, 0.0, 0.0])) assert output is not None @@ -277,7 +272,7 @@ def gripper_task(fake_ik: FakeIK) -> EEFTwistTask: return EEFTwistTask( "eef", EEFTwistTaskConfig( - joint_names=["arm/joint1", "arm/joint2", "arm/joint3"], + joint_names=["joint1", "joint2", "joint3"], control_ik=PinkControlIKConfig(robot_model=_fake_robot_model()), timeout=0.0, max_joint_delta_deg=15.0, diff --git a/dimos/control/tasks/teleop_task/test_teleop_task.py b/dimos/control/tasks/teleop_task/test_teleop_task.py index 3d39730c74..fc389aea38 100644 --- a/dimos/control/tasks/teleop_task/test_teleop_task.py +++ b/dimos/control/tasks/teleop_task/test_teleop_task.py @@ -72,7 +72,6 @@ def solve( def _robot(path: Path) -> RobotModelConfig: return RobotModelConfig( - name="arm", model_path=path, base_pose=PoseStamped(position=[0, 0, 0], orientation=[0, 0, 0, 1]), joint_names=["joint1", "joint2"], @@ -84,7 +83,6 @@ def _robot(path: Path) -> RobotModelConfig: tip_link="tool", ) ], - joint_name_mapping={"arm/joint1": "joint1", "arm/joint2": "joint2"}, home_joints=[0.0, 0.0], ) @@ -102,7 +100,7 @@ def _state( return CoordinatorState( joints=JointStateSnapshot( joint_positions={ - f"arm/joint{index + 1}": position for index, position in enumerate(positions) + f"joint{index + 1}": position for index, position in enumerate(positions) } ), t_now=t_now, @@ -146,7 +144,7 @@ def task(tmp_path: Path, fake_ik: _FakePinkIK) -> TeleopIKTask: return TeleopIKTask( "teleop_arm", TeleopIKTaskConfig( - joint_names=["arm/joint1", "arm/joint2"], + joint_names=["joint1", "joint2"], control_ik=_pink_config(tmp_path / "unused.urdf"), hand="right", min_dt=0.02, @@ -161,7 +159,7 @@ def gripper_task(tmp_path: Path, fake_ik: _FakePinkIK) -> TeleopIKTask: return TeleopIKTask( "teleop_arm", TeleopIKTaskConfig( - joint_names=["arm/joint1", "arm/joint2"], + joint_names=["joint1", "joint2"], control_ik=_pink_config(tmp_path / "unused.urdf"), hand="right", gripper_joint="arm/gripper", @@ -263,9 +261,9 @@ def test_gripper_claim_interpolation_and_hold_output( output = gripper_task.compute(_state(1.01, (0.4, 0.5))) - assert gripper_task.claim().joints == frozenset({"arm/joint1", "arm/joint2", "arm/gripper"}) + assert gripper_task.claim().joints == frozenset({"joint1", "joint2", "arm/gripper"}) assert output is not None - assert output.joint_names == ["arm/joint1", "arm/joint2", "arm/gripper"] + assert output.joint_names == ["joint1", "joint2", "arm/gripper"] assert output.positions == pytest.approx([0.4, 0.5, 0.6]) @@ -279,11 +277,11 @@ def test_pose_is_rejected_before_engage_and_after_release(task: TeleopIKTask) -> assert not task.on_cartesian_command(_delta(), 2.3) -def test_factory_requires_pink_configuration_and_matching_model(tmp_path: Path) -> None: +def test_factory_requires_pink_configuration_and_matching_model_dof(tmp_path: Path) -> None: legacy = TaskConfig( name="teleop", type="teleop_ik", - joint_names=["arm/joint1", "arm/joint2"], + joint_names=["joint1", "joint2"], params={"model_path": "legacy.xml", "ee_joint_id": 2, "hand": "right"}, ) with pytest.raises(ValueError, match="control_ik"): @@ -292,11 +290,23 @@ def test_factory_requires_pink_configuration_and_matching_model(tmp_path: Path) mismatched = TaskConfig( name="teleop", type="teleop_ik", - joint_names=["wrong/joint1", "wrong/joint2"], + joint_names=["joint1"], params={ "control_ik": {"robot_model": _robot(tmp_path / "unused.urdf")}, "hand": "right", }, ) - with pytest.raises(ValueError, match="task joints must match"): + with pytest.raises(ValueError, match="task joints must match model joints exactly"): create_task(mismatched, {}) + + wrong_names = TaskConfig( + name="teleop", + type="teleop_ik", + joint_names=["other1", "other2"], + params={ + "control_ik": {"robot_model": _robot(tmp_path / "unused.urdf")}, + "hand": "right", + }, + ) + with pytest.raises(ValueError, match="task joints must match model joints exactly"): + create_task(wrong_names, {}) diff --git a/dimos/e2e_tests/test_manipulation_planning_groups.py b/dimos/e2e_tests/test_manipulation_planning_groups.py index 390db63fde..41a4d3faf3 100644 --- a/dimos/e2e_tests/test_manipulation_planning_groups.py +++ b/dimos/e2e_tests/test_manipulation_planning_groups.py @@ -42,9 +42,8 @@ BLUEPRINT = "openarm-mock-planner-coordinator" -def _wait_for_robot_info( +def _wait_for_model_info( client: RPCClient, - robot_name: str, *, timeout: float = 120.0, ) -> dict[str, Any]: @@ -52,13 +51,13 @@ def _wait_for_robot_info( last_error: BaseException | None = None while time.time() < deadline: try: - info = client.get_robot_info(robot_name) + info = client.get_model_info() if info and info.get("planning_groups"): return cast("dict[str, Any]", info) except Exception as exc: last_error = exc time.sleep(0.5) - raise TimeoutError(f"Timed out waiting for {robot_name!r} robot info") from last_error + raise TimeoutError("Timed out waiting for model info") from last_error def _wait_for_trajectory_completion( @@ -94,55 +93,54 @@ def _wait_for_manipulation_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: try: - missing = tuple( - robot_name - for robot_name in robot_names - if client.get_current_joints(robot_name) is None - ) + if client.get_current_joints() is not None: + return except Exception: - # Robot metadata becomes visible while the planning world is still + # Model metadata becomes visible while the planning world is still # finalizing. Treat that readiness race like a missing joint state. - missing = robot_names - if not missing: - return + pass time.sleep(0.1) - raise TimeoutError(f"Timed out waiting for current joints from {missing}") + raise TimeoutError("Timed out waiting for current model joints") -def _prepare_for_planning(client: RPCClient, robot_names: tuple[str, ...]) -> None: +def _prepare_for_planning(client: RPCClient) -> 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 + _wait_for_current_joints(client) + # Model 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] +def _planning_group(info: dict[str, Any], group_id: str) -> PlanningGroup: + group = next( + group + for group in info["planning_groups"] + if (group.id if isinstance(group, PlanningGroup) else group["id"]) == group_id + ) if isinstance(group, PlanningGroup): - return group.id - group_id = group["id"] - assert isinstance(group_id, str) - return group_id + return group + return PlanningGroup(**group) -def _offset_target(client: RPCClient, robot_name: str, delta: float) -> JointState: - current = client.get_current_joints(robot_name) +def _offset_target( + client: RPCClient, info: dict[str, Any], group: PlanningGroup, delta: float +) -> JointState: + current = client.get_current_joints() assert current is not None - return JointState(position=[position + delta for position in current]) + positions = dict(zip(info["joint_names"], current, strict=True)) + return JointState( + name=list(group.joint_names), + position=[positions[name] + delta for name in group.joint_names], + ) def _start_openarm_mock_planner( @@ -163,15 +161,17 @@ def test_single_arm_plans_and_executes_through_control_coordinator( 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) + info = _wait_for_model_info(client) + left_group = _planning_group(info, "left_arm") tasks = coordinator_client.list_tasks() assert tasks == [DEFAULT_TRAJECTORY_TASK_NAME] - _prepare_for_planning(client, ("left_arm",)) + _prepare_for_planning(client) - planned = client.plan_to_joint_targets({left_id: _offset_target(client, "left_arm", 0.02)}) + planned = client.plan_to_joint_targets( + {left_group.id: _offset_target(client, info, left_group, 0.02)} + ) assert planned, client.get_error() assert client.has_planned_path() assert client.execute_plan() @@ -192,20 +192,19 @@ def test_dual_arm_plans_and_dispatches_both_arms_through_control_coordinator( 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) + info = _wait_for_model_info(client) + left_group = _planning_group(info, "left_arm") + right_group = _planning_group(info, "right_arm") tasks = coordinator_client.list_tasks() assert tasks == [DEFAULT_TRAJECTORY_TASK_NAME] - _prepare_for_planning(client, ("left_arm", "right_arm")) + _prepare_for_planning(client) planned = client.plan_to_joint_targets( { - left_id: _offset_target(client, "left_arm", 0.02), - right_id: _offset_target(client, "right_arm", -0.02), + left_group.id: _offset_target(client, info, left_group, 0.02), + right_group.id: _offset_target(client, info, right_group, -0.02), } ) assert planned, client.get_error() diff --git a/dimos/experimental/world_belief/xarm6_blueprint.py b/dimos/experimental/world_belief/xarm6_blueprint.py index 0b6dca1ed3..9ea0d78580 100644 --- a/dimos/experimental/world_belief/xarm6_blueprint.py +++ b/dimos/experimental/world_belief/xarm6_blueprint.py @@ -90,14 +90,11 @@ class _XArm6WorldBeliefCoordinator(ControlCoordinator): xarm6_worldbelief = autoconnect( # Provides wrist-camera FK/TF. ManipulationModule.blueprint( - robots=[ - make_xarm6_model_config( - name="arm", - add_gripper=False, - # Enables TF publication. - tf_extra_links=["link_base"], - ), - ], + model=make_xarm6_model_config( + add_gripper=False, + # Enables TF publication. + tf_extra_links=["link_base"], + ), ), RealSenseCamera.blueprint( width=640, diff --git a/dimos/manipulation/conftest.py b/dimos/manipulation/conftest.py index ded73ddf1a..b76d3a8e72 100644 --- a/dimos/manipulation/conftest.py +++ b/dimos/manipulation/conftest.py @@ -15,6 +15,7 @@ """Shared manipulation test fixtures.""" from collections.abc import Iterator +from pathlib import Path from typing import Any, Protocol, cast from unittest.mock import MagicMock @@ -28,6 +29,8 @@ TrajectoryExecutionStatus, ) from dimos.manipulation.manipulation_module import ManipulationModule +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition +from dimos.manipulation.planning.spec.config import RobotModelConfig class ModuleFactory(Protocol): @@ -48,18 +51,28 @@ def _mock_control_coordinator() -> MagicMock: return coordinator +def _test_model() -> RobotModelConfig: + return RobotModelConfig( + model_path=Path("/test/model.urdf"), + joint_names=["arm/j0"], + base_link="base", + planning_groups=[PlanningGroupDefinition("manipulator", ("arm/j0",), "base", "tool")], + ) + + @pytest.fixture def module_factory() -> Iterator[ModuleFactory]: """Create started modules and stop every instance during fixture teardown.""" modules: list[ManipulationModule] = [] def create(coordinator: ControlCoordinator | None = None) -> ManipulationModule: - module = ManipulationModule() + module = ManipulationModule(model=_test_model()) modules.append(module) module._control_coordinator = ( coordinator if coordinator is not None else _mock_control_coordinator() ) cast("Any", module).coordinator_joint_state = None + module._initialize_planning = MagicMock() module.start() return module diff --git a/dimos/manipulation/execution_manager.py b/dimos/manipulation/execution_manager.py index 4cab26178c..3b9b55dd80 100644 --- a/dimos/manipulation/execution_manager.py +++ b/dimos/manipulation/execution_manager.py @@ -16,10 +16,9 @@ from __future__ import annotations -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Sequence from enum import Enum, auto import threading -from types import MappingProxyType import attrs @@ -30,17 +29,12 @@ TrajectoryExecutionResult, TrajectoryExecutionStatus, ) -from dimos.manipulation.planning.spec.models import GeneratedPlan, RobotName +from dimos.manipulation.planning.spec.models import GeneratedPlan from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory from dimos.utils.logging_config import setup_logger logger = setup_logger() -_NON_EMPTY_STRING = attrs.validators.and_( - attrs.validators.instance_of(str), - attrs.validators.min_len(1), -) - class ExecutionOutcome(Enum): """Safety-aware outcome of dispatching planned execution.""" @@ -64,98 +58,6 @@ def accepted(self) -> bool: return self.outcome is ExecutionOutcome.ACCEPTED -def _to_model_joint_names(value: Sequence[str]) -> tuple[str, ...]: - return tuple(value) - - -def _to_immutable_joint_mapping(value: Mapping[str, str]) -> Mapping[str, str]: - return MappingProxyType(dict(value)) - - -@attrs.frozen(slots=False) -class ExecutionTarget: - """Immutable coordinator joint mapping for one robot.""" - - robot_name: RobotName = attrs.field(validator=_NON_EMPTY_STRING) - model_joint_names: tuple[str, ...] = attrs.field( - converter=_to_model_joint_names, - validator=attrs.validators.deep_iterable( - member_validator=attrs.validators.instance_of(str), - ), - ) - model_to_coordinator: Mapping[str, str] = attrs.field( - converter=_to_immutable_joint_mapping, - validator=attrs.validators.deep_mapping( - key_validator=attrs.validators.instance_of(str), - value_validator=attrs.validators.instance_of(str), - ), - repr=False, - ) - - @model_joint_names.validator - def _validate_model_joint_names( - self, - _attribute: attrs.Attribute[tuple[str, ...]], - value: tuple[str, ...], - ) -> None: - if not value or any(not name or "/" in name for name in value): - raise ValueError(f"Execution target '{self.robot_name}' has invalid local model joints") - if len(set(value)) != len(value): - raise ValueError( - f"Execution target '{self.robot_name}' has duplicate local model joints" - ) - - @model_to_coordinator.validator - def _validate_model_to_coordinator( - self, - _attribute: attrs.Attribute[Mapping[str, str]], - value: Mapping[str, str], - ) -> None: - if set(value) != set(self.model_joint_names): - raise ValueError(f"Execution target '{self.robot_name}' must resolve every model joint") - resolved_names = list(value.values()) - if any(not name for name in resolved_names) or len(set(resolved_names)) != len( - resolved_names - ): - raise ValueError( - f"Execution target '{self.robot_name}' has ambiguous coordinator joints" - ) - - @classmethod - def from_coordinator_mapping( - cls, - *, - robot_name: RobotName, - model_joint_names: Sequence[str], - # TODO: unify coordinator joint name with planner - coordinator_to_model: Mapping[str, str], - ) -> ExecutionTarget: - """Validate and invert a coordinator-to-model joint mapping.""" - local_names = tuple(model_joint_names) - known = set(local_names) - reverse: dict[str, str] = {} - for coordinator_name, model_name in coordinator_to_model.items(): - if model_name not in known: - raise ValueError( - f"Coordinator joint '{coordinator_name}' maps to unknown model joint " - f"'{model_name}' for '{robot_name}'" - ) - if model_name in reverse: - raise ValueError( - f"Multiple coordinator joints map to model joint '{model_name}' " - f"for '{robot_name}'" - ) - reverse[model_name] = coordinator_name - - return cls( - robot_name=robot_name, - model_joint_names=local_names, - model_to_coordinator={ - model_name: reverse.get(model_name, model_name) for model_name in local_names - }, - ) - - class _PlanRejectedError(Exception): """Expected rejection while mapping a generated plan.""" @@ -166,15 +68,12 @@ class PlanExecutionManager: def __init__( self, *, - targets: Iterable[ExecutionTarget], + joint_names: Sequence[str], coordinator: ControlCoordinator, ) -> None: - target_items = tuple(targets) - target_names = [target.robot_name for target in target_items] - if len(set(target_names)) != len(target_names): - raise ValueError("Execution targets must have unique robot names") - - self._targets = {target.robot_name: target for target in target_items} + self._joint_names = frozenset(joint_names) + if not self._joint_names or len(self._joint_names) != len(joint_names): + raise ValueError("Execution joint names must be non-empty and unique") self._coordinator = coordinator self._operation_lock = threading.Lock() @@ -228,33 +127,10 @@ def _prepare_trajectory(self, plan: GeneratedPlan) -> JointTrajectory: if not plan.is_success(): raise _PlanRejectedError("Generated plan status is not successful") - coordinator_names: list[str] = [] - for global_name in plan.trajectory.joint_names: - parts = global_name.split("/") - if len(parts) != 2 or not parts[0] or not parts[1]: - raise _PlanRejectedError( - f"Generated trajectory joint '{global_name}' is not globally named" - ) - robot_name, local_name = parts - target = self._targets.get(robot_name) - if target is None: - raise _PlanRejectedError( - f"Generated plan references unknown execution robot '{robot_name}'" - ) - coordinator_name = target.model_to_coordinator.get(local_name) - if coordinator_name is None: - raise _PlanRejectedError( - f"Generated trajectory joint '{global_name}' is not configured" - ) - coordinator_names.append(coordinator_name) - - if len(set(coordinator_names)) != len(coordinator_names): - raise _PlanRejectedError( - "Generated trajectory resolves to duplicate coordinator joints" - ) - - return JointTrajectory( - joint_names=coordinator_names, - points=plan.trajectory.points, - timestamp=plan.trajectory.timestamp, - ) + names = plan.trajectory.joint_names + unknown = [name for name in names if name not in self._joint_names] + if unknown: + raise _PlanRejectedError(f"Generated trajectory has unknown joints: {unknown}") + if len(set(names)) != len(names): + raise _PlanRejectedError("Generated trajectory has duplicate joints") + return plan.trajectory diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index bc5570fc12..c1e47af6bb 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -15,7 +15,7 @@ """Manipulation Module - Motion planning with ControlCoordinator execution. Base module providing core manipulation infrastructure: -- @rpc: Low-level building blocks (plan_to_pose, plan_to_joints, preview_path, execute) +- @rpc: Low-level building blocks (plan_to_pose, plan_to_joints, preview_plan, execute) - @skill (short-horizon): Single-step actions (move_to_pose, open_gripper, go_home, go_init) Subclass PickAndPlaceModule (pick_and_place_module.py) adds perception integration @@ -29,6 +29,7 @@ import math import threading import time +import traceback from typing import Any, Literal, TypeAlias from pydantic import Field @@ -42,19 +43,13 @@ from dimos.core.stream import In, Out from dimos.manipulation.execution_manager import ( ExecutionOutcome, - ExecutionTarget, PlanExecutionManager, ) -from dimos.manipulation.planning.factory import ( - KinematicsName, - WorldBackend, - create_planning_specs, - create_world, -) +from dimos.manipulation.planning.factory import WorldBackend, create_planning_specs, create_world from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupSelection from dimos.manipulation.planning.groups.utils import ( filter_joint_state_to_selected_joints, - joint_target_to_global_names, + normalize_joint_target, planning_group_id_from_selector, ) from dimos.manipulation.planning.kinematics.config import ( @@ -77,8 +72,6 @@ Obstacle, PlanningGroupID, PlanningResult, - RobotName, - WorldRobotID, ) from dimos.manipulation.planning.spec.protocols import ( KinematicsSpec, @@ -99,6 +92,7 @@ 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.Transform import Transform from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.sensor_msgs.JointState import JointState from dimos.msgs.tf2_msgs.TFMessage import TFMessage @@ -106,17 +100,10 @@ logger = setup_logger() -# Composite type aliases for readability (using semantic IDs from planning.spec) -RobotEntry: TypeAlias = tuple[WorldRobotID, RobotModelConfig] -"""(world_robot_id, config)""" - -RobotRegistry: TypeAlias = dict[RobotName, RobotEntry] -"""Maps robot_name -> RobotEntry""" - -RobotInfoValue: TypeAlias = ( +ModelInfoValue: TypeAlias = ( str | bool | float | list[str] | list[float] | list[PlanningGroup] | None ) -RobotInfoPayload: TypeAlias = dict[str, RobotInfoValue] +ModelInfoPayload: TypeAlias = dict[str, ModelInfoValue] ObstacleShape: TypeAlias = Literal["box", "sphere", "cylinder", "mesh"] @@ -141,7 +128,7 @@ class ManipulationState(Enum): class ManipulationModuleConfig(ModuleConfig): """Configuration for ManipulationModule.""" - robots: list[RobotModelConfig] = Field(default_factory=list) + model: RobotModelConfig planning_timeout: float = 10.0 world_backend: WorldBackend = "roboplan" visualization: ManipulationVisualizationConfig = Field( @@ -156,8 +143,6 @@ class ManipulationModuleConfig(ModuleConfig): ), ) kinematics: ManipulationKinematicsConfig = Field(default_factory=PinkKinematicsConfig) - # Deprecated: use kinematics.backend instead. - kinematics_name: KinematicsName | None = None # Floor plane Z height (meters). When set, a box obstacle is added at startup # to prevent the planner from routing trajectories below this height. # Set to None to disable. @@ -196,18 +181,14 @@ def __init__(self, **kwargs: Any) -> None: self._kinematics: KinematicsSpec | None = None self._trajectory_parametrizer: TrajectoryParametrizerSpec | None = None - # Robot registry: maps robot_name -> (world_robot_id, config) - self._robots: RobotRegistry = {} - - # Canonical generated plan for plan/preview/execute workflow. - # Robot-local paths and trajectories are derived from this plan on demand. + # Canonical generated plan for the plan/preview/execute workflow. self._last_plan: GeneratedPlan | None = None # Coordinator integration (initialized in start()) self._execution_manager: PlanExecutionManager - # Init joints: captured from first joint state per robot, used by go_init - self._init_joints: dict[RobotName, JointState] = {} + # Init joints captured from the first complete canonical state. + self._init_joints: JointState | None = None # TF publishing thread self._tf_stop_event = threading.Event() @@ -233,10 +214,6 @@ def start(self) -> None: def _initialize_planning(self) -> None: """Initialize world, planner, and trajectory generator.""" - if not self.config.robots: - logger.warning("No robots configured, planning disabled") - return - world = create_world( backend=self.config.world_backend, visualization=self.config.visualization, @@ -245,7 +222,6 @@ def _initialize_planning(self) -> None: world=world, world_backend=self.config.world_backend, planner=self.config.planner, - kinematics_name=self.config.kinematics_name, kinematics=self.config.kinematics, trajectory_parametrization=self.config.trajectory_parametrization, ) @@ -260,9 +236,7 @@ def _initialize_planning(self) -> None: manipulation_module=self, ) - for robot_config in self.config.robots: - robot_id = self._world_monitor.add_robot(robot_config) - self._robots[robot_config.name] = (robot_id, robot_config) + self._world_monitor.load_model(self.config.model) operator = ManipulationOperator(self, self._world_monitor) self._world_monitor.finalize(visualization, operator=operator) @@ -282,18 +256,18 @@ def _initialize_planning(self) -> None: dimensions=(0.6, 1.2, thickness), ) self._world_monitor.add_obstacle(floor_obs) - logger.info(f"Floor obstacle added at z={fz:.3f}") + logger.info("Floor obstacle added", z=fz) - for _, (robot_id, _) in self._robots.items(): - self._world_monitor.start_state_monitor(robot_id) + self._world_monitor.start_state_monitor() if self._world_monitor.visualization is not None: self._world_monitor.start_visualization_thread(rate_hz=10.0) if url := self._world_monitor.get_visualization_url(): - logger.info(f"Visualization: {url}") + logger.info("Visualization available", url=url) - # Start TF publishing thread if any robot has tf_extra_links - if any(c.tf_extra_links for _, c in self._robots.values()): + # Publish every pose-group tip plus any explicitly requested extra links. + has_pose_group = any(group.has_pose_target for group in self.config.model.planning_groups) + if has_pose_group or self.config.model.tf_extra_links: self._tf_stop_event.clear() self._tf_thread = threading.Thread( target=self._tf_publish_loop, name="ManipTFThread", daemon=True @@ -301,124 +275,65 @@ def _initialize_planning(self) -> None: self._tf_thread.start() logger.info("TF publishing thread started") - def _get_default_robot_name(self) -> RobotName | None: - """Get default robot name (first robot if only one, else None).""" - if len(self._robots) == 1: - return next(iter(self._robots.keys())) - return None - - def _get_robot( - self, robot_name: RobotName | None = None - ) -> tuple[RobotName, WorldRobotID, RobotModelConfig] | None: - """Get robot by name or default. - - Args: - robot_name: Robot name or None for default (if single robot) - - Returns: - (robot_name, robot_id, config) or None if not found - """ - if not robot_name: # None or empty string (LLMs often pass "") - robot_name = self._get_default_robot_name() - if robot_name is None: - logger.error("Multiple robots configured, must specify robot_name") - return None - - if robot_name not in self._robots: - logger.error(f"Unknown robot: {robot_name}") - return None - - robot_id, config = self._robots[robot_name] - return (robot_name, robot_id, config) - def _on_joint_state(self, msg: JointState) -> None: """Callback when joint state received from driver. - Splits the aggregated JointState by robot using each robot's - coordinator joint names, then routes to the correct monitor. + Validates and forwards the complete canonical model state. """ try: if self._world_monitor is None: return - # Build name → index map once for the whole message name_to_idx = {name: i for i, name in enumerate(msg.name)} - - for robot_name, (robot_id, config) in self._robots.items(): - coord_names = config.get_coordinator_joint_names() - indices = [name_to_idx.get(cn) for cn in coord_names] - if any(idx is None for idx in indices): - missing = [ - cn for cn, idx in zip(coord_names, indices, strict=False) if idx is None - ] - logger.warning(f"Skipping '{robot_name}': missing joints {missing}") - continue - - # Build per-robot sub-message (coordinator namespace) - sub_positions = [msg.position[idx] for idx in indices] # type: ignore[index] - sub_velocities = ( - [msg.velocity[idx] for idx in indices] # type: ignore[index] - if msg.velocity and len(msg.velocity) == len(msg.name) - else [] - ) - sub_msg = JointState( - name=list(coord_names), - position=sub_positions, - velocity=sub_velocities, - ) - - # Route to specific monitor - self._world_monitor.on_joint_state(sub_msg, robot_id=robot_id) - - # 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( - f"Init joints captured for '{robot_name}': " - f"[{', '.join(f'{j:.3f}' for j in sub_positions)}]" - ) + names = self.config.model.joint_names + missing = [name for name in names if name not in name_to_idx] + if missing: + logger.warning("Skipping incomplete model state", missing_joints=missing) + return + indices = [name_to_idx[name] for name in names] + state = JointState( + name=list(names), + position=[msg.position[index] for index in indices], + velocity=[msg.velocity[index] for index in indices] + if len(msg.velocity) == len(msg.name) + else [], + ) + self._world_monitor.on_joint_state(state) + if self._init_joints is None: + self._init_joints = state except Exception as e: - logger.error(f"Exception in _on_joint_state: {e}") - import traceback - + logger.error("Joint-state handling failed", error=str(e)) logger.error(traceback.format_exc()) def _tf_publish_loop(self) -> None: """Publish TF transforms at 10Hz for EE and extra links.""" - from dimos.msgs.geometry_msgs.Transform import Transform - period = 0.1 # 10Hz while not self._tf_stop_event.is_set(): try: if self._world_monitor is None: break transforms: list[Transform] = [] - for robot_id, config in self._robots.values(): - # Publish world → the unique pose-target group tip, when one exists. - group_id = self._world_monitor.planning_groups.primary_pose_group_id_for_robot( - config.name - ) - if group_id is not None: - group = self._world_monitor.planning_groups.get(group_id) - ee_pose = self._world_monitor.get_ee_pose(robot_id) - if ee_pose is not None and group.tip_link is not None: - ee_tf = Transform.from_pose(group.tip_link, ee_pose) - ee_tf.frame_id = "world" - transforms.append(ee_tf) - - # Publish world → each extra link - for link_name in config.tf_extra_links: - link_pose = self._world_monitor.get_link_pose(robot_id, link_name) - if link_pose is not None: - link_tf = Transform.from_pose(link_name, link_pose) - link_tf.frame_id = "world" - transforms.append(link_tf) + config = self.config.model + for group in self._world_monitor.planning_groups.list(): + if not group.has_pose_target or group.tip_link is None: + continue + ee_pose = self._world_monitor.get_group_ee_pose(group.id) + if ee_pose is not None and group.tip_link is not None: + ee_tf = Transform.from_pose(group.tip_link, ee_pose) + ee_tf.frame_id = "world" + transforms.append(ee_tf) + for link_name in config.tf_extra_links: + link_pose = self._world_monitor.get_link_pose(link_name) + if link_pose is not None: + link_tf = Transform.from_pose(link_name, link_pose) + link_tf.frame_id = "world" + transforms.append(link_tf) if transforms: self.tf.publish(TFMessage(*transforms)) except Exception as e: - logger.debug(f"TF publish error: {e}") + logger.warning("TF publish failed", error=str(e)) self._tf_stop_event.wait(period) @@ -511,70 +426,38 @@ def reset(self) -> SkillResult[ManipulationSkillError]: return SkillResult.ok("Reset to IDLE — ready for new commands") @rpc - def get_current_joints(self, robot_name: RobotName | None = None) -> list[float] | None: - """Get current joint positions. - - Args: - robot_name: Robot to query (required if multiple robots configured) - """ - if (robot := self._get_robot(robot_name)) and self._world_monitor: - state = self._world_monitor.get_current_joint_state(robot[1]) + def get_current_joints(self) -> list[float] | None: + """Get the complete canonical model joint positions.""" + if self._world_monitor: + state = self._world_monitor.get_current_joint_state() if state is not None: return list(state.position) return None @rpc - def get_ee_pose(self, robot_name: RobotName | None = None) -> Pose | None: - """Get current end-effector pose. - - Args: - robot_name: Robot to query (required if multiple robots configured) - """ - if (robot := self._get_robot(robot_name)) and self._world_monitor: + def get_ee_pose(self, group_id: PlanningGroupID | None = None) -> Pose | None: + """Get a planning group's current tip pose.""" + if self._world_monitor: try: - return self._world_monitor.get_ee_pose(robot[1], joint_state=None) + selected = group_id or self._require_unique_pose_group_id() + return self._world_monitor.get_group_ee_pose(selected) except ValueError as exc: - logger.warning("End-effector pose unavailable: %s", exc) + logger.warning("End-effector pose unavailable", error=str(exc)) return None return None @rpc - def is_collision_free(self, joints: list[float], robot_name: RobotName | None = None) -> bool: + def is_collision_free(self, joints: list[float]) -> bool: """Check if joint configuration is collision-free. Args: joints: Joint configuration to check - robot_name: Robot to check (required if multiple robots configured) """ - if (robot := self._get_robot(robot_name)) and self._world_monitor: - _, robot_id, config = robot - joint_state = JointState(name=config.joint_names, position=joints) - return self._world_monitor.is_state_valid(robot_id, joint_state) + if self._world_monitor: + joint_state = JointState(name=self.config.model.joint_names, position=joints) + return self._world_monitor.is_state_valid(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: - self._record_error("Planning not initialized") - return None - if (robot := self._get_robot(robot_name)) is None: - self._record_error("Robot not found or robot_name is required") - return None - with self._lock: - if self._state not in (ManipulationState.IDLE, ManipulationState.COMPLETED): - self._record_error(f"Cannot plan while state is {self._state.name}") - return None - self._planning_epoch += 1 - self._last_plan = None - self._state = ManipulationState.PLANNING - return robot[0], robot[1] - def _begin_group_planning(self) -> int | None: """Check state and begin planning for explicit planning-group APIs.""" if self._world_monitor is None: @@ -582,22 +465,21 @@ def _begin_group_planning(self) -> int | None: return None with self._lock: if self._state not in (ManipulationState.IDLE, ManipulationState.COMPLETED): - logger.warning(f"Cannot plan: state is {self._state.name}") + logger.warning("Cannot plan in current state", state=self._state.name) return None self._planning_epoch += 1 self._last_plan = None self._state = ManipulationState.PLANNING return self._planning_epoch - def _require_unique_pose_group_id_for_robot(self, robot_name: RobotName) -> PlanningGroupID: + def _require_unique_pose_group_id(self) -> 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) + group_id = self._world_monitor.planning_groups.primary_pose_group_id() if group_id is None: raise ValueError( - f"Robot '{robot_name}' has no pose-targetable planning group; " - "use an explicit planning group ID" + "Model has no unique pose-targetable planning group; use an explicit group ID" ) return group_id @@ -610,7 +492,7 @@ def _resolve_group_plan_start( assert self._world_monitor is not None try: selection = self._world_monitor.planning_groups.select(group_ids) - current = self._world_monitor.current_global_joint_state() + current = self._world_monitor.current_model_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}") @@ -668,7 +550,7 @@ def _plan_selected_path( ) return None - logger.info("Path: %d waypoints, groups=%s", len(result.path), group_ids) + logger.info("Path generated", waypoint_count=len(result.path), group_ids=group_ids) return self._store_generated_plan(group_ids, result, planning_epoch) def _record_error(self, message: str) -> bool: @@ -700,23 +582,10 @@ 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 - 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() + self._world_monitor.cancel_preview_animation() def _solve_ik_for_pose( self, - robot_id: WorldRobotID, pose: Pose, seed: JointState, check_collision: bool, @@ -732,7 +601,6 @@ def _solve_ik_for_pose( return self._kinematics.solve( world=self._world_monitor.world, - robot_id=robot_id, target_pose=target_pose, seed=seed, check_collision=check_collision, @@ -768,7 +636,7 @@ def inverse_kinematics( 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() + current = self._world_monitor.current_model_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) @@ -788,19 +656,15 @@ def inverse_kinematics( def inverse_kinematics_single( self, pose: Pose, - robot_name: RobotName | None = None, + group_id: PlanningGroupID | None = None, seed: JointState | None = None, check_collision: bool = True, ) -> IKResult: - """Solve IK for one robot's unique pose-targetable planning group.""" + """Solve IK for one selected or unambiguous pose-targetable 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) + selected_group_id = group_id or self._require_unique_pose_group_id() except ValueError as exc: return IKResult(status=IKStatus.NO_SOLUTION, message=str(exc)) target_pose = PoseStamped( @@ -809,14 +673,14 @@ def inverse_kinematics_single( orientation=pose.orientation, ) return self.inverse_kinematics( - {group_id: target_pose}, seed=seed, check_collision=check_collision + {selected_group_id: target_pose}, seed=seed, check_collision=check_collision ) @rpc def solve_ik( self, pose: Pose, - robot_name: RobotName | None = None, + group_id: PlanningGroupID | None = None, check_collision: bool = True, seed: JointState | None = None, ) -> IKResult: @@ -824,18 +688,13 @@ def solve_ik( Args: pose: Target end-effector pose - robot_name: Robot to solve for (required if multiple robots configured) + group_id: Planning group to solve for; omission requires one compatible group 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: self._record_error("Planning not initialized") return IKResult(status=IKStatus.NO_SOLUTION, message="Planning not initialized") - robot = self._get_robot(robot_name) - if robot is None: - self._record_error("Robot not found or robot_name is required") - return IKResult(status=IKStatus.NO_SOLUTION, message="Robot not found") - with self._lock: if self._state not in (ManipulationState.IDLE, ManipulationState.COMPLETED): self._record_error(f"Cannot solve IK while state is {self._state.name}") @@ -846,41 +705,36 @@ def solve_ik( self._state = ManipulationState.PLANNING result = self.inverse_kinematics_single( pose, - robot_name=robot_name, + group_id=group_id, seed=seed, check_collision=check_collision, ) self._state = ManipulationState.COMPLETED if result.is_success() else ManipulationState.IDLE if result.is_success(): - logger.info(f"IK solved, error: {result.position_error:.4f}m") + logger.info("IK solved", position_error=result.position_error) else: detail = f": {result.message}" if result.message else "" self._record_error(f"IK failed: {result.status.name}{detail}") return result @rpc - def plan_to_pose(self, pose: Pose, robot_name: RobotName | None = None) -> bool: + def plan_to_pose(self, pose: Pose, group_id: PlanningGroupID | None = None) -> bool: """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) + group_id: Planning group to use; omission requires one compatible group """ if self._kinematics is None or self._world_monitor is None: self._record_error("Planning not initialized") return False - robot = self._get_robot(robot_name) - if robot is None: - self._record_error("Robot not found or robot_name is required") - return False - selected_robot_name, _, _ = robot try: - group_id = self._require_unique_pose_group_id_for_robot(selected_robot_name) + selected_group_id = group_id or self._require_unique_pose_group_id() except ValueError as exc: - logger.warning("Pose planning unavailable: %s", exc) + logger.warning("Pose planning unavailable", error=str(exc)) self._record_error(str(exc)) return False - return self.plan_to_pose_targets({group_id: pose}) + return self.plan_to_pose_targets({selected_group_id: pose}) @rpc def plan_to_pose_targets( @@ -892,33 +746,22 @@ def plan_to_pose_targets( 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: + def plan_to_joints(self, joints: JointState, group_id: PlanningGroupID | None = None) -> bool: """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) + group_id: Planning group to use; omission requires exactly one group """ - 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]}" - ) + logger.info("Planning to joints", positions=joints.position) if self._world_monitor is None: self._record_error("Planning not initialized") 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, - ) + selected_group_id = group_id or self._world_monitor.planning_groups.default_group_id() + if selected_group_id is None: + logger.error("Model has no unique default planning group; select a group explicitly") return False - return self.plan_to_joint_targets({group_id: joints}) + return self.plan_to_joint_targets({selected_group_id: joints}) @rpc def plan_to_joint_targets( @@ -955,7 +798,7 @@ def generate_plan_to_joint_targets( 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) + target_global = normalize_joint_target(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}'") @@ -1003,7 +846,7 @@ def generate_plan_to_pose_targets( detail = f": {ik.message}" if ik.message else "" self._fail_planning_epoch(planning_epoch, f"IK failed: {ik.status.name}{detail}") return None - logger.info(f"IK solved, error: {ik.position_error:.4f}m") + logger.info("IK solved", position_error=ik.position_error) return self._plan_selected_path(group_ids, start, ik.joint_state, planning_epoch) @rpc @@ -1066,45 +909,17 @@ def generate_cartesian_plan( return None return self._store_generated_plan(group_ids, result, planning_epoch) - @rpc - def preview_path( - self, - duration: float | None = None, - robot_name: RobotName | None = None, - target_fps: float = 30.0, - ) -> bool: - """Compatibility wrapper for preview_plan(). - - Args: - 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. - """ - return self.preview_plan(None, duration, robot_name, target_fps) - @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.""" plan = plan or self._last_plan if plan is None or not plan.path: logger.warning("No generated plan to preview") return False - 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 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) @@ -1149,15 +964,6 @@ def clear_planned_path(self) -> bool: self._dismiss_preview(plan.group_ids) return True - @rpc - def list_robots(self) -> list[str]: - """List all configured robot names. - - Returns: - List of robot names - """ - return list(self._robots.keys()) - @rpc def list_planning_groups(self) -> list[PlanningGroup]: """Return all configured planning groups.""" @@ -1165,35 +971,17 @@ def list_planning_groups(self) -> list[PlanningGroup]: 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.""" + def get_current_joint_state(self) -> JointState | None: + """Return the complete canonical model joint state.""" 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) + return self._world_monitor.get_current_joint_state() @rpc - def get_robot_info(self, robot_name: RobotName | None = None) -> RobotInfoPayload | None: - """Get information about a robot. - - Args: - robot_name: Robot name (uses default if None) - - Returns: - Dict with robot info or None if not found - """ - robot = self._get_robot(robot_name) - if robot is None: - 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 [] - ) + def get_model_info(self) -> ModelInfoPayload: + """Get information about the configured logical robot model.""" + config = self.config.model + planning_groups = self.list_planning_groups() pose_tip_links = [ group.tip_link for group in planning_groups @@ -1202,65 +990,35 @@ def get_robot_info(self, robot_name: RobotName | None = None) -> RobotInfoPayloa end_effector_link = pose_tip_links[0] if len(pose_tip_links) == 1 else None return { - "name": config.name, - "world_robot_id": robot_id, "joint_names": config.joint_names, "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, - "has_joint_name_mapping": bool(config.joint_name_mapping), "home_joints": config.home_joints, "pre_grasp_offset": config.pre_grasp_offset, - "init_joints": list(init.position) - if (init := self._init_joints.get(robot_name)) + "init_joints": list(self._init_joints.position) + if self._init_joints is not None else None, } - def robot_items(self) -> list[tuple[RobotName, WorldRobotID, RobotModelConfig]]: - """Return configured robots for in-process visualization adapters.""" - return [(name, robot_id, config) for name, (robot_id, config) in self._robots.items()] - - def robot_id_for_name(self, robot_name: RobotName) -> WorldRobotID | None: - """Return the planning-world robot id for a configured robot name.""" - entry = self._robots.get(robot_name) - return entry[0] if entry is not None else None - - def robot_name_for_id(self, robot_id: WorldRobotID) -> RobotName | None: - """Return the configured robot name for a planning-world robot id.""" - for robot_name, (candidate_id, _) in self._robots.items(): - if candidate_id == robot_id: - return robot_name - return None - - def get_robot_config(self, robot_name: RobotName) -> RobotModelConfig | None: - """Return the robot model config for an in-process visualization adapter.""" - entry = self._robots.get(robot_name) - return entry[1] if entry is not None else None + def get_model_config(self) -> RobotModelConfig: + """Return the configured model for in-process visualization adapters.""" + return self.config.model @rpc - def get_init_joints(self, robot_name: RobotName | None = None) -> JointState | None: - """Get the init joint state (captured at startup or set manually). + def get_init_joints(self) -> JointState | None: + """Get the init joint state captured at startup or set manually.""" + return self._init_joints - Args: - robot_name: Robot name (uses default if None and only one robot) - """ - robot = self._get_robot(robot_name) - if robot is None: - return None - return self._init_joints.get(robot[0]) - - def evaluate_joint_target( - self, joints: JointState | None, robot_name: RobotName - ) -> TargetEvaluation: + def evaluate_joint_target(self, joints: JointState | None) -> TargetEvaluation: """Evaluate a joint target for visualization without planning a path.""" - robot_id = self.robot_id_for_name(robot_name) - if robot_id is None or self._world_monitor is None: + if self._world_monitor is None: return { "success": False, - "status": "NO_ROBOT", - "message": f"Unknown robot: {robot_name}", + "status": "UNAVAILABLE", + "message": "Planning is not initialized", "collision_free": False, "ee_pose": None, "joint_state": None, @@ -1275,27 +1033,18 @@ def evaluate_joint_target( "joint_state": None, } target = JointState(joints) - collision_free = self._world_monitor.is_state_valid(robot_id, target) + collision_free = self._world_monitor.is_state_valid(target) 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": self._world_monitor.get_ee_pose(robot_id, target), + "ee_pose": self._world_monitor.get_ee_pose(target), "joint_state": target, } - def evaluate_pose_target(self, pose: Pose, robot_name: RobotName) -> TargetEvaluation: + def evaluate_pose_target(self, pose: Pose) -> TargetEvaluation: """Evaluate a Cartesian target for visualization without planning a path.""" - robot_id = self.robot_id_for_name(robot_name) - if robot_id is None: - return { - "success": False, - "joint_state": None, - "status": "UNKNOWN_ROBOT", - "message": f"Unknown robot: {robot_name}", - "collision_free": False, - } if self._world_monitor is None or self._kinematics is None: return { "success": False, @@ -1304,7 +1053,7 @@ def evaluate_pose_target(self, pose: Pose, robot_name: RobotName) -> TargetEvalu "message": "Planning is not initialized or current state is unavailable", "collision_free": False, } - current = self._world_monitor.get_current_joint_state(robot_id) + current = self._world_monitor.get_current_joint_state() if current is None: return { "success": False, @@ -1313,10 +1062,10 @@ def evaluate_pose_target(self, pose: Pose, robot_name: RobotName) -> TargetEvalu "message": "Planning is not initialized or current state is unavailable", "collision_free": False, } - ik = self._solve_ik_for_pose(robot_id, pose, current, check_collision=True) + ik = self._solve_ik_for_pose(pose, current, check_collision=True) joint_state = JointState(ik.joint_state) if ik.is_success() and ik.joint_state else None collision_free = bool( - joint_state is not None and self._world_monitor.is_state_valid(robot_id, joint_state) + joint_state is not None and self._world_monitor.is_state_valid(joint_state) ) return { "success": joint_state is not None and collision_free, @@ -1329,59 +1078,32 @@ def evaluate_pose_target(self, pose: Pose, robot_name: RobotName) -> TargetEvalu } @rpc - def set_init_joints(self, joint_state: JointState, robot_name: RobotName | None = None) -> bool: + def set_init_joints(self, joint_state: JointState) -> bool: """Set the init joint state. Args: joint_state: New init joint state (names + positions) - robot_name: Robot name (uses default if None and only one robot) """ - robot = self._get_robot(robot_name) - if robot is None: - return False - self._init_joints[robot[0]] = joint_state - logger.info( - f"Init joints set for '{robot[0]}': " - f"[{', '.join(f'{j:.3f}' for j in joint_state.position)}]" - ) + self._init_joints = joint_state + logger.info("Init joints set", positions=joint_state.position) return True @rpc - def set_init_joints_to_current(self, robot_name: RobotName | None = None) -> bool: - """Set init joints to the current joint positions. - - Args: - robot_name: Robot to capture from (required if multiple robots configured) - """ - robot = self._get_robot(robot_name) - if robot is None: - return False - robot_name_resolved, robot_id, _ = robot + def set_init_joints_to_current(self) -> bool: + """Set init joints to the current joint positions.""" if self._world_monitor is None: return False - current = self._world_monitor.get_current_joint_state(robot_id) + current = self._world_monitor.get_current_joint_state() if current is None: logger.error("Cannot capture init joints — no current joint state") return False - self._init_joints[robot_name_resolved] = current - logger.info( - f"Init joints set to current for '{robot_name_resolved}': " - f"[{', '.join(f'{j:.3f}' for j in current.position)}]" - ) + self._init_joints = current return True def _initialize_execution(self) -> None: """Initialize coordinator access and planned execution policy.""" - targets = [ - ExecutionTarget.from_coordinator_mapping( - robot_name=config.name, - model_joint_names=config.joint_names, - coordinator_to_model=config.joint_name_mapping, - ) - for _, config in self._robots.values() - ] self._execution_manager = PlanExecutionManager( - targets=targets, + joint_names=self.config.model.joint_names, coordinator=self._control_coordinator, ) @@ -1449,7 +1171,7 @@ def add_obstacle( obstacle_type = _SHAPE_TO_OBSTACLE_TYPE.get(shape) if obstacle_type is None: - logger.warning(f"Unknown obstacle shape: {shape}") + logger.warning("Unknown obstacle shape", shape=shape) return "" # Validate mesh_path for mesh type @@ -1519,70 +1241,52 @@ def remove_obstacle(self, obstacle_id: str) -> bool: return False return self._world_monitor.remove_obstacle(obstacle_id) - def _get_gripper_hardware_id(self, robot_name: RobotName | None = None) -> str | None: - """Get gripper hardware ID for a robot.""" - robot = self._get_robot(robot_name) - if robot is None: - return None - _, _, config = robot + def _get_gripper_hardware_id(self) -> str | None: + """Get the configured legacy gripper hardware ID.""" + config = self.config.model if not config.gripper_hardware_id: - logger.warning(f"No gripper_hardware_id configured for '{config.name}'") + logger.warning("No gripper_hardware_id configured") return None return str(config.gripper_hardware_id) - def _set_gripper_position(self, position: float, robot_name: RobotName | None = None) -> bool: + def _set_gripper_position(self, position: float) -> bool: """Internal: set gripper position in meters.""" - hw_id = self._get_gripper_hardware_id(robot_name) + hw_id = self._get_gripper_hardware_id() if hw_id is None: return False return self._control_coordinator.set_gripper_position(hw_id, position) @rpc - def get_gripper(self, robot_name: RobotName | None = None) -> float | None: - """Get gripper position in meters. - - Args: - robot_name: Robot to query (required if multiple robots configured) - """ - hw_id = self._get_gripper_hardware_id(robot_name) + def get_gripper(self) -> float | None: + """Get gripper position in meters.""" + hw_id = self._get_gripper_hardware_id() if hw_id is None: return None result = self._control_coordinator.get_gripper_position(hw_id) return float(result) if result is not None else None @skill - def set_gripper( - self, position: float, robot_name: str | None = None - ) -> SkillResult[ManipulationSkillError]: + def set_gripper(self, position: float) -> SkillResult[ManipulationSkillError]: """Set gripper to a specific opening in meters. Args: position: Gripper opening in meters (0.0 = closed, 0.85 = fully open). - robot_name: Robot to control (only needed for multi-arm setups). """ - if self._set_gripper_position(position, robot_name): + if self._set_gripper_position(position): return SkillResult.ok(f"Gripper set to {position:.3f}m") return SkillResult.fail("GRIPPER_FAILED", "Failed to set gripper position") @skill - def open_gripper(self, robot_name: str | None = None) -> SkillResult[ManipulationSkillError]: - """Open the robot gripper fully. - - Args: - robot_name: Robot to control (only needed for multi-arm setups). - """ - if self._set_gripper_position(0.85, robot_name): + def open_gripper(self) -> SkillResult[ManipulationSkillError]: + """Open the robot gripper fully.""" + if self._set_gripper_position(0.85): return SkillResult.ok("Gripper opened") return SkillResult.fail("GRIPPER_FAILED", "Failed to open gripper") @skill - def close_gripper(self, robot_name: str | None = None) -> SkillResult[ManipulationSkillError]: - """Close the robot gripper fully. - - Args: - robot_name: Robot to control (only needed for multi-arm setups). - """ - if self._set_gripper_position(0.0, robot_name): + def close_gripper(self) -> SkillResult[ManipulationSkillError]: + """Close the robot gripper fully.""" + if self._set_gripper_position(0.0): return SkillResult.ok("Gripper closed") return SkillResult.fail("GRIPPER_FAILED", "Failed to close gripper") @@ -1593,40 +1297,41 @@ def _wait_for_trajectory_completion(self, timeout: float = 60.0) -> bool: return True wait_time = last_plan.trajectory.duration + 0.5 if wait_time > timeout: - logger.warning(f"Trajectory duration exceeds timeout of {timeout}s") + logger.warning("Trajectory duration exceeds timeout", timeout=timeout) return False time.sleep(wait_time) return True def _lift_if_low( - self, robot_name: RobotName | None = None, min_z: float = 0.05 + self, group_id: PlanningGroupID | None = None, min_z: float = 0.05 ) -> SkillResult[ManipulationSkillError]: """If the end-effector is below *min_z*, plan and execute a short lift.""" - ee = self.get_ee_pose(robot_name) + ee = self.get_ee_pose(group_id) if ee is None or ee.position.z >= min_z: return SkillResult.ok() lift_z = min_z + 0.05 - logger.info(f"EE z={ee.position.z:.3f} < {min_z}, lifting to z={lift_z:.3f}") + logger.info( + "Lifting low end effector", current_z=ee.position.z, minimum_z=min_z, target_z=lift_z + ) lift_pose = Pose(Vector3(ee.position.x, ee.position.y, lift_z), ee.orientation) - if not self.plan_to_pose(lift_pose, robot_name): + if not self.plan_to_pose(lift_pose, group_id): return SkillResult.fail( "PLANNING_FAILED", f"Failed to plan lift from z={ee.position.z:.3f}", ) - return self._preview_execute_wait(robot_name) + return self._preview_execute_wait() def _preview_execute_wait( - self, robot_name: RobotName | None = None, preview_duration: float = 0.5 + self, preview_duration: float = 0.5 ) -> SkillResult[ManipulationSkillError]: """Preview planned path, execute, and wait for completion. Args: - robot_name: Robot to operate on preview_duration: Duration to animate the preview in Meshcat (seconds) """ logger.info("Previewing trajectory...") - self.preview_path(preview_duration, robot_name) + self.preview_plan(duration=preview_duration) logger.info("Executing trajectory...") if not self.execute(): @@ -1638,28 +1343,26 @@ def _preview_execute_wait( return SkillResult.ok() @skill - def get_robot_state(self, robot_name: str | None = None) -> SkillResult[ManipulationSkillError]: - """Get current robot state: joint positions, end-effector pose, and gripper. - - Args: - robot_name: Robot to query (only needed for multi-arm setups). - """ + def get_robot_state( + self, group_id: PlanningGroupID | None = None + ) -> SkillResult[ManipulationSkillError]: + """Get model joints, a selected end-effector pose, and gripper state.""" lines: list[str] = [] - joints = self.get_current_joints(robot_name) + joints = self.get_current_joints() if joints is not None: lines.append(f"Joints: [{', '.join(f'{j:.3f}' for j in joints)}]") else: lines.append("Joints: unavailable (no state received)") - ee_pose = self.get_ee_pose(robot_name) + ee_pose = self.get_ee_pose(group_id) if ee_pose is not None: p = ee_pose.position lines.append(f"EE pose: ({p.x:.4f}, {p.y:.4f}, {p.z:.4f})") else: lines.append("EE pose: unavailable") - gripper_pos = self.get_gripper(robot_name) + gripper_pos = self.get_gripper() if gripper_pos is not None: lines.append(f"Gripper: {gripper_pos:.3f}m") else: @@ -1678,7 +1381,7 @@ def move_to_pose( roll: float | None = None, pitch: float | None = None, yaw: float | None = None, - robot_name: str | None = None, + group_id: PlanningGroupID | None = None, ) -> SkillResult[ManipulationSkillError]: """Move the robot end-effector to a target pose. @@ -1692,20 +1395,20 @@ def move_to_pose( roll: Target roll in radians (omit to keep current orientation). pitch: Target pitch in radians (omit to keep current orientation). yaw: Target yaw in radians (omit to keep current orientation). - robot_name: Robot to move (only needed for multi-arm setups). + group_id: Planning group to move; omission requires one compatible group. """ - logger.info(f"Planning motion to ({x:.3f}, {y:.3f}, {z:.3f})...") + logger.info("Planning motion to pose", x=x, y=y, z=z) # If no orientation specified, preserve the current EE orientation. # If partially specified, fill unspecified angles from current orientation. if roll is None and pitch is None and yaw is None: - current_pose = self.get_ee_pose(robot_name) + current_pose = self.get_ee_pose(group_id) if current_pose is not None: orientation = current_pose.orientation else: orientation = Quaternion(0, 0, 0, 1) # identity fallback else: - current_pose = self.get_ee_pose(robot_name) + current_pose = self.get_ee_pose(group_id) if current_pose is not None: current_euler = current_pose.orientation.to_euler() orientation = Quaternion.from_euler( @@ -1721,17 +1424,17 @@ def move_to_pose( pose = Pose(Vector3(x, y, z), orientation) # If EE is low, lift up first to clear obstacles - lift = self._lift_if_low(robot_name) + lift = self._lift_if_low(group_id) if not lift.is_success(): return lift - if not self.plan_to_pose(pose, robot_name): + if not self.plan_to_pose(pose, group_id): return SkillResult.fail( "PLANNING_FAILED", f"Pose ({x:.3f}, {y:.3f}, {z:.3f}) may be unreachable or in collision", ) - exec_result = self._preview_execute_wait(robot_name) + exec_result = self._preview_execute_wait() if not exec_result.is_success(): return exec_result @@ -1741,7 +1444,7 @@ def move_to_pose( def move_to_joints( self, joints: str, - robot_name: str | None = None, + group_id: PlanningGroupID | None = None, ) -> SkillResult[ManipulationSkillError]: """Move the robot to a target joint configuration. @@ -1749,7 +1452,7 @@ def move_to_joints( Args: joints: Comma-separated joint positions in radians, e.g. "0.1, -0.5, 1.2, 0.0, 0.3, -0.1". - robot_name: Robot to move (only needed for multi-arm setups). + group_id: Planning group to move; omission requires exactly one group. """ try: joint_values = [float(j.strip()) for j in joints.split(",")] @@ -1759,38 +1462,61 @@ def move_to_joints( f"Invalid joints format '{joints}'. Expected comma-separated floats.", ) - robot = self._get_robot(robot_name) - if robot is None: - return SkillResult.fail("ROBOT_NOT_FOUND", "Robot not found") - rname, _, config = robot - goal = JointState(name=config.joint_names, position=joint_values) - - logger.info(f"Planning motion to joints [{', '.join(f'{j:.3f}' for j in joint_values)}]...") - if not self.plan_to_joints(goal, rname): + if self._world_monitor is None: + return SkillResult.fail("PLANNING_FAILED", "Planning not initialized") + selected_group_id = group_id or self._world_monitor.planning_groups.default_group_id() + if selected_group_id is None: + return SkillResult.fail("INVALID_INPUT", "Select a planning group explicitly") + try: + self._world_monitor.planning_groups.get(selected_group_id) + except KeyError: + return SkillResult.fail("INVALID_INPUT", f"Unknown planning group: {selected_group_id}") + group = self._world_monitor.planning_groups.get(selected_group_id) + goal = JointState(name=list(group.joint_names), position=joint_values) + + logger.info("Planning motion to joints", positions=joint_values) + if not self.plan_to_joints(goal, selected_group_id): return SkillResult.fail( "PLANNING_FAILED", "Joint configuration may be unreachable or in collision", ) - exec_result = self._preview_execute_wait(robot_name) + exec_result = self._preview_execute_wait() if not exec_result.is_success(): return exec_result return SkillResult.ok("Reached target joint configuration") + def _model_positions_for_group( + self, positions: Sequence[float], group_id: PlanningGroupID | None + ) -> JointState | None: + return self._model_state_for_group( + JointState(name=self.config.model.joint_names, position=list(positions)), group_id + ) + + def _model_state_for_group( + self, state: JointState, group_id: PlanningGroupID | None + ) -> JointState | None: + if self._world_monitor is None: + return None + selected_group_id = group_id or self._world_monitor.planning_groups.default_group_id() + if selected_group_id is None: + return None + group = self._world_monitor.planning_groups.get(selected_group_id) + return filter_joint_state_to_selected_joints(state, group.joint_names) + @skill - def go_home(self, robot_name: str | None = None) -> SkillResult[ManipulationSkillError]: + def go_home( + self, group_id: PlanningGroupID | None = None + ) -> SkillResult[ManipulationSkillError]: """Move the robot to its home/observe joint configuration. Opens the gripper and moves to the predefined home position. Args: - robot_name: Robot to move (only needed for multi-arm setups). + group_id: Planning group to move; omission requires exactly one group. """ - robot = self._get_robot(robot_name) - if robot is None: - return SkillResult.fail("ROBOT_NOT_FOUND", "Robot not found") - rname, _, config = robot + config = self.config.model if config.home_joints is None: return SkillResult.fail( @@ -1799,74 +1525,81 @@ def go_home(self, robot_name: str | None = None) -> SkillResult[ManipulationSkil ) logger.info("Opening gripper...") - self._set_gripper_position(0.85, rname) + self._set_gripper_position(0.85) time.sleep(0.5) - goal = JointState(name=config.joint_names, position=config.home_joints) + goal = self._model_positions_for_group(config.home_joints, group_id) + if goal is None: + return SkillResult.fail("INVALID_INPUT", "Select a planning group explicitly") logger.info("Planning motion to home position...") - if not self.plan_to_joints(goal, rname): + if not self.plan_to_joints(goal, group_id): return SkillResult.fail("PLANNING_FAILED", "Failed to plan path to home position") - exec_result = self._preview_execute_wait(robot_name) + exec_result = self._preview_execute_wait() if not exec_result.is_success(): return exec_result return SkillResult.ok("Reached home position") @skill - def go_init(self, robot_name: str | None = None) -> SkillResult[ManipulationSkillError]: + def go_init( + self, group_id: PlanningGroupID | None = None + ) -> SkillResult[ManipulationSkillError]: """Move the robot to its init position (captured at startup or set manually). The init position is the joint configuration the robot was in when the module first received joint state. It can be changed with set_init_joints(). Args: - robot_name: Robot to move (only needed for multi-arm setups). + group_id: Planning group to move; omission requires exactly one group. """ - robot = self._get_robot(robot_name) - if robot is None: - return SkillResult.fail("ROBOT_NOT_FOUND", "Robot not found") - rname, robot_id, _ = robot - - init = self._init_joints.get(rname) + init = self._init_joints if init is None: return SkillResult.fail( "NOT_CONFIGURED", "No init joints captured — robot may not have reported joint state yet", ) + if self._world_monitor is None: + return SkillResult.fail( + "WORLD_MONITOR_UNAVAILABLE", "Planning world is not initialized" + ) + selected_group_id = group_id or self._world_monitor.planning_groups.default_group_id() + if selected_group_id is None: + return SkillResult.fail("INVALID_INPUT", "Select a planning group explicitly") + # Lift if EE is low before moving to init - lift = self._lift_if_low(robot_name) + lift = self._lift_if_low(selected_group_id) if not lift.is_success(): return lift # Move through a safe waypoint: 10cm above and 5cm in front of init pose. # This avoids direct paths through the workspace that could collide with objects. - if self._world_monitor is not None: - init_ee = self._world_monitor.get_ee_pose(robot_id, joint_state=init) - if init_ee is not None: - wp = Pose( - Vector3( - init_ee.position.x + 0.05, - init_ee.position.y, - init_ee.position.z + 0.10, - ), - init_ee.orientation, - ) - if self.plan_to_pose(wp, robot_name): - wp_result = self._preview_execute_wait(robot_name) - if not wp_result.is_success(): - return wp_result - else: - logger.warning("Safe waypoint unreachable, going directly to init") + init_ee = self._world_monitor.get_group_ee_pose(selected_group_id, init) + if init_ee is not None: + wp = Pose( + Vector3( + init_ee.position.x + 0.05, + init_ee.position.y, + init_ee.position.z + 0.10, + ), + init_ee.orientation, + ) + if self.plan_to_pose(wp, selected_group_id): + wp_result = self._preview_execute_wait() + if not wp_result.is_success(): + return wp_result + else: + logger.warning("Safe waypoint unreachable, going directly to init") - logger.info( - f"Planning motion to init position [{', '.join(f'{j:.3f}' for j in init.position)}]..." - ) - if not self.plan_to_joints(init, robot_name): + logger.info("Planning motion to init position", positions=init.position) + target = self._model_state_for_group(init, selected_group_id) + if target is None: + return SkillResult.fail("INVALID_INPUT", "Select a planning group explicitly") + if not self.plan_to_joints(target, selected_group_id): return SkillResult.fail("PLANNING_FAILED", "Failed to plan path to init position") - exec_result = self._preview_execute_wait(robot_name) + exec_result = self._preview_execute_wait() if not exec_result.is_success(): return exec_result diff --git a/dimos/manipulation/pick_and_place_module.py b/dimos/manipulation/pick_and_place_module.py index e00e049abc..0f7f86059f 100644 --- a/dimos/manipulation/pick_and_place_module.py +++ b/dimos/manipulation/pick_and_place_module.py @@ -34,6 +34,7 @@ ManipulationModule, ManipulationModuleConfig, ) +from dimos.manipulation.planning.spec.models import PlanningGroupID from dimos.manipulation.skill_errors import ManipulationSkillError from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.geometry_msgs.Quaternion import Quaternion @@ -338,25 +339,27 @@ def _resolve_object_position(self, object_name: str) -> tuple[float, float, floa return det.center.x, det.center.y, det.center.z @skill - def get_scene_info(self, robot_name: str | None = None) -> SkillResult[ManipulationSkillError]: + def get_scene_info( + self, group_id: PlanningGroupID | None = None + ) -> SkillResult[ManipulationSkillError]: """Get current robot state, detected objects, and scene information. Returns a summary of the robot's joint positions, end-effector pose, gripper state, detected objects, and obstacle count. Args: - robot_name: Robot to query (only needed for multi-arm setups). + group_id: Planning group used for the end-effector pose query. """ lines: list[str] = [] # Robot state - joints = self.get_current_joints(robot_name) + joints = self.get_current_joints() if joints is not None: lines.append(f"Joints: [{', '.join(f'{j:.3f}' for j in joints)}]") else: lines.append("Joints: unavailable (no state received)") - ee_pose = self.get_ee_pose(robot_name) + ee_pose = self.get_ee_pose(group_id) if ee_pose is not None: p = ee_pose.position lines.append(f"EE pose: ({p.x:.4f}, {p.y:.4f}, {p.z:.4f})") @@ -364,7 +367,7 @@ def get_scene_info(self, robot_name: str | None = None) -> SkillResult[Manipulat lines.append("EE pose: unavailable") # Gripper - gripper_pos = self.get_gripper(robot_name) + gripper_pos = self.get_gripper() if gripper_pos is not None: lines.append(f"Gripper: {gripper_pos:.3f}m") else: @@ -397,13 +400,11 @@ def get_scene_info(self, robot_name: str | None = None) -> SkillResult[Manipulat return SkillResult.ok("\n".join(lines)) @skill - def look(self, robot_name: str | None = None) -> SkillResult[ManipulationSkillError]: + def look(self) -> SkillResult[ManipulationSkillError]: """Quick check of what objects are visible from the current camera position. Does NOT move the arm. Returns objects currently detected in the camera view. - Args: - robot_name: Robot context (only needed for multi-arm setups). """ obstacles = self.refresh_obstacles(0.0) @@ -427,7 +428,7 @@ def look(self, robot_name: str | None = None) -> SkillResult[ManipulationSkillEr def scan_objects( self, min_duration: float = 0.0, - robot_name: str | None = None, + group_id: PlanningGroupID | None = None, ) -> SkillResult[ManipulationSkillError]: """Scan for objects — moves to init position first for a clear camera view, \ then refreshes perception obstacles. @@ -436,10 +437,10 @@ def scan_objects( Args: min_duration: Minimum time an object must be seen to be included. - robot_name: Robot context (only needed for multi-arm setups). + group_id: Planning group to move; omission requires exactly one group. """ # Go to init for a clear camera view - init_result = self.go_init(robot_name) + init_result = self.go_init(group_id) if not init_result.is_success(): return init_result @@ -467,7 +468,7 @@ def pick( self, object_name: str, object_id: str | None = None, - robot_name: str | None = None, + group_id: PlanningGroupID | None = None, ) -> SkillResult[ManipulationSkillError]: """Pick up an object by name using grasp planning and motion execution. @@ -477,12 +478,9 @@ def pick( Args: object_name: Name of the object to pick (e.g. "cup", "bottle", "can"). object_id: Optional unique object ID from perception for precise identification. - robot_name: Robot to use (only needed for multi-arm setups). + group_id: Planning group to use; omission requires exactly one group. """ - robot = self._get_robot(robot_name) - if robot is None: - return SkillResult.fail("ROBOT_NOT_FOUND", "Robot not found") - rname, _, config = robot + config = self.config.model pre_grasp_offset = config.pre_grasp_offset # 1. Generate grasps (uses already-cached detections — call scan_objects first) @@ -495,7 +493,7 @@ def pick( ) # Lift if EE is low before approaching - lift = self._lift_if_low(rname) + lift = self._lift_if_low(group_id) if not lift.is_success(): return lift @@ -509,38 +507,38 @@ def pick( pre_grasp_pose = self._compute_pre_grasp_pose(grasp_pose, offset) logger.info(f"Planning approach to pre-grasp (attempt {i + 1}/{max_attempts})...") - if not self.plan_to_pose(pre_grasp_pose, rname): + if not self.plan_to_pose(pre_grasp_pose, group_id): logger.info(f"Grasp candidate {i + 1} approach planning failed, trying next") continue # Try next candidate # 3. Open gripper before approach logger.info("Opening gripper...") - self._set_gripper_position(0.85, rname) + self._set_gripper_position(0.85) time.sleep(0.5) # 4. Execute approach to pre-grasp - exec_result = self._preview_execute_wait(rname) + exec_result = self._preview_execute_wait() if not exec_result.is_success(): return exec_result # 5. Move to grasp pose logger.info("Moving to grasp position...") - if not self.plan_to_pose(grasp_pose, rname): + if not self.plan_to_pose(grasp_pose, group_id): return SkillResult.fail("PLANNING_FAILED", "Grasp pose planning failed") - exec_result = self._preview_execute_wait(rname) + exec_result = self._preview_execute_wait() if not exec_result.is_success(): return exec_result # 6. Close gripper logger.info("Closing gripper...") - self._set_gripper_position(0.0, rname) + self._set_gripper_position(0.0) time.sleep(1.5) # Wait for gripper to close # 7. Retract to pre-grasp logger.info("Retracting with object...") - if not self.plan_to_pose(pre_grasp_pose, rname): + if not self.plan_to_pose(pre_grasp_pose, group_id): return SkillResult.fail("PLANNING_FAILED", "Retract planning failed") - exec_result = self._preview_execute_wait(rname) + exec_result = self._preview_execute_wait() if not exec_result.is_success(): return exec_result @@ -560,7 +558,7 @@ def place( x: float, y: float, z: float, - robot_name: str | None = None, + group_id: PlanningGroupID | None = None, ) -> SkillResult[ManipulationSkillError]: """Place a held object at the specified position. @@ -571,11 +569,11 @@ def place( x: Target X position in meters. y: Target Y position in meters. z: Target Z position in meters. - robot_name: Robot to use (only needed for multi-arm setups). + group_id: Planning group to use; omission requires exactly one group. """ xy_dist = (x**2 + y**2) ** 0.5 orientation = self._grasp_orientation(x, y, xy_dist) - return self._place_with_orientation(x, y, z, orientation, robot_name) + return self._place_with_orientation(x, y, z, orientation, group_id) def _place_with_orientation( self, @@ -583,13 +581,10 @@ def _place_with_orientation( y: float, z: float, orientation: Quaternion, - robot_name: str | None = None, + group_id: PlanningGroupID | None = None, ) -> SkillResult[ManipulationSkillError]: """Internal place with explicit orientation.""" - robot = self._get_robot(robot_name) - if robot is None: - return SkillResult.fail("ROBOT_NOT_FOUND", "Robot not found") - rname, _, config = robot + config = self.config.model pre_place_offset = config.pre_grasp_offset # Reduce pre-place height for far targets @@ -601,50 +596,52 @@ def _place_with_orientation( pre_place_pose = self._compute_pre_grasp_pose(place_pose, pre_place_offset) # Lift if EE is low before approaching - lift = self._lift_if_low(rname) + lift = self._lift_if_low(group_id) if not lift.is_success(): return lift # 1. Move to pre-place logger.info(f"Planning approach to place position ({x:.3f}, {y:.3f}, {z:.3f})...") - if not self.plan_to_pose(pre_place_pose, rname): + if not self.plan_to_pose(pre_place_pose, group_id): return SkillResult.fail("PLANNING_FAILED", "Pre-place approach planning failed") - exec_result = self._preview_execute_wait(rname) + exec_result = self._preview_execute_wait() if not exec_result.is_success(): return exec_result # 2. Lower to place position logger.info("Lowering to place position...") - if not self.plan_to_pose(place_pose, rname): + if not self.plan_to_pose(place_pose, group_id): return SkillResult.fail("PLANNING_FAILED", "Place pose planning failed") - exec_result = self._preview_execute_wait(rname) + exec_result = self._preview_execute_wait() if not exec_result.is_success(): return exec_result # 3. Release logger.info("Releasing object...") - self._set_gripper_position(0.85, rname) + self._set_gripper_position(0.85) time.sleep(1.0) # 4. Retract logger.info("Retracting...") - if not self.plan_to_pose(pre_place_pose, rname): + if not self.plan_to_pose(pre_place_pose, group_id): return SkillResult.fail("PLANNING_FAILED", "Retract planning failed") - exec_result = self._preview_execute_wait(rname) + exec_result = self._preview_execute_wait() if not exec_result.is_success(): return exec_result return SkillResult.ok(f"Place complete — object released at ({x:.3f}, {y:.3f}, {z:.3f})") @skill - def place_back(self, robot_name: str | None = None) -> SkillResult[ManipulationSkillError]: + def place_back( + self, group_id: PlanningGroupID | None = None + ) -> SkillResult[ManipulationSkillError]: """Place the held object back at its original pick position. Uses the position stored from the last successful pick operation. Args: - robot_name: Robot to use (only needed for multi-arm setups). + group_id: Planning group to use; omission requires exactly one group. """ if self._last_pick_pose is None: return SkillResult.fail( @@ -655,14 +652,14 @@ def place_back(self, robot_name: str | None = None) -> SkillResult[ManipulationS p = self._last_pick_pose.position o = self._last_pick_pose.orientation logger.info(f"Placing back at original position ({p.x:.3f}, {p.y:.3f}, {p.z:.3f})...") - return self._place_with_orientation(p.x, p.y, p.z, o, robot_name) + return self._place_with_orientation(p.x, p.y, p.z, o, group_id) @skill def drop_on( self, target_object_name: str, z_offset: float = 0.1, - robot_name: str | None = None, + group_id: PlanningGroupID | None = None, ) -> SkillResult[ManipulationSkillError]: """Drop a held object on top of a detected object. @@ -672,7 +669,7 @@ def drop_on( Args: target_object_name: Name of the target object to drop onto (e.g. "cup", "bowl"). z_offset: Height above the target object's center to release (meters). - robot_name: Robot to use (only needed for multi-arm setups). + group_id: Planning group to use; omission requires exactly one group. """ pos = self._resolve_object_position(target_object_name) if pos is None: @@ -685,7 +682,7 @@ def drop_on( logger.info( f"Dropping on '{target_object_name}' at corrected position ({x:.3f}, {y:.3f}, {z:.3f})" ) - return self.place(x, y, z, robot_name) + return self.place(x, y, z, group_id) @skill def pick_and_place( @@ -695,7 +692,7 @@ def pick_and_place( place_y: float, place_z: float, object_id: str | None = None, - robot_name: str | None = None, + group_id: PlanningGroupID | None = None, ) -> SkillResult[ManipulationSkillError]: """Pick up an object and place it at a target location. @@ -707,7 +704,7 @@ def pick_and_place( place_y: Target Y position to place the object (meters). place_z: Target Z position to place the object (meters). object_id: Optional unique object ID from perception. - robot_name: Robot to use (only needed for multi-arm setups). + group_id: Planning group to use; omission requires exactly one group. """ logger.info( f"Starting pick and place: pick '{object_name}' → place at " @@ -715,12 +712,12 @@ def pick_and_place( ) # Pick phase - pick_result = self.pick(object_name, object_id, robot_name) + pick_result = self.pick(object_name, object_id, group_id) if not pick_result.is_success(): return pick_result # Place phase - return self.place(place_x, place_y, place_z, robot_name) + return self.place(place_x, place_y, place_z, group_id) @rpc def stop(self) -> None: diff --git a/dimos/manipulation/planning/examples/manipulation_client.py b/dimos/manipulation/planning/examples/manipulation_client.py index de815371ec..60b19162a8 100644 --- a/dimos/manipulation/planning/examples/manipulation_client.py +++ b/dimos/manipulation/planning/examples/manipulation_client.py @@ -37,8 +37,7 @@ execute() Execute planned trajectory via coordinator home() Move to home position url() Get Meshcat visualization URL - robots() List configured robots - info(robot) Get robot config details + info() Get model configuration details gripper(pos) Set gripper position (0.0=closed, 0.85=open) add_box(name,x,y,z) Add box obstacle add_sphere(name,x,y,z) Add sphere obstacle @@ -69,14 +68,14 @@ _client = RPCClient(None, ManipulationModule) -def joints(robot_name: str | None = None) -> list[float] | None: +def joints() -> list[float] | None: """Get current joint positions.""" - return _client.get_current_joints(robot_name) + return _client.get_current_joints() -def ee(robot_name: str | None = None) -> Pose | None: +def ee(group_id: str | None = None) -> Pose | None: """Get end-effector pose.""" - return _client.get_ee_pose(robot_name) + return _client.get_ee_pose(group_id) def state() -> str: @@ -84,10 +83,10 @@ def state() -> str: return _client.get_state() -def plan(target_joints: list[float], robot_name: str | None = None) -> bool: +def plan(target_joints: list[float], group_id: str | None = None) -> bool: """Plan to joint configuration. e.g. plan([0.1]*7)""" js = JointState(position=target_joints) - return _client.plan_to_joints(js, robot_name) + return _client.plan_to_joints(js, group_id) def groups() -> list[PlanningGroup]: @@ -112,21 +111,20 @@ def _make_target_pose( roll: float | None = None, pitch: float | None = None, yaw: float | None = None, - robot_name: str | None = None, + group_id: str | None = None, ) -> Pose: """Create a target pose, preserving current orientation if rpy is not given.""" if roll is not None or pitch is not None or yaw is not None: orientation = Quaternion.from_euler(Vector3(x=roll or 0, y=pitch or 0, z=yaw or 0)) else: # Preserve current EE orientation - current = _client.get_ee_pose(robot_name) + current = _client.get_ee_pose(group_id) orientation = current.orientation if current else Quaternion(0, 0, 0, 1) return Pose(position=Vector3(x=x, y=y, z=z), orientation=orientation) def _make_seed_joint_state( seed_joints: list[float] | JointState | None, - robot_name: str | None, ) -> JointState | None: """Create a seed JointState for IK from explicit joints, if provided.""" if seed_joints is None: @@ -134,7 +132,7 @@ def _make_seed_joint_state( if isinstance(seed_joints, JointState): return seed_joints - info = _client.get_robot_info(robot_name) or {} + info = _client.get_model_info() or {} joint_names = info.get("joint_names", []) if len(joint_names) != len(seed_joints): joint_names = [] @@ -148,7 +146,7 @@ def ik_pose( roll: float | None = None, pitch: float | None = None, yaw: float | None = None, - robot_name: str | None = None, + group_id: str | None = None, seed_joints: list[float] | JointState | None = None, ) -> IKResult: """Solve IK for a Cartesian pose without path planning. @@ -160,13 +158,13 @@ def ik_pose( roll: Optional target roll. Preserves current orientation if omitted. 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. + group_id: Planning group to solve; omission requires one compatible group. 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.inverse_kinematics_single(target, robot_name, seed) + target = _make_target_pose(x, y, z, roll, pitch, yaw, group_id) + seed = _make_seed_joint_state(seed_joints) + return _client.inverse_kinematics_single(target, group_id, seed) def ik_group_pose( @@ -196,11 +194,11 @@ def plan_pose( roll: float | None = None, pitch: float | None = None, yaw: float | None = None, - robot_name: str | None = None, + group_id: str | None = None, ) -> bool: """Plan to Cartesian pose. Preserves current orientation if rpy not given.""" - target = _make_target_pose(x, y, z, roll, pitch, yaw, robot_name) - return _client.plan_to_pose(target, robot_name) + target = _make_target_pose(x, y, z, roll, pitch, yaw, group_id) + return _client.plan_to_pose(target, group_id) def plan_group_pose( @@ -219,10 +217,9 @@ def plan_group_pose( def preview( duration: float | None = None, - robot_name: str | None = None, ) -> bool: """Preview the last generated plan in the visualizer.""" - return _client.preview_plan(None, duration, robot_name) + return _client.preview_plan(None, duration) def execute() -> bool: @@ -230,12 +227,12 @@ def execute() -> bool: return _client.execute() -def home(robot_name: str | None = None) -> bool: +def home(group_id: str | None = None) -> bool: """Plan and execute move to home position.""" from dimos.msgs.sensor_msgs.JointState import JointState - 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) + home_joints = _client.get_model_info().get("home_joints", [0.0] * 7) + success = _client.plan_to_joints(JointState(position=home_joints), group_id) if success: return _client.execute() return False @@ -246,19 +243,14 @@ def url() -> str | None: return _client.get_visualization_url() -def robots() -> list[str]: - """List configured robots.""" - return _client.list_robots() +def info() -> dict[str, Any]: + """Get configured model details.""" + return _client.get_model_info() -def info(robot_name: str | None = None) -> dict[str, Any] | None: - """Get robot config details.""" - return _client.get_robot_info(robot_name) - - -def gripper(position: float, robot_name: str | None = None) -> str: +def gripper(position: float) -> str: """Set gripper position (0.0=closed, 0.85=open).""" - return _client.set_gripper(position, robot_name) + return _client.set_gripper(position) def add_box( @@ -357,9 +349,9 @@ def remove(obstacle_id: str) -> bool: return _client.remove_obstacle(obstacle_id) -def collision_free(target_joints: list[float], robot_name: str | None = None) -> bool: +def collision_free(target_joints: list[float]) -> bool: """Check if a joint configuration is collision-free.""" - return _client.is_collision_free(target_joints, robot_name) + return _client.is_collision_free(target_joints) def commands() -> None: diff --git a/dimos/manipulation/planning/factory.py b/dimos/manipulation/planning/factory.py index 453307c77f..6fecc499fd 100644 --- a/dimos/manipulation/planning/factory.py +++ b/dimos/manipulation/planning/factory.py @@ -216,15 +216,12 @@ def create_planning_specs( world: WorldSpec, world_backend: str = "roboplan", planner: ManipulationPlannerConfig | None = None, - kinematics_name: str | None = None, kinematics: ManipulationKinematicsConfig | None = None, trajectory_parametrization: TrajectoryParametrizationConfig | None = None, ) -> PlanningSpecs: """Create planning specs around an already-created world.""" from dimos.manipulation.planning.monitor.world_monitor import WorldMonitor - if kinematics_name is not None: - kinematics = kinematics_config_from_name(kinematics_name) if kinematics is None: kinematics = kinematics_config_from_name(DEFAULT_KINEMATICS_NAME) if planner is None: @@ -259,20 +256,18 @@ def create_planning_stack( world_backend: str = "roboplan", visualization: ManipulationVisualizationConfig | None = None, planner: ManipulationPlannerConfig | None = None, - kinematics_name: str | None = None, kinematics: ManipulationKinematicsConfig | None = None, -) -> tuple[WorldSpec, KinematicsSpec, PlannerSpec, str]: - """Create complete planning stack. Returns (world, kinematics, planner, robot_id).""" +) -> tuple[WorldSpec, KinematicsSpec, PlannerSpec]: + """Create and finalize a complete single-model planning stack.""" world = create_world(backend=world_backend, visualization=visualization) planning_specs = create_planning_specs( world=world, world_backend=world_backend, planner=planner, - kinematics_name=kinematics_name, kinematics=kinematics, ) - robot_id = world.add_robot(robot_config) + world.load_model(robot_config) world.finalize() - return world, planning_specs.kinematics, planning_specs.planner, robot_id + return world, planning_specs.kinematics, planning_specs.planner diff --git a/dimos/manipulation/planning/groups/discovery.py b/dimos/manipulation/planning/groups/discovery.py index 5cf596c245..5a927602b3 100644 --- a/dimos/manipulation/planning/groups/discovery.py +++ b/dimos/manipulation/planning/groups/discovery.py @@ -35,7 +35,6 @@ class PlanningGroupDiscoveryError(ValueError): def discover_planning_group_definitions( *, - robot_name: str, model_path: Path, model: ModelDescription, controllable_joint_names: list[str], @@ -57,7 +56,7 @@ def discover_planning_group_definitions( return groups logger.warning( f"No supported planning groups found in SRDF {resolved_srdf_path} " - f"for robot {robot_name}; trying fallback generation" + "for the configured model; trying fallback generation" ) return [ diff --git a/dimos/manipulation/planning/groups/identifiers.py b/dimos/manipulation/planning/groups/identifiers.py index 6cf19511f8..1b5094318f 100644 --- a/dimos/manipulation/planning/groups/identifiers.py +++ b/dimos/manipulation/planning/groups/identifiers.py @@ -12,103 +12,28 @@ # 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 +"""Canonical planning-group and joint-name validation.""" from collections.abc import Sequence -from dimos.manipulation.planning.spec.models import ( - GlobalJointName, - LocalModelJointName, - PlanningGroupID, - RobotName, -) +from dimos.manipulation.planning.spec.models import JointName, PlanningGroupID -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_group_id(group_id: PlanningGroupID) -> None: + """Validate a stable, unprefixed planning-group name.""" + if not group_id or "/" in group_id: + raise ValueError(f"Invalid planning group ID: {group_id!r}") -def assert_valid_local_joint_name(local_joint_name: LocalModelJointName) -> None: - """Validate an exact model joint name before adding the legacy robot scope.""" - if not local_joint_name or local_joint_name.startswith("/") or local_joint_name.endswith("/"): - raise ValueError(f"Invalid model joint name: {local_joint_name!r}") - if any(not part for part in local_joint_name.split("/")): - raise ValueError(f"Invalid model joint name: {local_joint_name!r}") +def assert_valid_joint_name(joint_name: JointName) -> None: + """Validate an exact canonical model joint name.""" + if not joint_name or joint_name.startswith("/") or joint_name.endswith("/"): + raise ValueError(f"Invalid canonical joint name: {joint_name!r}") + if any(not part for part in joint_name.split("/")): + raise ValueError(f"Invalid canonical joint name: {joint_name!r}") -def assert_local_joint_names(names: Sequence[LocalModelJointName]) -> None: - """Validate that names are local model joint names, not global joint names.""" +def assert_valid_joint_names(names: Sequence[JointName]) -> None: + """Validate canonical 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 a robot scope followed by a canonical model name.""" - parts = name.split("/") - return len(parts) >= 2 and all(parts) - - -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 + assert_valid_joint_name(name) diff --git a/dimos/manipulation/planning/groups/models.py b/dimos/manipulation/planning/groups/models.py index 599a9ab85d..0ee63bcc78 100644 --- a/dimos/manipulation/planning/groups/models.py +++ b/dimos/manipulation/planning/groups/models.py @@ -20,10 +20,8 @@ from typing import Literal, TypeAlias from dimos.manipulation.planning.spec.models import ( - GlobalJointName, - LocalModelJointName, + JointName, PlanningGroupID, - RobotName, ) PlanningGroupSource: TypeAlias = Literal["srdf", "fallback"] @@ -38,7 +36,7 @@ class PlanningGroupDefinition: """ name: str - joint_names: tuple[LocalModelJointName, ...] + joint_names: tuple[JointName, ...] base_link: str tip_link: str | None = None source: PlanningGroupSource = "srdf" @@ -53,15 +51,11 @@ def has_pose_target(self) -> bool: 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. + A planning group exposes one stable name and canonical model joint names. """ id: PlanningGroupID - robot_name: RobotName - group_name: str - joint_names: tuple[GlobalJointName, ...] - local_joint_names: tuple[LocalModelJointName, ...] + joint_names: tuple[JointName, ...] base_link: str tip_link: str | None = None source: PlanningGroupSource = "srdf" @@ -82,23 +76,19 @@ class PlanningGroupSelection: groups: tuple[PlanningGroup, ...] group_ids: tuple[PlanningGroupID, ...] - joint_names: tuple[GlobalJointName, ...] - robot_names: tuple[RobotName, ...] + joint_names: tuple[JointName, ...] @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] = [] + """Build a selection, rejecting overlapping selected joints.""" + seen_joints: dict[JointName, PlanningGroupID] = {} + joint_names: list[JointName] = [] 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 " + "Selected planning groups overlap on joint " f"{joint_name}: {previous_group_id} and {group.id}" ) seen_joints[joint_name] = group.id @@ -108,5 +98,4 @@ def from_groups(cls, groups: tuple[PlanningGroup, ...]) -> PlanningGroupSelectio 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 index fb982b3562..bdd253c6bc 100644 --- a/dimos/manipulation/planning/groups/registry.py +++ b/dimos/manipulation/planning/groups/registry.py @@ -19,57 +19,41 @@ 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 +from dimos.manipulation.planning.spec.models import PlanningGroupID if TYPE_CHECKING: from dimos.manipulation.planning.spec.config import RobotModelConfig class PlanningGroupRegistry: - """Registry of public planning groups derived from robot configs.""" + """Registry of public planning groups derived from one model config.""" - def __init__(self, robot_configs: Iterable[RobotModelConfig] = ()) -> None: + def __init__(self, model_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) + for config in model_configs: + self.add_model(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] = [] + def add_model(self, config: RobotModelConfig) -> None: + """Register all planning groups declared by the model config.""" + if self._groups: + raise ValueError("A model is already registered") for definition in config.planning_groups: - group_id = make_planning_group_id(config.name, definition.name) + group_id = 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, + 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) + return tuple(self._groups.values()) def get(self, group_id: PlanningGroupID) -> PlanningGroup: """Return one planning group by public ID.""" @@ -84,34 +68,19 @@ def select(self, group_ids: Iterable[PlanningGroupID]) -> PlanningGroupSelection 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 default_group_id(self) -> PlanningGroupID | None: + """Return the sole group ID when selection is unambiguous.""" + groups = self.list() + return groups[0].id if len(groups) == 1 else 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 - ] + def primary_pose_group_id(self) -> PlanningGroupID | None: + """Return the unique pose-targetable group ID.""" + pose_groups = [group for group in self.list() 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; " + f"Model 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 index 22a7be40b6..fd84f8cc4a 100644 --- a/dimos/manipulation/planning/groups/test_planning_groups.py +++ b/dimos/manipulation/planning/groups/test_planning_groups.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,9 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for planning groups.""" - -from __future__ import annotations +"""Tests for canonical, single-model planning groups.""" from pathlib import Path @@ -27,19 +25,15 @@ 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, + normalize_joint_target, 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 @@ -61,691 +55,136 @@ def _serial_model(*joint_types: str) -> ModelDescription: ) -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 _model_with_branched_prismatic_gripper() -> ModelDescription: - return ModelDescription( - joints=[ - JointDescription( - name="arm_joint", - type="revolute", - parent_link="base", - child_link="wrist", - ), - JointDescription( - name="left_finger_joint", - type="prismatic", - parent_link="wrist", - child_link="left_finger", - ), - JointDescription( - name="right_finger_joint", - type="prismatic", - parent_link="wrist", - child_link="right_finger", - ), - ], - root_link="base", - links=["base", "wrist", "left_finger", "right_finger"], - ) - - -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: +def _config(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", - ) + model_path=Path("/tmp/model.urdf"), + joint_names=["left/j1", "left/j2", "right/j1"], + planning_groups=groups + or [ + PlanningGroupDefinition("left_arm", ("left/j1", "left/j2"), "base", "left/tool"), + PlanningGroupDefinition("right_arm", ("right/j1",), "base", "right/tool"), ], ) def test_parse_srdf_chain_group(tmp_path: Path) -> None: - model = _serial_model("revolute", "revolute", "revolute") - srdf_path = _write_srdf( - tmp_path, - "", + path = tmp_path / "model.srdf" + path.write_text( + "" + "" ) - 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, + path, + model=_serial_model("revolute", "revolute"), controllable_joint_names=["joint1", "joint2"], ) + assert groups == [ + PlanningGroupDefinition("arm", ("joint1", "joint2"), "link0", "link2", "srdf") + ] - 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") +def test_fallback_generation_and_branch_rejection() -> None: group = generate_fallback_planning_group( - model=model, + model=_serial_model("revolute", "revolute", "prismatic"), controllable_joint_names=["joint1", "joint2", "joint3"], ) - + assert group.name == FALLBACK_PLANNING_GROUP_NAME assert group.joint_names == ("joint1", "joint2") - assert group.tip_link == "link2" - assert group.source == "fallback" - -def test_fallback_excludes_branched_terminal_prismatic_gripper_joints() -> None: - group = generate_fallback_planning_group( - model=_model_with_branched_prismatic_gripper(), - controllable_joint_names=[ - "arm_joint", - "left_finger_joint", - "right_finger_joint", + branching = ModelDescription( + joints=[ + JointDescription("left", "revolute", "base", "left_link"), + JointDescription("right", "revolute", "base", "right_link"), ], + root_link="base", + links=["base", "left_link", "right_link"], ) - - assert group.joint_names == ("arm_joint",) - assert group.base_link == "base" - assert group.tip_link == "wrist" - - -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"], + model=branching, controllable_joint_names=["left", "right"] ) -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_path=tmp_path / "model.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_uses_stable_unprefixed_ids_and_canonical_joints() -> None: + registry = PlanningGroupRegistry([_config()]) + assert [group.id for group in registry.list()] == ["left_arm", "right_arm"] + assert registry.get("left_arm").joint_names == ("left/j1", "left/j2") + assert registry.default_group_id() is None + with pytest.raises(ValueError, match="explicit planning group ID"): + registry.primary_pose_group_id() -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: +def test_registry_default_requires_exactly_one_compatible_group() -> 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, + [_config([PlanningGroupDefinition("arm", ("left/j1",), "base", "tool")])] ) + assert registry.default_group_id() == "arm" + assert registry.primary_pose_group_id() == "arm" - 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_selection_preserves_order_and_rejects_overlap() -> None: + registry = PlanningGroupRegistry([_config()]) + selection = registry.select(("right_arm", "left_arm")) + assert selection.group_ids == ("right_arm", "left_arm") + assert selection.joint_names == ("right/j1", "left/j1", "left/j2") -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( + overlap = _config( [ - _robot_config( - planning_groups=[ - PlanningGroupDefinition("arm", ("joint1", "joint2"), "base", "tool"), - PlanningGroupDefinition("gripper", ("joint3",), "tool"), - ] - ) + PlanningGroupDefinition("one", ("left/j1",), "base"), + PlanningGroupDefinition("two", ("left/j1",), "base"), ] ) - - 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] + PlanningGroupRegistry([overlap]).select(("one", "two")) -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]}) - ) +def test_normalize_joint_target_accepts_exact_or_unnamed_target() -> None: + group = PlanningGroup("left_arm", ("left/j1", "left/j2"), "base", "tool") + named = normalize_joint_target( + group, + JointState(name=["left/j1", "left/j2"], position=[1.0, 2.0]), + ) + unnamed = normalize_joint_target(group, JointState(position=[3.0, 4.0])) + assert named.name == ["left/j1", "left/j2"] + assert unnamed.name == ["left/j1", "left/j2"] 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]}) + normalize_joint_target( + group, + JointState(name=["j1", "j2"], position=[1.0, 2.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]}) +def test_state_projection_requires_exact_canonical_names() -> None: + state = JointState(name=["right/j1", "left/j2", "left/j1"], position=[3.0, 2.0, 1.0]) + projected = filter_joint_state_to_selected_joints(state, ("left/j1", "right/j1")) + assert projected.name == ["left/j1", "right/j1"] + assert projected.position == [1.0, 3.0] + assert joint_state_to_ordered_positions( + state, joint_names=("left/j1", "left/j2", "right/j1") + ).tolist() == [1.0, 2.0, 3.0] + with pytest.raises(ValueError, match="missing"): + filter_joint_state_to_selected_joints(state, ("missing",)) - 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_keeps_joint_mapping_without_model_wide_tip() -> 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 not hasattr(config, "end_effector_link") - 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_accepts_canonical_slash_joint_names() -> None: - config = RobotModelConfig( - name="arm", - model_path=Path("robot.urdf"), - joint_names=["left/j1", "right/j1"], - planning_groups=[ - PlanningGroupDefinition( - name="left_arm", - joint_names=("left/j1",), - base_link="base", - tip_link="left_tool", - ), - PlanningGroupDefinition( - name="right_arm", - joint_names=("right/j1",), - base_link="base", - tip_link="right_tool", - ), - ], - ) - - assert config.joint_names == ["left/j1", "right/j1"] +def test_selector_accepts_id_or_group() -> None: + group = PlanningGroup("arm", ("left/j1",), "base", "tool") + assert planning_group_id_from_selector("arm") == "arm" + assert planning_group_id_from_selector(group) == "arm" -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, - ) +def test_model_config_rejects_obsolete_name_and_mapping_fields() -> None: + base = {"model_path": Path("/tmp/model.urdf"), "joint_names": ["joint1"]} + with pytest.raises(ValueError): + RobotModelConfig(**base, name="arm") + with pytest.raises(ValueError): + RobotModelConfig(**base, joint_name_mapping={"arm/joint1": "joint1"}) diff --git a/dimos/manipulation/planning/groups/utils.py b/dimos/manipulation/planning/groups/utils.py index 5d9161e5bd..f7e16b1fbd 100644 --- a/dimos/manipulation/planning/groups/utils.py +++ b/dimos/manipulation/planning/groups/utils.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,91 +12,39 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Shared helpers for planning-group selectors and joint-state projection.""" +"""Planning-group selection and canonical joint-state projection.""" -from collections.abc import Mapping, Sequence +from collections.abc import 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.manipulation.planning.spec.models import PlanningGroupID 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 + return selector.id if isinstance(selector, PlanningGroup) else selector def filter_joint_state_to_selected_joints( - joint_state: JointState, - global_joint_names: Sequence[GlobalJointName], - local_joint_names: Sequence[LocalModelJointName] = (), + joint_state: JointState, joint_names: Sequence[str] ) -> 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") - + """Project a canonical joint state to selected joints in the requested order.""" 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) - + missing = [name for name in joint_names if name not in positions_by_name] if missing: - raise ValueError(f"IK result is missing selected joints: {missing}") - - return JointState({"name": list(global_joint_names), "position": selected_positions}) - + raise ValueError(f"Joint state is missing selected joints: {missing}") + return JointState( + name=list(joint_names), + position=[float(positions_by_name[name]) for name in joint_names], + ) -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. - """ +def normalize_joint_target(group: PlanningGroup, target: JointState) -> JointState: + """Normalize one group target to canonical group joint order.""" if not target.name: if len(target.position) != len(group.joint_names): raise ValueError( @@ -104,133 +52,42 @@ def joint_target_to_global_names( 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) + positions = dict(zip(target.name, target.position, strict=True)) + missing = set(group.joint_names) - positions.keys() + extra = positions.keys() - set(group.joint_names) if missing: - raise ValueError(f"Target for '{group.id}' is missing joints: {missing}") - - extra = set(target_names) - set(expected_names) + raise ValueError(f"Target for '{group.id}' is missing joints: {sorted(missing)}") 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) + return JointState( + name=list(group.joint_names), + position=[float(positions[name]) for name in group.joint_names], ) - 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], + joint_state: JointState, *, joint_names: Sequence[str] ) -> NDArray[np.float64]: - """Convert a JointState to an array ordered by local robot joint names.""" + """Normalize an unnamed or canonically named state to model joint order.""" 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" + f"Joint state has {len(joint_state.position)} positions, " + f"expected {len(joint_names)}" ) - - 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] + return np.asarray(joint_state.position, dtype=np.float64) + if len(joint_state.name) != len(joint_state.position): + raise ValueError("Joint state names and positions must have the same length") + positions = dict(zip(joint_state.name, joint_state.position, strict=True)) + missing = [name for name in joint_names if name not in positions] + extra = set(positions) - set(joint_names) if missing: - raise ValueError(f"JointState missing joints: {missing}") - return np.asarray([name_to_pos[name] for name in joint_names], dtype=np.float64) + raise ValueError(f"Joint state is missing joints: {missing}") + if extra: + raise ValueError(f"Joint state has unknown joints: {sorted(extra)}") + return np.asarray([positions[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 91da25986c..92ad481ac5 100644 --- a/dimos/manipulation/planning/kinematics/drake_optimization_ik.py +++ b/dimos/manipulation/planning/kinematics/drake_optimization_ik.py @@ -17,7 +17,7 @@ from __future__ import annotations from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast import numpy as np @@ -25,10 +25,10 @@ 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, + unique_pose_target_frame as _unique_pose_target_frame, ) 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 from dimos.manipulation.planning.spec.protocols import WorldSpec from dimos.manipulation.planning.utils.kinematics_utils import compute_pose_error from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped @@ -39,6 +39,8 @@ if TYPE_CHECKING: from numpy.typing import NDArray + from dimos.manipulation.planning.world.drake_world import DrakeWorld + try: from pydrake.math import RigidTransform, RotationMatrix from pydrake.multibody.inverse_kinematics import ( @@ -77,7 +79,6 @@ def _validate_world(self, world: WorldSpec) -> IKResult | None: def solve( self, world: WorldSpec, - robot_id: WorldRobotID, target_pose: PoseStamped, seed: JointState | None = None, position_tolerance: float = 0.001, @@ -90,7 +91,7 @@ def solve( if error is not None: return error - target_frame_name = _unique_pose_target_frame_for_robot(world, robot_id) + target_frame_name = _unique_pose_target_frame(world) if target_frame_name is None: return _create_failure_result( IKStatus.UNSUPPORTED, @@ -104,12 +105,12 @@ def solve( ).to_matrix() # Get joint limits - lower_limits, upper_limits = world.get_joint_limits(robot_id) + lower_limits, upper_limits = world.get_joint_limits() # Get seed from current state if not provided if seed is None: with world.scratch_context() as ctx: - seed = world.get_joint_state(ctx, robot_id) + seed = world.get_joint_state(ctx) # Extract joint names and seed positions joint_names = seed.name @@ -132,7 +133,6 @@ def solve( # Solve IK result = self._solve_single( world=world, - robot_id=robot_id, target_transform=target_transform, seed=current_seed, joint_names=joint_names, @@ -146,7 +146,7 @@ def solve( if result.is_success() and result.joint_state is not None: # Check collision if requested if check_collision: - if not world.check_config_collision_free(robot_id, result.joint_state): + if not world.check_config_collision_free(result.joint_state): continue # Try another seed # Check error @@ -200,7 +200,7 @@ def solve_pose_targets( "DrakeOptimizationIK requires a pose-targetable planning group", ) - lower_limits, upper_limits = world.get_joint_limits(request.robot_id) + lower_limits, upper_limits = world.get_joint_limits() target_matrix = Transform( translation=request.target_pose.position, rotation=request.target_pose.orientation, @@ -226,7 +226,6 @@ def solve_pose_targets( result = self._solve_single( world=world, - robot_id=request.robot_id, target_transform=target_transform, seed=current_seed, joint_names=request.joint_names, @@ -239,9 +238,7 @@ def solve_pose_targets( ) 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 - ): + if check_collision and not world.check_config_collision_free(result.joint_state): continue total_error = result.position_error + result.orientation_error if total_error < best_error: @@ -263,7 +260,6 @@ def solve_pose_targets( def _solve_single( self, world: WorldSpec, - robot_id: WorldRobotID, target_transform: RigidTransform, seed: NDArray[np.float64], joint_names: list[str], @@ -274,16 +270,14 @@ def _solve_single( 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] - plant = world.plant # type: ignore[attr-defined] + drake_world = cast("DrakeWorld", world) + plant = drake_world.plant + joint_indices = drake_world.get_model_joint_indices() # Create IK problem ik = InverseKinematics(plant) - target_frame = plant.GetBodyByName( - target_frame_name, robot_data.model_instance - ).body_frame() + target_frame = drake_world.get_body_frame(target_frame_name) # Add position constraint ik.AddPositionConstraint( @@ -308,12 +302,12 @@ def _solve_single( q = ik.q() for local_index, value in (locked_joint_positions or {}).items(): - joint_idx = robot_data.joint_indices[local_index] + joint_idx = 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): + for i, joint_idx in enumerate(joint_indices): full_seed[joint_idx] = seed[i] prog.SetInitialGuess(q, full_seed) @@ -328,7 +322,7 @@ def _solve_single( # Extract solution for this robot's joints full_solution = result.GetSolution(q) - joint_solution = np.array([full_solution[idx] for idx in robot_data.joint_indices]) + joint_solution = np.array([full_solution[idx] for idx in joint_indices]) # Clip to limits joint_solution = np.clip(joint_solution, lower_limits, upper_limits) @@ -336,8 +330,8 @@ def _solve_single( # Compute actual error using FK 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_matrix = world.get_link_pose(ctx, robot_id, target_frame_name) + world.set_joint_state(ctx, solution_state) + actual_matrix = world.get_link_pose(ctx, target_frame_name) position_error, orientation_error = compute_pose_error( actual_matrix, diff --git a/dimos/manipulation/planning/kinematics/jacobian_ik.py b/dimos/manipulation/planning/kinematics/jacobian_ik.py index 4c4e16207a..110b13ec46 100644 --- a/dimos/manipulation/planning/kinematics/jacobian_ik.py +++ b/dimos/manipulation/planning/kinematics/jacobian_ik.py @@ -34,7 +34,7 @@ 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.models import IKResult from dimos.manipulation.planning.spec.protocols import WorldSpec from dimos.manipulation.planning.utils.kinematics_utils import ( check_singularity, @@ -73,7 +73,7 @@ class JacobianIK: Example: ik = JacobianIK(damping=0.01) result = ik.solve_iterative( - world, robot_id, + world, target_pose=target, seed=current_joints, ) @@ -101,7 +101,6 @@ def __init__( def solve( self, world: WorldSpec, - robot_id: WorldRobotID, target_pose: PoseStamped, seed: JointState | None = None, position_tolerance: float = 0.001, @@ -116,7 +115,6 @@ def solve( Args: world: World for FK/collision checking - robot_id: Robot to solve IK for target_pose: Target end-effector pose seed: Initial guess (uses current state if None) position_tolerance: Required position accuracy (meters) @@ -130,12 +128,12 @@ def solve( if not world.is_finalized: return _create_failure_result(IKStatus.NO_SOLUTION, "World must be finalized before IK") - lower_limits, upper_limits = world.get_joint_limits(robot_id) + lower_limits, upper_limits = world.get_joint_limits() # Get seed from current state if not provided if seed is None: with world.scratch_context() as ctx: - seed = world.get_joint_state(ctx, robot_id) + seed = world.get_joint_state(ctx) # Extract joint names for creating random seeds joint_names = seed.name @@ -157,7 +155,6 @@ def solve( # Solve iterative IK result = self.solve_iterative( world=world, - robot_id=robot_id, target_pose=target_pose, seed=current_seed, max_iterations=self._max_iterations, @@ -168,7 +165,7 @@ def solve( if result.is_success() and result.joint_state is not None: # Check collision if requested if check_collision: - if not world.check_config_collision_free(robot_id, result.joint_state): + if not world.check_config_collision_free(result.joint_state): continue # Try another seed # Check error @@ -223,7 +220,6 @@ def solve_pose_targets( ) 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), @@ -243,14 +239,13 @@ def solve_pose_targets( full_state = JointState( {"name": request.joint_names, "position": full_positions.tolist()} ) - if not world.check_config_collision_free(request.robot_id, full_state): + if not world.check_config_collision_free(full_state): return _create_failure_result(IKStatus.COLLISION, "IK solution is in collision") return result def solve_iterative( self, world: WorldSpec, - robot_id: WorldRobotID, target_pose: PoseStamped, seed: JointState, max_iterations: int = 100, @@ -266,7 +261,6 @@ def solve_iterative( Args: world: World for FK/Jacobian computation - robot_id: Robot to solve IK for target_pose: Target end-effector pose seed: Initial joint configuration max_iterations: Maximum iterations before giving up @@ -287,7 +281,7 @@ def solve_iterative( 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) + lower_limits, upper_limits = world.get_joint_limits() for iteration in range(max_iterations): with world.scratch_context() as ctx: @@ -295,10 +289,10 @@ def solve_iterative( current_state = JointState( {"name": joint_names, "position": current_joints.tolist()} ) - world.set_joint_state(ctx, robot_id, current_state) + world.set_joint_state(ctx, current_state) if group is None: - current_pose = pose_to_matrix(world.get_ee_pose(ctx, robot_id)) + current_pose = pose_to_matrix(world.get_ee_pose(ctx)) else: current_pose = pose_to_matrix(world.get_group_ee_pose(ctx, group.id)) @@ -316,7 +310,7 @@ def solve_iterative( ) if group is None: - jacobian = world.get_jacobian(ctx, robot_id) + jacobian = world.get_jacobian(ctx) else: jacobian = world.get_group_jacobian(ctx, group.id) @@ -352,9 +346,9 @@ def solve_iterative( # Compute final error with world.scratch_context() as ctx: final_state = JointState({"name": joint_names, "position": current_joints.tolist()}) - world.set_joint_state(ctx, robot_id, final_state) + world.set_joint_state(ctx, final_state) if group is None: - final_pose = pose_to_matrix(world.get_ee_pose(ctx, robot_id)) + final_pose = pose_to_matrix(world.get_ee_pose(ctx)) 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) @@ -368,7 +362,6 @@ def solve_iterative( def solve_differential( self, world: WorldSpec, - robot_id: WorldRobotID, current_joints: JointState, twist: Twist, dt: float, @@ -380,7 +373,6 @@ def solve_differential( Args: world: World for Jacobian computation - robot_id: Robot to compute for current_joints: Current joint configuration twist: Desired end-effector twist (linear + angular velocity) dt: Time step (not used, but kept for interface compatibility) @@ -403,8 +395,8 @@ def solve_differential( joint_names = current_joints.name with world.scratch_context() as ctx: - world.set_joint_state(ctx, robot_id, current_joints) - J = world.get_jacobian(ctx, robot_id) + world.set_joint_state(ctx, current_joints) + J = world.get_jacobian(ctx) # Check for singularity if check_singularity(J, threshold=self._singularity_threshold): @@ -418,7 +410,7 @@ def solve_differential( q_dot = J_pinv @ twist_array # Apply velocity limits if available - config = world.get_robot_config(robot_id) + config = world.get_model_config() if config.velocity_limits is not None: velocity_limits = np.array(config.velocity_limits) # Only consider joints with non-zero velocity limits @@ -433,7 +425,6 @@ def solve_differential( def solve_differential_position_only( self, world: WorldSpec, - robot_id: WorldRobotID, current_joints: JointState, linear_velocity: Vector3, ) -> JointState | None: @@ -444,7 +435,6 @@ def solve_differential_position_only( Args: world: World for Jacobian computation - robot_id: Robot to compute for current_joints: Current joint configuration linear_velocity: Desired linear velocity @@ -458,8 +448,8 @@ def solve_differential_position_only( joint_names = current_joints.name with world.scratch_context() as ctx: - world.set_joint_state(ctx, robot_id, current_joints) - J = world.get_jacobian(ctx, robot_id) + world.set_joint_state(ctx, current_joints) + J = world.get_jacobian(ctx) # Extract linear part (first 3 rows) J_linear = J[:3, :] diff --git a/dimos/manipulation/planning/kinematics/pink_ik.py b/dimos/manipulation/planning/kinematics/pink_ik.py index 3456c2e27f..866821dd67 100644 --- a/dimos/manipulation/planning/kinematics/pink_ik.py +++ b/dimos/manipulation/planning/kinematics/pink_ik.py @@ -17,7 +17,7 @@ from __future__ import annotations from collections.abc import Mapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, replace import importlib from pathlib import Path from types import ModuleType @@ -28,14 +28,12 @@ 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, + unique_pose_target_frame as _unique_pose_target_frame, ) from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import IKStatus -from dimos.manipulation.planning.spec.models import IKResult, RobotName, WorldRobotID +from dimos.manipulation.planning.spec.models import IKResult 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 @@ -106,12 +104,11 @@ def __init__( config_values.update(overrides) self.config = PinkKinematicsConfig(**config_values) self._modules = _load_optional_dependencies(self.config.solver) - self._robot_contexts: dict[tuple[str, str], _PinkRobotContext] = {} + self._model_context: _PinkRobotContext | None = None def solve( self, world: WorldSpec, - robot_id: WorldRobotID, target_pose: PoseStamped, seed: JointState | None = None, position_tolerance: float = 0.001, @@ -123,7 +120,7 @@ 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) + target_frame_name = _unique_pose_target_frame(world) if target_frame_name is None: return _failure( IKStatus.NO_SOLUTION, @@ -131,16 +128,16 @@ def solve( ) try: - robot_context = self._get_robot_context(world, robot_id, target_frame_name) + robot_context = self._get_model_context(world, target_frame_name) except (FileNotFoundError, ImportError, ValueError) as exc: return _failure(IKStatus.NO_SOLUTION, f"Pink IK model setup failed: {exc}") if seed is None: with world.scratch_context() as ctx: - seed = world.get_joint_state(ctx, robot_id) + seed = world.get_joint_state(ctx) - lower_limits, upper_limits = world.get_joint_limits(robot_id) - target_model = self._target_in_model_frame(world.get_robot_config(robot_id), target_pose) + lower_limits, upper_limits = world.get_joint_limits() + target_model = self._target_in_model_frame(world.get_model_config(), target_pose) fallback_result: IKResult | None = None @@ -166,9 +163,7 @@ def solve( fallback_result = result continue - if check_collision and not world.check_config_collision_free( - robot_id, result.joint_state - ): + if check_collision and not world.check_config_collision_free(result.joint_state): fallback_result = _collision_failure(result) continue @@ -211,61 +206,42 @@ def solve_pose_targets( 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) + config = world.get_model_config() 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) + selected_indices = [joint_names.index(name) for name in selection.joint_names] + seed_positions = _seed_positions_with_world_fallback(world, joint_names, seed) + lower_limits, upper_limits = world.get_joint_limits() 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}") + targets = [ + ( + self._get_model_context(world, group.tip_link or ""), + self._target_in_model_frame(config, target), + ) + for group, target in pose_targets.items() + ] + except (FileNotFoundError, ImportError, ValueError) as exc: + return _failure(IKStatus.NO_SOLUTION, f"Pink IK model setup failed: {exc}") - fallback_result: IKResult | None = None + if not targets: + positions = seed_positions + result = _success(joint_names, positions, 0.0, 0.0, 0) + else: + fallback: IKResult | None = None + result = _failure(IKStatus.NO_SOLUTION, "Pink IK did not produce a solution") for attempt in range(max_attempts): - current_positions = seed_positions.copy() - if attempt > 0: - current_positions[selected_indices] = np.random.uniform( + current = seed_positions.copy() + if attempt: + current[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( + q0 = self._q_from_dimos_positions(targets[0][0], current) + result = ( + self._solve_single( robot_context=targets[0][0], target_model=targets[0][1], seed_q=q0, @@ -275,8 +251,8 @@ def solve_pose_targets( orientation_tolerance=orientation_tolerance, locked_joint_positions=locked_positions, ) - else: - result = self._solve_multi( + if len(targets) == 1 + else self._solve_multi( targets=targets, seed_q=q0, lower_limits=lower_limits, @@ -285,79 +261,32 @@ def solve_pose_targets( 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: + ) + except (ValueError, RuntimeError) 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 + if result.is_success() and result.joint_state is not None: + break + fallback = fallback or result else: - if fallback_result is not None: - return fallback_result - return _failure( - IKStatus.NO_SOLUTION, f"Pink IK failed after {max_attempts} attempts" - ) + return fallback or result - 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( + assert result.joint_state is not None + if check_collision and not world.check_config_collision_free(result.joint_state): + return _collision_failure(result) + positions_by_name = dict( + zip(result.joint_state.name, result.joint_state.position, strict=True) + ) + return IKResult( status=IKStatus.SUCCESS, joint_state=JointState( - { - "name": selected_names, - "position": selected_positions, - } + name=list(selection.joint_names), + position=[positions_by_name[name] for name in selection.joint_names], ), - position_error=max_position_error, - orientation_error=max_orientation_error, - iterations=iterations, - message="Pink IK solution found", + position_error=result.position_error, + orientation_error=result.orientation_error, + iterations=result.iterations, + message=result.message, ) - 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, @@ -522,18 +451,16 @@ def _solve_single( message="Pink IK did not converge within the iteration budget", ) - 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), frame_name - ) - return self._robot_contexts[cache_key] + def _get_model_context(self, world: WorldSpec, frame_name: str) -> _PinkRobotContext: + if self._model_context is None: + self._model_context = self._build_robot_context(world.get_model_config(), frame_name) + if self._model_context.frame_name == frame_name: + return self._model_context + return replace( + self._model_context, + frame_id=_get_frame_id(self._model_context.model, frame_name), + frame_name=frame_name, + ) def _build_robot_context(self, config: RobotModelConfig, frame_name: str) -> _PinkRobotContext: pinocchio = self._modules.pinocchio @@ -668,7 +595,7 @@ def _build_joint_mapping(model: Any, config: RobotModelConfig) -> _JointMapping: model_joint_names: list[str] = [] for dimos_name in config.joint_names: - model_joint_name = config.get_urdf_joint_name(dimos_name) + model_joint_name = dimos_name joint_id = _get_joint_id(model, model_joint_name) joint = model.joints[joint_id] nq = int(getattr(joint, "nq", 1)) @@ -766,23 +693,6 @@ 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], diff --git a/dimos/manipulation/planning/kinematics/test_drake_optimization_ik_selection.py b/dimos/manipulation/planning/kinematics/test_drake_optimization_ik_selection.py index afac663e72..4c717e3224 100644 --- a/dimos/manipulation/planning/kinematics/test_drake_optimization_ik_selection.py +++ b/dimos/manipulation/planning/kinematics/test_drake_optimization_ik_selection.py @@ -31,38 +31,32 @@ 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"], + joint_names=["arm/base", "arm/shoulder", "arm/elbow", "arm/wrist"], ) self.current_state = JointState( - {"name": ["base", "shoulder", "elbow", "wrist"], "position": [1.0, 2.0, 3.0, 4.0]} + { + "name": ["arm/base", "arm/shoulder", "arm/elbow", "arm/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 + def get_model_config(self) -> RobotModelConfig: return self.config - def get_joint_limits(self, robot_id: str) -> tuple[np.ndarray, np.ndarray]: - assert robot_id == self.robot_id + def get_joint_limits(self) -> tuple[np.ndarray, np.ndarray]: 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 + def get_joint_state(self, ctx: object) -> JointState: return self.current_state - def check_config_collision_free(self, robot_id: str, joint_state: JointState) -> bool: - assert robot_id == self.robot_id + def check_config_collision_free(self, joint_state: JointState) -> bool: self.collision_checked_state = joint_state return True @@ -73,11 +67,8 @@ def test_solve_pose_targets_uses_group_tip_locks_seed_fallback_and_filters(monke world = FakeWorld() group = PlanningGroup( - id="arm/reach", - robot_name="arm", - group_name="reach", + id="reach", joint_names=("arm/shoulder", "arm/wrist"), - local_joint_names=("shoulder", "wrist"), base_link="base_link", tip_link="group_tip_link", ) @@ -89,7 +80,7 @@ def fake_solve_single(self, **kwargs) -> IKResult: status=IKStatus.SUCCESS, joint_state=JointState( { - "name": ["base", "shoulder", "elbow", "wrist"], + "name": ["arm/base", "arm/shoulder", "arm/elbow", "arm/wrist"], "position": [10.0, 20.0, 30.0, 40.0], } ), @@ -103,7 +94,7 @@ def fake_solve_single(self, **kwargs) -> IKResult: result = DrakeOptimizationIK().solve_pose_targets( world=world, # type: ignore[arg-type] pose_targets={group: PoseStamped()}, - seed=JointState({"name": ["shoulder"], "position": [22.0]}), + seed=JointState({"name": ["arm/shoulder"], "position": [22.0]}), check_collision=False, max_attempts=1, ) @@ -198,27 +189,27 @@ def GetSolution(self, q: list[str]) -> np.ndarray: class FakeDrakeWorld: def __init__(self) -> None: self.plant = FakePlant() - self._robots = {"robot-instance": _FakeRobotData()} - self.link_pose_calls: list[tuple[str, str]] = [] + self.link_pose_calls: list[str] = [] self.set_joint_state_calls: list[JointState] = [] + def get_body_frame(self, link_name: str) -> str: + return self.plant.GetBodyByName(link_name, "model-instance").body_frame() + + def get_model_joint_indices(self) -> list[int]: + return [1, 3, 4] + @contextmanager def scratch_context(self): yield "ctx" - def set_joint_state(self, ctx: str, robot_id: str, joint_state: JointState) -> None: + def set_joint_state(self, ctx: 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)) + def get_link_pose(self, ctx: str, target_frame_name: str) -> np.ndarray: + self.link_pose_calls.append(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) @@ -230,7 +221,6 @@ def test_solve_single_uses_target_frame_for_constraints_error_and_joint_locks(mo 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"], @@ -247,7 +237,7 @@ def test_solve_single_uses_target_frame_for_constraints_error_and_joint_locks(mo 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 world.link_pose_calls == ["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 index 15d966b7a4..d19a7bf9fa 100644 --- a/dimos/manipulation/planning/kinematics/test_jacobian_ik_selection.py +++ b/dimos/manipulation/planning/kinematics/test_jacobian_ik_selection.py @@ -37,11 +37,8 @@ def _pose(x: float = 0.0) -> PoseStamped: def _group(tip_link: str | None = "tool") -> PlanningGroup: return PlanningGroup( - id="arm/manipulator", - robot_name="arm", - group_name="manipulator", + id="manipulator", joint_names=("arm/joint_a", "arm/joint_b"), - local_joint_names=("joint_a", "joint_b"), base_link="base", tip_link=tip_link, ) @@ -53,40 +50,36 @@ class _World: 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"], + joint_names=["arm/joint_a", "arm/joint_b", "arm/gripper"], base_link="base", planning_groups=[ PlanningGroupDefinition( name="manipulator", - joint_names=("joint_a", "joint_b"), + joint_names=("arm/joint_a", "arm/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: + def get_model_config(self) -> RobotModelConfig: return self.config - def get_joint_limits(self, robot_id: str) -> tuple[np.ndarray, np.ndarray]: + def get_joint_limits(self) -> 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 get_joint_state(self, ctx: object) -> JointState: + return JointState( + {"name": ["arm/joint_a", "arm/joint_b", "arm/gripper"], "position": [0.0, 0.0, 0.9]} + ) - def set_joint_state(self, ctx: object, robot_id: str, joint_state: JointState) -> None: + def set_joint_state(self, ctx: object, joint_state: JointState) -> None: self.last_state = joint_state def get_group_ee_pose(self, ctx: object, group_id: str) -> PoseStamped: @@ -97,15 +90,7 @@ 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: + def check_config_collision_free(self, joint_state: JointState) -> bool: return True @@ -125,8 +110,6 @@ def test_solve_pose_targets_filters_to_group_and_uses_group_world_methods() -> N 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: @@ -144,4 +127,4 @@ 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 + assert "no tip" in result.message diff --git a/dimos/manipulation/planning/kinematics/test_pink_ik.py b/dimos/manipulation/planning/kinematics/test_pink_ik.py index d0324edb56..65f63e7ef9 100644 --- a/dimos/manipulation/planning/kinematics/test_pink_ik.py +++ b/dimos/manipulation/planning/kinematics/test_pink_ik.py @@ -167,7 +167,6 @@ def solve_ik( def _robot_config() -> RobotModelConfig: return RobotModelConfig( - name="arm", 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"], @@ -210,121 +209,49 @@ def __init__(self, collision_free: bool = True) -> None: 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"), + "manipulator": PlanningGroup( + id="manipulator", + 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",), + "no_tip": PlanningGroup( + id="no_tip", + 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",), + "wrist": PlanningGroup( + id="wrist", + 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: + def get_model_config(self) -> RobotModelConfig: return self.config def scratch_context(self) -> nullcontext[None]: return nullcontext(None) - def get_joint_state(self, ctx: object, robot_id: str) -> JointState: + def get_joint_state(self, ctx: object) -> JointState: 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]: + def get_joint_limits(self) -> 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: + def check_config_collision_free(self, joint_state: JointState) -> bool: return self.collision_free - def set_joint_state(self, ctx: object, robot_id: str, joint_state: JointState) -> None: + def set_joint_state(self, ctx: object, joint_state: JointState) -> None: self.joint_state = joint_state - def is_collision_free(self, ctx: object, robot_id: str) -> bool: + def is_collision_free(self, ctx: object) -> 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, ) -> None: @@ -447,11 +374,10 @@ 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", "tool"): context} + ik._model_context = context result = ik.solve( world=cast("Any", _FakeWorld(collision_free=False)), - robot_id="robot", target_pose=PoseStamped( position=Vector3(0.1, 0.0, 0.0), orientation=Quaternion(0.0, 0.0, 0.0, 1.0), @@ -467,7 +393,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", "tool"): context} + ik._model_context = context calls = 0 def fake_solve_single(**_: object) -> IKResult: @@ -494,7 +420,6 @@ def fake_solve_single(**_: object) -> IKResult: result = ik.solve( world=cast("Any", _FakeWorld(collision_free=True)), - robot_id="robot", target_pose=PoseStamped( position=Vector3(0.1, 0.0, 0.0), orientation=Quaternion(0.0, 0.0, 0.0, 1.0), @@ -507,7 +432,9 @@ def fake_solve_single(**_: object) -> IKResult: assert result.status == IKStatus.SUCCESS -def test_robot_context_cache_key_includes_tip_frame(mocker: MockerFixture, tmp_path: Path) -> None: +def test_robot_context_is_built_once_and_reused_for_every_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) @@ -517,12 +444,15 @@ def test_robot_context_cache_key_includes_tip_frame(mocker: MockerFixture, tmp_p world = _FakeWorld() world.config.model_path = model_path ik = PinkIK(PinkIKConfig(max_iterations=1)) + build_context = mocker.spy(ik, "_build_robot_context") - first = ik._get_robot_context(cast("Any", world), "robot", "tool") - second = ik._get_robot_context(cast("Any", world), "robot", "base") + first = ik._get_model_context(cast("Any", world), "tool") + second = ik._get_model_context(cast("Any", world), "base") assert first is not second - assert set(ik._robot_contexts) == {("robot", "tool"), ("robot", "base")} + assert first.model is second.model + assert first.data is second.data + build_context.assert_called_once_with(world.config, "tool") def test_build_robot_context_rejects_base_link_not_model_root( @@ -548,7 +478,7 @@ def test_solve_pose_targets_uses_group_tip_and_filters_group_joints( ) -> None: ik = _pink_ik(mocker, converge=True) context = _context() - get_context = mocker.patch.object(ik, "_get_robot_context", return_value=context) + get_context = mocker.patch.object(ik, "_get_model_context", return_value=context) mocker.patch.object( ik, "_solve_single", @@ -564,22 +494,20 @@ def test_solve_pose_targets_uses_group_tip_and_filters_group_joints( result = ik.solve_pose_targets( world=cast("Any", world), pose_targets={ - world.groups["arm/manipulator"]: PoseStamped( + world.groups["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]} - ), + seed=JointState({"name": ["joint_a", "joint_b", "joint_c"], "position": [0.0, 0.0, 0.0]}), max_attempts=1, ) - get_context.assert_called_once_with(cast("Any", world), "robot", "tool") + get_context.assert_called_once_with(cast("Any", world), "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.name == ["joint_a", "joint_b"] assert result.joint_state.position == [0.1, 0.2] - assert world.joint_state_calls == 0 + assert world.joint_state_calls == 1 def test_solve_pose_targets_rejects_group_without_tip(mocker: MockerFixture) -> None: @@ -589,7 +517,7 @@ def test_solve_pose_targets_rejects_group_without_tip(mocker: MockerFixture) -> result = ik.solve_pose_targets( world=cast("Any", world), pose_targets={ - world.groups["arm/no_tip"]: PoseStamped( + world.groups["no_tip"]: PoseStamped( position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) ) }, @@ -601,7 +529,7 @@ def test_solve_pose_targets_rejects_group_without_tip(mocker: MockerFixture) -> 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, "_get_model_context", return_value=_context()) mocker.patch.object( ik, "_solve_single", @@ -617,11 +545,11 @@ def test_solve_pose_targets_partial_seed_reads_world_state(mocker: MockerFixture result = ik.solve_pose_targets( world=cast("Any", world), pose_targets={ - world.groups["arm/manipulator"]: PoseStamped( + world.groups["manipulator"]: PoseStamped( position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) ) }, - seed=JointState({"name": ["arm/joint_a"], "position": [0.0]}), + seed=JointState({"name": ["joint_a"], "position": [0.0]}), max_attempts=1, ) @@ -632,7 +560,7 @@ def test_solve_pose_targets_partial_seed_reads_world_state(mocker: MockerFixture 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()) + mocker.patch.object(ik, "_get_model_context", return_value=_context()) solve_multi = mocker.patch.object( ik, "_solve_multi", @@ -649,84 +577,24 @@ def test_solve_pose_targets_multi_target_uses_multi_frame_solve(mocker: MockerFi result = ik.solve_pose_targets( world=cast("Any", world), pose_targets={ - world.groups["arm/manipulator"]: PoseStamped( + world.groups["manipulator"]: PoseStamped( position=Vector3(), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) ), - world.groups["arm/wrist"]: PoseStamped( + world.groups["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]} - ), + seed=JointState({"name": ["joint_a", "joint_b", "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.name == ["joint_a", "joint_b", "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: @@ -736,14 +604,12 @@ def test_solve_pose_targets_auxiliary_only_retains_seed_selection_order( 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]} - ), + auxiliary_groups=[world.groups["no_tip"], world.groups["manipulator"]], + seed=JointState({"name": ["joint_a", "joint_b", "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.name == ["joint_c", "joint_a", "joint_b"] assert result.joint_state.position == [0.3, 0.1, 0.2] - assert world.joint_state_calls == 0 + assert world.joint_state_calls == 1 diff --git a/dimos/manipulation/planning/kinematics/utils.py b/dimos/manipulation/planning/kinematics/utils.py index 3d0fa0509c..f2dd7919c5 100644 --- a/dimos/manipulation/planning/kinematics/utils.py +++ b/dimos/manipulation/planning/kinematics/utils.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,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Shared IK-only helpers for planning-group-scoped kinematics backends.""" +"""Shared helpers for planning-group-scoped kinematics backends.""" from collections.abc import Mapping, Sequence from dataclasses import dataclass @@ -23,7 +23,7 @@ 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.models import IKResult from dimos.manipulation.planning.spec.protocols import WorldSpec from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.JointState import JointState @@ -33,71 +33,42 @@ 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 +def unique_pose_target_frame(world: WorldSpec) -> str | None: + frames = [ + group.tip_link + for group in world.get_model_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 + unique = list(dict.fromkeys(frames)) + return unique[0] if len(unique) == 1 else None def seed_positions_with_world_fallback( - world: WorldSpec, - robot_id: WorldRobotID, - robot_name: RobotName, - local_joint_names: list[str], - seed: JointState | None, + world: WorldSpec, joint_names: list[str], seed: JointState | None ) -> NDArray[np.float64]: - """Return full robot positions, reading world only for absent seed joints.""" + with world.scratch_context() as ctx: + current = world.get_joint_state(ctx) + fallback = positions_by_name(current, joint_names) 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 + return fallback + if not seed.name: + return positions_by_name(seed, joint_names) + known = set(joint_names) + for name, position in zip(seed.name, seed.position, strict=True): + if name not in known: + raise ValueError(f"Unrecognized seed joint '{name}'") + fallback[joint_names.index(name)] = position + return fallback def resolve_single_pose_target_request( world: WorldSpec, - pose_targets: dict[PlanningGroup, PoseStamped] | Mapping[PlanningGroup, PoseStamped], + pose_targets: Mapping[PlanningGroup, PoseStamped], auxiliary_groups: Sequence[PlanningGroup], seed: JointState | None, backend_name: str, @@ -109,96 +80,34 @@ def resolve_single_pose_target_request( 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", - ) - + group = next(iter(pose_targets)) + if not group.has_pose_target: + return None, _failure(IKStatus.UNSUPPORTED, f"Planning group '{group.id}' has no tip") 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] + joint_names = list(world.get_model_config().joint_names) + seed_positions = seed_positions_with_world_fallback(world, joint_names, seed) + indices = [joint_names.index(name) for name in group.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) + return SinglePoseTargetRequest( + group=group, + target_pose=pose_targets[group], + joint_names=joint_names, + seed_positions=seed_positions, + group_indices=indices, + ), None + + +def positions_by_name(state: JointState, joint_names: list[str]) -> NDArray[np.float64]: + if not state.name: + if len(state.position) != len(joint_names): + raise ValueError("JointState position count does not match model joints") + return np.asarray(state.position, dtype=np.float64) + positions = dict(zip(state.name, state.position, strict=True)) + missing = [name for name in joint_names if name not in positions] 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 + return np.asarray([positions[name] for name in joint_names], dtype=np.float64) def filter_result_to_group(result: IKResult, group: PlanningGroup) -> IKResult: @@ -208,15 +117,10 @@ def filter_result_to_group(result: IKResult, group: PlanningGroup) -> IKResult: 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, + result.joint_state, selection.joint_names ), position_error=result.position_error, orientation_error=result.orientation_error, @@ -225,12 +129,5 @@ def filter_result_to_selection(result: IKResult, selection: PlanningGroupSelecti ) -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/robot_state_monitor.py b/dimos/manipulation/planning/monitor/robot_state_monitor.py index 49c6b56366..837ed36de7 100644 --- a/dimos/manipulation/planning/monitor/robot_state_monitor.py +++ b/dimos/manipulation/planning/monitor/robot_state_monitor.py @@ -15,11 +15,11 @@ """ Robot State Monitor -Per-robot monitor that tracks joint state and syncs it to a WorldSpec instance. +Monitor that tracks joint state and syncs it to a WorldSpec instance. This is the WorldSpec-based replacement for StateMonitor. Example: - monitor = RobotStateMonitor(world, lock, robot_id, joint_names) + monitor = RobotStateMonitor(world, lock, joint_names) monitor.start() monitor.on_joint_state(joint_state_msg) # Called by subscriber """ @@ -67,9 +67,7 @@ def __init__( self, world: WorldSpec, lock: threading.RLock, - robot_id: str, joint_names: list[str], - joint_name_mapping: dict[str, str] | None = None, timeout: float = 1.0, ) -> None: """Create a world state monitor. @@ -77,24 +75,14 @@ def __init__( Args: world: WorldSpec instance to sync state to lock: Shared lock for thread-safe access - robot_id: ID of the robot to monitor - joint_names: Ordered list of joint names for this robot (URDF names) - joint_name_mapping: Maps coordinator joint names to URDF joint names. - Example: {"left/joint1": "joint1"} means messages with "left/joint1" - will be mapped to URDF "joint1". If None, names must match exactly. + joint_names: Ordered list of canonical model joint names timeout: Timeout for waiting for initial state (seconds) """ self._world = world self._lock = lock - self._robot_id = robot_id self._joint_names = joint_names self._timeout = timeout - # Joint name mapping: coordinator name -> URDF name - self._joint_name_mapping = joint_name_mapping or {} - # Build reverse mapping: URDF name -> coordinator name - self._reverse_mapping = {v: k for k, v in self._joint_name_mapping.items()} - # Latest state self._latest_positions: NDArray[np.float64] | None = None self._latest_velocities: NDArray[np.float64] | None = None @@ -103,28 +91,20 @@ def __init__( # Running state self._running = False - # Callbacks: (robot_id, joint_state) called on each state update - self._state_callbacks: list[Callable[[str, JointState], None]] = [] + self._state_callbacks: list[Callable[[JointState], None]] = [] def start(self) -> None: """Start the state monitor.""" self._running = True - logger.info(f"World state monitor started for robot '{self._robot_id}'") def stop(self) -> None: """Stop the state monitor.""" self._running = False - logger.info(f"World state monitor stopped for robot '{self._robot_id}'") def is_running(self) -> bool: """Check if monitor is running.""" return self._running - @property - def robot_id(self) -> str: - """Get the robot ID being monitored.""" - return self._robot_id - def on_joint_state(self, msg: JointState) -> None: """Handle incoming joint state message. @@ -170,14 +150,14 @@ def on_joint_state(self, msg: JointState) -> None: name=self._joint_names, position=positions.tolist(), ) - self._world.sync_from_joint_state(self._robot_id, joint_state) + self._world.sync_from_joint_state(joint_state) except Exception as e: logger.error(f"Failed to sync joint state to live context: {e}") # Call registered callbacks for callback in self._state_callbacks: try: - callback(self._robot_id, joint_state) + callback(joint_state) except Exception as e: logger.error(f"State callback error: {e}") @@ -190,29 +170,20 @@ def on_joint_state(self, msg: JointState) -> None: def _extract_positions(self, msg: JointState) -> NDArray[np.float64] | None: """Extract positions for our joints from JointState message. - Handles joint name translation from coordinator namespace to URDF namespace. - If joint_name_mapping is set, message names are looked up via the reverse mapping. - Args: msg: JointState message (may use coordinator joint names) Returns: Array of joint positions or None if any joint is missing """ - # Build name->index map from message (coordinator names) + # Build name->index map from canonical message names. name_to_idx = {name: i for i, name in enumerate(msg.name)} positions = [] - for urdf_joint_name in self._joint_names: - # Try direct match first (when no mapping or names already match) - if urdf_joint_name in name_to_idx: - idx = name_to_idx[urdf_joint_name] - else: - # Try reverse mapping: URDF name -> coordinator name -> msg index - orch_name = self._reverse_mapping.get(urdf_joint_name) - if orch_name is None or orch_name not in name_to_idx: - return None # Missing joint - idx = name_to_idx[orch_name] + for joint_name in self._joint_names: + if joint_name not in name_to_idx: + return None + idx = name_to_idx[joint_name] if idx >= len(msg.position): return None # Position not available @@ -223,7 +194,7 @@ def _extract_positions(self, msg: JointState) -> NDArray[np.float64] | None: def _extract_velocities(self, msg: JointState) -> NDArray[np.float64] | None: """Extract velocities for our joints. - Uses same name translation as _extract_positions. + Uses the same canonical-name lookup as _extract_positions. """ if not msg.velocity or len(msg.velocity) == 0: return None @@ -231,16 +202,10 @@ def _extract_velocities(self, msg: JointState) -> NDArray[np.float64] | None: name_to_idx = {name: i for i, name in enumerate(msg.name)} velocities = [] - for urdf_joint_name in self._joint_names: - # Try direct match first - if urdf_joint_name in name_to_idx: - idx = name_to_idx[urdf_joint_name] - else: - # Try reverse mapping - orch_name = self._reverse_mapping.get(urdf_joint_name) - if orch_name is None or orch_name not in name_to_idx: - return None - idx = name_to_idx[orch_name] + for joint_name in self._joint_names: + if joint_name not in name_to_idx: + return None + idx = name_to_idx[joint_name] if idx >= len(msg.velocity): return None @@ -313,18 +278,18 @@ def is_state_stale(self, max_age: float = 1.0) -> bool: def add_state_callback( self, - callback: Callable[[str, JointState], None], + callback: Callable[[JointState], None], ) -> None: """Add callback for state updates. Args: - callback: Function called with (robot_id, joint_state) on each update + callback: Function called with the canonical joint state """ self._state_callbacks.append(callback) def remove_state_callback( self, - callback: Callable[[str, JointState], None], + callback: Callable[[JointState], None], ) -> None: """Remove a state callback.""" if callback in self._state_callbacks: diff --git a/dimos/manipulation/planning/monitor/test_world_monitor.py b/dimos/manipulation/planning/monitor/test_world_monitor.py index 9423aa3b4d..746e5c6e64 100644 --- a/dimos/manipulation/planning/monitor/test_world_monitor.py +++ b/dimos/manipulation/planning/monitor/test_world_monitor.py @@ -42,26 +42,6 @@ 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" @@ -73,21 +53,16 @@ def __exit__(self, exc_type: object, exc: object, traceback: object) -> bool: class FakeWorld: def __init__(self) -> None: self.calls: list[tuple[Any, ...]] = [] - self.configs: dict[str, RobotModelConfig] = {} - - def add_robot(self, config): - self.calls.append(("add_robot", config)) - robot_id = f"robot-{len(self.configs) + 1}" - self.configs[robot_id] = config - return robot_id + self.config: RobotModelConfig | None = None - def get_robot_ids(self): - return [] + def load_model(self, config): + self.calls.append(("load_model", config)) + self.config = config - def get_robot_config(self, robot_id): - return None + def get_model_config(self): + return self.config - def get_joint_limits(self, robot_id): + def get_joint_limits(self): return ([], []) def add_obstacle(self, obstacle): @@ -125,39 +100,43 @@ def scratch_context(self): self.calls.append(("scratch_context", None)) return _ScratchContext() - def sync_from_joint_state(self, robot_id, joint_state) -> None: + def sync_from_joint_state(self, 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)) + def set_joint_state(self, ctx, joint_state) -> None: + self.calls.append(("set_joint_state", ctx, joint_state)) return None - def get_joint_state(self, ctx, robot_id): - return None + def get_joint_state(self, ctx): + if self.config is None: + return JointState() + return JointState( + name=self.config.joint_names, position=[0.0] * len(self.config.joint_names) + ) - def is_collision_free(self, ctx, robot_id): + def is_collision_free(self, ctx): return True - def get_min_distance(self, ctx, robot_id): + def get_min_distance(self, ctx): return 0.0 - def check_config_collision_free(self, robot_id, joint_state): + def check_config_collision_free(self, joint_state): return True - def check_edge_collision_free(self, robot_id, start, end, step_size: float = 0.05): + def check_edge_collision_free(self, start, end, step_size: float = 0.05): return True - def get_ee_pose(self, ctx, robot_id): + def get_ee_pose(self, ctx): 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): + def get_link_pose(self, ctx, link_name): return [] - def get_jacobian(self, ctx, robot_id): + def get_jacobian(self, ctx): return [] def get_group_jacobian(self, ctx, group_id): @@ -223,7 +202,6 @@ def clear_vis_obstacles(self) -> None: def _robot_config() -> RobotModelConfig: return RobotModelConfig( - name="arm", model_path=Path("/tmp/arm.urdf"), base_pose=PoseStamped(position=Vector3(), orientation=Quaternion([0, 0, 0, 1])), joint_names=["j1", "j2"], @@ -256,16 +234,17 @@ def _three_joint_reordered_group_config() -> RobotModelConfig: ) -def test_world_monitor_add_robot_records_scene_without_visualization_probe() -> None: +def test_world_monitor_load_model_records_scene_without_visualization_probe() -> None: fake_world = FakeWorld() fake_viz = FakeViz() monitor = world_monitor_module.WorldMonitor(world=fake_world, visualization=fake_viz) # type: ignore[arg-type] - monitor.add_robot(_robot_config()) - assert fake_world.calls[0][0] == "add_robot" + config = _robot_config() + monitor.load_model(config) + assert fake_world.calls[0][0] == "load_model" assert fake_viz.calls == [] - assert monitor.planning_scene_info().robots["robot-1"].name == "arm" + assert monitor.planning_scene_info().model is config def test_world_monitor_syncs_planning_scene_to_visualization() -> None: @@ -273,25 +252,26 @@ def test_world_monitor_syncs_planning_scene_to_visualization() -> None: fake_viz = FakeViz() monitor = world_monitor_module.WorldMonitor(world=fake_world, visualization=fake_viz) # type: ignore[arg-type] - monitor.add_robot(_robot_config()) + config = _robot_config() + monitor.load_model(config) operator = object() monitor.finalize(fake_viz, operator=operator) # type: ignore[arg-type] monitor.add_obstacle(object()) # type: ignore[arg-type] - assert [call[0] for call in fake_world.calls] == ["add_robot", "finalize", "add_obstacle"] + assert [call[0] for call in fake_world.calls] == ["load_model", "finalize", "add_obstacle"] assert fake_viz.calls[0][0] == "initialize" session = fake_viz.calls[0][1] assert session.operator is operator scene = session.scene assert isinstance(scene, PlanningSceneInfo) - assert scene.robots["robot-1"].name == "arm" - assert scene.planning_groups[0].id == "arm/manipulator" + assert scene.model is config + assert scene.planning_groups[0].id == "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=[]) + trajectory = JointTrajectory(joint_names=["j1"], points=[]) assert isinstance(fake_viz, VisualizationSpec) monitor.cancel_preview_animation() @@ -373,12 +353,12 @@ def test_create_planning_specs_wraps_existing_world(mocker: MockerFixture) -> No 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()) + monitor.load_model(_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 + assert [group.id for group in monitor.planning_groups.list()] == ["manipulator"] + with pytest.raises(ValueError, match="already loaded"): + monitor.load_model(_robot_config()) + assert [call[0] for call in fake_world.calls].count("load_model") == 1 def test_world_monitor_invalid_duplicate_group_config_does_not_mutate_backend() -> None: @@ -396,9 +376,9 @@ def test_world_monitor_invalid_duplicate_group_config_does_not_mutate_backend() ) with pytest.raises(ValueError, match="already registered"): - monitor.add_robot(invalid_config) + monitor.load_model(invalid_config) - assert [call[0] for call in fake_world.calls].count("add_robot") == 0 + assert [call[0] for call in fake_world.calls].count("load_model") == 0 def test_world_monitor_invalid_group_joint_name_does_not_mutate_backend() -> None: @@ -416,159 +396,132 @@ def test_world_monitor_invalid_group_joint_name_does_not_mutate_backend() -> Non ) with pytest.raises(ValueError, match="unknown model joints"): - monitor.add_robot(invalid_config) + monitor.load_model(invalid_config) - assert [call[0] for call in fake_world.calls].count("add_robot") == 0 + assert [call[0] for call in fake_world.calls].count("load_model") == 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] + monitor.load_model(_three_joint_reordered_group_config()) + monitor.start_state_monitor() + monitor.on_joint_state(JointState(name=["j1", "j2", "j3"], position=[0.1, 0.2, 0.3])) - state = monitor.current_group_joint_state("arm/manipulator") + state = monitor.current_group_joint_state("manipulator") - assert state.name == ["arm/j2", "arm/j1"] + assert state.name == ["j2", "j1"] assert state.position == [0.2, 0.1] -def test_current_global_joint_state_skips_stale_robots_and_preserves_state_order() -> None: +def test_current_model_joint_state_rejects_stale_state(mocker) -> 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" - ) - ], - ) - ) + monitor.load_model(_three_joint_reordered_group_config()) + monitor.start_state_monitor() + monitor.on_joint_state(JointState(name=["j1", "j2", "j3"], position=[1.0, 2.0, 3.0])) + mocker.patch.object(monitor, "is_state_stale", return_value=True) - state = monitor.current_global_joint_state(max_age=0.5) + state = monitor.current_model_joint_state(max_age=0.5) - assert state.name == ["arm/j1", "arm/j2", "arm/j3"] - assert state.position == [0.1, 0.2, 0.3] + assert state.name == [] + assert state.position == [] -def test_current_group_joint_state_rejects_stale_or_unavailable_state() -> None: +def test_current_group_joint_state_rejects_stale_state(mocker) -> 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] + stale_monitor.load_model(_three_joint_reordered_group_config()) + stale_monitor.start_state_monitor() + stale_monitor.on_joint_state(JointState(name=["j1", "j2", "j3"], position=[0.1, 0.2, 0.3])) + mocker.patch.object(stale_monitor, "is_state_stale", return_value=True) 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") + stale_monitor.current_group_joint_state("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] + monitor.load_model(_three_joint_reordered_group_config()) + monitor.start_state_monitor() + monitor.on_joint_state(JointState(name=["j1", "j2", "j3"], position=[0.1, 0.2, 0.3])) - pose = monitor.get_group_ee_pose("arm/manipulator") + pose = monitor.get_group_ee_pose("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 set_calls[0][2].name == ["j1", "j2", "j3"] + assert set_calls[0][2].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: +def test_group_ee_pose_without_joint_state_rejects_stale_state(mocker) -> 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] + stale_monitor.load_model(_three_joint_reordered_group_config()) + stale_monitor.start_state_monitor() + stale_monitor.on_joint_state(JointState(name=["j1", "j2", "j3"], position=[0.1, 0.2, 0.3])) + mocker.patch.object(stale_monitor, "is_state_stale", return_value=True) 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") + stale_monitor.get_group_ee_pose("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()) + monitor.load_model(_three_joint_reordered_group_config()) pose = monitor.get_group_ee_pose( - "arm/manipulator", + "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 set_calls[0][2].name == ["j1", "j2", "j3"] + assert set_calls[0][2].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()) + monitor.load_model(_three_joint_reordered_group_config()) pose = monitor.get_group_ee_pose( - "arm/manipulator", + "manipulator", JointState(name=["j1", "j2", "j3"], position=[0.9, 0.8, 0.3]), ) jacobian = monitor.get_group_jacobian( - "arm/manipulator", + "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 set_calls[0][2].name == ["j1", "j2", "j3"] + assert set_calls[0][2].position == [0.9, 0.8, 0.3] + assert set_calls[1][2].name == ["j1", "j2", "j3"] + assert set_calls[1][2].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 + assert ("get_group_ee_pose", "scratch", "manipulator") in fake_world.calls + assert ("get_group_jacobian", "scratch", "manipulator") in fake_world.calls -def test_legacy_wrappers_fail_for_no_pose_and_ambiguous_pose_groups() -> None: +def test_convenience_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( + monitor.load_model( _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])) + with pytest.raises(ValueError, match="no unique pose-targetable"): + monitor.get_ee_pose(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( + monitor2.load_model( _robot_config_with_groups( [ PlanningGroupDefinition( @@ -581,7 +534,7 @@ def test_legacy_wrappers_fail_for_no_pose_and_ambiguous_pose_groups() -> None: ) ) with pytest.raises(ValueError, match="pose-targetable planning groups"): - monitor2.get_jacobian(ambiguous_id, JointState(name=["j1", "j2"], position=[0.0, 0.0])) + monitor2.get_jacobian(JointState(name=["j1", "j2"], position=[0.0, 0.0])) def test_world_monitor_obstacle_mutations_cover_failure_and_visualization_errors( diff --git a/dimos/manipulation/planning/monitor/world_monitor.py b/dimos/manipulation/planning/monitor/world_monitor.py index fd5bb6c67d..e0a3cb4ab9 100644 --- a/dimos/manipulation/planning/monitor/world_monitor.py +++ b/dimos/manipulation/planning/monitor/world_monitor.py @@ -16,16 +16,12 @@ 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.identifiers import assert_valid_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 @@ -37,6 +33,7 @@ ) from dimos.manipulation.planning.spec.protocols import VisualizationSpec, WorldSpec from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Quaternion import Quaternion from dimos.msgs.sensor_msgs.JointState import JointState from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory from dimos.utils.logging_config import setup_logger @@ -53,8 +50,6 @@ JointPath, Obstacle, PlanningGroupID, - RobotName, - WorldRobotID, ) from dimos.msgs.vision_msgs.Detection3D import Detection3D from dimos.perception.experimental.object import Object @@ -76,31 +71,25 @@ def __init__( # 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._model_config: RobotModelConfig | None = None self._planning_groups = PlanningGroupRegistry() - self._state_monitors: dict[WorldRobotID, RobotStateMonitor] = {} + self._state_monitor: RobotStateMonitor | None = None self._obstacle_monitor: WorldObstacleMonitor | None = None self._viz_thread: threading.Thread | None = None self._viz_stop_event = threading.Event() self._viz_rate_hz: float = 10.0 - # Robot Management + # Model Management - def add_robot(self, config: RobotModelConfig) -> WorldRobotID: - """Add a robot. Returns robot_id.""" + def load_model(self, config: RobotModelConfig) -> None: + """Load the one logical robot model.""" with self._lock: - if config.name in self._robot_ids_by_name: - raise ValueError(f"Robot name '{config.name}' is already registered") + if self._model_config is not None: + raise ValueError("A model is already loaded") 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 + self._world.load_model(config) + self._model_config = config + self._planning_groups.add_model(config) @property def planning_groups(self) -> PlanningGroupRegistry: @@ -111,26 +100,21 @@ 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), + model=self.get_model_config(), planning_groups=tuple(self._planning_groups.list()), ) - def get_robot_ids(self) -> list[WorldRobotID]: - """Get all robot IDs.""" - with self._lock: - return self._world.get_robot_ids() - - def get_robot_config(self, robot_id: WorldRobotID) -> RobotModelConfig: - """Get robot configuration.""" + def get_model_config(self) -> RobotModelConfig: + """Get the configured model.""" with self._lock: - return self._world.get_robot_config(robot_id) + if self._model_config is None: + raise RuntimeError("Model is not loaded") + return self._model_config - def get_joint_limits( - self, robot_id: WorldRobotID - ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: - """Get joint limits for a robot.""" + def get_joint_limits(self) -> tuple[NDArray[np.float64], NDArray[np.float64]]: + """Get model joint limits.""" with self._lock: - return self._world.get_joint_limits(robot_id) + return self._world.get_joint_limits() # Obstacle Management @@ -194,42 +178,20 @@ def clear_obstacles(self) -> None: # Monitor Control - def start_state_monitor( - self, - robot_id: WorldRobotID, - joint_names: list[str] | None = None, - joint_name_mapping: dict[str, str] | None = None, - ) -> None: - """Start monitoring joint states. Uses config defaults if args are None.""" + def start_state_monitor(self) -> None: + """Start monitoring canonical model joint state.""" with self._lock: - if robot_id in self._state_monitors: - logger.warning(f"State monitor for '{robot_id}' already started") + if self._state_monitor is not None: + logger.warning("State monitor already started") return - - # Get config for defaults - config = self._world.get_robot_config(robot_id) - - # Get joint names from config if not provided - if joint_names is None: - if robot_id in self._robot_joints: - joint_names = self._robot_joints[robot_id] - else: - joint_names = config.joint_names - - # Get joint name mapping from config if not provided - if joint_name_mapping is None and config.joint_name_mapping: - joint_name_mapping = config.joint_name_mapping - + config = self.get_model_config() monitor = RobotStateMonitor( world=self._world, lock=self._lock, - robot_id=robot_id, - joint_names=joint_names, - joint_name_mapping=joint_name_mapping, + joint_names=config.joint_names, ) monitor.start() - self._state_monitors[robot_id] = monitor - logger.info(f"State monitor started for '{robot_id}'") + self._state_monitor = monitor def start_obstacle_monitor(self) -> None: """Start monitoring obstacle updates.""" @@ -250,9 +212,9 @@ def stop_all_monitors(self) -> None: self.stop_visualization_thread() with self._lock: - for _robot_id, monitor in self._state_monitors.items(): - monitor.stop() - self._state_monitors.clear() + if self._state_monitor is not None: + self._state_monitor.stop() + self._state_monitor = None if self._obstacle_monitor is not None: self._obstacle_monitor.stop() @@ -268,18 +230,11 @@ def stop_all_monitors(self) -> None: # Message Handlers - def on_joint_state(self, msg: JointState, robot_id: WorldRobotID | None = None) -> None: - """Handle joint state message. Broadcasts to all monitors if robot_id is None.""" + def on_joint_state(self, msg: JointState) -> None: + """Handle a canonical model joint-state message.""" try: - if robot_id is not None: - if robot_id in self._state_monitors: - self._state_monitors[robot_id].on_joint_state(msg) - else: - logger.warning(f"No state monitor for robot_id: {robot_id}") - else: - # Broadcast to all monitors - for monitor in self._state_monitors.values(): - monitor.on_joint_state(msg) + if self._state_monitor is not None: + self._state_monitor.on_joint_state(msg) except Exception as e: logger.error(f"[WorldMonitor] Exception in on_joint_state: {e}") import traceback @@ -345,16 +300,14 @@ def list_added_obstacles(self) -> list[dict[str, Any]]: # State Access - def get_current_joint_state(self, robot_id: WorldRobotID) -> JointState | None: + def get_current_joint_state(self) -> JointState | None: """Get current joint state. Returns None if not yet received.""" - # Try state monitor first for positions - if robot_id in self._state_monitors: - positions = self._state_monitors[robot_id].get_current_positions() - velocities = self._state_monitors[robot_id].get_current_velocities() + if self._state_monitor is not None: + positions = self._state_monitor.get_current_positions() + velocities = self._state_monitor.get_current_velocities() if positions is not None: - joint_names = self._robot_joints.get(robot_id, []) return JointState( - name=joint_names, + name=self.get_model_config().joint_names, position=positions.tolist(), velocity=velocities.tolist() if velocities is not None else [], ) @@ -362,47 +315,35 @@ def get_current_joint_state(self, robot_id: WorldRobotID) -> JointState | None: # Fall back to world's live context with self._lock: 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) + return self._world.get_joint_state(ctx) + + def current_model_joint_state(self, max_age: float = 1.0) -> JointState: + """Return the fresh canonical model state.""" + if self._state_monitor is not None and self.is_state_stale(max_age): + return JointState() + return self.get_current_joint_state() or JointState() 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 self._state_monitor is not None and self.is_state_stale(max_age): + raise ValueError("Current model state is stale") + state = self.get_current_joint_state() 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 - ) + raise ValueError("Current model state is unavailable") + return filter_joint_state_to_selected_joints(state, group.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() model_joint_names = set(config.joint_names) for definition in config.planning_groups: - group_id = make_planning_group_id(config.name, definition.name) + group_id = definition.name + assert_valid_group_id(group_id) 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) unknown_joint_names = set(definition.joint_names) - model_joint_names if unknown_joint_names: raise ValueError( @@ -411,25 +352,26 @@ def _validate_planning_group_config(self, config: RobotModelConfig) -> None: ) seen_group_names.add(definition.name) - def get_current_velocities(self, robot_id: WorldRobotID) -> JointState | None: + def get_current_velocities(self) -> JointState | None: """Get current joint velocities as JointState. Returns None if not available.""" - if robot_id in self._state_monitors: - velocities = self._state_monitors[robot_id].get_current_velocities() + if self._state_monitor is not None: + velocities = self._state_monitor.get_current_velocities() if velocities is not None: - joint_names = self._robot_joints.get(robot_id, []) - return JointState(name=joint_names, velocity=velocities.tolist()) + return JointState( + name=self.get_model_config().joint_names, velocity=velocities.tolist() + ) return None - def wait_for_state(self, robot_id: WorldRobotID, timeout: float = 1.0) -> bool: + def wait_for_state(self, timeout: float = 1.0) -> bool: """Wait until state is received. Returns False on timeout.""" - if robot_id in self._state_monitors: - return self._state_monitors[robot_id].wait_for_state(timeout) + if self._state_monitor is not None: + return self._state_monitor.wait_for_state(timeout) return False - def is_state_stale(self, robot_id: WorldRobotID, max_age: float = 1.0) -> bool: + def is_state_stale(self, max_age: float = 1.0) -> bool: """Check if state is stale.""" - if robot_id in self._state_monitors: - return self._state_monitors[robot_id].is_state_stale(max_age) + if self._state_monitor is not None: + return self._state_monitor.is_state_stale(max_age) return True # Context Management @@ -446,17 +388,14 @@ def get_live_context(self) -> Any: # Collision Checking - def is_state_valid(self, robot_id: WorldRobotID, joint_state: JointState) -> bool: + def is_state_valid(self, joint_state: JointState) -> bool: """Check if configuration is collision-free.""" - return self._world.check_config_collision_free(robot_id, joint_state) + return self._world.check_config_collision_free(joint_state) - def is_path_valid( - self, robot_id: WorldRobotID, path: JointPath, step_size: float = 0.05 - ) -> bool: + def is_path_valid(self, path: JointPath, step_size: float = 0.05) -> bool: """Check if path is collision-free with interpolation. Args: - robot_id: Robot to check path: List of JointState waypoints step_size: Max step size for interpolation (radians) @@ -464,70 +403,63 @@ def is_path_valid( True if entire path is collision-free """ if len(path) < 2: - return len(path) == 0 or self._world.check_config_collision_free(robot_id, path[0]) + return len(path) == 0 or self._world.check_config_collision_free(path[0]) # Check each edge for i in range(len(path) - 1): - if not self._world.check_edge_collision_free(robot_id, path[i], path[i + 1], step_size): + if not self._world.check_edge_collision_free(path[i], path[i + 1], step_size): return False return True - def get_min_distance(self, robot_id: WorldRobotID) -> float: + def get_min_distance(self) -> float: """Get minimum distance to obstacles for current state.""" with self._world.scratch_context() as ctx: - return self._world.get_min_distance(ctx, robot_id) + return self._world.get_min_distance(ctx) # Kinematics - def get_ee_pose( - self, robot_id: WorldRobotID, joint_state: JointState | None = None - ) -> PoseStamped: + def get_ee_pose(self, 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) + group_id = self._planning_groups.primary_pose_group_id() if group_id is None: - raise ValueError(f"Robot '{robot_name}' has no pose-targetable planning group") + raise ValueError("Model has no unique 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] + self._planning_groups.get(group_id) with self._world.scratch_context() as ctx: 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 self._state_monitor is not None and self.is_state_stale(): + raise ValueError("Current model state is stale") + joint_state = self.get_current_joint_state() 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) + raise ValueError("Current model state is unavailable") + self._world.set_joint_state(ctx, joint_state) 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 + self, link_name: str, joint_state: JointState | None = None ) -> PoseStamped | None: """Get arbitrary link pose as PoseStamped. Args: - robot_id: Robot to query link_name: Name of the link in the URDF joint_state: Joint state to use (uses current if None) """ - from dimos.msgs.geometry_msgs.Quaternion import Quaternion - with self._world.scratch_context() as ctx: if joint_state is None: - joint_state = self.get_current_joint_state(robot_id) + joint_state = self.get_current_joint_state() if joint_state is not None: - self._world.set_joint_state(ctx, robot_id, joint_state) + self._world.set_joint_state(ctx, joint_state) try: - mat = self._world.get_link_pose(ctx, robot_id, link_name) + mat = self._world.get_link_pose(ctx, link_name) except KeyError: - logger.warning(f"Link '{link_name}' not found in robot '{robot_id}'") + logger.warning(f"Link '{link_name}' not found in model") return None pos = mat[:3, 3] @@ -539,22 +471,20 @@ def get_link_pose( orientation=[float(quat.x), float(quat.y), float(quat.z), float(quat.w)], ) - def get_jacobian(self, robot_id: WorldRobotID, joint_state: JointState) -> NDArray[np.float64]: + def get_jacobian(self, 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) + group_id = self._planning_groups.primary_pose_group_id() if group_id is None: - raise ValueError(f"Robot '{robot_name}' has no pose-targetable planning group") + raise ValueError("Model has no unique 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] + self._planning_groups.get(group_id) with self._world.scratch_context() as ctx: - self._world.set_joint_state(ctx, robot_id, joint_state) + self._world.set_joint_state(ctx, joint_state) return self._world.get_group_jacobian(ctx, group_id) # Lifecycle @@ -573,7 +503,7 @@ def finalize( if attached_visualization is not None: session = VisualizationSession( scene=PlanningSceneInfo( - robots=dict(self._robot_configs), + model=self.get_model_config(), planning_groups=tuple(self._planning_groups.list()), ), operator=operator, @@ -598,14 +528,7 @@ def get_visualization_url(self) -> str | None: 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) + return VisualizationStateFrame(joint_state=self.get_current_joint_state()) def update_visualization_state(self) -> None: """Push current state to visualization.""" @@ -613,41 +536,20 @@ def update_visualization_state(self) -> None: with self._visualization_lock: self._visualization.update_state(self.visualization_state_frame()) - def cancel_preview_animation(self, robot_ids: Sequence[WorldRobotID] | None = None) -> None: + def cancel_preview_animation(self) -> None: """Cancel active visualization preview animation.""" if self._visualization is not None: - if robot_ids is None: - self._visualization.cancel_preview_animation() - else: - self._visualization.cancel_preview_animation(robot_ids) + self._visualization.cancel_preview_animation() def animate_trajectory( self, trajectory: JointTrajectory, duration: float | None = None ) -> None: """Animate a raw generated-plan trajectory if visualization is available.""" if self._visualization is not None: - 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() + 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.""" if self._viz_thread is not None and self._viz_thread.is_alive(): @@ -704,9 +606,9 @@ def visualization(self) -> VisualizationSpec | None: """Get optional visualization backend.""" return self._visualization - def get_state_monitor(self, robot_id: str) -> RobotStateMonitor | None: - """Get state monitor for a robot (may be None).""" - return self._state_monitors.get(robot_id) + def get_state_monitor(self) -> RobotStateMonitor | None: + """Get the state monitor (may be None).""" + return self._state_monitor @property def obstacle_monitor(self) -> WorldObstacleMonitor | None: diff --git a/dimos/manipulation/planning/planners/roboplan_planner.py b/dimos/manipulation/planning/planners/roboplan_planner.py index f6e7492280..9e7bfd7ca5 100644 --- a/dimos/manipulation/planning/planners/roboplan_planner.py +++ b/dimos/manipulation/planning/planners/roboplan_planner.py @@ -44,7 +44,6 @@ CartesianTarget, PlanningGroupID, PlanningResult, - WorldRobotID, ) from dimos.manipulation.planning.utils.path_utils import compute_path_length from dimos.manipulation.planning.world.roboplan_model import ( @@ -83,43 +82,41 @@ def __init__(self, world: WorldSpec, config: RoboPlanPlannerConfig) -> None: def plan_joint_path( self, world: WorldSpec, - robot_id: WorldRobotID, start: JointState, goal: JointState, timeout: float = 10.0, ) -> PlanningResult: - """Plan using the legacy robot-scoped local-name contract.""" + """Plan a path for the configured model's canonical joints.""" if world is not self._world: return PlanningResult( status=PlanningStatus.NO_SOLUTION, message="RoboPlan-native planner requires its RoboPlanWorld instance", ) try: - q_start = self._world._joint_state_to_q(robot_id, start) + q_start = self._world.ordered_joint_positions(start) except ValueError as exc: return PlanningResult(status=PlanningStatus.INVALID_START, message=str(exc)) try: - q_goal = self._world._joint_state_to_q(robot_id, goal) + q_goal = self._world.ordered_joint_positions(goal) except ValueError as exc: return PlanningResult(status=PlanningStatus.INVALID_GOAL, message=str(exc)) - if not self._world._is_ready(): + if not self._world.is_ready(): return PlanningResult( status=PlanningStatus.INVALID_START, message="RoboPlan planning scene is not ready: authoritative state is incomplete", ) - robot = self._world._get_robot(robot_id) - group = self._world._legacy_group(robot.config.name) + config = self._world.get_model_config() + group = self._world.all_planning_group() with self._world.scratch_context() as ctx: self._world.set_joint_state( ctx, - robot_id, - JointState(name=list(robot.config.joint_names), position=q_start.tolist()), + JointState(name=list(config.joint_names), position=q_start.tolist()), ) return self._plan_group( ctx, group, - dict(zip(robot.config.joint_names, q_start, strict=True)), - dict(zip(robot.config.joint_names, q_goal, strict=True)), + dict(zip(config.joint_names, q_start, strict=True)), + dict(zip(config.joint_names, q_goal, strict=True)), timeout, 5000, ) @@ -144,7 +141,7 @@ def plan_selected_joint_path( status=PlanningStatus.INVALID_GOAL, message="No planning groups selected", ) - group = self._world._require_model().groups.get(frozenset(selection.group_ids)) + group = self._world.planning_group(selection.group_ids) if group is None: return PlanningResult( status=PlanningStatus.UNSUPPORTED, @@ -195,7 +192,7 @@ def plan_cartesian_path( except ValueError as exc: return PlanningResult(status=PlanningStatus.INVALID_START, message=str(exc)) - group = self._world._require_model().groups.get(frozenset(selection.group_ids)) + group = self._world.planning_group(selection.group_ids) if group is None: return PlanningResult( status=PlanningStatus.UNSUPPORTED, @@ -256,7 +253,7 @@ def _normalize_selection_start( start: JointState, ) -> JointState: """Validate readiness and normalize the request's authoritative start.""" - if not self._world._is_ready(): + if not self._world.is_ready(): raise ValueError( "RoboPlan planning scene is not ready: authoritative state is incomplete" ) @@ -345,13 +342,12 @@ def _validate_cartesian_request( def _apply_selected_state(self, ctx: RoboPlanContext, state: JointState) -> None: positions = dict(zip(state.name, state.position, strict=True)) - for robot_id, robot in self._world._robots.items(): - q = ctx.q_by_robot[robot_id].copy() - for index, local_name in enumerate(robot.config.joint_names): - global_name = f"{robot.config.name}/{local_name}" - if global_name in positions: - q[index] = positions[global_name] - ctx.q_by_robot[robot_id] = q + config = self._world.get_model_config() + q = ctx.q.copy() + for index, name in enumerate(config.joint_names): + if name in positions: + q[index] = positions[name] + ctx.q = q def _build_cartesian_path( self, @@ -359,7 +355,6 @@ def _build_cartesian_path( selection: PlanningGroupSelection, targets: Mapping[PlanningGroupID, CartesianTarget], ) -> Any: - model = self._world._require_model() base_frames: list[str] = [] tip_frames: list[str] = [] waypoint_paths: list[list[NDArray[np.float64]]] = [] @@ -379,7 +374,7 @@ def _build_cartesian_path( f"Cartesian target for '{group.id}' must begin at its current TCP pose" ) base_frames.append(ROBOPLAN_WORLD_FRAME) - tip_frames.append(model.native_link(group.robot_name, group.tip_link)) + tip_frames.append(self._world.native_link_name(group.tip_link)) waypoint_paths.append(target_matrices) return roboplan_core.CartesianPath(base_frames, tip_frames, waypoint_paths) @@ -512,13 +507,11 @@ def _combined_path_collision_free( position=(q_start + fraction * (q_end - q_start)).tolist(), ) self._apply_selected_state(ctx, sample) - first_robot_id = next(iter(self._world._robots)) - if not self._world.is_collision_free(ctx, first_robot_id): + if not self._world.is_collision_free(ctx): return False if len(path) == 1: self._apply_selected_state(ctx, path[0]) - first_robot_id = next(iter(self._world._robots)) - return self._world.is_collision_free(ctx, first_robot_id) + return self._world.is_collision_free(ctx) return True def _plan_group( diff --git a/dimos/manipulation/planning/planners/rrt_planner.py b/dimos/manipulation/planning/planners/rrt_planner.py index 6c3b0586bc..c13f1c88b8 100644 --- a/dimos/manipulation/planning/planners/rrt_planner.py +++ b/dimos/manipulation/planning/planners/rrt_planner.py @@ -39,7 +39,6 @@ JointPath, PlanningGroupID, PlanningResult, - WorldRobotID, ) from dimos.manipulation.planning.spec.protocols import WorldSpec from dimos.manipulation.planning.utils.path_utils import compute_path_length @@ -93,7 +92,6 @@ def __init__( def plan_joint_path( self, world: WorldSpec, - robot_id: WorldRobotID, start: JointState, goal: JointState, timeout: float = 10.0, @@ -107,14 +105,14 @@ def plan_joint_path( q_goal = np.array(goal.position, dtype=np.float64) joint_names = start.name # Store for converting back to JointState - error = self._validate_inputs(world, robot_id, start, goal) + error = self._validate_inputs(world, start, goal) if error is not None: return error - if world.check_edge_collision_free(robot_id, start, goal, self._collision_step_size): + if world.check_edge_collision_free(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) + lower, upper = world.get_joint_limits() start_tree = [TreeNode(config=q_start.copy())] goal_tree = [TreeNode(config=q_goal.copy())] trees_swapped = False @@ -129,14 +127,11 @@ def plan_joint_path( ) sample = np.random.uniform(lower, upper) - extended = self._extend_tree( - world, robot_id, start_tree, sample, self._step_size, joint_names - ) + extended = self._extend_tree(world, start_tree, sample, self._step_size, joint_names) if extended is not None: connected = self._connect_tree( world, - robot_id, goal_tree, extended.config, self._connect_step_size, @@ -146,7 +141,7 @@ def plan_joint_path( path = self._extract_path(extended, connected, joint_names) if trees_swapped: path = list(reversed(path)) - path = self._simplify_path(world, robot_id, path) + path = self._simplify_path(world, path) return _create_success_result(path, time.time() - start_time, iteration + 1) start_tree, goal_tree = goal_tree, start_tree @@ -174,9 +169,9 @@ def plan_selected_joint_path( ) -> 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. + The search space uses the selected canonical-joint order. Collision checks + project candidates into the full model state while holding unselected joints + at their current values. """ start_time = time.time() if not world.is_finalized: @@ -359,7 +354,6 @@ def _connect_selected_tree( def _validate_inputs( self, world: WorldSpec, - robot_id: WorldRobotID, start: JointState, goal: JointState, ) -> PlanningResult | None: @@ -371,29 +365,22 @@ def _validate_inputs( "World must be finalized before planning", ) - # Check robot exists - if robot_id not in world.get_robot_ids(): - return _create_failure_result( - PlanningStatus.NO_SOLUTION, - f"Robot '{robot_id}' not found", - ) - # Check start validity using context-free method - if not world.check_config_collision_free(robot_id, start): + if not world.check_config_collision_free(start): return _create_failure_result( PlanningStatus.COLLISION_AT_START, "Start configuration is in collision", ) # Check goal validity using context-free method - if not world.check_config_collision_free(robot_id, goal): + if not world.check_config_collision_free(goal): return _create_failure_result( PlanningStatus.COLLISION_AT_GOAL, "Goal configuration is in collision", ) # Check limits with small tolerance for driver floating-point drift - lower, upper = world.get_joint_limits(robot_id) + lower, upper = world.get_joint_limits() q_start = np.array(start.position, dtype=np.float64) q_goal = np.array(goal.position, dtype=np.float64) limit_eps = 1e-3 # ~0.06 degrees @@ -415,7 +402,6 @@ def _validate_inputs( def _extend_tree( self, world: WorldSpec, - robot_id: WorldRobotID, tree: list[TreeNode], target: NDArray[np.float64], step_size: float, @@ -437,9 +423,7 @@ def _extend_tree( # 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()}) - if world.check_edge_collision_free( - robot_id, start_state, end_state, self._collision_step_size - ): + if world.check_edge_collision_free(start_state, end_state, self._collision_step_size): new_node = TreeNode(config=new_config, parent=nearest) nearest.children.append(new_node) tree.append(new_node) @@ -450,7 +434,6 @@ def _extend_tree( def _connect_tree( self, world: WorldSpec, - robot_id: WorldRobotID, tree: list[TreeNode], target: NDArray[np.float64], step_size: float, @@ -459,7 +442,7 @@ def _connect_tree( """Try to connect tree to target, returns connected node if successful.""" # Keep extending toward target while True: - result = self._extend_tree(world, robot_id, tree, target, step_size, joint_names) + result = self._extend_tree(world, tree, target, step_size, joint_names) if result is None: return None # Extension failed @@ -491,7 +474,6 @@ def _extract_path( def _simplify_path( self, world: WorldSpec, - robot_id: WorldRobotID, path: JointPath, max_iterations: int = 100, ) -> JointPath: @@ -512,7 +494,7 @@ def _simplify_path( # Check if direct connection is valid using context-free method # path elements are already JointState if world.check_edge_collision_free( - robot_id, simplified[i], simplified[j], self._collision_step_size + simplified[i], simplified[j], self._collision_step_size ): # Remove intermediate waypoints simplified = simplified[: i + 1] + simplified[j:] diff --git a/dimos/manipulation/planning/planners/selected_joint_space.py b/dimos/manipulation/planning/planners/selected_joint_space.py index 89d6efbadc..de783b463e 100644 --- a/dimos/manipulation/planning/planners/selected_joint_space.py +++ b/dimos/manipulation/planning/planners/selected_joint_space.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Selected planning-group joint-space projection helpers.""" +"""Projection between selected planning groups and one full-model joint vector.""" from __future__ import annotations @@ -21,92 +21,54 @@ 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.models import JointPath 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.""" + """One full-model baseline plus projections for selected canonical joints.""" - 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 + model_joint_names: tuple[str, ...] + selected_joint_names: tuple[str, ...] + base_positions: NDArray[np.float64] + lower_limits: NDArray[np.float64] + upper_limits: NDArray[np.float64] @classmethod - def from_world( - cls, - world: WorldSpec, - selection: PlanningGroupSelection, - ) -> SelectedJointSpace: + def from_world(cls, world: WorldSpec, selection: PlanningGroupSelection) -> SelectedJointSpace: + config = world.get_model_config() + with world.scratch_context() as ctx: + current = world.get_joint_state(ctx) + base = _ordered_positions(current, config.joint_names, "Current state") + lower, upper = world.get_joint_limits() + indices = [config.joint_names.index(name) for name in selection.joint_names] return cls( - robot_projections=_build_robot_projections(world, selection), - selected_joint_names=list(selection.joint_names), + model_joint_names=tuple(config.joint_names), + selected_joint_names=selection.joint_names, + base_positions=base, + lower_limits=np.asarray(lower, dtype=np.float64)[indices], + upper_limits=np.asarray(upper, dtype=np.float64)[indices], ) 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) + return self.lower_limits.copy(), self.upper_limits.copy() + + def project_config(self, selected_positions: NDArray[np.float64]) -> JointState: + if len(selected_positions) != len(self.selected_joint_names): + raise ValueError("Selected position count does not match selected joints") + positions = self.base_positions.copy() + indices = {name: i for i, name in enumerate(self.model_joint_names)} + for name, value in zip(self.selected_joint_names, selected_positions, strict=True): + positions[indices[name]] = value + return JointState(name=list(self.model_joint_names), position=positions.tolist()) def config_collision_free( - self, - world: WorldSpec, - selected_positions: NDArray[np.float64], + 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 - ) + return world.check_config_collision_free(self.project_config(selected_positions)) def edge_collision_free( self, @@ -117,34 +79,10 @@ def edge_collision_free( ) -> 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) + return all( + self.config_collision_free(world, start + (step / steps) * (end - start)) + for step in range(steps + 1) ) - 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, @@ -153,167 +91,50 @@ def simplify_path( 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): + if self.edge_collision_free( + world, + np.asarray(simplified[i].position), + np.asarray(simplified[j].position), + collision_step_size, + ): simplified = simplified[: i + 1] + simplified[j:] return simplified def normalize_selection_target( - selection: PlanningGroupSelection, - target: JointState, - label: str, + selection: PlanningGroupSelection, target: JointState, label: str ) -> JointState: - """Normalize a selected-joint target to global selection order.""" - selected_global_names = list(selection.joint_names) + """Normalize a target to canonical selection order.""" + names = list(selection.joint_names) if not target.name: - if len(target.position) != len(selected_global_names): + if len(target.position) != len(names): raise ValueError( - f"{label} target has {len(target.position)} positions, " - f"expected {len(selected_global_names)}" + f"{label} target has {len(target.position)} positions, expected {len(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] + return JointState(name=names, position=list(target.position)) + positions = dict(zip(target.name, target.position, strict=True)) + missing = [name for name in names if name not in positions] + extra = sorted(set(positions) - set(names)) 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 + return JointState(name=names, position=[float(positions[name]) for name in names]) -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") +def _ordered_positions(state: JointState, names: list[str], label: str) -> NDArray[np.float64]: + if not state.name: + if len(state.position) != len(names): + raise ValueError(f"{label} position count does not match model joints") + return np.asarray(state.position, dtype=np.float64) + positions = dict(zip(state.name, state.position, strict=True)) + missing = [name for name in names if name not in positions] + if missing: + raise ValueError(f"{label} is missing joints: {missing}") + return np.asarray([positions[name] for name in names], dtype=np.float64) diff --git a/dimos/manipulation/planning/planners/test_rrt_planner_selection.py b/dimos/manipulation/planning/planners/test_rrt_planner_selection.py index 3d401c8af9..b03c95636b 100644 --- a/dimos/manipulation/planning/planners/test_rrt_planner_selection.py +++ b/dimos/manipulation/planning/planners/test_rrt_planner_selection.py @@ -44,11 +44,8 @@ def _pose() -> PoseStamped: def _group(name: str, joints: tuple[str, ...]) -> PlanningGroup: return PlanningGroup( - id=f"arm/{name}", - robot_name="arm", - group_name=name, + id=name, joint_names=tuple(f"arm/{joint}" for joint in joints), - local_joint_names=joints, base_link="base", tip_link="tool", ) @@ -61,35 +58,37 @@ 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"], + joint_names=["arm/joint_a", "arm/joint_b", "arm/gripper"], base_link="base", planning_groups=[ - PlanningGroupDefinition("arm", ("joint_a", "joint_b"), "base", "tool") + PlanningGroupDefinition("arm", ("arm/joint_a", "arm/joint_b"), "base", "tool") ], ) - def get_robot_ids(self) -> list[str]: - return ["robot"] - - def get_robot_config(self, robot_id: str) -> RobotModelConfig: + def get_model_config(self) -> 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_state(self, ctx: object) -> JointState: + return JointState( + {"name": ["arm/joint_a", "arm/joint_b", "arm/gripper"], "position": self.current} + ) - def get_joint_limits(self, robot_id: str) -> tuple[np.ndarray, np.ndarray]: + def get_joint_limits(self) -> 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: + def set_joint_state(self, ctx: object, joint_state: JointState) -> None: self.projected_states.append(joint_state) - def is_collision_free(self, ctx: object, robot_id: str) -> bool: + def is_collision_free(self, ctx: object) -> bool: + return True + + def check_config_collision_free(self, joint_state: JointState) -> bool: + self.projected_states.append(joint_state) return True @@ -108,12 +107,6 @@ def is_collision_free(self, ctx: object, robot_id: str) -> bool: [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( @@ -152,7 +145,7 @@ def test_plan_selected_joint_path_normalizes_target_forms( JointState({"name": ["arm/joint_a", "joint_b"], "position": [0.0, 0.0]}), JointState({"position": [0.0, 0.0]}), PlanningStatus.INVALID_START, - "mixes", + "missing", ), ], ) @@ -181,7 +174,7 @@ def test_plan_selected_joint_path_rejects_local_names_for_multi_group_selection( ) assert result.status == PlanningStatus.INVALID_START - assert "multi-group" in result.message + assert "missing" in result.message def test_plan_selected_joint_path_direct_edge_projects_full_state_with_unselected_joints() -> None: @@ -197,7 +190,10 @@ def test_plan_selected_joint_path_direct_edge_projects_full_state_with_unselecte 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.name == ["arm/joint_a", "arm/joint_b", "arm/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 971fc43251..d52ca3d400 100644 --- a/dimos/manipulation/planning/spec/config.py +++ b/dimos/manipulation/planning/spec/config.py @@ -21,7 +21,7 @@ from pydantic import Field from dimos.core.module import ModuleConfig -from dimos.manipulation.planning.groups.identifiers import assert_valid_robot_name +from dimos.manipulation.planning.groups.identifiers import assert_valid_joint_names from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped @@ -30,7 +30,6 @@ class RobotModelConfig(ModuleConfig): """Configuration for adding a robot to the world. Attributes: - name: Human-readable robot name model_path: Path to robot model file (.urdf, .xacro, or .xml/MJCF) srdf_path: Optional path to SRDF file containing planning group definitions base_pose: Placement transform. This is the canonical world placement for @@ -50,12 +49,8 @@ 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 model joint names. - This is retained for current coordinator/monitor integrations while planning - APIs move toward globally scoped joint names. """ - name: str model_path: Path srdf_path: Path | None = None base_pose: PoseStamped = Field(default_factory=PoseStamped) @@ -72,8 +67,6 @@ class RobotModelConfig(ModuleConfig): # Motion constraints for trajectory generation max_velocity: float = 1.0 max_acceleration: float = 2.0 - # Coordinator integration - joint_name_mapping: dict[str, str] = Field(default_factory=dict) gripper_hardware_id: str | None = None # TF publishing for extra links (e.g., camera mount) tf_extra_links: list[str] = Field(default_factory=list) @@ -84,27 +77,8 @@ class RobotModelConfig(ModuleConfig): def model_post_init(self, __context: object) -> None: """Validate configuration-level naming constraints.""" - assert_valid_robot_name(self.name) + assert_valid_joint_names(self.joint_names) if any(not name for name in self.joint_names): raise ValueError("RobotModelConfig.joint_names must contain non-empty names") if len(self.joint_names) != len(set(self.joint_names)): - raise ValueError( - f"RobotModelConfig '{self.name}' contains duplicate canonical joint names" - ) - - def get_urdf_joint_name(self, coordinator_name: str) -> str: - """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 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 - - 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(joint_name) for joint_name in self.joint_names] + raise ValueError("RobotModelConfig contains duplicate canonical joint names") diff --git a/dimos/manipulation/planning/spec/models.py b/dimos/manipulation/planning/spec/models.py index a3b6f0da21..7afbccb965 100644 --- a/dimos/manipulation/planning/spec/models.py +++ b/dimos/manipulation/planning/spec/models.py @@ -16,7 +16,7 @@ from __future__ import annotations -from collections.abc import Mapping, Sequence +from collections.abc import Sequence from dataclasses import dataclass, field from typing import TYPE_CHECKING, TypeAlias @@ -38,20 +38,11 @@ from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory -RobotName: TypeAlias = str -"""User-facing robot name (e.g., 'left_arm', 'right_arm')""" - -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.""" +"""Stable public planning-group name.""" -GlobalJointName: TypeAlias = str -"""Public joint name of the form {robot_name}/{local_joint_name}.""" +JointName: TypeAlias = str +"""Canonical joint name used by model, planning, state, and execution.""" JointPath: TypeAlias = "list[JointState]" """List of joint states forming a path (each waypoint has names + positions)""" @@ -72,8 +63,8 @@ class PlanningSceneInfo: backend handles, mutable world contexts, GUI state, or execution state. """ - robots: Mapping[WorldRobotID, RobotModelConfig] - """Robot model configurations keyed by world robot ID.""" + model: RobotModelConfig + """The configured logical robot model.""" planning_groups: tuple[PlanningGroup, ...] = () """Resolved immutable planning groups for the initialized scene.""" @@ -92,7 +83,7 @@ class VisualizationSession: class VisualizationStateFrame: """Pushed current joint states for visualization backends.""" - joint_states: Mapping[WorldRobotID, JointState] + joint_state: JointState | None Jacobian: TypeAlias = "NDArray[np.float64]" diff --git a/dimos/manipulation/planning/spec/protocols.py b/dimos/manipulation/planning/spec/protocols.py index 1223d555a2..720bebc655 100644 --- a/dimos/manipulation/planning/spec/protocols.py +++ b/dimos/manipulation/planning/spec/protocols.py @@ -41,7 +41,6 @@ PlanningResult, VisualizationSession, VisualizationStateFrame, - WorldRobotID, ) from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.JointState import JointState @@ -66,23 +65,17 @@ class WorldSpec(Protocol): - DrakeWorld: Uses Drake's MultibodyPlant and SceneGraph """ - # Robot Management - def add_robot(self, config: RobotModelConfig) -> WorldRobotID: - """Add a robot to the world. Returns unique robot ID.""" + # Model Management + def load_model(self, config: RobotModelConfig) -> None: + """Load the logical robot model.""" ... - def get_robot_ids(self) -> list[WorldRobotID]: - """Get all robot IDs.""" + def get_model_config(self) -> RobotModelConfig: + """Get the logical robot model configuration.""" ... - def get_robot_config(self, robot_id: WorldRobotID) -> RobotModelConfig: - """Get robot configuration.""" - ... - - def get_joint_limits( - self, robot_id: WorldRobotID - ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: # lower limits, upper limits - """Get joint limits (lower, upper) for a robot.""" + def get_joint_limits(self) -> tuple[NDArray[np.float64], NDArray[np.float64]]: + """Get model joint limits (lower, upper).""" ... # Obstacle Management @@ -129,36 +122,35 @@ def scratch_context(self) -> AbstractContextManager[Any]: """Get a scratch context for planning (thread-safe clone).""" ... - def sync_from_joint_state(self, robot_id: WorldRobotID, joint_state: JointState) -> None: + def sync_from_joint_state(self, joint_state: JointState) -> None: """Sync live context from joint state message.""" ... # State Operations (require context) - def set_joint_state(self, ctx: Any, robot_id: WorldRobotID, joint_state: JointState) -> None: + def set_joint_state(self, ctx: Any, joint_state: JointState) -> None: """Set robot joint state in a context.""" ... - def get_joint_state(self, ctx: Any, robot_id: WorldRobotID) -> JointState: + def get_joint_state(self, ctx: Any) -> JointState: """Get robot joint state from a context.""" ... # Collision Checking (require context) - def is_collision_free(self, ctx: Any, robot_id: WorldRobotID) -> bool: + def is_collision_free(self, ctx: Any) -> bool: """Check if robot configuration is collision-free.""" ... - def get_min_distance(self, ctx: Any, robot_id: WorldRobotID) -> float: + def get_min_distance(self, ctx: Any) -> float: """Get minimum distance to obstacles (negative if collision).""" ... # Collision Checking (context-free, for planning) - def check_config_collision_free(self, robot_id: WorldRobotID, joint_state: JointState) -> bool: + def check_config_collision_free(self, joint_state: JointState) -> bool: """Check if a joint state is collision-free (manages context internally).""" ... def check_edge_collision_free( self, - robot_id: WorldRobotID, start: JointState, end: JointState, step_size: float = 0.05, @@ -167,17 +159,15 @@ def check_edge_collision_free( ... # Forward Kinematics (require context) - def get_ee_pose(self, ctx: Any, robot_id: WorldRobotID) -> PoseStamped: + def get_ee_pose(self, ctx: Any) -> PoseStamped: """Get end-effector pose.""" ... - def get_link_pose( - self, ctx: Any, robot_id: WorldRobotID, link_name: str - ) -> NDArray[np.float64]: + def get_link_pose(self, ctx: Any, link_name: str) -> NDArray[np.float64]: """Get link pose as 4x4 homogeneous transform.""" ... - def get_jacobian(self, ctx: Any, robot_id: WorldRobotID) -> NDArray[np.float64]: + def get_jacobian(self, ctx: Any) -> NDArray[np.float64]: """Get end-effector Jacobian (6 x n_joints).""" ... @@ -230,7 +220,7 @@ def get_visualization_url(self) -> str | None: ... def update_state(self, frame: VisualizationStateFrame) -> None: - """Receive current joint states keyed by initialized world robot ID.""" + """Receive the current model joint state.""" ... def animate_trajectory( @@ -239,7 +229,7 @@ def animate_trajectory( """Animate a raw globally named trajectory.""" ... - def cancel_preview_animation(self, robot_ids: Sequence[WorldRobotID] | None = None) -> None: + def cancel_preview_animation(self) -> None: """Cancel an active preview animation without waiting for its renderer to finish.""" ... @@ -255,7 +245,6 @@ class KinematicsSpec(Protocol): def solve( self, world: WorldSpec, - robot_id: WorldRobotID, target_pose: PoseStamped, seed: JointState | None = None, position_tolerance: float = 0.001, @@ -302,7 +291,6 @@ class PlannerSpec(Protocol): def plan_joint_path( self, world: WorldSpec, - robot_id: WorldRobotID, start: JointState, goal: JointState, timeout: float = 10.0, diff --git a/dimos/manipulation/planning/spec/test_model_validation.py b/dimos/manipulation/planning/spec/test_model_validation.py index 0311d1eccf..d4aafcb705 100644 --- a/dimos/manipulation/planning/spec/test_model_validation.py +++ b/dimos/manipulation/planning/spec/test_model_validation.py @@ -61,7 +61,6 @@ def _write_slash_model(path: Path) -> None: def _config(path: Path) -> RobotModelConfig: return RobotModelConfig( - name="robot", model_path=path, joint_names=["left/j1", "right/j1"], base_link="world", diff --git a/dimos/manipulation/planning/spec/validation.py b/dimos/manipulation/planning/spec/validation.py index 63abb59c77..14eb07f6ba 100644 --- a/dimos/manipulation/planning/spec/validation.py +++ b/dimos/manipulation/planning/spec/validation.py @@ -42,7 +42,7 @@ def validate_robot_model_config(config: RobotModelConfig) -> ModelDescription: """Validate one prepared robot model and its canonical planning groups.""" if not config.joint_names: - raise ValueError(f"RobotModelConfig '{config.name}' contains no controllable joints") + raise ValueError("RobotModelConfig contains no controllable joints") try: model = parse_model( config.model_path, @@ -50,86 +50,68 @@ def validate_robot_model_config(config: RobotModelConfig) -> ModelDescription: xacro_args=config.xacro_args, ) except (ET.ParseError, OSError, ValueError) as exc: - raise ValueError( - f"RobotModelConfig '{config.name}' has an invalid model asset: {exc}" - ) from exc + raise ValueError(f"RobotModelConfig has an invalid model asset: {exc}") from exc if config.srdf_path is not None: - _validate_srdf(config.name, config.srdf_path) + _validate_srdf(config.srdf_path) duplicate_model_joints = _duplicates(joint.name for joint in model.joints) if duplicate_model_joints: raise ValueError( - f"RobotModelConfig '{config.name}' model contains duplicate joint names: " - f"{duplicate_model_joints}" + f"RobotModelConfig model contains duplicate joint names: {duplicate_model_joints}" ) duplicate_links = _duplicates(model.links) if duplicate_links: - raise ValueError( - f"RobotModelConfig '{config.name}' model contains duplicate link names: " - f"{duplicate_links}" - ) + raise ValueError(f"RobotModelConfig model contains duplicate link names: {duplicate_links}") model_joint_names = {joint.name for joint in model.joints} missing_joints = sorted(set(config.joint_names) - model_joint_names) if missing_joints: raise ValueError( - f"RobotModelConfig '{config.name}' configured joints are missing from the model: " - f"{missing_joints}" + f"RobotModelConfig configured joints are missing from the model: {missing_joints}" ) if config.base_link not in model.links: - raise ValueError( - f"RobotModelConfig '{config.name}' base link '{config.base_link}' is missing" - ) + raise ValueError(f"RobotModelConfig base link '{config.base_link}' is missing") duplicate_group_names = _duplicates(group.name for group in config.planning_groups) if duplicate_group_names: raise ValueError( - f"RobotModelConfig '{config.name}' contains duplicate planning groups: " - f"{duplicate_group_names}" + f"RobotModelConfig contains duplicate planning groups: {duplicate_group_names}" ) controllable = set(config.joint_names) for group in config.planning_groups: if not group.name: - raise ValueError( - f"RobotModelConfig '{config.name}' contains an empty planning-group name" - ) + raise ValueError("RobotModelConfig contains an empty planning-group name") if not group.joint_names: - raise ValueError( - f"Planning group '{group.name}' for '{config.name}' contains no joints" - ) + raise ValueError(f"Planning group '{group.name}' contains no joints") duplicate_group_joints = _duplicates(group.joint_names) if duplicate_group_joints: raise ValueError( - f"Planning group '{group.name}' for '{config.name}' contains duplicate joints: " - f"{duplicate_group_joints}" + f"Planning group '{group.name}' contains duplicate joints: {duplicate_group_joints}" ) unknown_group_joints = sorted(set(group.joint_names) - controllable) if unknown_group_joints: raise ValueError( - f"Planning group '{group.name}' for '{config.name}' references joints outside " + f"Planning group '{group.name}' references joints outside " f"the controllable model set: {unknown_group_joints}" ) for role, link_name in (("base", group.base_link), ("tip", group.tip_link)): if link_name is not None and link_name not in model.links: raise ValueError( - f"Planning group '{group.name}' for '{config.name}' has missing " - f"{role} link '{link_name}'" + f"Planning group '{group.name}' has missing {role} link '{link_name}'" ) return model -def _validate_srdf(robot_name: str, srdf_path: Path) -> None: +def _validate_srdf(srdf_path: Path) -> None: if not srdf_path.exists(): - raise ValueError(f"RobotModelConfig '{robot_name}' SRDF file is missing: {srdf_path}") + raise ValueError(f"RobotModelConfig SRDF file is missing: {srdf_path}") try: root = ET.parse(srdf_path).getroot() except (ET.ParseError, OSError) as exc: - raise ValueError(f"RobotModelConfig '{robot_name}' has an invalid SRDF: {exc}") from exc + raise ValueError(f"RobotModelConfig has an invalid SRDF: {exc}") from exc if root.tag != "robot": - raise ValueError( - f"RobotModelConfig '{robot_name}' SRDF root must be , got <{root.tag}>" - ) + raise ValueError(f"RobotModelConfig SRDF root must be , got <{root.tag}>") def _duplicates(values: Iterable[str]) -> list[str]: diff --git a/dimos/manipulation/planning/trajectory_generator/simple_parametrizer.py b/dimos/manipulation/planning/trajectory_generator/simple_parametrizer.py index f939029693..e174b09a60 100644 --- a/dimos/manipulation/planning/trajectory_generator/simple_parametrizer.py +++ b/dimos/manipulation/planning/trajectory_generator/simple_parametrizer.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Compatibility trajectory parametrizer using segmented trapezoids.""" +"""Trajectory parametrizer using segmented trapezoids.""" import math @@ -79,28 +79,19 @@ def _selected_limits( world: WorldSpec, selection: PlanningGroupSelection, ) -> tuple[tuple[float, ...], tuple[float, ...]]: - configs = {} - for robot_id in world.get_robot_ids(): - config = world.get_robot_config(robot_id) - configs[config.name] = config + config = world.get_model_config() velocities: list[float] = [] accelerations: list[float] = [] - for global_name in selection.joint_names: - if "/" not in global_name: - raise TrajectoryParametrizationError(f"Joint '{global_name}' is not globally named") - robot_name, local_name = global_name.split("/", 1) - selected_config = configs.get(robot_name) - if selected_config is None: - raise TrajectoryParametrizationError(f"Unknown robot for joint '{global_name}'") - if local_name not in selected_config.joint_names: - raise TrajectoryParametrizationError(f"Unknown local joint '{global_name}'") - velocity = float(selected_config.max_velocity) - acceleration = float(selected_config.max_acceleration) + for joint_name in selection.joint_names: + if joint_name not in config.joint_names: + raise TrajectoryParametrizationError(f"Unknown model joint '{joint_name}'") + velocity = float(config.max_velocity) + acceleration = float(config.max_acceleration) if not math.isfinite(velocity) or velocity <= 0.0: - raise TrajectoryParametrizationError(f"Invalid velocity limit for '{global_name}'") + raise TrajectoryParametrizationError(f"Invalid velocity limit for '{joint_name}'") if not math.isfinite(acceleration) or acceleration <= 0.0: raise TrajectoryParametrizationError( - f"Invalid acceleration limit for '{global_name}'" + f"Invalid acceleration limit for '{joint_name}'" ) velocities.append(velocity) accelerations.append(acceleration) diff --git a/dimos/manipulation/planning/trajectory_generator/test_parametrizer.py b/dimos/manipulation/planning/trajectory_generator/test_parametrizer.py index 8d49c58eec..cb303cde22 100644 --- a/dimos/manipulation/planning/trajectory_generator/test_parametrizer.py +++ b/dimos/manipulation/planning/trajectory_generator/test_parametrizer.py @@ -54,11 +54,8 @@ def _selection() -> PlanningGroupSelection: return PlanningGroupSelection.from_groups( ( PlanningGroup( - id="arm/group", - robot_name="arm", - group_name="group", + id="group", joint_names=("arm/a", "arm/b"), - local_joint_names=("a", "b"), base_link="base", ), ) diff --git a/dimos/manipulation/planning/trajectory_generator/test_roboplan_toppra_parametrizer.py b/dimos/manipulation/planning/trajectory_generator/test_roboplan_toppra_parametrizer.py index d225f831b0..85bcada6f5 100644 --- a/dimos/manipulation/planning/trajectory_generator/test_roboplan_toppra_parametrizer.py +++ b/dimos/manipulation/planning/trajectory_generator/test_roboplan_toppra_parametrizer.py @@ -71,7 +71,7 @@ def parametrization_model(self): def _model(*, unbounded_acceleration: bool = False) -> RoboPlanModel: group = RoboPlanGroup( - group_ids=("left/arm", "right/arm"), + group_ids=("left_arm", "right_arm"), name="composite", native_names=("native_b", "native_a"), public_names=("right/b", "left/a"), @@ -79,9 +79,8 @@ def _model(*, unbounded_acceleration: bool = False) -> RoboPlanModel: return RoboPlanModel( scene=_Scene(unbounded_acceleration=unbounded_acceleration), groups={frozenset(group.group_ids): group}, - legacy_group_ids={}, - native_joint_by_global={}, - native_link_by_robot={}, + native_joints={}, + native_links={}, all_group=group, ) @@ -95,19 +94,13 @@ def _selection_and_result( } groups_by_name = { "left/a": PlanningGroup( - id="left/arm", - robot_name="left", - group_name="arm", + id="left_arm", joint_names=("left/a",), - local_joint_names=("a",), base_link="base", ), "right/b": PlanningGroup( - id="right/arm", - robot_name="right", - group_name="arm", + id="right_arm", joint_names=("right/b",), - local_joint_names=("b",), base_link="base", ), } @@ -268,7 +261,7 @@ def test_roboplan_parametrizer_reports_missing_generated_group() -> None: with pytest.raises( TrajectoryParametrizationError, - match=r"RoboPlan has no generated group for \['left/arm', 'right/arm'\]", + match=r"RoboPlan has no generated group for \['left_arm', 'right_arm'\]", ): RoboPlanTOPPRAParametrizer(RoboPlanTOPPRAParametrizationConfig()).materialize_plan( _World(model), selection, result diff --git a/dimos/manipulation/planning/trajectory_generator/test_simple_parametrizer.py b/dimos/manipulation/planning/trajectory_generator/test_simple_parametrizer.py index 37c353610b..24d0dd236a 100644 --- a/dimos/manipulation/planning/trajectory_generator/test_simple_parametrizer.py +++ b/dimos/manipulation/planning/trajectory_generator/test_simple_parametrizer.py @@ -45,11 +45,8 @@ def _selection() -> PlanningGroupSelection: return PlanningGroupSelection.from_groups( ( PlanningGroup( - id="arm/manipulator", - robot_name="arm", - group_name="manipulator", + id="manipulator", joint_names=("arm/a", "arm/b"), - local_joint_names=("a", "b"), base_link="base", tip_link="tip", ), @@ -59,17 +56,15 @@ def _selection() -> PlanningGroupSelection: def _world(*, velocity: float = 2.0, acceleration: float = 6.0) -> WorldSpec: config = RobotModelConfig( - name="arm", model_path=Path("/robot.urdf"), base_pose=PoseStamped(), - joint_names=["a", "b"], + joint_names=["arm/a", "arm/b"], base_link="base", max_velocity=velocity, max_acceleration=acceleration, ) world = MagicMock(spec=WorldSpec) - world.get_robot_ids.return_value = ["arm-id"] - world.get_robot_config.return_value = config + world.get_model_config.return_value = config return world @@ -102,7 +97,7 @@ def test_simple_parametrizer_materializes_segmented_trapezoid_plan() -> None: speed_scale=0.5, ) - assert plan.group_ids == ("arm/manipulator",) + assert plan.group_ids == ("manipulator",) assert plan.trajectory.joint_names == ["arm/a", "arm/b"] assert len(plan.trajectory.points) == 9 assert plan.trajectory.points[0].positions == [0.0, 0.0] diff --git a/dimos/manipulation/planning/utils/kinematics_utils.py b/dimos/manipulation/planning/utils/kinematics_utils.py index 02e885f1ae..29e02f6e6d 100644 --- a/dimos/manipulation/planning/utils/kinematics_utils.py +++ b/dimos/manipulation/planning/utils/kinematics_utils.py @@ -60,7 +60,7 @@ def damped_pseudoinverse( n x 6 pseudoinverse matrix Example: - J = world.get_jacobian(ctx, robot_id) + J = world.get_jacobian(ctx) J_pinv = damped_pseudoinverse(J, damping=0.01) q_dot = J_pinv @ twist """ @@ -88,7 +88,7 @@ def check_singularity( True if near singularity (manipulability < threshold) Example: - J = world.get_jacobian(ctx, robot_id) + J = world.get_jacobian(ctx) if check_singularity(J, threshold=0.001): logger.warning("Near singularity, using damped IK") """ @@ -114,7 +114,7 @@ def get_manipulability(J: Jacobian) -> float: Manipulability measure (non-negative) Example: - J = world.get_jacobian(ctx, robot_id) + J = world.get_jacobian(ctx) w = get_manipulability(J) print(f"Manipulability: {w:.4f}") """ @@ -140,7 +140,7 @@ def compute_pose_error( Tuple of (position_error, orientation_error) in meters and radians Example: - current = world.get_ee_pose(ctx, robot_id) + current = world.get_ee_pose(ctx) pos_err, ori_err = compute_pose_error(current, target) converged = pos_err < 0.001 and ori_err < 0.01 """ diff --git a/dimos/manipulation/planning/utils/path_utils.py b/dimos/manipulation/planning/utils/path_utils.py index dd5de1a0a4..0a9235b0ed 100644 --- a/dimos/manipulation/planning/utils/path_utils.py +++ b/dimos/manipulation/planning/utils/path_utils.py @@ -37,7 +37,7 @@ if TYPE_CHECKING: from numpy.typing import NDArray - from dimos.manipulation.planning.spec.models import JointPath, WorldRobotID + from dimos.manipulation.planning.spec.models import JointPath from dimos.manipulation.planning.spec.protocols import WorldSpec @@ -59,7 +59,7 @@ def interpolate_path( Example: # After planning, interpolate for smoother execution - raw_path = planner.plan_joint_path(world, robot_id, start, goal).path + raw_path = planner.plan_joint_path(world, start, goal).path smooth_path = interpolate_path(raw_path, resolution=0.02) """ if len(path) <= 1: @@ -109,7 +109,7 @@ def interpolate_segment( # Check collision along a segment segment = interpolate_segment(start_state, end_state, step_size=0.02) for state in segment: - if not world.check_config_collision_free(robot_id, state): + if not world.check_config_collision_free(state): return False """ q_start = np.array(start.position, dtype=np.float64) @@ -135,7 +135,6 @@ def interpolate_segment( def simplify_path( world: WorldSpec, - robot_id: WorldRobotID, path: JointPath, max_iterations: int = 100, collision_step_size: float = 0.02, @@ -148,7 +147,6 @@ def simplify_path( Args: world: World for collision checking - robot_id: Which robot path: Original path (list of JointState waypoints) max_iterations: Maximum shortcutting attempts collision_step_size: Step size for collision checking along shortcuts @@ -157,8 +155,8 @@ def simplify_path( Simplified path with fewer waypoints Example: - raw_path = planner.plan_joint_path(world, robot_id, start, goal).path - simplified = simplify_path(world, robot_id, raw_path) + raw_path = planner.plan_joint_path(world, start, goal).path + simplified = simplify_path(world, raw_path) """ if len(path) <= 2: return list(path) @@ -174,9 +172,7 @@ def simplify_path( j = np.random.randint(i + 2, len(simplified)) # Check if direct connection is valid using context-free API - if world.check_edge_collision_free( - robot_id, simplified[i], simplified[j], collision_step_size - ): + if world.check_edge_collision_free(simplified[i], simplified[j], collision_step_size): # Remove intermediate waypoints simplified = simplified[: i + 1] + simplified[j:] diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index 7ee76d52b1..75a7cb7024 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -16,7 +16,6 @@ from __future__ import annotations -from collections.abc import Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from copy import deepcopy @@ -28,15 +27,12 @@ import numpy as np -from dimos.manipulation.planning.groups.identifiers import ( - make_global_joint_names, - make_planning_group_id, -) +from dimos.manipulation.planning.groups.identifiers import assert_valid_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 Obstacle, PlanningGroupID, WorldRobotID +from dimos.manipulation.planning.spec.models import Obstacle, PlanningGroupID from dimos.manipulation.planning.spec.protocols import VisualizationSpec, WorldSpec from dimos.manipulation.planning.spec.validation import ( validate_obstacle, @@ -103,7 +99,6 @@ class _RobotData: """Internal data for tracking a robot in the world.""" - robot_id: WorldRobotID config: RobotModelConfig model_instance: Any # ModelInstanceIndex joint_indices: list[int] # Indices into plant's position vector @@ -185,8 +180,8 @@ def __init__(self, time_step: float = 0.0, enable_viz: bool = False) -> None: self._builder, time_step=time_step ) self._parser = Parser(self._plant) - # Enable auto-renaming to avoid conflicts when adding multiple robots - # with the same URDF (e.g., 4 XArm6 arms all have model name "UF_ROBOT") + # The visualization preview loads a second copy of the configured model. + # Auto-renaming prevents its internal model name from colliding with the live copy. self._parser.SetAutoRenaming(True) # Visualization — wrapped to enforce Drake's thread affinity @@ -199,9 +194,8 @@ def __init__(self, time_step: float = 0.0, enable_viz: bool = False) -> None: self._obstacles_model_instance = self._plant.AddModelInstance("obstacles") # Tracking data - self._robots: dict[WorldRobotID, _RobotData] = {} + self._model: _RobotData | None = None self._obstacles: dict[str, _ObstacleData] = {} - self._robot_counter = 0 self._obstacle_counter = 0 # Built diagram and contexts (created after finalize) @@ -211,28 +205,21 @@ def __init__(self, time_step: float = 0.0, enable_viz: bool = False) -> 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 - def add_robot(self, config: RobotModelConfig) -> WorldRobotID: - """Add a robot to the world. Returns robot_id. - - Same model_path + base_pose reuses the model instance (e.g. two arms in one URDF). - """ + def load_model(self, config: RobotModelConfig) -> None: + """Load the one logical robot model.""" if self._finalized: 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") + if self._model is not None: + raise ValueError("A model is already loaded") validate_robot_model_config(config) self._validate_planning_group_config(config) - self._robot_counter += 1 - robot_id = f"robot_{self._robot_counter}" - model_instance = self._load_model(config) self._weld_base_if_needed(config, model_instance) @@ -256,8 +243,7 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: preview_model_instance = self._load_model(config) self._weld_base_if_needed(config, preview_model_instance) - self._robots[robot_id] = _RobotData( - robot_id=robot_id, + self._model = _RobotData( config=config, model_instance=model_instance, joint_indices=[], @@ -265,8 +251,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 def _load_model(self, config: RobotModelConfig) -> Any: """Load robot model (URDF/xacro/MJCF) and return model instance.""" @@ -292,11 +276,9 @@ def _load_model(self, config: RobotModelConfig) -> Any: for pkg_name, pkg_path in config.package_paths.items(): self._parser.package_map().Add(pkg_name, Path(pkg_path)) else: - self._parser.package_map().Add( - f"{config.name}_description", prepared_path_obj.parent - ) + self._parser.package_map().Add("robot_description", prepared_path_obj.parent) - logger.info(f"Using prepared model: {prepared_path_obj}") + logger.info("Using prepared model", model_path=str(prepared_path_obj)) model_instances = self._parser.AddModels(prepared_path_obj) if not model_instances: @@ -375,24 +357,31 @@ def _validate_joints(self, config: RobotModelConfig, model_instance: Any) -> Non except RuntimeError: raise ValueError(f"Joint '{joint_name}' not found in URDF") - def get_robot_ids(self) -> list[WorldRobotID]: - """Get all robot IDs in the world.""" - return list(self._robots.keys()) + def get_model_config(self) -> RobotModelConfig: + """Get the logical robot model configuration.""" + return self._require_model().config - def get_robot_config(self, robot_id: WorldRobotID) -> RobotModelConfig: - """Get robot configuration by ID.""" - if robot_id not in self._robots: - raise KeyError(f"Robot '{robot_id}' not found") - return self._robots[robot_id].config + def get_body_frame(self, link_name: str) -> Any: + """Return a configured model link frame for Drake-native planning backends.""" + robot_data = self._require_model() + return self._plant.GetBodyByName(link_name, robot_data.model_instance).body_frame() + + def get_model_joint_indices(self) -> list[int]: + """Return Drake position indices in canonical model-joint order.""" + return list(self._require_model().joint_indices) + + def _require_model(self) -> _RobotData: + if self._model is None: + raise RuntimeError("Model is not loaded") + return self._model @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) + assert_valid_group_id(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 @@ -400,13 +389,9 @@ 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)) + if definition.name == group_id: return PlanningGroup( group_id, - config.name, - definition.name, - joint_names, definition.joint_names, definition.base_link, definition.tip_link, @@ -415,12 +400,7 @@ def _planning_group_from_config( 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}") + return self._planning_group_from_config(self._require_model().config, group_id) @staticmethod def _primary_pose_group_id_for_config(config: RobotModelConfig) -> PlanningGroupID | None: @@ -428,17 +408,13 @@ def _primary_pose_group_id_for_config(config: RobotModelConfig) -> PlanningGroup 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) + raise ValueError("Model has multiple pose groups") + return pose_groups[0].name - def get_joint_limits( - self, robot_id: WorldRobotID - ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: + def get_joint_limits(self) -> tuple[NDArray[np.float64], NDArray[np.float64]]: """Get joint limits (lower, upper) in radians.""" - if robot_id not in self._robots: - raise KeyError(f"Robot '{robot_id}' not found") - - config = self._robots[robot_id].config + robot_data = self._require_model() + config = robot_data.config if config.joint_limits_lower is not None and config.joint_limits_upper is not None: return ( @@ -448,7 +424,6 @@ def get_joint_limits( # Query Drake plant if finalized (limits from URDF/MJCF) if self._finalized: - robot_data = self._robots[robot_id] lower = [] upper = [] for joint_name in config.joint_names: @@ -486,7 +461,7 @@ def add_obstacle(self, obstacle: Obstacle) -> str | None: # Check for duplicate in our tracking if obstacle_id in self._obstacles: - logger.debug(f"Obstacle '{obstacle_id}' already exists, skipping") + logger.debug("Obstacle already exists", obstacle_id=obstacle_id) return None snapshot = deepcopy(obstacle) @@ -498,12 +473,16 @@ def add_obstacle(self, obstacle: Obstacle) -> str | None: geometry_id=geometry_id, source_id=self._obstacle_source_id, ) - logger.debug(f"Added obstacle '{obstacle_id}': {obstacle.obstacle_type.value}") + logger.debug( + "Added obstacle", + obstacle_id=obstacle_id, + obstacle_type=obstacle.obstacle_type.value, + ) except RuntimeError as e: # Handle case where geometry name already exists in SceneGraph # (can happen with concurrent access) if "already been used" in str(e): - logger.debug(f"Obstacle '{obstacle_id}' already in SceneGraph, skipping") + logger.debug("Obstacle already in SceneGraph", obstacle_id=obstacle_id) return None else: raise @@ -667,7 +646,7 @@ def remove_obstacle(self, obstacle_id: str) -> bool: obstacle_data = self._obstacles[obstacle_id] self._remove_obstacle_geometry(obstacle_data) del self._obstacles[obstacle_id] - logger.debug(f"Removed obstacle '{obstacle_id}'") + logger.debug("Removed obstacle", obstacle_id=obstacle_id) return True def update_obstacle(self, obstacle: Obstacle) -> bool: @@ -717,9 +696,10 @@ def _set_preview_colors(self) -> None: source_id: Any = self._plant.get_source_id() preview_color = Rgba(1.0, 0.8, 0.0, 0.4) - for robot_data in self._robots.values(): + robot_data = self._model + if robot_data is not None: if robot_data.preview_model_instance is None: - continue + return for body_idx in self._plant.GetBodyIndices(robot_data.preview_model_instance): body = self._plant.get_body(body_idx) for geom_id in self._plant.GetVisualGeometriesForBody(body): @@ -731,9 +711,10 @@ def _remove_preview_collision_roles(self) -> None: """Remove proximity (collision) role from all preview robot geometries.""" source_id: Any = self._plant.get_source_id() # SourceId - for robot_data in self._robots.values(): + robot_data = self._model + if robot_data is not None: if robot_data.preview_model_instance is None: - continue + return for body_idx in self._plant.GetBodyIndices(robot_data.preview_model_instance): body = self._plant.get_body(body_idx) for geom_id in self._plant.GetCollisionGeometriesForBody(body): @@ -752,8 +733,9 @@ def finalize(self) -> None: # Finalize plant self._plant.Finalize() - # Compute joint indices for each robot (live + preview) - for robot_id, robot_data in self._robots.items(): + robot_data = self._require_model() + # Compute joint indices for the model (live + preview) + if robot_data is not None: joint_indices: list[int] = [] for joint_name in robot_data.config.joint_names: joint = self._plant.GetJointByName(joint_name, robot_data.model_instance) @@ -761,7 +743,7 @@ def finalize(self) -> None: num_positions = joint.num_positions() joint_indices.extend(range(start_idx, start_idx + num_positions)) robot_data.joint_indices = joint_indices - logger.debug(f"Robot '{robot_id}' joint indices: {joint_indices}") + logger.debug("Computed model joint indices", joint_indices=joint_indices) # Compute preview joint indices if robot_data.preview_model_instance is not None: @@ -774,7 +756,7 @@ def finalize(self) -> None: num_positions = joint.num_positions() preview_indices.extend(range(start_idx, start_idx + num_positions)) robot_data.preview_joint_indices = preview_indices - logger.debug(f"Robot '{robot_id}' preview joint indices: {preview_indices}") + logger.debug("Computed preview joint indices", joint_indices=preview_indices) # Setup collision filters self._setup_collision_filters() @@ -812,20 +794,17 @@ def finalize(self) -> None: ) # Set home pose for robots that have one configured - for robot_data in self._robots.values(): - if robot_data.config.home_joints is not None: - home = np.array(robot_data.config.home_joints, dtype=np.float64) - self._set_positions_internal(self._plant_context, robot_data.robot_id, home) + if robot_data.config.home_joints is not None: + home = np.array(robot_data.config.home_joints, dtype=np.float64) + self._set_positions_internal(self._plant_context, home) self._finalized = True - logger.info(f"World finalized with {len(self._robots)} robots") # Initial visualization publish (routed to Meshcat thread) if self._meshcat_visualizer is not None: self._publish_visualization() # Hide all preview robots initially - for robot_id in self._robots: - self._set_preview_visibility(robot_id, False) + self._set_preview_visibility(False) @property def is_finalized(self) -> bool: @@ -843,7 +822,8 @@ def _require_finalized(self) -> None: def _setup_collision_filters(self) -> None: """Filter collisions between adjacent links and user-specified pairs.""" - for robot_data in self._robots.values(): + robot_data = self._require_model() + if robot_data is not None: # Filter parent-child pairs (adjacent links always "collide") for joint_idx in self._plant.GetJointIndices(robot_data.model_instance): joint = self._plant.get_joint(joint_idx) @@ -858,7 +838,9 @@ def _setup_collision_filters(self) -> None: body2 = self._plant.GetBodyByName(name2, robot_data.model_instance) self._exclude_body_pair(body1, body2) except RuntimeError: - logger.warning(f"Collision exclusion: link not found: {name1} or {name2}") + logger.warning( + "Collision exclusion link not found", first_link=name1, second_link=name2 + ) logger.info("Collision filters applied") @@ -896,18 +878,18 @@ def scratch_context(self) -> Generator[Context, None, None]: # Copy live robot states so inter-robot collision checking works if self._plant_context is not None: plant_ctx = self._diagram.GetMutableSubsystemContext(self._plant, ctx) - for robot_data in self._robots.values(): - try: - positions = self._plant.GetPositions( - self._plant_context, robot_data.model_instance - ) - self._plant.SetPositions(plant_ctx, robot_data.model_instance, positions) - except RuntimeError: - pass # Robot not yet synced + robot_data = self._require_model() + try: + positions = self._plant.GetPositions( + self._plant_context, robot_data.model_instance + ) + self._plant.SetPositions(plant_ctx, robot_data.model_instance, positions) + except RuntimeError: + pass # Model not yet synced yield ctx - def sync_from_joint_state(self, robot_id: WorldRobotID, joint_state: JointState) -> None: + def sync_from_joint_state(self, joint_state: JointState) -> None: """Sync live context from driver's joint state message. Called by StateMonitor when new JointState arrives. @@ -915,11 +897,11 @@ 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 - positions = self._joint_state_to_q(robot_id, joint_state) + positions = self._joint_state_to_q(joint_state) with self._lock: self._require_usable() - self._set_positions_internal(self._plant_context, robot_id, positions) + self._set_positions_internal(self._plant_context, positions) # NOTE: ForcedPublish is intentionally NOT called here. # Calling ForcedPublish from the LCM callback thread blocks message processing. @@ -927,24 +909,17 @@ def sync_from_joint_state(self, robot_id: WorldRobotID, joint_state: JointState) # State Operations (context-based) - def set_joint_state( - self, ctx: Context, robot_id: WorldRobotID, joint_state: JointState - ) -> None: + def set_joint_state(self, ctx: Context, joint_state: JointState) -> None: """Set robot joint state in given context.""" with self._lock: self._require_finalized() - positions = self._joint_state_to_q(robot_id, joint_state) + positions = self._joint_state_to_q(joint_state) plant_ctx = self._diagram.GetMutableSubsystemContext(self._plant, ctx) - self._set_positions_internal(plant_ctx, robot_id, positions) + self._set_positions_internal(plant_ctx, positions) - def _set_positions_internal( - self, plant_ctx: Context, robot_id: WorldRobotID, positions: NDArray[np.float64] - ) -> None: + def _set_positions_internal(self, plant_ctx: Context, positions: NDArray[np.float64]) -> None: """Internal: Set positions in a plant context.""" - if robot_id not in self._robots: - raise KeyError(f"Robot '{robot_id}' not found") - - robot_data = self._robots[robot_id] + robot_data = self._require_model() full_positions = self._plant.GetPositions(plant_ctx).copy() for i, joint_idx in enumerate(robot_data.joint_indices): @@ -952,35 +927,19 @@ 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] + def _joint_state_to_q(self, joint_state: JointState) -> NDArray[np.float64]: + """Normalize a canonical JointState to model joint order.""" + robot_data = self._require_model() 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: + def get_joint_state(self, ctx: Context) -> JointState: """Get robot joint state from given context.""" with self._lock: self._require_finalized() - if robot_id not in self._robots: - raise KeyError(f"Robot '{robot_id}' not found") - robot_data = self._robots[robot_id] + robot_data = self._require_model() plant_ctx = self._diagram.GetSubsystemContext(self._plant, ctx) full_positions = self._plant.GetPositions(plant_ctx) positions = [float(full_positions[idx]) for idx in robot_data.joint_indices] @@ -988,17 +947,16 @@ def get_joint_state(self, ctx: Context, robot_id: WorldRobotID) -> JointState: # Collision Checking (context-based) - def is_collision_free(self, ctx: Context, robot_id: WorldRobotID) -> bool: + def is_collision_free(self, ctx: Context) -> bool: """Check if current configuration in context is collision-free.""" with self._lock: self._require_finalized() - if robot_id not in self._robots: - raise KeyError(f"Robot '{robot_id}' not found") + self._require_model() scene_graph_ctx = self._diagram.GetSubsystemContext(self._scene_graph, ctx) query_object = self._scene_graph.get_query_output_port().Eval(scene_graph_ctx) return not query_object.HasCollisions() - def get_min_distance(self, ctx: Context, robot_id: WorldRobotID) -> float: + def get_min_distance(self, ctx: Context) -> float: """Get minimum signed distance (positive = clearance, negative = penetration).""" with self._lock: self._require_finalized() @@ -1011,7 +969,7 @@ def get_min_distance(self, ctx: Context, robot_id: WorldRobotID) -> float: # Collision Checking (context-free, for planning) - def check_config_collision_free(self, robot_id: WorldRobotID, joint_state: JointState) -> bool: + def check_config_collision_free(self, joint_state: JointState) -> bool: """Check if a joint state is collision-free (manages context internally). This is a convenience method for planners that don't need to manage contexts. @@ -1019,12 +977,11 @@ def check_config_collision_free(self, robot_id: WorldRobotID, joint_state: Joint with self._lock: self._require_finalized() with self.scratch_context() as ctx: - self.set_joint_state(ctx, robot_id, joint_state) - return self.is_collision_free(ctx, robot_id) + self.set_joint_state(ctx, joint_state) + return self.is_collision_free(ctx) def check_edge_collision_free( self, - robot_id: WorldRobotID, start: JointState, end: JointState, step_size: float = 0.05, @@ -1041,30 +998,26 @@ def check_edge_collision_free( q_end = np.array(end.position, dtype=np.float64) dist = float(np.linalg.norm(q_end - q_start)) if dist < 1e-8: - return self.check_config_collision_free(robot_id, start) + return self.check_config_collision_free(start) n_steps = max(2, int(np.ceil(dist / step_size)) + 1) with self.scratch_context() as ctx: for i in range(n_steps): t = i / (n_steps - 1) q = q_start + t * (q_end - q_start) interp_state = JointState(name=start.name, position=q.tolist()) - self.set_joint_state(ctx, robot_id, interp_state) - if not self.is_collision_free(ctx, robot_id): + self.set_joint_state(ctx, interp_state) + if not self.is_collision_free(ctx): return False return True # Forward Kinematics (context-based) - def get_ee_pose(self, ctx: Context, robot_id: WorldRobotID) -> PoseStamped: + def get_ee_pose(self, ctx: Context) -> PoseStamped: """Get end-effector pose.""" - if robot_id not in self._robots: - raise KeyError(f"Robot '{robot_id}' not found") - robot_data = self._robots[robot_id] + robot_data = self._require_model() 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" - ) + raise ValueError("Model 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: @@ -1077,7 +1030,7 @@ def _get_group_ee_pose(self, ctx: Context, group_id: PlanningGroupID) -> PoseSta 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)] + robot_data = self._require_model() plant_ctx = self._diagram.GetSubsystemContext(self._plant, ctx) ee_body = self._plant.GetBodyByName(group.tip_link, robot_data.model_instance) @@ -1093,46 +1046,35 @@ def _get_group_ee_pose(self, ctx: Context, group_id: PlanningGroupID) -> PoseSta orientation=[float(quat.x()), float(quat.y()), float(quat.z()), float(quat.w())], ) - def get_link_pose( - self, ctx: Context, robot_id: WorldRobotID, link_name: str - ) -> NDArray[np.float64]: + def get_link_pose(self, ctx: Context, link_name: str) -> NDArray[np.float64]: """Get link pose as 4x4 transform.""" with self._lock: self._require_finalized() - return self._get_link_pose(ctx, robot_id, link_name) - - def _get_link_pose( - self, ctx: Context, robot_id: WorldRobotID, link_name: str - ) -> NDArray[np.float64]: - if robot_id not in self._robots: - raise KeyError(f"Robot '{robot_id}' not found") + return self._get_link_pose(ctx, link_name) - robot_data = self._robots[robot_id] + def _get_link_pose(self, ctx: Context, link_name: str) -> NDArray[np.float64]: + robot_data = self._require_model() plant_ctx = self._diagram.GetSubsystemContext(self._plant, ctx) try: body = self._plant.GetBodyByName(link_name, robot_data.model_instance) except RuntimeError: - raise KeyError(f"Link '{link_name}' not found in robot '{robot_id}'") + raise KeyError(f"Link '{link_name}' not found in model") X_WL = self._plant.EvalBodyPoseInWorld(plant_ctx, body) result = X_WL.GetAsMatrix4() return result # type: ignore[no-any-return] - def get_jacobian(self, ctx: Context, robot_id: WorldRobotID) -> NDArray[np.float64]: + def get_jacobian(self, ctx: Context) -> NDArray[np.float64]: """Get geometric Jacobian (6 x n_joints). Rows: [vx, vy, vz, wx, wy, wz] (linear, then angular) """ - if robot_id not in self._robots: - raise KeyError(f"Robot '{robot_id}' not found") - robot_data = self._robots[robot_id] + robot_data = self._require_model() 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" - ) + raise ValueError("Model 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]: @@ -1145,7 +1087,7 @@ def _get_group_jacobian(self, ctx: Context, group_id: PlanningGroupID) -> NDArra 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)] + robot_data = self._require_model() plant_ctx = self._diagram.GetSubsystemContext(self._plant, ctx) tip_frame = self._plant.GetBodyByName( group.tip_link, robot_data.model_instance @@ -1167,14 +1109,14 @@ def _get_group_jacobian(self, ctx: Context, group_id: PlanningGroupID) -> NDArra ) missing = [ joint_name - for joint_name in group.local_joint_names + for joint_name in group.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] + group_joint_indices = [joint_indices_by_name[name] for name in group.joint_names] n_joints = len(group_joint_indices) J_robot = np.zeros((6, n_joints)) @@ -1234,11 +1176,9 @@ 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: - """Set preview robot positions in a plant context.""" - robot_data = self._robots.get(robot_id) + def _set_preview_positions(self, plant_ctx: Context, positions: NDArray[np.float64]) -> None: + """Set preview model positions in a plant context.""" + robot_data = self._model if robot_data is None or robot_data.preview_model_instance is None: return @@ -1247,110 +1187,73 @@ def _set_preview_positions( full_positions[idx] = positions[i] self._plant.SetPositions(plant_ctx, full_positions) - def _set_preview_visibility(self, robot_id: WorldRobotID, visible: bool) -> None: - """Set one preview robot's Meshcat visibility.""" + def _set_preview_visibility(self, visible: bool) -> None: + """Set preview model Meshcat visibility.""" if self._meshcat is None: return - robot_data = self._robots.get(robot_id) + robot_data = self._model 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", visible) - def cancel_preview_animation(self, robot_ids: Sequence[WorldRobotID] | None = None) -> None: + def cancel_preview_animation(self) -> 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 + self._set_preview_visibility(False) + + def _trajectory_indices(self, trajectory: JointTrajectory) -> list[tuple[int, str]]: + known = set(self._require_model().config.joint_names) + unknown = [name for name in trajectory.joint_names if name not in known] + if unknown: + raise ValueError(f"trajectory references unknown joints: {unknown}") + return list(enumerate(trajectory.joint_names)) def animate_trajectory( self, trajectory: JointTrajectory, duration: float | None = None ) -> None: - """Render raw globally named trajectory on its stored shared clock.""" + """Render a canonical trajectory on its stored shared clock.""" if self._meshcat is None or len(trajectory.points) < 2: return import time - robot_indices = self._robot_trajectory_indices(trajectory) - robot_ids = list(robot_indices) + trajectory_indices = self._trajectory_indices(trajectory) 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, - ) + generation = self._preview_animation_generation + robot_data = self._require_model() + self._set_preview_visibility(True) + baseline = np.array(self.get_joint_state(self._live_context).position, dtype=np.float64) + joint_positions_by_name = 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: + if self._preview_animation_generation != generation: 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) + positions = baseline.copy() + for trajectory_index, joint_name in trajectory_indices: + positions[joint_positions_by_name[joint_name]] = point.positions[ + trajectory_index + ] + self._set_preview_positions(self._plant_context, positions) self._publish_visualization() if frame_index < len(trajectory.points) - 1: next_time = trajectory.points[frame_index + 1].time_from_start @@ -1358,12 +1261,8 @@ def animate_trajectory( previous_time = next_time finally: with self._lock: - 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) + if self._preview_animation_generation == generation: + self._set_preview_visibility(False) def close(self) -> None: """Shut down the viz thread.""" diff --git a/dimos/manipulation/planning/world/roboplan_model.py b/dimos/manipulation/planning/world/roboplan_model.py index c7ea33ce84..03a7fd5cec 100644 --- a/dimos/manipulation/planning/world/roboplan_model.py +++ b/dimos/manipulation/planning/world/roboplan_model.py @@ -18,7 +18,6 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass, replace -from itertools import combinations from pathlib import Path from typing import Any, Protocol import xml.etree.ElementTree as ET @@ -29,13 +28,12 @@ from dimos.manipulation.planning.groups.models import PlanningGroup from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry from dimos.manipulation.planning.spec.config import RobotModelConfig -from dimos.manipulation.planning.spec.models import PlanningGroupID, RobotName +from dimos.manipulation.planning.spec.models import PlanningGroupID from dimos.manipulation.planning.utils.mesh_utils import prepare_urdf_for_drake from dimos.utils.transform_utils import pose_to_matrix ROBOPLAN_WORLD_FRAME = "dimos_world" -_MAX_COMPOSITE_GROUPS = 64 _ROOT_LINK = ROBOPLAN_WORLD_FRAME _ROOT_JOINT = "dimos_world_joint" _FREE_ROOTS = {"world", "map", _ROOT_LINK} @@ -50,6 +48,8 @@ "parent_frame_id", "child_frame_id", ) +_MODEL_KEY = "model" +_MODEL_NAME = "dimos_model" class _BuildRobot(Protocol): @@ -88,16 +88,15 @@ class RoboPlanModel: scene: Any groups: Mapping[frozenset[PlanningGroupID], RoboPlanGroup] - legacy_group_ids: Mapping[RobotName, PlanningGroupID] - native_joint_by_global: Mapping[str, str] - native_link_by_robot: Mapping[RobotName, Mapping[str, str]] + native_joints: Mapping[str, str] + native_links: Mapping[str, str] all_group: RoboPlanGroup - def native_joint(self, robot_name: RobotName, local_name: str) -> str: - return self.native_joint_by_global[f"{robot_name}/{local_name}"] + def native_joint(self, canonical_name: str) -> str: + return self.native_joints[canonical_name] - def native_link(self, robot_name: RobotName, local_name: str) -> str: - return self.native_link_by_robot[robot_name][local_name] + def native_link(self, canonical_name: str) -> str: + return self.native_links[canonical_name] @dataclass(frozen=True) @@ -111,18 +110,16 @@ class _NameMap: @dataclass(frozen=True) class _Composed: xml: str - maps: Mapping[RobotName, _NameMap] + maps: Mapping[str, _NameMap] adjacent_links: tuple[tuple[str, str], ...] def build_roboplan_model( - robots: Sequence[_BuildRobot], + robot: _BuildRobot, registry: PlanningGroupRegistry, scene_factory: _SceneFactory, ) -> RoboPlanModel: """Build one composite scene transactionally.""" - if not robots: - raise ValueError("RoboPlanWorld requires at least one robot") prepared = [ ( robot, @@ -135,18 +132,13 @@ def build_roboplan_model( ) ), ) - for robot in robots ] - composite = len(robots) > 1 - composed = _compose(prepared, composite) - groups, legacy_ids, all_group = _groups(robots, registry, composed.maps, composite) - model_name = "dimos_composite" if composite else robots[0].config.name - srdf = _srdf(model_name, robots, groups, composed) - package_paths = list( - dict.fromkeys(str(path) for robot in robots for path in robot.config.package_paths.values()) - ) + composed = _compose(prepared, False) + groups, all_group = _groups(robot, registry, composed.maps) + srdf = _srdf(_MODEL_NAME, [robot], groups, composed) + package_paths = [str(path) for path in robot.config.package_paths.values()] scene = scene_factory( - name=model_name, + name=_MODEL_NAME, urdf=composed.xml, srdf=srdf, package_paths=package_paths, @@ -154,17 +146,12 @@ def build_roboplan_model( groups = _validate_group_order(scene, groups) all_group = groups[frozenset(all_group.group_ids)] _apply_collision_exclusions(scene, srdf) - native_joint_by_global = { - f"{robot.config.name}/{local}": composed.maps[robot.config.name].joints[local] - for robot in robots - for local in robot.config.joint_names - } + mapping = composed.maps[_MODEL_KEY] return RoboPlanModel( scene=scene, groups=groups, - legacy_group_ids=legacy_ids, - native_joint_by_global=native_joint_by_global, - native_link_by_robot={name: mapping.links for name, mapping in composed.maps.items()}, + native_joints={name: mapping.joints[name] for name in robot.config.joint_names}, + native_links=mapping.links, all_group=all_group, ) @@ -172,18 +159,18 @@ def build_roboplan_model( def _compose(prepared: Sequence[tuple[_BuildRobot, Path]], composite: bool) -> _Composed: result = ET.Element( "robot", - {"name": "dimos_composite" if composite else prepared[0][0].config.name}, + {"name": "dimos_composite" if composite else _MODEL_NAME}, ) ET.SubElement(result, "link", {"name": _ROOT_LINK}) - maps: dict[RobotName, _NameMap] = {} + maps: dict[str, _NameMap] = {} used_names: set[str] = {_ROOT_LINK} for robot, path in prepared: config = robot.config root = ET.parse(path).getroot() if _tag(root.tag) != "robot": - raise ValueError(f"Prepared model for '{config.name}' is not a URDF robot") + raise ValueError("Prepared model is not a URDF robot") _add_missing_acceleration_limits(root) - mapping = _name_map(root, config.name, composite) + mapping = _name_map(root, _MODEL_KEY, composite) mapped_names = { value for table in (mapping.links, mapping.joints, mapping.materials, mapping.frames) @@ -194,13 +181,11 @@ def _compose(prepared: Sequence[tuple[_BuildRobot, Path]], composite: bool) -> _ raise ValueError(f"Duplicate composed model names: {sorted(duplicates)}") used_names.update(mapped_names) if config.base_link not in mapping.links: - raise ValueError(f"Robot '{config.name}' base link '{config.base_link}' is missing") + raise ValueError(f"Model base link '{config.base_link}' is missing") missing_joints = set(config.joint_names) - set(mapping.joints) if missing_joints: - raise ValueError( - f"Robot '{config.name}' configured joints are missing: {sorted(missing_joints)}" - ) - authored_root = _authored_root(root, config.base_link, config.name) + raise ValueError(f"Configured model joints are missing: {sorted(missing_joints)}") + authored_root = _authored_root(root, config.base_link, _MODEL_KEY) authored_parent = ( _joint_link(authored_root, "parent") if authored_root is not None else None ) @@ -216,15 +201,15 @@ def _compose(prepared: Sequence[tuple[_BuildRobot, Path]], composite: bool) -> _ copied = ET.fromstring(ET.tostring(element, encoding="unicode")) _rewrite(copied, mapping) result.append(copied) - attachment_name = _qualified(config.name, _ROOT_JOINT) if composite else _ROOT_JOINT + attachment_name = _qualified(_MODEL_KEY, _ROOT_JOINT) if composite else _ROOT_JOINT if attachment_name in used_names: - raise ValueError(f"Robot '{config.name}' collides with its synthetic attachment name") + raise ValueError("Model collides with its synthetic attachment name") used_names.add(attachment_name) joint = ET.SubElement(result, "joint", {"name": attachment_name, "type": "fixed"}) ET.SubElement(joint, "parent", {"link": _ROOT_LINK}) ET.SubElement(joint, "child", {"link": mapping.links[config.base_link]}) ET.SubElement(joint, "origin", _pose_attributes(config.base_pose)) - maps[config.name] = mapping + maps[_MODEL_KEY] = mapping adjacent: list[tuple[str, str]] = [] for joint in result: if _tag(joint.tag) != "joint": @@ -248,10 +233,10 @@ def _add_missing_acceleration_limits(root: ET.Element) -> None: limit.set("acceleration", str(_DEFAULT_ACCELERATION_LIMIT)) -def _name_map(root: ET.Element, robot_name: RobotName, prefix: bool) -> _NameMap: +def _name_map(root: ET.Element, model_key: str, prefix: bool) -> _NameMap: def names(tag: str) -> dict[str, str]: return { - name: _qualified(robot_name, name) if prefix else name + name: _qualified(model_key, name) if prefix else name for element in root.iter() if _tag(element.tag) == tag if (name := element.get("name")) @@ -260,7 +245,7 @@ def names(tag: str) -> dict[str, str]: return _NameMap(names("link"), names("joint"), names("material"), names("frame")) -def _authored_root(root: ET.Element, base_link: str, robot_name: RobotName) -> ET.Element | None: +def _authored_root(root: ET.Element, base_link: str, model_key: str) -> ET.Element | None: links = {element.get("name") for element in root if _tag(element.tag) == "link"} roots = {"world", "map", root.get("name", "")} & links | {"world", "map"} matches = [ @@ -271,11 +256,11 @@ def _authored_root(root: ET.Element, base_link: str, robot_name: RobotName) -> E if _joint_link(joint, "child") == base_link ] if len(matches) > 1: - raise ValueError(f"Robot '{robot_name}' has ambiguous world attachment") + raise ValueError(f"Model '{model_key}' has ambiguous world attachment") if not matches: return None if matches[0].get("type") != "fixed": - raise ValueError(f"Robot '{robot_name}' world attachment must be fixed") + raise ValueError(f"Model '{model_key}' world attachment must be fixed") return matches[0] @@ -308,82 +293,42 @@ def _rewrite(element: ET.Element, mapping: _NameMap) -> None: def _groups( - robots: Sequence[_BuildRobot], + robot: _BuildRobot, registry: PlanningGroupRegistry, - maps: Mapping[RobotName, _NameMap], - composite: bool, + maps: Mapping[str, _NameMap], ) -> tuple[ dict[frozenset[PlanningGroupID], RoboPlanGroup], - dict[RobotName, PlanningGroupID], RoboPlanGroup, ]: groups: dict[frozenset[PlanningGroupID], RoboPlanGroup] = {} - legacy_ids: dict[RobotName, PlanningGroupID] = {} - for robot in robots: - config = robot.config - group_id = f"{config.name}/__roboplan_legacy__" - legacy_ids[config.name] = group_id - legacy_group = RoboPlanGroup( - (group_id,), - f"_dimos_legacy__{_safe(config.name)}" if composite else config.name, - tuple(maps[config.name].joints[name] for name in config.joint_names), - tuple(config.joint_names), - ) - groups[frozenset(legacy_group.group_ids)] = legacy_group configured = registry.list() for group in configured: - layout = _group_layout((group,), maps, composite) + layout = _group_layout((group,), maps) groups[frozenset(layout.group_ids)] = layout - generated = 0 - for size in range(2, len(configured) + 1): - for selected in combinations(configured, size): - if len({group.robot_name for group in selected}) < 2: - continue - if len({name for group in selected for name in group.joint_names}) != sum( - len(group.joint_names) for group in selected - ): - continue - generated += 1 - if generated > _MAX_COMPOSITE_GROUPS: - raise ValueError( - f"RoboPlan composite planning groups exceed {_MAX_COMPOSITE_GROUPS}" - ) - layout = _group_layout(selected, maps, True) - groups[frozenset(layout.group_ids)] = layout all_id = "__dimos_all_configured__" + config = robot.config all_group = RoboPlanGroup( (all_id,), all_id, - tuple( - maps[robot.config.name].joints[name] - for robot in robots - for name in robot.config.joint_names - ), - tuple( - f"{robot.config.name}/{name}" for robot in robots for name in robot.config.joint_names - ), + tuple(maps[_MODEL_KEY].joints[name] for name in config.joint_names), + tuple(config.joint_names), ) groups[frozenset(all_group.group_ids)] = all_group names = [group.name for group in groups.values()] if len(names) != len(set(names)): raise ValueError("Generated RoboPlan planning-group names are not unique") - return groups, legacy_ids, all_group + return groups, all_group def _group_layout( selected: Sequence[PlanningGroup], - maps: Mapping[RobotName, _NameMap], - composite: bool, + maps: Mapping[str, _NameMap], ) -> RoboPlanGroup: ids = tuple(group.id for group in selected) return RoboPlanGroup( ids, - _composite_group_name(ids) if composite else selected[0].group_name, - tuple( - maps[group.robot_name].joints[local] - for group in selected - for local in group.local_joint_names - ), + _composite_group_name(ids) if len(selected) > 1 else selected[0].id, + tuple(maps[_MODEL_KEY].joints[name] for group in selected for name in group.joint_names), tuple(name for group in selected for name in group.joint_names), ) @@ -402,7 +347,7 @@ def _srdf( pairs = {tuple(sorted(pair)) for pair in composed.adjacent_links if _ROOT_LINK not in pair} for robot in robots: config = robot.config - mapping = composed.maps[config.name] + mapping = composed.maps[_MODEL_KEY] configured = list(config.collision_exclusion_pairs) if config.srdf_path is not None: configured.extend(_source_exclusions(config.srdf_path)) @@ -413,8 +358,7 @@ def _srdf( continue if not first_exists or not second_exists: raise ValueError( - f"Robot '{config.name}' collision exclusion references unknown links: " - f"{first} <-> {second}" + f"Model collision exclusion references unknown links: {first} <-> {second}" ) pairs.add(tuple(sorted((mapping.links[first], mapping.links[second])))) lines.extend( @@ -495,8 +439,8 @@ def _joint_link(joint: ET.Element | None, tag: str) -> str | None: ) -def _qualified(robot_name: RobotName, local_name: str) -> str: - return f"{_safe(robot_name)}__{_safe(local_name)}" +def _qualified(model_key: str, local_name: str) -> str: + return f"{_safe(model_key)}__{_safe(local_name)}" def _safe(value: str) -> str: diff --git a/dimos/manipulation/planning/world/roboplan_world.py b/dimos/manipulation/planning/world/roboplan_world.py index 9dc8d60c7b..afb78c0feb 100644 --- a/dimos/manipulation/planning/world/roboplan_world.py +++ b/dimos/manipulation/planning/world/roboplan_world.py @@ -21,6 +21,7 @@ from __future__ import annotations +from collections.abc import Sequence from contextlib import contextmanager from copy import deepcopy from dataclasses import dataclass, field, replace @@ -45,8 +46,6 @@ from dimos.manipulation.planning.spec.models import ( Obstacle, PlanningGroupID, - RobotName, - WorldRobotID, ) from dimos.manipulation.planning.spec.validation import ( validate_obstacle, @@ -73,7 +72,6 @@ @dataclass class _RoboPlanRobotData: - robot_id: WorldRobotID config: RobotModelConfig lower_limits: NDArray[np.float64] | None = None upper_limits: NDArray[np.float64] | None = None @@ -83,7 +81,7 @@ class _RoboPlanRobotData: class RoboPlanContext: """DimOS context wrapper for RoboPlan world state.""" - q_by_robot: dict[WorldRobotID, NDArray[np.float64]] = field(default_factory=dict) + q: NDArray[np.float64] = field(default_factory=lambda: np.empty(0, dtype=np.float64)) class RoboPlanWorld: @@ -96,57 +94,63 @@ def __init__(self, enable_viz: bool = False, **_: object) -> None: if enable_viz: logger.warning("RoboPlanWorld does not currently provide manipulation visualization") - self._robots: dict[WorldRobotID, _RoboPlanRobotData] = {} + self._model_data: _RoboPlanRobotData | None = None self._planning_groups = PlanningGroupRegistry() self._obstacles: dict[str, Obstacle] = {} - self._authoritative_robot_ids: set[WorldRobotID] = set() - self._robot_counter = 0 + self._has_authoritative_state = False self._finalized = False self._usable = True self._live_context = RoboPlanContext() self._state_lock = RLock() self._lock = RLock() - # Robot Management + # Model Management - def add_robot(self, config: RobotModelConfig) -> WorldRobotID: - """Register a robot for the scene built by :meth:`finalize`.""" + def load_model(self, config: RobotModelConfig) -> None: + """Register the logical robot model for :meth:`finalize`.""" if self._finalized: raise RuntimeError("Cannot add robot after world is finalized") - if any(data.config.name == config.name for data in self._robots.values()): - raise ValueError(f"Robot name '{config.name}' is already registered") + if self._model_data is not None: + raise ValueError("A model is already loaded") validate_robot_model_config(config) self._validate_planning_group_config(config) self._validate_robot_config(config) - self._robot_counter += 1 - robot_id = f"robot_{self._robot_counter}" - self._robots[robot_id] = _RoboPlanRobotData( - robot_id=robot_id, - config=config, - ) - self._planning_groups.add_robot(config) - self._live_context.q_by_robot[robot_id] = np.zeros( - len(config.joint_names), dtype=np.float64 - ) - return robot_id - - def get_robot_ids(self) -> list[WorldRobotID]: - """Get all robot IDs in the world.""" - return list(self._robots.keys()) + self._model_data = _RoboPlanRobotData(config=config) + self._planning_groups.add_model(config) + self._live_context.q = np.zeros(len(config.joint_names), dtype=np.float64) - def get_robot_config(self, robot_id: WorldRobotID) -> RobotModelConfig: - """Get robot configuration by ID.""" - return self._get_robot(robot_id).config + def get_model_config(self) -> RobotModelConfig: + """Get the logical robot model configuration.""" + return self._get_model_data().config - def get_joint_limits( - self, robot_id: WorldRobotID - ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: + def get_joint_limits(self) -> tuple[NDArray[np.float64], NDArray[np.float64]]: """Get joint limits in DimOS joint order.""" - robot = self._get_robot(robot_id) + robot = self._get_model_data() if robot.lower_limits is None or robot.upper_limits is None: raise RuntimeError("Joint limits are available after RoboPlan finalization") return robot.lower_limits.copy(), robot.upper_limits.copy() + def ordered_joint_positions(self, joint_state: JointState) -> NDArray[np.float64]: + """Return a canonical joint state in configured model order.""" + return self._joint_state_to_q(joint_state) + + def is_ready(self) -> bool: + """Return whether authoritative state is available for planning.""" + with self._state_lock: + return self._model_data is not None and self._has_authoritative_state + + def planning_group(self, group_ids: Sequence[PlanningGroupID]) -> RoboPlanGroup | None: + """Return the native group generated for a public group selection.""" + return self._require_model().groups.get(frozenset(group_ids)) + + def all_planning_group(self) -> RoboPlanGroup: + """Return the generated group spanning every canonical model joint.""" + return self._require_model().all_group + + def native_link_name(self, canonical_name: str) -> str: + """Return the backend link name for a canonical model link.""" + return self._require_model().native_link(canonical_name) + # Obstacle Management def add_obstacle(self, obstacle: Obstacle) -> str | None: @@ -237,18 +241,18 @@ def finalize(self) -> None: if self._finalized: return model = build_roboplan_model( - list(self._robots.values()), + self._get_model_data(), self._planning_groups, roboplan_core.Scene, ) self._model = model self._scene = model.scene try: - for robot in self._robots.values(): - group = self._legacy_group(robot.config.name) - lower, upper = self._extract_joint_limits(robot.config, group) - robot.lower_limits = lower - robot.upper_limits = upper + robot = self._get_model_data() + group = model.all_group + lower, upper = self._extract_joint_limits(robot.config, group) + robot.lower_limits = lower + robot.upper_limits = upper for obstacle_id, obstacle in self._obstacles.items(): self._add_obstacle_to_scene(obstacle, obstacle_id) except BaseException: @@ -274,50 +278,41 @@ def scratch_context(self) -> Generator[RoboPlanContext, None, None]: """Create a per-consumer context with independent collision scratch.""" with self._state_lock: self._require_finalized() - ctx = RoboPlanContext( - q_by_robot={ - robot_id: q.copy() for robot_id, q in self._live_context.q_by_robot.items() - } - ) + ctx = RoboPlanContext(q=self._live_context.q.copy()) yield ctx - def sync_from_joint_state(self, robot_id: WorldRobotID, joint_state: JointState) -> None: + def sync_from_joint_state(self, joint_state: JointState) -> None: """Sync live context from a driver joint-state message.""" if not self._finalized: return - q = self._joint_state_to_q(robot_id, joint_state) + q = self._joint_state_to_q(joint_state) with self._state_lock: - self._live_context.q_by_robot[robot_id] = q - self._authoritative_robot_ids.add(robot_id) + self._live_context.q = q + self._has_authoritative_state = True # State Operations - def set_joint_state( - self, ctx: RoboPlanContext, robot_id: WorldRobotID, joint_state: JointState - ) -> None: + def set_joint_state(self, ctx: RoboPlanContext, joint_state: JointState) -> None: """Set robot joint state in a context.""" self._require_finalized() - ctx.q_by_robot[robot_id] = self._joint_state_to_q(robot_id, joint_state) + ctx.q = self._joint_state_to_q(joint_state) - def get_joint_state(self, ctx: RoboPlanContext, robot_id: WorldRobotID) -> JointState: + def get_joint_state(self, ctx: RoboPlanContext) -> JointState: """Get robot joint state from a context.""" - robot = self._get_robot(robot_id) - q = ctx.q_by_robot.get(robot_id) - if q is None: + robot = self._get_model_data() + q = ctx.q + if not len(q): q = np.zeros(len(robot.config.joint_names), dtype=np.float64) return JointState(name=robot.config.joint_names, position=q.astype(float).tolist()) # Collision Checking - def is_collision_free(self, ctx: RoboPlanContext, robot_id: WorldRobotID) -> bool: + def is_collision_free(self, ctx: RoboPlanContext) -> bool: """Check if the robot configuration in a context is collision-free.""" self._require_finalized() - q = ctx.q_by_robot.get(robot_id) - if q is None: - raise KeyError(f"Robot '{robot_id}' not found in context") - return not self._has_collisions(ctx, robot_id, q) + return not self._has_collisions(ctx, ctx.q) - def get_min_distance(self, ctx: RoboPlanContext, robot_id: WorldRobotID) -> float: + def get_min_distance(self, ctx: RoboPlanContext) -> float: """Get minimum signed distance. RoboPlan signed-distance semantics are not verified yet, so do not return @@ -325,34 +320,33 @@ def get_min_distance(self, ctx: RoboPlanContext, robot_id: WorldRobotID) -> floa """ raise NotImplementedError("RoboPlanWorld.get_min_distance is not implemented") - def check_config_collision_free(self, robot_id: WorldRobotID, joint_state: JointState) -> bool: + def check_config_collision_free(self, joint_state: JointState) -> bool: """Check a joint state using a scratch collision context.""" with self.scratch_context() as ctx: - self.set_joint_state(ctx, robot_id, joint_state) - return self.is_collision_free(ctx, robot_id) + self.set_joint_state(ctx, joint_state) + return self.is_collision_free(ctx) def check_edge_collision_free( self, - robot_id: WorldRobotID, start: JointState, end: JointState, step_size: float = 0.05, ) -> bool: """Check if an interpolated edge is collision-free.""" self._require_finalized() - q_start = self._joint_state_to_q(robot_id, start) - q_end = self._joint_state_to_q(robot_id, end) + q_start = self._joint_state_to_q(start) + q_end = self._joint_state_to_q(end) with self.scratch_context() as ctx: - return not self._call_path_collision_checker(ctx, robot_id, q_start, q_end, step_size) + return not self._call_path_collision_checker(ctx, q_start, q_end, step_size) # Forward Kinematics - def get_ee_pose(self, ctx: RoboPlanContext, robot_id: WorldRobotID) -> PoseStamped: + def get_ee_pose(self, ctx: RoboPlanContext) -> PoseStamped: """Get end-effector pose if RoboPlan exposes FK.""" - robot = self._get_robot(robot_id) + robot = self._get_model_data() 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") + raise ValueError("Model 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: @@ -360,7 +354,7 @@ def get_group_ee_pose(self, ctx: RoboPlanContext, group_id: PlanningGroupID) -> 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) + mat = self.get_link_pose(ctx, group.tip_link) pose = matrix_to_pose(mat) return PoseStamped( frame_id="world", @@ -373,31 +367,26 @@ def get_group_ee_pose(self, ctx: RoboPlanContext, group_id: PlanningGroupID) -> ], ) - def get_link_pose( - self, ctx: RoboPlanContext, robot_id: WorldRobotID, link_name: str - ) -> NDArray[np.float64]: + def get_link_pose(self, ctx: RoboPlanContext, link_name: str) -> NDArray[np.float64]: """Get link pose as a 4x4 homogeneous transform.""" - q = ctx.q_by_robot.get(robot_id) - if q is None: - raise KeyError(f"Robot '{robot_id}' not found in context") + q = ctx.q scene = self._require_scene() - robot = self._get_robot(robot_id) with self._lock: - scene_q = self._full_scene_q(ctx, overlay=(robot_id, q)) + scene_q = self._full_scene_q(ctx, overlay=q) scene.setJointPositions(scene_q) result = scene.forwardKinematics( scene_q, - self._require_model().native_link(robot.config.name, link_name), + self._require_model().native_link(link_name), "", ) return np.asarray(result, dtype=np.float64) - def get_jacobian(self, ctx: RoboPlanContext, robot_id: WorldRobotID) -> NDArray[np.float64]: + def get_jacobian(self, ctx: RoboPlanContext) -> NDArray[np.float64]: """Get end-effector Jacobian if RoboPlan exposes a compatible API.""" - robot = self._get_robot(robot_id) + robot = self._get_model_data() 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") + raise ValueError("Model has no pose-targetable planning group") return self.get_group_jacobian(ctx, group_id) def get_group_jacobian( @@ -407,8 +396,6 @@ def get_group_jacobian( 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) scene = self._require_scene() model = self._require_model() with self._lock: @@ -416,7 +403,7 @@ def get_group_jacobian( scene.setJointPositions(scene_q) result = scene.computeFrameJacobian( scene_q, - model.native_link(robot.config.name, group.tip_link), + model.native_link(group.tip_link), True, ) arr = np.asarray(result, dtype=np.float64) @@ -424,9 +411,7 @@ def get_group_jacobian( raise ValueError(f"Unexpected RoboPlan Jacobian shape: {arr.shape}; expected 6 x n") scene_joint_order = list(scene.getJointNames()) if arr.shape[1] == len(scene_joint_order): - native_names = [ - model.native_joint(group.robot_name, name) for name in group.local_joint_names - ] + native_names = [model.native_joint(name) for name in group.joint_names] return arr[:, [scene_joint_order.index(name) for name in native_names]] raise ValueError( f"Unexpected RoboPlan Jacobian shape: {arr.shape}; cannot project group '{group_id}'" @@ -447,10 +432,27 @@ def _extract_joint_limits( lower = np.asarray(config.joint_limits_lower, dtype=np.float64) upper = np.asarray(config.joint_limits_upper, dtype=np.float64) else: - lower, upper = self._require_scene().getPositionLimitVectors(group.name, False) + scene = self._require_scene() + lower, upper = scene.getPositionLimitVectors(group.name, False) lower = np.asarray(lower, dtype=np.float64) upper = np.asarray(upper, dtype=np.float64) - by_name = dict(zip(group.public_names, zip(lower, upper, strict=True), strict=True)) + native_names = tuple(scene.getJointGroupInfo(group.name).joint_names) + native_to_canonical = { + native: canonical + for canonical, native in self._require_model().native_joints.items() + } + try: + canonical_names = tuple(native_to_canonical[name] for name in native_names) + except KeyError as exc: + raise ValueError( + f"RoboPlan joint-limit group contains unknown native joint '{exc.args[0]}'" + ) from exc + if set(canonical_names) != set(config.joint_names): + raise ValueError( + "RoboPlan joint-limit group does not match the composed model: " + f"{sorted(canonical_names)} != {sorted(config.joint_names)}" + ) + by_name = dict(zip(canonical_names, zip(lower, upper, strict=True), strict=True)) lower = np.asarray([by_name[name][0] for name in config.joint_names]) upper = np.asarray([by_name[name][1] for name in config.joint_names]) if len(lower) != len(config.joint_names) or len(upper) != len(config.joint_names): @@ -467,30 +469,18 @@ def _planning_group_from_id(self, group_id: PlanningGroupID) -> PlanningGroup: return self._planning_groups.get(group_id) def _primary_pose_group_id_for_config(self, config: RobotModelConfig) -> PlanningGroupID | None: - return self._planning_groups.primary_pose_group_id_for_robot(config.name) + return self._planning_groups.primary_pose_group_id() - 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 _get_model_data(self) -> _RoboPlanRobotData: + if self._model_data is None: + raise RuntimeError("Model is not loaded") + return self._model_data - 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) + def _joint_state_to_q(self, joint_state: JointState) -> NDArray[np.float64]: + robot = self._get_model_data() 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: @@ -523,61 +513,48 @@ def parametrization_model(self) -> Generator[RoboPlanModel, None, None]: def _full_scene_q( self, ctx: RoboPlanContext, - overlay: tuple[WorldRobotID, NDArray[np.float64]] | None = None, + overlay: NDArray[np.float64] | None = None, ) -> NDArray[np.float64]: scene = self._require_scene() group = self._require_model().all_group - positions = self._current_global_positions(ctx, overlay) + positions = self._current_positions(ctx, overlay) q = np.asarray([positions[name] for name in group.public_names], dtype=np.float64) return np.asarray(scene.toFullJointPositions(group.name, q), dtype=np.float64) - def _current_global_positions( + def _current_positions( self, ctx: RoboPlanContext | None = None, - overlay: tuple[WorldRobotID, NDArray[np.float64]] | None = None, + overlay: NDArray[np.float64] | None = None, ) -> dict[str, float]: context = ctx if ctx is not None else self._live_context - positions: dict[str, float] = {} - for robot_id, robot in self._robots.items(): - q = ( - overlay[1] - if overlay is not None and overlay[0] == robot_id - else context.q_by_robot.get(robot_id) - ) - if q is None or len(q) != len(robot.config.joint_names): - raise RuntimeError(f"Missing authoritative state for robot '{robot_id}'") - positions.update( - { - f"{robot.config.name}/{name}": float(value) - for name, value in zip(robot.config.joint_names, q, strict=True) - } - ) - return positions + robot = self._get_model_data() + q = overlay if overlay is not None else context.q + if len(q) != len(robot.config.joint_names): + raise RuntimeError("Missing authoritative model state") + return dict(zip(robot.config.joint_names, map(float, q), strict=True)) def _has_collisions( self, ctx: RoboPlanContext, - robot_id: WorldRobotID, q: NDArray[np.float64], ) -> bool: with self._lock: scene = self._require_scene() - scene_q = self._full_scene_q(ctx, overlay=(robot_id, q)) + scene_q = self._full_scene_q(ctx, overlay=q) scene.setJointPositions(scene_q) return bool(scene.hasCollisions(scene_q)) def _call_path_collision_checker( self, ctx: RoboPlanContext, - robot_id: WorldRobotID, q_start: NDArray[np.float64], q_end: NDArray[np.float64], step_size: float, ) -> bool: with self._lock: scene = self._require_scene() - scene_q_start = self._full_scene_q(ctx, overlay=(robot_id, q_start)) - scene_q_end = self._full_scene_q(ctx, overlay=(robot_id, q_end)) + scene_q_start = self._full_scene_q(ctx, overlay=q_start) + scene_q_end = self._full_scene_q(ctx, overlay=q_end) scene.setJointPositions(scene_q_start) return bool( roboplan_core.hasCollisionsAlongPath( @@ -647,12 +624,3 @@ def _require_dimensions(self, obstacle: Obstacle, n_dims: int) -> None: f"{obstacle.obstacle_type.name} obstacle requires {n_dims} dimensions, " f"got {len(obstacle.dimensions)}" ) - - def _legacy_group(self, robot_name: RobotName) -> RoboPlanGroup: - model = self._require_model() - group_id = model.legacy_group_ids[robot_name] - return model.groups[frozenset((group_id,))] - - def _is_ready(self) -> bool: - with self._state_lock: - return bool(self._robots) and self._authoritative_robot_ids == set(self._robots) diff --git a/dimos/manipulation/planning/world/test_drake_world_planning_groups.py b/dimos/manipulation/planning/world/test_drake_world_planning_groups.py index 63f7ce36cf..4785ebacd8 100644 --- a/dimos/manipulation/planning/world/test_drake_world_planning_groups.py +++ b/dimos/manipulation/planning/world/test_drake_world_planning_groups.py @@ -120,7 +120,6 @@ 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"], @@ -142,10 +141,9 @@ def test_drake_loads_canonical_slash_names_natively(tmp_path: Path) -> None: urdf = tmp_path / "canonical.urdf" _write_canonical_urdf(urdf) config = RobotModelConfig( - name="robot", model_path=urdf, joint_names=["left/j1"], - base_link="world", + base_link="left/base", planning_groups=[ PlanningGroupDefinition( name="left_arm", @@ -157,10 +155,10 @@ def test_drake_loads_canonical_slash_names_natively(tmp_path: Path) -> None: ) world = DrakeWorld() - robot_id = world.add_robot(config) + world.load_model(config) world.finalize() - assert world.get_robot_config(robot_id).joint_names == ["left/j1"] + assert world.get_model_config().joint_names == ["left/j1"] def test_drake_config_group_helpers_resolve_groups_without_drake_runtime(tmp_path: Path) -> None: @@ -168,12 +166,11 @@ def test_drake_config_group_helpers_resolve_groups_without_drake_runtime(tmp_pat _write_urdf(urdf) config = _config(urdf, [_arm_group("joint2", "joint1", name="wrist")]) - group = DrakeWorld._planning_group_from_config(config, "arm/wrist") + group = DrakeWorld._planning_group_from_config(config, "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 DrakeWorld._primary_pose_group_id_for_config(config) == "wrist" + assert group.id == "wrist" + assert group.joint_names == ("joint2", "joint1") assert group.tip_link == "tool0" @@ -196,7 +193,7 @@ def test_drake_config_group_helpers_validate_duplicate_and_ambiguous_groups( 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") + DrakeWorld._planning_group_from_config(ambiguous, "missing") @requires_drake @@ -206,7 +203,7 @@ def test_drake_obstacle_ids_are_world_owned_and_invalid_insertions_are_rejected( urdf = tmp_path / "robot.urdf" _write_urdf(urdf) world = DrakeWorld() - robot_id = world.add_robot(_config(urdf, [_arm_group("joint1", "joint2")])) + world.load_model(_config(urdf, [_arm_group("joint1", "joint2")])) obstacle = Obstacle( name="box", @@ -231,7 +228,7 @@ def test_drake_obstacle_ids_are_world_owned_and_invalid_insertions_are_rejected( assert world.add_obstacle(unnamed) is None assert world.add_obstacle(obstacle) == "box" joint_state = JointState(name=["joint1", "joint2"], position=[0.0, 0.0]) - assert world.check_config_collision_free(robot_id, joint_state) + assert world.check_config_collision_free(joint_state) original_geometry_id = world._obstacles["box"].geometry_id assert world.add_obstacle(obstacle) is None assert world.remove_obstacle("missing") is False @@ -241,7 +238,7 @@ def test_drake_obstacle_ids_are_world_owned_and_invalid_insertions_are_rejected( assert world.update_obstacle_pose("box", moved_pose) assert world._obstacles["box"].geometry_id != original_geometry_id assert world.get_obstacles()[0].pose.position.x == pytest.approx(0.0) - assert not world.check_config_collision_free(robot_id, joint_state) + assert not world.check_config_collision_free(joint_state) replacement = replace( obstacle, @@ -250,7 +247,7 @@ def test_drake_obstacle_ids_are_world_owned_and_invalid_insertions_are_rejected( color=(0.0, 1.0, 0.0, 1.0), ) assert world.update_obstacle(replacement) - assert world.check_config_collision_free(robot_id, joint_state) + assert world.check_config_collision_free(joint_state) replacement.dimensions = (9.0,) retrieved = world.get_obstacles()[0] retrieved.dimensions = (8.0,) @@ -264,7 +261,7 @@ def test_drake_obstacle_replacement_failure_invalidates_world( urdf = tmp_path / "robot.urdf" _write_urdf(urdf) world = DrakeWorld() - world.add_robot(_config(urdf, [_arm_group("joint1", "joint2")])) + world.load_model(_config(urdf, [_arm_group("joint1", "joint2")])) world.finalize() obstacle = Obstacle( name="box", @@ -286,23 +283,21 @@ def test_drake_obstacle_replacement_failure_invalidates_world( @requires_drake -def test_drake_group_fk_uses_tip_link_and_legacy_unique_pose_group(tmp_path: Path) -> None: +def test_drake_group_fk_uses_tip_link_and_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.load_model(_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]}) - ) + world.set_joint_state(ctx, 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) + group_pose = world.get_group_ee_pose(ctx, "arm") + default_pose = world.get_ee_pose(ctx) 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) + assert default_pose.position.x == pytest.approx(group_pose.position.x) + assert world.get_jacobian(ctx).shape == (6, 2) @requires_drake @@ -312,9 +307,8 @@ def test_drake_applies_config_base_pose_when_urdf_has_world_base_joint( urdf = tmp_path / "robot_with_world.urdf" _write_urdf_with_world_base_joint(urdf) world = DrakeWorld(enable_viz=False) - left_id = world.add_robot( + world.load_model( RobotModelConfig( - name="left_arm", model_path=urdf, base_pose=PoseStamped(position=[0, 0.5, 0], orientation=[0, 0, 0, 1]), joint_names=["joint1", "joint2"], @@ -322,25 +316,12 @@ def test_drake_applies_config_base_pose_when_urdf_has_world_base_joint( 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") + base_pose = world.get_link_pose(ctx, "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]) + assert base_pose[1, 3] == pytest.approx(0.5) @requires_drake @@ -348,7 +329,7 @@ def test_drake_group_jacobian_shape_and_group_local_order(tmp_path: Path) -> Non urdf = tmp_path / "robot.urdf" _write_urdf(urdf) world = DrakeWorld() - robot_id = world.add_robot( + world.load_model( _config( urdf, [ @@ -359,12 +340,10 @@ def test_drake_group_jacobian_shape_and_group_local_order(tmp_path: Path) -> Non ) world.finalize() ctx = world.get_live_context() - world.set_joint_state( - ctx, robot_id, JointState({"name": ["joint1", "joint2"], "position": [0.0, 0.0]}) - ) + world.set_joint_state(ctx, 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") + forward_jacobian = world.get_group_jacobian(ctx, "wrist_forward") + reverse_jacobian = world.get_group_jacobian(ctx, "wrist_reverse") assert reverse_jacobian.shape == (6, 2) np.testing.assert_allclose(reverse_jacobian[:, 0], forward_jacobian[:, 1]) @@ -372,17 +351,17 @@ def test_drake_group_jacobian_shape_and_group_local_order(tmp_path: Path) -> Non @requires_drake -def test_drake_legacy_wrappers_fail_at_call_time_for_no_or_ambiguous_pose(tmp_path: Path) -> None: +def test_drake_default_pose_methods_fail_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.load_model(_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) + no_pose.get_ee_pose(no_pose.get_live_context()) ambiguous = DrakeWorld() - ambiguous_id = ambiguous.add_robot( + ambiguous.load_model( _config( urdf, [ @@ -393,193 +372,138 @@ def test_drake_legacy_wrappers_fail_at_call_time_for_no_or_ambiguous_pose(tmp_pa ) ambiguous.finalize() with pytest.raises(ValueError, match="multiple pose"): - ambiguous.get_jacobian(ambiguous.get_live_context(), ambiguous_id) + ambiguous.get_jacobian(ambiguous.get_live_context()) @requires_drake -def test_drake_group_jacobian_rejects_non_controllable_group_joints(tmp_path: Path) -> None: +def test_drake_load_rejects_group_joints_outside_controllable_set(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") + with pytest.raises(ValueError, match="outside the controllable model set"): + world.load_model(_config(urdf, [_arm_group("joint1", "joint2")], joints=["joint1"])) @requires_drake -def test_drake_animate_trajectory_projects_all_robots_on_shared_ticks( - tmp_path: Path, monkeypatch +def test_drake_animate_trajectory_projects_selected_joints_on_shared_ticks( + tmp_path: Path, monkeypatch: pytest.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.load_model(_config(urdf, [_arm_group("joint1")])) 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, ...]] = [] + world.set_joint_state( + world.get_live_context(), + JointState(name=["joint1", "joint2"], position=[0.1, 0.2]), + ) + updates: list[list[float]] = [] + visibility: list[bool] = [] 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,)), + lambda _ctx, positions: updates.append(positions.tolist()), ) + monkeypatch.setattr(world, "_set_preview_visibility", visibility.append) 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]), - ] + + world.animate_trajectory(_trajectory(["joint1"], [1.0], [3.0]), duration=2.0) + + assert visibility == [True, False] + assert updates == [[1.0, 0.2], [3.0, 0.2]] assert sleeps == [2.0] @requires_drake def test_drake_animate_trajectory_validates_before_visibility_and_cleans_up( - tmp_path: Path, monkeypatch + tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: urdf = tmp_path / "robot.urdf" _write_urdf(urdf) world = DrakeWorld() - robot_id = world.add_robot(_config(urdf, [_arm_group("joint1")])) + world.load_model(_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,)), - ) + visibility: list[bool] = [] + monkeypatch.setattr(world, "_set_preview_visibility", visibility.append) 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]) + with pytest.raises(ValueError, match="unknown joints"): + world.animate_trajectory(_trajectory(["unknown/joint1"], [0.0], [1.0])) + assert visibility == [] 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,)] + world.animate_trajectory(_trajectory(["joint1"], [0.0], [1.0])) + assert visibility == [True, False] @requires_drake -def test_drake_cancel_preview_hides_ghosts_before_animation_resumes( - tmp_path: Path, monkeypatch +def test_drake_cancel_preview_hides_model_before_animation_resumes( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: urdf = tmp_path / "robot.urdf" _write_urdf(urdf) world = DrakeWorld() - robot_id = world.add_robot(_config(urdf, [_arm_group("joint1")])) + world.load_model(_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,)), - ) + visibility: list[bool] = [] + snapshots_during_sleep: list[list[bool]] = [] + monkeypatch.setattr(world, "_set_preview_visibility", visibility.append) 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)) + snapshots_during_sleep.append(list(visibility)) monkeypatch.setattr("time.sleep", cancel_during_sleep) - world.animate_trajectory(_trajectory(["arm/joint1"], [0.0], [1.0])) + world.animate_trajectory(_trajectory(["joint1"], [0.0], [1.0])) - assert hidden_snapshots_during_sleep == [[(robot_id,)]] - assert hidden[0] == (robot_id,) + assert snapshots_during_sleep == [[True, False]] @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 +def test_drake_animate_trajectory_cancellation_stops_stale_frames( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: urdf = tmp_path / "robot.urdf" _write_urdf(urdf) world = DrakeWorld() - robot_id = world.add_robot(_config(urdf, [_arm_group("joint1")])) + world.load_model(_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, ...]] = [] + visibility: list[bool] = [] 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, + lambda _ctx, positions: updates.append(positions.tolist()), ) + monkeypatch.setattr(world, "_set_preview_visibility", visibility.append) 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])) + world.animate_trajectory(_trajectory(["joint1"], [1.0], [2.0])) assert updates == [[1.0, 0.0]] - assert hidden[0] == (robot_id,) + assert visibility == [True, False] diff --git a/dimos/manipulation/test_execution_manager.py b/dimos/manipulation/test_execution_manager.py index be12f9053e..0ca01367d9 100644 --- a/dimos/manipulation/test_execution_manager.py +++ b/dimos/manipulation/test_execution_manager.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,9 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from __future__ import annotations +"""Tests for direct canonical trajectory execution.""" -from threading import Event, Thread from unittest.mock import MagicMock import pytest @@ -26,12 +25,7 @@ TrajectoryExecutionResult, TrajectoryExecutionStatus, ) -from dimos.manipulation.execution_manager import ( - ExecutionDispatchResult, - ExecutionOutcome, - ExecutionTarget, - PlanExecutionManager, -) +from dimos.manipulation.execution_manager import ExecutionOutcome, PlanExecutionManager from dimos.manipulation.planning.spec.enums import PlanningStatus from dimos.manipulation.planning.spec.models import GeneratedPlan from dimos.msgs.sensor_msgs.JointState import JointState @@ -39,47 +33,18 @@ from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint -def _target( - robot_name: str = "arm", - *, - model_joint_names: tuple[str, ...] = ("j1", "j2"), - coordinator_to_model: dict[str, str] | None = None, -) -> ExecutionTarget: - return ExecutionTarget.from_coordinator_mapping( - robot_name=robot_name, - model_joint_names=model_joint_names, - coordinator_to_model=coordinator_to_model or {}, - ) - - def _plan( - joint_names: tuple[str, ...] = ("arm/j1", "arm/j2"), - *, + names: tuple[str, ...] = ("left/j1", "right/j1"), status: PlanningStatus = PlanningStatus.SUCCESS, ) -> GeneratedPlan: - width = len(joint_names) - trajectory = JointTrajectory( - joint_names=list(joint_names), - points=[ - TrajectoryPoint( - positions=[0.0] * width, - velocities=[0.0] * width, - time_from_start=0.0, - ), - TrajectoryPoint( - positions=[1.0] * width, - velocities=[0.0] * width, - time_from_start=1.0, - ), - ], - ) + points = [ + TrajectoryPoint(positions=[0.0] * len(names), time_from_start=0.0), + TrajectoryPoint(positions=[1.0] * len(names), time_from_start=1.0), + ] return GeneratedPlan( - group_ids=("arm/manipulator",), - trajectory=trajectory, - path=[ - JointState(name=list(joint_names), position=[0.0] * width), - JointState(name=list(joint_names), position=[1.0] * width), - ], + group_ids=("both_arms",), + trajectory=JointTrajectory(joint_names=list(names), points=points), + path=[JointState(name=list(names), position=point.positions) for point in points], status=status, ) @@ -95,261 +60,64 @@ def _coordinator() -> MagicMock: return coordinator -def _manager( - *targets: ExecutionTarget, coordinator: MagicMock | None = None -) -> PlanExecutionManager: +def _manager(coordinator: MagicMock | None = None) -> PlanExecutionManager: return PlanExecutionManager( - targets=targets or (_target(),), - coordinator=coordinator if coordinator is not None else _coordinator(), - ) - - -def test_execution_target_inverts_coordinator_mapping() -> None: - target = _target(coordinator_to_model={"hardware/a": "j1", "hardware/b": "j2"}) - - assert dict(target.model_to_coordinator) == { - "j1": "hardware/a", - "j2": "hardware/b", - } - - -@pytest.mark.parametrize( - ("model_joint_names", "mapping", "message"), - [ - ((), {}, "invalid local model joints"), - (("j1", "j1"), {}, "duplicate local model joints"), - (("arm/j1",), {}, "invalid local model joints"), - (("j1",), {"hardware/a": "missing"}, "unknown model joint"), - ( - ("j1", "j2"), - {"hardware/a": "j1", "hardware/b": "j1"}, - "Multiple coordinator joints", - ), - ], -) -def test_execution_target_rejects_invalid_configuration( - model_joint_names: tuple[str, ...], - mapping: dict[str, str], - message: str, -) -> None: - with pytest.raises(ValueError, match=message): - _target(model_joint_names=model_joint_names, coordinator_to_model=mapping) - - -def test_manager_requires_unique_robot_targets() -> None: - with pytest.raises(ValueError, match="unique robot names"): - _manager(_target(), _target()) - - -def test_execute_maps_all_robots_into_one_trajectory() -> None: - coordinator = _coordinator() - manager = _manager( - _target( - "left", - model_joint_names=("j1",), - coordinator_to_model={"left_hw/j1": "j1"}, - ), - _target( - "right", - model_joint_names=("j1",), - coordinator_to_model={"right_hw/j1": "j1"}, - ), - coordinator=coordinator, + joint_names=("left/j1", "left/j2", "right/j1"), + coordinator=coordinator or _coordinator(), ) - plan = _plan(("left/j1", "right/j1")) - result = manager.execute(plan) - assert result.outcome is ExecutionOutcome.ACCEPTED - coordinator.execute_trajectory.assert_called_once() - trajectory = coordinator.execute_trajectory.call_args.args[0] - assert trajectory.joint_names == ["left_hw/j1", "right_hw/j1"] - assert trajectory.points == plan.trajectory.points - assert trajectory.timestamp == plan.trajectory.timestamp +def test_manager_rejects_empty_or_duplicate_model_joint_names() -> None: + with pytest.raises(ValueError, match="non-empty and unique"): + PlanExecutionManager(joint_names=(), coordinator=_coordinator()) + with pytest.raises(ValueError, match="non-empty and unique"): + PlanExecutionManager(joint_names=("j1", "j1"), coordinator=_coordinator()) -def test_execute_preserves_single_robot_subset() -> None: +def test_execute_forwards_same_canonical_trajectory_object_unchanged() -> None: coordinator = _coordinator() - manager = _manager( - _target("left", model_joint_names=("j1", "j2")), - _target("right", model_joint_names=("j1",)), - coordinator=coordinator, - ) - - result = manager.execute(_plan(("left/j2",))) - + plan = _plan() + result = _manager(coordinator).execute(plan) assert result.accepted - trajectory = coordinator.execute_trajectory.call_args.args[0] - assert trajectory.joint_names == ["j2"] - assert trajectory.points[0].positions == [0.0] + assert coordinator.execute_trajectory.call_args.args[0] is plan.trajectory @pytest.mark.parametrize( ("plan", "message"), [ (_plan(status=PlanningStatus.NO_SOLUTION), "status is not successful"), - (_plan(("not-global",)), "not globally named"), - (_plan(("unknown/j1",)), "unknown execution robot"), - (_plan(("arm/missing",)), "is not configured"), + (_plan(("unknown",)), "unknown joints"), ], ) -def test_execute_rejects_unmappable_plan_before_rpc( - plan: GeneratedPlan, - message: str, -) -> None: +def test_execute_rejects_invalid_plan_before_rpc(plan: GeneratedPlan, message: str) -> None: coordinator = _coordinator() - manager = _manager(coordinator=coordinator) - - result = manager.execute(plan) - + result = _manager(coordinator).execute(plan) assert result.outcome is ExecutionOutcome.REJECTED assert message in result.message coordinator.execute_trajectory.assert_not_called() -def test_execute_rejects_cross_robot_mapping_collision() -> None: +def test_execute_preserves_coordinator_rejection() -> None: coordinator = _coordinator() - manager = _manager( - _target( - "left", - model_joint_names=("j1",), - coordinator_to_model={"shared/j1": "j1"}, - ), - _target( - "right", - model_joint_names=("j1",), - coordinator_to_model={"shared/j1": "j1"}, - ), - coordinator=coordinator, + rejection = TrajectoryExecutionResult( + TrajectoryExecutionStatus.INVALID_TRAJECTORY, "specific rejection" ) - - result = manager.execute(_plan(("left/j1", "right/j1"))) - + coordinator.execute_trajectory.return_value = rejection + result = _manager(coordinator).execute(_plan()) assert result.outcome is ExecutionOutcome.REJECTED - assert "duplicate coordinator joints" in result.message - coordinator.execute_trajectory.assert_not_called() - - -@pytest.mark.parametrize( - "status", - [ - TrajectoryExecutionStatus.NO_TRAJECTORY_TASK, - TrajectoryExecutionStatus.INVALID_TRAJECTORY, - TrajectoryExecutionStatus.START_STATE_UNAVAILABLE, - TrajectoryExecutionStatus.START_STATE_MISMATCH, - ], -) -def test_execute_preserves_coordinator_rejection(status: TrajectoryExecutionStatus) -> None: - coordinator = _coordinator() - coordinator_result = TrajectoryExecutionResult(status, "specific rejection") - coordinator.execute_trajectory.return_value = coordinator_result - manager = _manager(coordinator=coordinator) - - result = manager.execute(_plan()) - - assert result.outcome is ExecutionOutcome.REJECTED - assert result.message == "specific rejection" - assert result.coordinator_result is coordinator_result + assert result.coordinator_result is rejection def test_execute_rpc_failure_is_uncertain() -> None: coordinator = _coordinator() coordinator.execute_trajectory.side_effect = TimeoutError("timed out") - - result = _manager(coordinator=coordinator).execute(_plan()) - + result = _manager(coordinator).execute(_plan()) assert result.outcome is ExecutionOutcome.UNCERTAIN assert "timed out" in result.message -@pytest.mark.parametrize( - ("status", "safe", "cancelled"), - [ - ( - TrajectoryCancellationStatus.CANCELLED, - True, - True, - ), - ( - TrajectoryCancellationStatus.ALREADY_STOPPED, - True, - False, - ), - ( - TrajectoryCancellationStatus.NO_TRAJECTORY_TASK, - True, - False, - ), - ], -) -def test_cancel_preserves_coordinator_semantics( - status: TrajectoryCancellationStatus, - safe: bool, - cancelled: bool, -) -> None: +def test_cancel_forwards_to_coordinator() -> None: coordinator = _coordinator() - coordinator_result = TrajectoryCancellationResult(status, "cancel result") - coordinator.cancel_trajectory.return_value = coordinator_result - - result = _manager(coordinator=coordinator).cancel() - - assert result is coordinator_result - assert result.safe is safe - assert result.cancelled is cancelled - - -def test_cancel_rpc_failure_is_uncertain() -> None: - coordinator = _coordinator() - coordinator.cancel_trajectory.side_effect = TimeoutError("timed out") - - result = _manager(coordinator=coordinator).cancel() - - assert result.status is TrajectoryCancellationStatus.UNCERTAIN - assert not result.safe - assert not result.cancelled - assert "timed out" in result.message - - -def test_cancel_waits_for_in_flight_execute_then_cancels() -> None: - coordinator = _coordinator() - execute_started = Event() - release_execute = Event() - - def execute_trajectory(_trajectory: JointTrajectory) -> TrajectoryExecutionResult: - execute_started.set() - if not release_execute.wait(timeout=1.0): - raise TimeoutError("test did not release execute RPC") - return TrajectoryExecutionResult(TrajectoryExecutionStatus.ACCEPTED) - - coordinator.execute_trajectory.side_effect = execute_trajectory - manager = _manager(coordinator=coordinator) - execute_results: list[ExecutionDispatchResult] = [] - cancel_results: list[TrajectoryCancellationResult] = [] - - execute_thread = Thread(target=lambda: execute_results.append(manager.execute(_plan()))) - cancel_thread = Thread(target=lambda: cancel_results.append(manager.cancel())) - execute_thread.start() - execute_was_started = execute_started.wait(timeout=1.0) - cancel_was_started = False - cancel_called_before_release = False - try: - if execute_was_started: - cancel_thread.start() - cancel_was_started = True - cancel_called_before_release = coordinator.cancel_trajectory.called - finally: - release_execute.set() - execute_thread.join(timeout=1.0) - if cancel_was_started: - cancel_thread.join(timeout=1.0) - - assert execute_was_started - assert not execute_thread.is_alive() - assert cancel_was_started - assert not cancel_thread.is_alive() - assert not cancel_called_before_release - assert len(execute_results) == 1 - assert execute_results[0].outcome is ExecutionOutcome.ACCEPTED - assert len(cancel_results) == 1 - assert cancel_results[0].status is TrajectoryCancellationStatus.ALREADY_STOPPED + result = _manager(coordinator).cancel() + assert result.status is TrajectoryCancellationStatus.ALREADY_STOPPED coordinator.cancel_trajectory.assert_called_once_with() diff --git a/dimos/manipulation/test_generated_plan_materialization.py b/dimos/manipulation/test_generated_plan_materialization.py index 00e0943add..4c32e5b2ea 100644 --- a/dimos/manipulation/test_generated_plan_materialization.py +++ b/dimos/manipulation/test_generated_plan_materialization.py @@ -70,20 +70,28 @@ def generate(self, waypoints: list[list[float]]) -> JointTrajectory: ) -def _robot(name: str, joints: list[str], velocity: float, acceleration: float) -> RobotModelConfig: +def _model() -> RobotModelConfig: return RobotModelConfig( - name=name, model_path=Path("/robot.urdf"), base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), - joint_names=joints, + joint_names=["left/a", "left/b", "right/c"], base_link="base", planning_groups=[ PlanningGroupDefinition( - name="group", joint_names=tuple(reversed(joints)), base_link="base", tip_link="tip" - ) + name="left_arm", + joint_names=("left/b", "left/a"), + base_link="base", + tip_link="left_tip", + ), + PlanningGroupDefinition( + name="right_arm", + joint_names=("right/c",), + base_link="base", + tip_link="right_tip", + ), ], - max_velocity=velocity, - max_acceleration=acceleration, + max_velocity=3.0, + max_acceleration=4.0, ) @@ -95,21 +103,12 @@ def _module(monkeypatch: pytest.MonkeyPatch, module_factory): "simple_parametrizer.JointTrajectoryGenerator", RecordingGenerator, ) - left = _robot("left", ["a", "b"], 1.0, 2.0) - right = _robot("right", ["c"], 3.0, 4.0) + model = _model() module = module_factory() - module._robots = { - "left": ("left_id", left), - "right": ("right_id", right), - } module._world_monitor = MagicMock() module._world_monitor.world = MagicMock() - module._world_monitor.world.get_robot_ids.return_value = ["left_id", "right_id"] - module._world_monitor.world.get_robot_config.side_effect = { - "left_id": left, - "right_id": right, - }.__getitem__ - module._world_monitor.planning_groups = PlanningGroupRegistry([left, right]) + module._world_monitor.world.get_model_config.return_value = model + module._world_monitor.planning_groups = PlanningGroupRegistry([model]) module._planner = MagicMock() module._trajectory_parametrizer = SimpleTrapezoidParametrizer( SimpleTrapezoidParametrizationConfig() @@ -134,9 +133,9 @@ def test_materializes_once_with_reordered_groups_heterogeneous_limits_and_distin status=PlanningStatus.SUCCESS, path=path ) - assert module._plan_selected_path(("left/group", "right/group"), path[0], path[-1], 1) + assert module._plan_selected_path(("left_arm", "right_arm"), path[0], path[-1], 1) 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]) + assert RecordingGenerator.limits == ([3.0, 3.0, 3.0], [4.0, 4.0, 4.0]) assert module._last_plan is not None assert module._last_plan.path is not module._last_plan.trajectory.points assert module._last_plan.trajectory.joint_names == names @@ -153,7 +152,7 @@ def test_cartesian_plan_preserves_planner_timestamps_and_velocities(monkeypatch, JointState(name=names, position=[0.0, 0.0], velocity=[0.0, 0.0]), JointState(name=names, position=[0.2, 0.1], velocity=[0.4, 0.2]), ] - module._world_monitor.current_global_joint_state.return_value = start + module._world_monitor.current_model_joint_state.return_value = start module._planner.plan_cartesian_path.return_value = PlanningResult( status=PlanningStatus.SUCCESS, path=path, @@ -162,7 +161,7 @@ def test_cartesian_plan_preserves_planner_timestamps_and_velocities(monkeypatch, success = module.plan_cartesian_targets( { - "left/group": ( + "left_arm": ( Transform.identity(), Transform(translation=Vector3(0.01, 0.0, 0.0)), ) @@ -196,7 +195,7 @@ def test_zero_generation_after_caching_for_status_and_completion(monkeypatch, mo module._planner.plan_selected_joint_path.return_value = PlanningResult( status=PlanningStatus.SUCCESS, path=path ) - assert module._plan_selected_path(("left/group",), path[0], path[-1], 1) + assert module._plan_selected_path(("left_arm",), path[0], path[-1], 1) RecordingGenerator.calls = [] module._wait_for_trajectory_completion(timeout=0.0) diff --git a/dimos/manipulation/test_manipulation_module.py b/dimos/manipulation/test_manipulation_module.py index 8ff5b204f1..09e9c62e8f 100644 --- a/dimos/manipulation/test_manipulation_module.py +++ b/dimos/manipulation/test_manipulation_module.py @@ -67,7 +67,6 @@ def _get_xarm7_config() -> RobotModelConfig: """Create XArm7 robot config for testing.""" desc_path = get_data("xarm_description") return RobotModelConfig( - name="test_arm", 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"], @@ -85,15 +84,6 @@ def _get_xarm7_config() -> RobotModelConfig: auto_convert_meshes=True, max_velocity=1.0, max_acceleration=2.0, - joint_name_mapping={ - "arm/joint1": "joint1", - "arm/joint2": "joint2", - "arm/joint3": "joint3", - "arm/joint4": "joint4", - "arm/joint5": "joint5", - "arm/joint6": "joint6", - "arm/joint7": "joint7", - }, ) @@ -107,13 +97,13 @@ def joint_state_zeros(): """Create a JointState message with zeros for XArm7.""" return JointState( name=[ - "arm/joint1", - "arm/joint2", - "arm/joint3", - "arm/joint4", - "arm/joint5", - "arm/joint6", - "arm/joint7", + "joint1", + "joint2", + "joint3", + "joint4", + "joint5", + "joint6", + "joint7", ], position=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], velocity=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], @@ -132,7 +122,7 @@ def module(xarm7_config): TrajectoryCancellationStatus.ALREADY_STOPPED ) mod = ManipulationModule( - robots=[xarm7_config], + model=xarm7_config, planning_timeout=10.0, world_backend="drake", planner=RRTConnectPlannerConfig(), @@ -157,7 +147,7 @@ def test_module_initialization(self, module): assert module._world_monitor is not None assert module._planner is not None assert module._kinematics is not None - assert "test_arm" in module._robots + assert module.get_model_config() == module.config.model def test_joint_state_sync(self, module, joint_state_zeros): """Test joint state synchronization to Drake world.""" @@ -189,20 +179,18 @@ def test_plan_to_joints(self, module, joint_state_zeros): assert module._last_plan is not None assert len(module._last_plan.trajectory.points) > 1 assert module._last_plan.trajectory.duration > 0 - assert module._last_plan.group_ids == ("test_arm/manipulator",) + assert module._last_plan.group_ids == ("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)} - ) + success = module.plan_to_joint_targets({"manipulator": JointState(position=[0.05] * 7)}) assert success is True assert module._state == ManipulationState.COMPLETED assert module._last_plan is not None - assert module._last_plan.group_ids == ("test_arm/manipulator",) + assert module._last_plan.group_ids == ("manipulator",) assert module.has_planned_path() is True assert module._last_plan.trajectory.points @@ -222,21 +210,18 @@ def test_add_and_remove_obstacle(self, module, joint_state_zeros): removed = module.remove_obstacle(obstacle_id) assert removed is True - def test_robot_info(self, module): - """Test getting robot information.""" - info = module.get_robot_info() + def test_model_info(self, module): + """Test getting model information.""" + info = module.get_model_info() - assert info is not None - assert info["name"] == "test_arm" assert len(info["joint_names"]) == 7 assert info["end_effector_link"] == "link7" - assert info["has_joint_name_mapping"] is True groups = info["planning_groups"] assert len(groups) == 1 - assert groups[0].id == "test_arm/manipulator" + assert groups[0].id == "manipulator" all_groups = module.list_planning_groups() - assert [group.id for group in all_groups] == ["test_arm/manipulator"] + assert [group.id for group in all_groups] == ["manipulator"] def test_ee_pose(self, module, joint_state_zeros): """Test getting end-effector pose.""" @@ -249,19 +234,18 @@ 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_trajectory_uses_canonical_names(self, module, joint_state_zeros): + """Test that execution preserves canonical model joint names.""" module._on_joint_state(joint_state_zeros) success = module.plan_to_joints(JointState(position=[0.05] * 7)) assert success is True assert module._last_plan is not None - robot_config = module._robots["test_arm"][1] assert module.execute() is True trajectory = module._control_coordinator.execute_trajectory.call_args.args[0] - assert trajectory.joint_names == list(robot_config.joint_name_mapping.keys()) + assert trajectory.joint_names == module.config.model.joint_names @pytest.mark.skipif(not _drake_available(), reason="Drake not installed") @@ -286,9 +270,7 @@ def test_execute_with_mock_coordinator(self, module, joint_state_zeros): trajectory = module._control_coordinator.execute_trajectory.call_args.args[0] assert len(trajectory.points) > 1 - # Joint names should be translated - robot_config = module._robots["test_arm"][1] - assert trajectory.joint_names == list(robot_config.joint_name_mapping.keys()) + assert trajectory.joint_names == module.config.model.joint_names def test_execute_rejected_by_coordinator(self, module, joint_state_zeros): """Test handling of coordinator rejection.""" diff --git a/dimos/manipulation/test_manipulation_monitor_preview.py b/dimos/manipulation/test_manipulation_monitor_preview.py index ae69ddb25f..0f2f796073 100644 --- a/dimos/manipulation/test_manipulation_monitor_preview.py +++ b/dimos/manipulation/test_manipulation_monitor_preview.py @@ -38,33 +38,26 @@ @pytest.fixture -def robot_config_with_mapping() -> RobotModelConfig: - """Create a robot config with joint name mapping.""" +def canonical_model_config() -> RobotModelConfig: + """Create a model whose joint names match coordinator-facing names.""" return RobotModelConfig( - name="left_arm", model_path=Path("/path/to/robot.urdf"), base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), - joint_names=["joint1", "joint2", "joint3"], + joint_names=["left/joint1", "left/joint2", "left/joint3"], base_link="link_base", planning_groups=[ PlanningGroupDefinition( name="manipulator", - joint_names=("joint1", "joint2", "joint3"), + joint_names=("left/joint1", "left/joint2", "left/joint3"), base_link="link_base", tip_link="link_tcp", ) ], - joint_name_mapping={ - "left/joint1": "joint1", - "left/joint2": "joint2", - "left/joint3": "joint3", - }, ) 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"], @@ -82,9 +75,7 @@ def _install_generated_plan( config: RobotModelConfig, *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)} + """Install a canonical generated plan and monitor state.""" module._world_monitor = MagicMock() module._world_monitor.planning_groups = PlanningGroupRegistry([config]) module._world_monitor.get_current_joint_state.return_value = JointState( @@ -93,7 +84,7 @@ def _install_generated_plan( ) module._last_plan = GeneratedPlan( trajectory=JointTrajectory( - joint_names=global_joint_names, + joint_names=config.joint_names, points=[ TrajectoryPoint( time_from_start=float(index), @@ -103,11 +94,11 @@ def _install_generated_plan( for index, point in enumerate(points) ], ), - group_ids=(f"{config.name}/manipulator",), + group_ids=("manipulator",), status=PlanningStatus.SUCCESS, path=[ JointState( - name=global_joint_names, + name=config.joint_names, position=list(point), ) for point in points @@ -115,17 +106,11 @@ def _install_generated_plan( ) -def _make_module_with_monitor( - module_factory, - *configs: RobotModelConfig, -) -> ManipulationModule: - """Create a ManipulationModule with a mocked world monitor and robots configured.""" +def _make_module_with_monitor(module_factory) -> ManipulationModule: + """Create a ManipulationModule with a mocked world monitor.""" module = module_factory() module._world_monitor = MagicMock() - module._init_joints = {} - for config in configs: - robot_id = f"robot_{config.name}" - module._robots[config.name] = (robot_id, config) + module._init_joints = None return module @@ -193,51 +178,47 @@ def close(self) -> None: class TestOnJointState: - """Test _on_joint_state routing, splitting, and init capture.""" + """Test complete canonical model-state routing and init capture.""" - def test_routes_positions_to_monitor(self, robot_config_with_mapping, module_factory): - """Joint positions from aggregated message are routed to the correct monitor.""" - module = _make_module_with_monitor(module_factory, robot_config_with_mapping) + def test_routes_canonical_state_in_model_order(self, canonical_model_config, module_factory): + module = _make_module_with_monitor(module_factory) + module.config.model = canonical_model_config msg = JointState( - name=["left/joint1", "left/joint2", "left/joint3"], - position=[0.1, 0.2, 0.3], - velocity=[1.0, 2.0, 3.0], + name=["left/joint3", "unrelated", "left/joint1", "left/joint2"], + position=[0.3, 9.0, 0.1, 0.2], + velocity=[3.0, 9.0, 1.0, 2.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" + state = module._world_monitor.on_joint_state.call_args.args[0] + assert state.name == canonical_model_config.joint_names + assert state.position == [0.1, 0.2, 0.3] + assert state.velocity == [1.0, 2.0, 3.0] - def test_skips_robot_with_missing_joints(self, robot_config_with_mapping, module_factory): - """Robots whose joints are absent from the message are skipped.""" - module = _make_module_with_monitor(module_factory, robot_config_with_mapping) + def test_skips_incomplete_model_state(self, canonical_model_config, module_factory): + module = _make_module_with_monitor(module_factory) + module.config.model = canonical_model_config - # Message has none of left_arm's joints msg = JointState( - name=["right/joint1", "right/joint2"], + name=["left/joint1", "left/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, module_factory): - """First joint state is stored as init joints; subsequent calls don't overwrite.""" - module = _make_module_with_monitor(module_factory, robot_config_with_mapping) + def test_captures_init_joints_once(self, canonical_model_config, module_factory): + module = _make_module_with_monitor(module_factory) + module.config.model = canonical_model_config 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] + assert module._init_joints is not None + assert module._init_joints.position == [0.1, 0.2, 0.3] # Second call should NOT overwrite second_msg = JointState( @@ -245,61 +226,12 @@ def test_captures_init_joints_on_first_call(self, robot_config_with_mapping, mod 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, module_factory): - """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"], - base_link="base", - planning_groups=[ - PlanningGroupDefinition( - name="manipulator", joint_names=("j1", "j2"), base_link="base", tip_link="ee" - ) - ], - joint_name_mapping={"left/j1": "j1", "left/j2": "j2"}, - ) - right_config = RobotModelConfig( - name="right", - model_path=Path("/path/to/robot.urdf"), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), - joint_names=["j1", "j2"], - base_link="base", - planning_groups=[ - PlanningGroupDefinition( - name="manipulator", joint_names=("j1", "j2"), base_link="base", tip_link="ee" - ) - ], - joint_name_mapping={"right/j1": "j1", "right/j2": "j2"}, - ) - module = _make_module_with_monitor(module_factory, left_config, right_config) + assert module._init_joints.position == [0.1, 0.2, 0.3] - 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, module_factory): + def test_no_monitor_returns_early(self, canonical_model_config, module_factory): """When world_monitor is None, _on_joint_state returns without error.""" module = module_factory() - module._robots = {"left_arm": ("id", robot_config_with_mapping)} + module.config.model = canonical_model_config module._world_monitor = None # Should not raise @@ -315,11 +247,14 @@ def test_visualization_routing_and_stop_all_monitors(self): viz = FakeVisualization() monitor = _make_world_monitor_with_viz(viz) state_monitor = MagicMock() + state_monitor.get_current_positions.return_value = None obstacle_monitor = MagicMock() - monitor._state_monitors = {"robot": state_monitor} + monitor._state_monitor = state_monitor monitor._obstacle_monitor = obstacle_monitor monitor._viz_thread = MagicMock() monitor._viz_thread.is_alive.return_value = False + monitor._world.get_live_context.return_value = object() + monitor._world.get_joint_state.return_value = JointState() assert monitor.get_visualization_url() == "123" monitor.update_visualization_state() @@ -327,7 +262,7 @@ def test_visualization_routing_and_stop_all_monitors(self): path = _make_path([1.0], [2.0], [3.0]) plan = GeneratedPlan( trajectory=JointTrajectory(), - group_ids=("robot/group",), + group_ids=("manipulator",), status=PlanningStatus.SUCCESS, path=path, ) @@ -357,7 +292,7 @@ def test_visualization_none_is_noop(self): class TestManipulationPreview: def test_clear_planned_path_invalidates_before_dismissing_preview(self, module_factory): module = module_factory() - plan = GeneratedPlan(trajectory=JointTrajectory(), group_ids=("arm/manipulator",), path=[]) + plan = GeneratedPlan(trajectory=JointTrajectory(), group_ids=("manipulator",), path=[]) module._last_plan = plan module._world_monitor = MagicMock() plan_during_dismissal: list[GeneratedPlan | None] = [] @@ -374,7 +309,7 @@ def test_clear_planned_path_invalidates_before_dismissing_preview(self, module_f def test_clear_planned_path_clears_without_a_world_monitor(self, module_factory): module = module_factory() module._last_plan = GeneratedPlan( - trajectory=JointTrajectory(), group_ids=("arm/manipulator",), path=[] + trajectory=JointTrajectory(), group_ids=("manipulator",), path=[] ) assert module.clear_planned_path() is True @@ -383,13 +318,13 @@ def test_clear_planned_path_clears_without_a_world_monitor(self, module_factory) def test_dismiss_preview_noop_without_monitor(self, module_factory): module = module_factory() - module._dismiss_preview(["arm/manipulator"]) + module._dismiss_preview(["manipulator"]) def test_dismiss_preview_routes_to_monitor(self, module_factory): module = module_factory() module._world_monitor = MagicMock() - module._dismiss_preview(["arm/manipulator"]) + module._dismiss_preview(["manipulator"]) module._world_monitor.cancel_preview_animation.assert_called_once_with() @@ -403,37 +338,3 @@ def test_preview_routes_one_complete_plan_with_default_duration(self, module_fac module._world_monitor.animate_trajectory.assert_called_once_with( module._last_plan.trajectory, None ) - - def test_preview_robot_name_validates_affectedness_without_trimming(self, module_factory): - module = module_factory() - left = _one_joint_config("left") - right = _one_joint_config("right") - module._robots = { - "left": ("left_id", left), - "right": ("right_id", right), - } - module._world_monitor = MagicMock() - module._world_monitor.planning_groups = PlanningGroupRegistry([left, right]) - module._last_plan = GeneratedPlan( - trajectory=JointTrajectory(), - group_ids=("left/manipulator", "right/manipulator"), - status=PlanningStatus.SUCCESS, - path=[ - JointState(name=["left/j0", "right/j0"], position=[0.0, 0.0]), - JointState(name=["left/j0", "right/j0"], position=[1.0, 1.0]), - ], - ) - - assert module.preview_plan(duration=2.5, robot_name="left") is True - - module._world_monitor.animate_trajectory.assert_called_once_with( - module._last_plan.trajectory, 2.5 - ) - - def test_preview_rejects_unaffected_compatibility_robot(self, module_factory): - module = module_factory() - config = _one_joint_config() - _install_generated_plan(module, config, [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 c84841c9df..976b74d3ba 100644 --- a/dimos/manipulation/test_manipulation_unit.py +++ b/dimos/manipulation/test_manipulation_unit.py @@ -17,7 +17,7 @@ from __future__ import annotations from pathlib import Path -from unittest.mock import ANY, MagicMock +from unittest.mock import ANY, MagicMock, call import pytest from pytest_mock import MockerFixture @@ -75,7 +75,6 @@ def _control_coordinator( def robot_config(): """Create a robot config for testing.""" return RobotModelConfig( - name="test_arm", model_path=Path("/path/to/robot.urdf"), base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), joint_names=["joint1", "joint2", "joint3"], @@ -93,43 +92,31 @@ def robot_config(): ) -@pytest.fixture -def robot_config_with_mapping(): - """Create a robot config with joint name mapping (dual-arm scenario).""" +def _one_joint_config() -> RobotModelConfig: return RobotModelConfig( - name="left_arm", - model_path=Path("/path/to/robot.urdf"), + model_path=Path("/path"), base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), - joint_names=["joint1", "joint2", "joint3"], - base_link="link_base", + joint_names=["j0"], + base_link="base_link", planning_groups=[ PlanningGroupDefinition( - name="manipulator", - joint_names=("joint1", "joint2", "joint3"), - base_link="link_base", - tip_link="link_tcp", + name="manipulator", joint_names=("j0",), base_link="base_link", tip_link="ee" ) ], - joint_name_mapping={ - "left/joint1": "joint1", - "left/joint2": "joint2", - "left/joint3": "joint3", - }, ) -def _one_joint_config(name: str = "arm") -> RobotModelConfig: +def _bimanual_config() -> RobotModelConfig: return RobotModelConfig( - name=name, - model_path=Path("/path"), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), - joint_names=["j0"], - base_link="base_link", + model_path=Path("/path/to/bimanual.urdf"), + joint_names=["left/j1", "right/j1"], + base_link="base", planning_groups=[ - PlanningGroupDefinition( - name="manipulator", joint_names=("j0",), base_link="base_link", tip_link="ee" - ) + PlanningGroupDefinition("left_arm", ("left/j1",), "base", "left/tool"), + PlanningGroupDefinition("right_arm", ("right/j1",), "base", "right/tool"), + PlanningGroupDefinition("both_arms", ("left/j1", "right/j1"), "base"), ], + home_joints=[0.0, 0.0], ) @@ -138,22 +125,17 @@ def _install_generated_plan( config: RobotModelConfig, *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)} + """Install a canonical generated plan and current model state.""" + module.config.model = config module._world_monitor = MagicMock() module._world_monitor.planning_groups = PlanningGroupRegistry([config]) - module._world_monitor.get_current_joint_state.return_value = JointState( + module._world_monitor.current_model_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], - ) module._last_plan = GeneratedPlan( trajectory=JointTrajectory( - joint_names=global_joint_names, + joint_names=config.joint_names, points=[ TrajectoryPoint( time_from_start=float(index), @@ -163,11 +145,11 @@ def _install_generated_plan( for index, point in enumerate(points) ], ), - group_ids=(f"{config.name}/manipulator",), + group_ids=("manipulator",), status=PlanningStatus.SUCCESS, path=[ JointState( - name=global_joint_names, + name=config.joint_names, position=list(point), ) for point in points @@ -310,7 +292,7 @@ def test_cancel_hides_active_plan_preview(self, module_factory): module = module_factory() module._state = ManipulationState.EXECUTING module._last_plan = GeneratedPlan( - trajectory=JointTrajectory(), group_ids=("arm/manipulator",), path=[] + trajectory=JointTrajectory(), group_ids=("manipulator",), path=[] ) module._world_monitor = MagicMock() @@ -363,7 +345,7 @@ def test_motion_speed_applies_to_future_plans_only(self, module_factory): module = module_factory() accepted = GeneratedPlan( trajectory=JointTrajectory(), - group_ids=("arm/manipulator",), + group_ids=("manipulator",), path=[JointState(name=["arm/j0"], position=[0.0])], ) module._last_plan = accepted @@ -381,60 +363,11 @@ def test_motion_speed_rejects_invalid_values(self, module_factory, invalid: floa assert module.get_motion_speed() == pytest.approx(0.5) assert "motion speed scale" in module.get_error() - def test_begin_planning_state_checks(self, robot_config, module_factory): - """_begin_planning only allowed from IDLE or COMPLETED.""" - module = module_factory() - module._world_monitor = MagicMock() - module._robots = {"test_arm": ("robot_id", robot_config)} - - # From IDLE - OK - module._state = ManipulationState.IDLE - assert module._begin_planning() == ("test_arm", "robot_id") - assert module._state == ManipulationState.PLANNING - - # 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 TestRobotSelection: - """Test robot selection logic.""" - - def test_single_robot_default(self, robot_config, module_factory): - """Single robot is used by default.""" - module = module_factory() - module._robots = {"arm": ("id", robot_config)} - - result = module._get_robot() - assert result is not None - assert result[0] == "arm" - - def test_multiple_robots_require_name(self, robot_config, module_factory): - """Multiple robots require explicit name.""" - module = module_factory() - module._robots = { - "left": ("id1", robot_config), - "right": ("id2", robot_config), - } - - # No name - fails - assert module._get_robot() is None - - # With name - works - result = module._get_robot("left") - assert result is not None - assert result[0] == "left" - class PlanningInitializationHarness: def __init__(self, mocker: MockerFixture) -> None: self.mock_world = MagicMock() self.mock_world_monitor = MagicMock(spec=WorldMonitor) - self.mock_world_monitor.add_robot.return_value = "robot_id" self.planning_specs = MagicMock( world_monitor=self.mock_world_monitor, planner=MagicMock(), @@ -460,17 +393,18 @@ def planning_initialization(mocker: MockerFixture) -> PlanningInitializationHarn class TestPlanningInitialization: """Test planning backend configuration wiring.""" - def test_default_kinematics_config_uses_pink(self) -> None: + def test_default_kinematics_config_uses_pink(self, robot_config) -> None: """Pink IK is the default solver for manipulation modules.""" - config = ManipulationModuleConfig() + config = ManipulationModuleConfig(model=robot_config) assert isinstance(config.kinematics, PinkKinematicsConfig) def test_start_eagerly_initializes_planning_and_execution( self, mocker: MockerFixture, + robot_config, ) -> None: - module = ManipulationModule() + module = ManipulationModule(model=robot_config) module.coordinator_joint_state = None initialize_planning = mocker.patch.object(module, "_initialize_planning") initialize_execution = mocker.patch.object(module, "_initialize_execution") @@ -489,54 +423,30 @@ def test_kinematics_config_is_passed_to_factory( module = module_factory() kinematics = PinkKinematicsConfig(max_iterations=100, dt=0.02) module.config = ManipulationModuleConfig( - robots=[robot_config], + model=robot_config, kinematics=kinematics, ) - module._initialize_planning() + ManipulationModule._initialize_planning(module) planning_initialization.mock_planning_specs.assert_called_once_with( world=planning_initialization.mock_world, world_backend="roboplan", planner=module.config.planner, - kinematics_name=None, kinematics=kinematics, trajectory_parametrization=ANY, ) - def test_legacy_kinematics_name_still_selects_backend( - self, - robot_config, - planning_initialization: PlanningInitializationHarness, - module_factory, - ): - """The old kinematics_name field remains a compatibility shim.""" - module = module_factory() - module.config = ManipulationModuleConfig( - robots=[robot_config], - kinematics_name="pink", - ) - - module._initialize_planning() - - planning_initialization.mock_planning_specs.assert_called_once_with( - world=planning_initialization.mock_world, - world_backend="roboplan", - planner=module.config.planner, - kinematics_name="pink", - kinematics=module.config.kinematics, - trajectory_parametrization=ANY, - ) - - def test_nested_kinematics_config_parses_cli_override_shape(self) -> None: + def test_nested_kinematics_config_parses_cli_override_shape(self, robot_config) -> None: """Pydantic parses the nested shape used by dynamic CLI overrides.""" config = ManipulationModuleConfig( + model=robot_config, kinematics={ "backend": "pink", "max_iterations": "100", "dt": "0.02", "posture_cost": "0.0", - } + }, ) assert isinstance(config.kinematics, PinkKinematicsConfig) @@ -547,16 +457,17 @@ def test_nested_kinematics_config_parses_cli_override_shape(self) -> None: def test_solve_ik_rpc_calls_configured_backend(self, robot_config, module_factory): """solve_ik returns the backend IKResult without path planning.""" module = module_factory() - module._robots = {"test_arm": ("robot_id", robot_config)} + module.config.model = robot_config module._world_monitor = MagicMock() module._world_monitor.world = MagicMock() + module._world_monitor.world.get_model_config.return_value = robot_config module._world_monitor.planning_groups = PlanningGroupRegistry([robot_config]) current = JointState(name=robot_config.joint_names, position=[0.0, 0.0, 0.0]) current_global = JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + name=["joint1", "joint2", "joint3"], position=[0.0, 0.0, 0.0], ) - module._world_monitor.current_global_joint_state.return_value = current_global + module._world_monitor.current_model_joint_state.return_value = current_global expected = IKResult( status=IKStatus.SUCCESS, joint_state=JointState(name=robot_config.joint_names, position=[0.1, 0.2, 0.3]), @@ -581,17 +492,18 @@ def test_solve_ik_rpc_calls_configured_backend(self, robot_config, module_factor assert kwargs["seed"].position == current.position assert kwargs["check_collision"] is True [(group, target_pose)] = kwargs["pose_targets"].items() - assert group.id == "test_arm/manipulator" + assert group.id == "manipulator" assert target_pose.frame_id == "world" assert target_pose.position.x == 0.45 def test_solve_ik_rpc_returns_failure_without_joint_state(self, robot_config, module_factory): """solve_ik reports a failed IKResult when no seed state is available.""" module = module_factory() - module._robots = {"test_arm": ("robot_id", robot_config)} + module.config.model = robot_config + module.config.model = robot_config module._world_monitor = MagicMock() module._world_monitor.planning_groups = PlanningGroupRegistry([robot_config]) - module._world_monitor.current_global_joint_state.return_value = JointState( + module._world_monitor.current_model_joint_state.return_value = JointState( name=[], position=[] ) module._kinematics = MagicMock() @@ -609,11 +521,12 @@ def test_solve_ik_rpc_accepts_explicit_seed_without_current_state( ): """solve_ik succeeds with an explicit seed when no current state is available.""" module = module_factory() - module._robots = {"test_arm": ("robot_id", robot_config)} + module.config.model = robot_config module._world_monitor = MagicMock() module._world_monitor.world = MagicMock() + module._world_monitor.world.get_model_config.return_value = robot_config module._world_monitor.planning_groups = PlanningGroupRegistry([robot_config]) - module._world_monitor.get_current_joint_state.return_value = None + module._world_monitor.current_model_joint_state.return_value = None 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._kinematics = MagicMock() @@ -625,64 +538,47 @@ def test_solve_ik_rpc_accepts_explicit_seed_without_current_state( assert result is expected _, kwargs = module._kinematics.solve_pose_targets.call_args assert kwargs["seed"] is explicit_seed - module._world_monitor.current_global_joint_state.assert_not_called() + module._world_monitor.current_model_joint_state.assert_not_called() class TestPlanningGroupApis: """Test explicit planning-group API behavior.""" - def test_list_planning_groups_and_robot_info_include_groups(self, robot_config, module_factory): + def test_list_planning_groups_and_model_info_include_groups(self, robot_config, module_factory): module = module_factory() + module.config.model = robot_config registry = PlanningGroupRegistry([robot_config]) - module._robots = {"test_arm": ("robot_id", robot_config)} module._world_monitor = MagicMock() module._world_monitor.planning_groups = registry - module._init_joints = {} + module._init_joints = None groups = module.list_planning_groups() - info = module.get_robot_info() + info = module.get_model_info() - assert [group.id for group in groups] == ["test_arm/manipulator"] - assert info is not None + assert [group.id for group in groups] == ["manipulator"] assert info["planning_groups"] == groups assert info["end_effector_link"] == "link_tcp" - assert info["has_joint_name_mapping"] is False - def test_plan_to_joint_targets_stores_generated_plan_and_legacy_caches( - self, robot_config, module_factory - ): + def test_plan_to_joint_targets_stores_generated_plan(self, robot_config, module_factory): module = module_factory() + module.config.model = robot_config registry = PlanningGroupRegistry([robot_config]) - module._robots = {"test_arm": ("robot_id", robot_config)} _enable_simple_parametrization(module) module._world_monitor = MagicMock() module._world_monitor.world = MagicMock() - module._world_monitor.world.get_robot_ids.return_value = ["robot_id"] - module._world_monitor.world.get_robot_config.return_value = robot_config + module._world_monitor.world.get_model_config.return_value = robot_config module._world_monitor.planning_groups = registry - module._world_monitor.current_global_joint_state.return_value = JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.0, 0.0, 0.0], - ) - module._world_monitor.get_current_joint_state.return_value = JointState( - name=robot_config.joint_names, - position=[0.0, 0.0, 0.0], - ) - module._world_monitor.current_global_joint_state.return_value = JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.0, 0.0, 0.0], - ) - module._world_monitor.current_global_joint_state.return_value = JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + module._world_monitor.current_model_joint_state.return_value = JointState( + name=["joint1", "joint2", "joint3"], position=[0.0, 0.0, 0.0], ) result_path = [ JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + name=["joint1", "joint2", "joint3"], position=[0.0, 0.0, 0.0], ), JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + name=["joint1", "joint2", "joint3"], position=[0.1, 0.2, 0.3], ), ] @@ -698,7 +594,7 @@ def test_plan_to_joint_targets_stores_generated_plan_and_legacy_caches( success = module.plan_to_joint_targets( { - "test_arm/manipulator": JointState( + "manipulator": JointState( name=robot_config.joint_names, position=[0.1, 0.2, 0.3], ) @@ -707,21 +603,21 @@ def test_plan_to_joint_targets_stores_generated_plan_and_legacy_caches( assert success is True assert module._last_plan is not None - assert module._last_plan.group_ids == ("test_arm/manipulator",) + assert module._last_plan.group_ids == ("manipulator",) assert module._last_plan.path == result_path assert module._last_plan.trajectory.points[-1].positions == [0.1, 0.2, 0.3] module._planner.plan_selected_joint_path.assert_called_once() _, kwargs = module._planner.plan_selected_joint_path.call_args - assert kwargs["selection"].group_ids == ("test_arm/manipulator",) + assert kwargs["selection"].group_ids == ("manipulator",) assert kwargs["goal"].name == [ - "test_arm/joint1", - "test_arm/joint2", - "test_arm/joint3", + "joint1", + "joint2", + "joint3", ] success = module.plan_to_joint_targets( { - "test_arm/manipulator": JointState( + "manipulator": JointState( name=robot_config.joint_names, position=[0.1, 0.2, 0.3], ) @@ -735,28 +631,19 @@ def test_plan_to_pose_targets_uses_group_ik_and_selected_path( self, robot_config, module_factory ): module = module_factory() + module.config.model = robot_config registry = PlanningGroupRegistry([robot_config]) - module._robots = {"test_arm": ("robot_id", robot_config)} _enable_simple_parametrization(module) module._world_monitor = MagicMock() module._world_monitor.world = MagicMock() - module._world_monitor.world.get_robot_ids.return_value = ["robot_id"] - module._world_monitor.world.get_robot_config.return_value = robot_config + module._world_monitor.world.get_model_config.return_value = robot_config module._world_monitor.planning_groups = registry - module._world_monitor.current_global_joint_state.return_value = JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], - position=[0.0, 0.0, 0.0], - ) - module._world_monitor.get_current_joint_state.return_value = JointState( - name=robot_config.joint_names, - position=[0.0, 0.0, 0.0], - ) - module._world_monitor.current_global_joint_state.return_value = JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + module._world_monitor.current_model_joint_state.return_value = JointState( + name=["joint1", "joint2", "joint3"], position=[0.0, 0.0, 0.0], ) ik_goal = JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + name=["joint1", "joint2", "joint3"], position=[0.1, 0.2, 0.3], ) module._kinematics = MagicMock() @@ -769,7 +656,7 @@ def test_plan_to_pose_targets_uses_group_ik_and_selected_path( status=PlanningStatus.SUCCESS, path=[ JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + name=["joint1", "joint2", "joint3"], position=[0.0, 0.0, 0.0], ), ik_goal, @@ -777,50 +664,50 @@ def test_plan_to_pose_targets_uses_group_ik_and_selected_path( ) pose = Pose(position=Vector3(x=0.45, y=0.0, z=0.25), orientation=Quaternion()) - success = module.plan_to_pose_targets({"test_arm/manipulator": pose}) + success = module.plan_to_pose_targets({"manipulator": pose}) assert success is True module._kinematics.solve_pose_targets.assert_called_once() _, ik_kwargs = module._kinematics.solve_pose_targets.call_args target_groups = list(ik_kwargs["pose_targets"].keys()) - assert [group.id for group in target_groups] == ["test_arm/manipulator"] + assert [group.id for group in target_groups] == ["manipulator"] target_pose = ik_kwargs["pose_targets"][target_groups[0]] assert target_pose.position.x == 0.45 assert ik_kwargs["seed"].name == [ - "test_arm/joint1", - "test_arm/joint2", - "test_arm/joint3", + "joint1", + "joint2", + "joint3", ] _, planner_kwargs = module._planner.plan_selected_joint_path.call_args assert planner_kwargs["goal"] is ik_goal def test_failed_plan_materialization_clears_generated_plan(self, robot_config, module_factory): module = module_factory() + module.config.model = robot_config registry = PlanningGroupRegistry([robot_config]) - module._robots = {"test_arm": ("robot_id", robot_config)} module._world_monitor = MagicMock() module._world_monitor.world = MagicMock() module._world_monitor.planning_groups = registry - module._world_monitor.current_global_joint_state.return_value = JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + module._world_monitor.current_model_joint_state.return_value = JointState( + name=["joint1", "joint2", "joint3"], position=[0.0, 0.0, 0.0], ) - module._world_monitor.get_current_joint_state.return_value = None + module._world_monitor.current_model_joint_state.return_value = None module._last_plan = GeneratedPlan( trajectory=_generated_plan_trajectory( - ["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + ["joint1", "joint2", "joint3"], [0.0, 0.0, 0.0], [0.1, 0.2, 0.3], ), - group_ids=("test_arm/manipulator",), + group_ids=("manipulator",), status=PlanningStatus.SUCCESS, path=[ JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + name=["joint1", "joint2", "joint3"], position=[0.0, 0.0, 0.0], ), JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + name=["joint1", "joint2", "joint3"], position=[0.2, 0.2, 0.2], ), ], @@ -830,14 +717,14 @@ def test_failed_plan_materialization_clears_generated_plan(self, robot_config, m status=PlanningStatus.SUCCESS, path=[ JointState( - name=["test_arm/joint1", "test_arm/joint2", "test_arm/joint3"], + name=["joint1", "joint2", "joint3"], position=[0.1, 0.2, 0.3], ) ], ) success = module.plan_to_joint_targets( - {"test_arm/manipulator": JointState(position=[0.1, 0.2, 0.3])} + {"manipulator": JointState(position=[0.1, 0.2, 0.3])} ) assert success is False @@ -845,41 +732,32 @@ def test_failed_plan_materialization_clears_generated_plan(self, robot_config, m assert module._last_plan is None assert module.has_planned_path() is False - def test_execute_plan_dispatches_selected_subsets_once_with_shared_clock_and_mapping( - self, module_factory - ): - left = RobotModelConfig( - name="left", + def test_execute_plan_forwards_one_canonical_trajectory_unchanged(self, module_factory): + model = RobotModelConfig( model_path=Path("/path"), base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), - joint_names=["j0", "j1"], + joint_names=["left/j1", "right/k0"], planning_groups=[ PlanningGroupDefinition( - name="wrist", joint_names=("j1",), base_link="base", tip_link="ee" - ) - ], - joint_name_mapping={"left_coord_j1": "j1"}, - ) - right = RobotModelConfig( - name="right", - model_path=Path("/path"), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), - joint_names=["k0", "k1"], - planning_groups=[ + name="left_arm", + joint_names=("left/j1",), + base_link="base", + tip_link="left/ee", + ), PlanningGroupDefinition( - name="elbow", joint_names=("k0",), base_link="base", tip_link="ee" - ) + name="right_arm", + joint_names=("right/k0",), + base_link="base", + tip_link="right/ee", + ), ], ) module = module_factory() - module._robots = { - "left": ("left_id", left), - "right": ("right_id", right), - } + module.config.model = model module._world_monitor = MagicMock() - module._world_monitor.planning_groups = PlanningGroupRegistry([left, right]) + module._world_monitor.planning_groups = PlanningGroupRegistry([model]) module._last_plan = GeneratedPlan( - group_ids=("left/wrist", "right/elbow"), + group_ids=("left_arm", "right_arm"), status=PlanningStatus.SUCCESS, path=[ JointState(name=["left/j1", "right/k0"], position=[0.0, 1.0]), @@ -904,11 +782,7 @@ def test_execute_plan_dispatches_selected_subsets_once_with_shared_clock_and_map mock_coordinator = _control_coordinator() module._control_coordinator = mock_coordinator module._initialize_execution() - module._world_monitor.get_current_joint_state.side_effect = lambda robot_id: { - "left_id": JointState(name=["j0", "j1"], position=[9.0, 0.0]), - "right_id": JointState(name=["k0", "k1"], position=[1.0, 9.0]), - }[robot_id] - module._world_monitor.current_global_joint_state.return_value = JointState( + module._world_monitor.current_model_joint_state.return_value = JointState( name=["left/j1", "right/k0"], position=[0.0, 1.0] ) @@ -916,16 +790,76 @@ def test_execute_plan_dispatches_selected_subsets_once_with_shared_clock_and_map mock_coordinator.execute_trajectory.assert_called_once() payload = mock_coordinator.execute_trajectory.call_args.args[0] - assert payload.joint_names == ["left_coord_j1", "k0"] + assert payload is module._last_plan.trajectory + assert payload.joint_names == ["left/j1", "right/k0"] assert [point.time_from_start for point in payload.points] == [0.0, 2.5] assert [point.positions for point in payload.points] == [[0.0, 1.0], [0.5, 1.5]] assert [point.velocities for point in payload.points] == [[0.0, 0.0], [0.2, 0.4]] + def test_go_init_uses_selected_group_for_safe_waypoint_fk( + self, module_factory, mocker: MockerFixture + ) -> None: + model = _bimanual_config() + module = module_factory() + module.config.model = model + module._init_joints = JointState(name=model.joint_names, position=[0.1, -0.1]) + module._world_monitor = MagicMock(spec=WorldMonitor) + module._world_monitor.planning_groups = PlanningGroupRegistry([model]) + module._world_monitor.get_group_ee_pose.return_value = PoseStamped( + position=Vector3(0.4, 0.2, 0.3), orientation=Quaternion(0.0, 0.0, 0.0, 1.0) + ) + module._world_monitor.get_ee_pose.side_effect = AssertionError( + "bimanual init must not use model-wide end-effector lookup" + ) + mocker.patch.object(module, "_lift_if_low", return_value=MagicMock(is_success=lambda: True)) + mocker.patch.object( + module, "_preview_execute_wait", return_value=MagicMock(is_success=lambda: True) + ) + plan_to_pose = mocker.patch.object(module, "plan_to_pose", return_value=True) + plan_to_joints = mocker.patch.object(module, "plan_to_joints", return_value=True) + + result = module.go_init("left_arm") + + assert result.is_success() + module._world_monitor.get_group_ee_pose.assert_called_once_with( + "left_arm", module._init_joints + ) + plan_to_pose.assert_called_once() + plan_to_joints.assert_called_once() + + def test_tf_loop_publishes_every_pose_group_for_bimanual_model( + self, module_factory, mocker: MockerFixture + ) -> None: + model = _bimanual_config() + module = module_factory() + module.config.model = model + module._world_monitor = MagicMock(spec=WorldMonitor) + module._world_monitor.planning_groups = PlanningGroupRegistry([model]) + module._world_monitor.get_group_ee_pose.side_effect = [ + PoseStamped(position=Vector3(0.4, 0.2, 0.3)), + PoseStamped(position=Vector3(0.4, -0.2, 0.3)), + ] + publish = mocker.patch.object(module.tf, "publish") + + def stop_after_first_iteration(_period: float) -> bool: + module._tf_stop_event.set() + return True + + mocker.patch.object(module._tf_stop_event, "wait", side_effect=stop_after_first_iteration) + module._tf_stop_event.clear() + + module._tf_publish_loop() + + assert module._world_monitor.get_group_ee_pose.call_args_list == [ + call("left_arm"), + call("right_arm"), + ] + publish.assert_called_once() + def test_pose_wrappers_fail_safely_without_unique_pose_group( self, robot_config, module_factory ): 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, @@ -939,7 +873,7 @@ def test_pose_wrappers_fail_safely_without_unique_pose_group( ], ) module = module_factory() - module._robots = {"test_arm": ("robot_id", no_pose_config)} + module.config.model = no_pose_config module._world_monitor = MagicMock() module._world_monitor.planning_groups = PlanningGroupRegistry([no_pose_config]) module._world_monitor.get_ee_pose.side_effect = ValueError("no pose group") @@ -951,13 +885,12 @@ def test_pose_wrappers_fail_safely_without_unique_pose_group( assert module.plan_to_pose(pose) is False result = module.inverse_kinematics_single(pose) assert result.status == IKStatus.NO_SOLUTION - assert "no pose-targetable planning group" in result.message + assert "no unique pose-targetable planning group" in result.message def test_pose_wrappers_fail_safely_with_multiple_pose_groups( self, robot_config, module_factory ): 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, @@ -978,7 +911,7 @@ def test_pose_wrappers_fail_safely_with_multiple_pose_groups( ], ) module = module_factory() - module._robots = {"test_arm": ("robot_id", multi_pose_config)} + module.config.model = multi_pose_config module._world_monitor = MagicMock() module._world_monitor.planning_groups = PlanningGroupRegistry([multi_pose_config]) module._world_monitor.get_ee_pose.side_effect = ValueError("multiple pose groups") @@ -995,12 +928,12 @@ def test_pose_wrappers_fail_safely_with_multiple_pose_groups( def test_solve_ik_preserves_backend_failure_detail(self, robot_config, module_factory): """IK diagnostics include the backend's human-readable failure message.""" module = module_factory() - module._robots = {"test_arm": ("robot_id", robot_config)} + module.config.model = robot_config module._world_monitor = MagicMock() module._world_monitor.world = MagicMock() module._world_monitor.planning_groups = PlanningGroupRegistry([robot_config]) - module._world_monitor.current_global_joint_state.return_value = JointState( - name=[f"test_arm/{name}" for name in robot_config.joint_names], + module._world_monitor.current_model_joint_state.return_value = JointState( + name=robot_config.joint_names, position=[0.0, 0.0, 0.0], ) module._kinematics = MagicMock() @@ -1019,21 +952,18 @@ class TestPlanningDiagnostics: def test_planner_failure_preserves_backend_detail(self, robot_config, module_factory): """Planning diagnostics include the backend message.""" module = module_factory() + module.config.model = robot_config module._world_monitor = MagicMock() module._world_monitor.planning_groups = PlanningGroupRegistry([robot_config]) - module._world_monitor.current_global_joint_state.return_value = JointState( - name=[f"test_arm/{name}" for name in robot_config.joint_names], + module._world_monitor.current_model_joint_state.return_value = JointState( + name=robot_config.joint_names, position=[0.0, 0.0, 0.0], ) module._planner = MagicMock() module._planner.plan_selected_joint_path.return_value = PlanningResult( status=PlanningStatus.TIMEOUT, message="planner timed out" ) - - module._robots = {"test_arm": ("robot_id", robot_config)} - assert not module.plan_to_joints( - JointState(position=[1.0, 1.0, 1.0]), robot_name="test_arm" - ) + assert not module.plan_to_joints(JointState(position=[1.0, 1.0, 1.0])) assert module.get_error() == "Planning failed: TIMEOUT: planner timed out" assert module._state == ManipulationState.FAULT @@ -1045,23 +975,6 @@ class TestExecute: def test_execute_requires_trajectory(self, robot_config, module_factory): """Execute fails without planned trajectory.""" module = module_factory() - module._robots = {"test_arm": ("id", robot_config)} assert module.execute() is False assert module._state == ManipulationState.IDLE - - -class TestRobotModelConfigMapping: - """Test RobotModelConfig joint name mapping helpers.""" - - def test_bidirectional_mapping(self, robot_config_with_mapping): - """Test URDF <-> coordinator name translation.""" - config = robot_config_with_mapping - - # Coordinator -> URDF - assert config.get_urdf_joint_name("left/joint1") == "joint1" - assert config.get_urdf_joint_name("unknown") == "unknown" - - # URDF -> Coordinator - assert config.get_coordinator_joint_name("joint1") == "left/joint1" - assert config.get_coordinator_joint_name("unknown") == "unknown" diff --git a/dimos/manipulation/test_plan_execution.py b/dimos/manipulation/test_plan_execution.py index cc8de4296b..220393799a 100644 --- a/dimos/manipulation/test_plan_execution.py +++ b/dimos/manipulation/test_plan_execution.py @@ -14,7 +14,6 @@ """Tests for ManipulationModule plan-execution result projection.""" -from pathlib import Path from unittest.mock import MagicMock from dimos.control.coordinator import ControlCoordinator @@ -25,8 +24,6 @@ TrajectoryExecutionStatus, ) from dimos.manipulation.manipulation_module import ManipulationModule, ManipulationState -from dimos.manipulation.planning.groups.models import PlanningGroupDefinition -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.msgs.sensor_msgs.JointState import JointState @@ -52,7 +49,7 @@ def _plan(final_position: float = 1.0) -> GeneratedPlan: ], ) return GeneratedPlan( - group_ids=("arm/manipulator",), + group_ids=("manipulator",), trajectory=trajectory, path=[ JointState(name=names, position=[0.0]), @@ -67,22 +64,6 @@ def _module_with_coordinator( module_factory, ) -> ManipulationModule: module = module_factory(coordinator) - config = RobotModelConfig( - name="arm", - model_path=Path("/path/to/robot.urdf"), - joint_names=["j0"], - base_link="base", - planning_groups=[ - PlanningGroupDefinition( - name="manipulator", - joint_names=("j0",), - base_link="base", - tip_link="tool", - ) - ], - ) - module._robots = {"arm": ("arm_id", config)} - module._initialize_execution() return module diff --git a/dimos/manipulation/test_planning_factory.py b/dimos/manipulation/test_planning_factory.py index 8bf988b6e4..a73a24c6fb 100644 --- a/dimos/manipulation/test_planning_factory.py +++ b/dimos/manipulation/test_planning_factory.py @@ -17,16 +17,19 @@ from __future__ import annotations from collections.abc import Callable, Generator +import inspect from pathlib import Path import sys from types import ModuleType from typing import Any from unittest.mock import ANY +from pydantic import ValidationError import pytest from pytest_mock import MockerFixture -from dimos.manipulation.manipulation_module import ManipulationModule +from dimos.manipulation.manipulation_module import ManipulationModule, ManipulationModuleConfig +from dimos.manipulation.pick_and_place_module import PickAndPlaceModule from dimos.manipulation.planning.factory import ( create_kinematics, create_planner, @@ -72,7 +75,6 @@ def _make(**kwargs: Any) -> ManipulationModule: @pytest.fixture def robot_config() -> RobotModelConfig: return RobotModelConfig( - name="arm", model_path=Path("/path/to/robot.urdf"), base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), # type: ignore[call-arg] joint_names=["joint1", "joint2"], @@ -263,7 +265,6 @@ def test_create_planning_stack_defaults_to_roboplan( mocker: MockerFixture, robot_config: RobotModelConfig ) -> None: world = mocker.MagicMock() - world.add_robot.return_value = "robot-id" kinematics = mocker.MagicMock(name="kinematics") planner = mocker.MagicMock(name="planner") @@ -286,7 +287,7 @@ def test_create_planning_stack_defaults_to_roboplan( result = create_planning_stack(robot_config) - assert result == (world, kinematics, planner, "robot-id") + assert result == (world, kinematics, planner) mock_world.assert_called_once_with(backend="roboplan", visualization=None) mock_kinematics.assert_called_once_with(config=PinkKinematicsConfig()) mock_planner.assert_called_once_with( @@ -294,25 +295,26 @@ def test_create_planning_stack_defaults_to_roboplan( world=world, world_backend="roboplan", ) - world.add_robot.assert_called_once_with(robot_config) + world.load_model.assert_called_once_with(robot_config) world.finalize.assert_called_once() -def test_start_with_no_robots_skips_planning( - mocker: MockerFixture, make_module: Callable[..., ManipulationModule] -) -> None: - module = make_module(robots=[]) - create_world_mock = mocker.patch("dimos.manipulation.manipulation_module.create_world") - create_planning_specs_mock = mocker.patch( - "dimos.manipulation.manipulation_module.create_planning_specs" - ) +def test_configuration_requires_one_model_and_rejects_robots() -> None: + with pytest.raises(ValidationError, match="model"): + ManipulationModuleConfig() + with pytest.raises(ValidationError, match="robots"): + ManipulationModuleConfig.model_validate({"robots": []}) - module._initialize_planning() - assert module._robots == {} - assert module._world_monitor is None - create_world_mock.assert_not_called() - create_planning_specs_mock.assert_not_called() +def test_public_manipulation_surface_has_no_robot_selectors_or_listing_apis() -> None: + forbidden_parameters = {"robot_name", "robot_id", "hardware_id"} + for module_type in (ManipulationModule, PickAndPlaceModule): + for _, method in inspect.getmembers(module_type, inspect.isfunction): + if getattr(method, "__rpc__", False): + assert forbidden_parameters.isdisjoint(inspect.signature(method).parameters) + + for obsolete_method in ("get_robot_info", "list_robots", "preview_path"): + assert not hasattr(ManipulationModule, obsolete_method) def test_start_uses_configured_planner_and_kinematics( @@ -322,13 +324,12 @@ def test_start_uses_configured_planner_and_kinematics( ) -> None: planner_config = RRTConnectPlannerConfig() module = make_module( - robots=[robot_config], + model=robot_config, planner=planner_config, kinematics=JacobianKinematicsConfig(), ) world = mocker.MagicMock(name="world") world_monitor = mocker.MagicMock() - world_monitor.add_robot.return_value = "robot-id" planner = mocker.MagicMock(name="planner") kinematics = mocker.MagicMock(name="kinematics") planning_specs = mocker.MagicMock( @@ -353,10 +354,9 @@ def test_start_uses_configured_planner_and_kinematics( world=world, world_backend="roboplan", planner=planner_config, - kinematics_name=None, kinematics=module.config.kinematics, trajectory_parametrization=ANY, ) assert module._planner is planner assert module._kinematics is kinematics - assert module._robots["arm"][0] == "robot-id" + world_monitor.load_model.assert_called_once_with(robot_config) diff --git a/dimos/manipulation/test_roboplan.py b/dimos/manipulation/test_roboplan.py index b3adb4f894..3516b02b08 100644 --- a/dimos/manipulation/test_roboplan.py +++ b/dimos/manipulation/test_roboplan.py @@ -241,7 +241,7 @@ def getPositionLimitVectors( def getJointGroupInfo(self, name: str) -> FakeJointGroupInfo: names = ( self.joint_group_joint_names - if self.joint_group_joint_names is not None and name == self.constructor_kwargs["name"] + if self.joint_group_joint_names is not None and name == "__dimos_all_configured__" else self.groups[name] ) return FakeJointGroupInfo(list(names)) @@ -464,7 +464,6 @@ def robot_config(tmp_path: Path) -> RobotModelConfig: """ ) return RobotModelConfig( - name="arm", model_path=model_path, base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), # type: ignore[call-arg] joint_names=["joint1", "joint2"], @@ -486,14 +485,13 @@ def _make_world( fake_roboplan: None, robot_config: RobotModelConfig, planner_config: RoboPlanPlannerConfig | None = None, -) -> tuple[Any, str]: +) -> Any: module = _import_roboplan_world(fake_roboplan) world = module.RoboPlanWorld() - robot_id = world.add_robot(robot_config) + world.load_model(robot_config) world.finalize() world.sync_from_joint_state( - robot_id, JointState( name=list(robot_config.joint_names), position=[0.0] * len(robot_config.joint_names), @@ -504,7 +502,7 @@ def _make_world( world, planner_config or RoboPlanPlannerConfig(), ) - return world, robot_id + return world def test_roboplan_loads_canonical_slash_names_natively( @@ -527,7 +525,6 @@ def test_roboplan_loads_canonical_slash_names_natively( """ ) config = RobotModelConfig( - name="robot", model_path=model_path, joint_names=["left/j1"], base_link="world", @@ -543,40 +540,12 @@ def test_roboplan_loads_canonical_slash_names_natively( joint_limits_upper=[1.0], ) - world, robot_id = _make_world(fake_roboplan, config) + world = _make_world(fake_roboplan, config) - assert world.get_robot_config(robot_id).joint_names == ["left/j1"] + assert world.get_model_config().joint_names == ["left/j1"] assert "left/j1" in world._scene.constructor_kwargs["urdf"] -def _make_two_robot_world( - fake_roboplan: None, robot_config: RobotModelConfig -) -> tuple[Any, str, str, RobotModelConfig]: - module = _import_roboplan_world(fake_roboplan) - second_config = robot_config.model_copy( - update={ - "name": "right", - "base_pose": PoseStamped( - position=Vector3(1, 0, 0), - orientation=Quaternion(), - ), - } - ) - world = module.RoboPlanWorld() - first_id = world.add_robot(robot_config) - second_id = world.add_robot(second_config) - world.finalize() - world.sync_from_joint_state( - first_id, - JointState(name=list(robot_config.joint_names), position=[0.0, 0.0]), - ) - world.sync_from_joint_state( - second_id, - JointState(name=list(second_config.joint_names), position=[0.0, 0.0]), - ) - return world, first_id, second_id, second_config - - def _selection( configs: tuple[RobotModelConfig, ...], *group_ids: str, @@ -648,17 +617,17 @@ def test_roboplan_planner_rejects_non_roboplan_world(fake_roboplan: None) -> Non def test_robot_registration_finalization_and_joint_limits( fake_roboplan: None, robot_config: RobotModelConfig ) -> None: - world, robot_id = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) - assert world.get_robot_ids() == [robot_id] - assert world.get_robot_config(robot_id) is robot_config - assert world._scene.constructor_kwargs["name"] == "arm" - assert ET.fromstring(world._scene.constructor_kwargs["urdf"]).get("name") == "arm" - assert world._scene.constructor_kwargs["srdf"].startswith('') + assert [world.get_model_config()] == [robot_config] + assert world.get_model_config() is robot_config + assert world._scene.constructor_kwargs["name"] == "dimos_model" + assert ET.fromstring(world._scene.constructor_kwargs["urdf"]).get("name") == "dimos_model" + assert world._scene.constructor_kwargs["srdf"].startswith('') assert ( 'disable_collisions link1="base" link2="link1"' in world._scene.constructor_kwargs["srdf"] ) - lower, upper = world.get_joint_limits(robot_id) + lower, upper = world.get_joint_limits() np.testing.assert_allclose(lower, [-1.0, -2.0]) np.testing.assert_allclose(upper, [1.0, 2.0]) @@ -675,9 +644,9 @@ def test_scene_joint_limits_are_reordered_to_configured_joint_order( monkeypatch.setattr(FakeScene, "position_limits_lower", [-2.0, -1.0]) monkeypatch.setattr(FakeScene, "position_limits_upper", [2.0, 1.0]) - world, robot_id = _make_world(fake_roboplan, config) + world = _make_world(fake_roboplan, config) - lower, upper = world.get_joint_limits(robot_id) + lower, upper = world.get_joint_limits() np.testing.assert_allclose(lower, [-1.0, -2.0]) np.testing.assert_allclose(upper, [1.0, 2.0]) @@ -697,72 +666,27 @@ def test_scene_joint_limits_validate_joint_names( def test_context_cloning_and_joint_state_round_trip( fake_roboplan: None, robot_config: RobotModelConfig ) -> None: - world, robot_id = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) live_state = JointState(name=["joint1", "joint2"], position=[0.1, 0.2]) - world.sync_from_joint_state(robot_id, live_state) + world.sync_from_joint_state(live_state) with world.scratch_context() as scratch: - scratch_state = world.get_joint_state(scratch, robot_id) + scratch_state = world.get_joint_state(scratch) assert scratch_state.name == ["joint1", "joint2"] assert scratch_state.position == [0.1, 0.2] - world.set_joint_state( - scratch, robot_id, JointState(name=["joint1", "joint2"], position=[0.3, 0.4]) - ) + world.set_joint_state(scratch, JointState(name=["joint1", "joint2"], position=[0.3, 0.4])) - live_round_trip = world.get_joint_state(world.get_live_context(), robot_id) + live_round_trip = world.get_joint_state(world.get_live_context()) assert live_round_trip.position == [0.1, 0.2] -def test_joint_name_mapping_is_applied_to_input_states( - 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.sync_from_joint_state( - robot_id, JointState(name=["arm/j1", "arm/j2"], position=[0.2, 0.3]) - ) - - live_round_trip = world.get_joint_state(world.get_live_context(), robot_id) - 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.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) - - 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, mocker: MockerFixture, ) -> None: - world, _ = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) add_box = mocker.patch.object( FakeScene, "addBoxGeometry", @@ -800,7 +724,7 @@ def test_obstacle_operations_require_finalization( ) -> None: module = _import_roboplan_world(fake_roboplan) world = module.RoboPlanWorld() - world.add_robot(robot_config) + world.load_model(robot_config) obstacle = Obstacle( name="box", obstacle_type=ObstacleType.BOX, @@ -826,7 +750,7 @@ def test_failed_obstacle_add_rolls_back_and_can_be_retried( robot_config: RobotModelConfig, mocker: MockerFixture, ) -> None: - world, _ = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) obstacle = Obstacle( name="retryable", obstacle_type=ObstacleType.BOX, @@ -854,7 +778,7 @@ def test_concurrent_remove_waits_for_obstacle_add( robot_config: RobotModelConfig, mocker: MockerFixture, ) -> None: - world, _ = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) obstacle = Obstacle( name="concurrent", obstacle_type=ObstacleType.BOX, @@ -901,7 +825,7 @@ def blocking_add(*args: Any, **kwargs: Any) -> None: def test_obstacle_ids_are_world_owned_and_invalid_insertions_are_rejected( fake_roboplan: None, robot_config: RobotModelConfig ) -> None: - world, _ = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) unnamed = Obstacle( name="", @@ -922,7 +846,7 @@ def test_obstacle_ids_are_world_owned_and_invalid_insertions_are_rejected( def test_complete_update_rejects_invalid_obstacle_values( fake_roboplan: None, robot_config: RobotModelConfig ) -> None: - world, _ = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) valid = Obstacle( name="shape", obstacle_type=ObstacleType.BOX, @@ -959,7 +883,7 @@ def test_complete_update_rejects_invalid_obstacle_values( def test_complete_replacement_and_defensive_obstacle_snapshots( fake_roboplan: None, robot_config: RobotModelConfig ) -> None: - world, _ = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) original = Obstacle( name="shape", obstacle_type=ObstacleType.BOX, @@ -998,7 +922,7 @@ def test_collision_query_blocks_during_obstacle_replacement( robot_config: RobotModelConfig, monkeypatch: pytest.MonkeyPatch, ) -> None: - world, robot_id = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) obstacle = Obstacle( name="box", obstacle_type=ObstacleType.BOX, @@ -1025,7 +949,6 @@ def blocking_remove(obstacle_id: str) -> None: query_thread = threading.Thread( target=lambda: ( world.check_config_collision_free( - robot_id, JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), ), query_finished.set(), @@ -1044,7 +967,7 @@ def test_obstacle_replacement_blocks_during_collision_query( robot_config: RobotModelConfig, monkeypatch: pytest.MonkeyPatch, ) -> None: - world, robot_id = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) obstacle = Obstacle( name="box", obstacle_type=ObstacleType.BOX, @@ -1065,7 +988,6 @@ def blocking_query(q: np.ndarray) -> bool: monkeypatch.setattr(world._scene, "hasCollisions", blocking_query) query_thread = threading.Thread( target=lambda: world.check_config_collision_free( - robot_id, JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), ) ) @@ -1090,7 +1012,7 @@ def test_native_update_failure_invalidates_world( robot_config: RobotModelConfig, monkeypatch: pytest.MonkeyPatch, ) -> None: - world, robot_id = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) obstacle = Obstacle( name="box", obstacle_type=ObstacleType.BOX, @@ -1110,7 +1032,6 @@ def test_native_update_failure_invalidates_world( world.get_obstacles() with pytest.raises(RuntimeError, match="invalid"): world.check_config_collision_free( - robot_id, JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), ) @@ -1120,7 +1041,7 @@ def test_native_pose_update_failure_invalidates_world( robot_config: RobotModelConfig, monkeypatch: pytest.MonkeyPatch, ) -> None: - world, _ = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) obstacle = Obstacle( name="box", obstacle_type=ObstacleType.BOX, @@ -1143,37 +1064,37 @@ def test_native_pose_update_failure_invalidates_world( def test_collision_config_and_edge_checks( fake_roboplan: None, robot_config: RobotModelConfig ) -> None: - world, robot_id = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) safe = JointState(name=["joint1", "joint2"], position=[0.1, 0.2]) colliding = JointState(name=["joint1", "joint2"], position=[0.95, 0.2]) - assert world.check_config_collision_free(robot_id, safe) - assert not world.check_config_collision_free(robot_id, colliding) - assert not world.check_edge_collision_free(robot_id, safe, colliding, step_size=0.05) + assert world.check_config_collision_free(safe) + assert not world.check_config_collision_free(colliding) + assert not world.check_edge_collision_free(safe, colliding, step_size=0.05) def test_collision_check_uses_scene_queries( fake_roboplan: None, robot_config: RobotModelConfig ) -> None: - world, robot_id = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) safe = JointState(name=["joint1", "joint2"], position=[0.1, 0.2]) colliding = JointState(name=["joint1", "joint2"], position=[0.95, 0.2]) - assert world.check_config_collision_free(robot_id, safe) - assert not world.check_config_collision_free(robot_id, colliding) + assert world.check_config_collision_free(safe) + assert not world.check_config_collision_free(colliding) def test_generic_rrt_planner_uses_roboplan_world_collision_checks( fake_roboplan: None, robot_config: RobotModelConfig ) -> None: - world, robot_id = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) planner = RRTConnectPlanner(step_size=0.5, connect_step_size=0.5, goal_tolerance=10.0) start = JointState(name=["joint1", "joint2"], position=[0.0, 0.0]) goal = JointState(name=["joint1", "joint2"], position=[0.2, 0.1]) - result = planner.plan_joint_path(world, robot_id, start, goal, timeout=1.0, max_iterations=3) + result = planner.plan_joint_path(world, start, goal, timeout=1.0, max_iterations=3) assert result.status == PlanningStatus.SUCCESS assert len(result.path) >= 2 @@ -1184,7 +1105,7 @@ def test_generic_planner_allows_update_between_collision_checks( robot_config: RobotModelConfig, monkeypatch: pytest.MonkeyPatch, ) -> None: - world, robot_id = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) obstacle = Obstacle( name="box", obstacle_type=ObstacleType.BOX, @@ -1196,9 +1117,9 @@ def test_generic_planner_allows_update_between_collision_checks( checks = 0 updated = False - def checking_with_interleaved_update(robot: str, state: JointState) -> bool: + def checking_with_interleaved_update(state: JointState) -> bool: nonlocal checks, updated - result = original_check(robot, state) + result = original_check(state) checks += 1 if checks == 1: updated = world.update_obstacle(replace(obstacle, dimensions=(1.0, 1.0, 1.0))) @@ -1208,7 +1129,6 @@ def checking_with_interleaved_update(robot: str, state: JointState) -> bool: planner = RRTConnectPlanner(step_size=0.5, connect_step_size=0.5, goal_tolerance=10.0) result = planner.plan_joint_path( world, - robot_id, JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), JointState(name=["joint1", "joint2"], position=[0.2, 0.1]), timeout=1.0, @@ -1223,17 +1143,15 @@ def checking_with_interleaved_update(robot: str, state: JointState) -> bool: def test_fk_jacobian_and_explicit_min_distance_unsupported( fake_roboplan: None, robot_config: RobotModelConfig ) -> None: - world, robot_id = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) ctx = world.get_live_context() - world.set_joint_state( - ctx, robot_id, JointState(name=["joint1", "joint2"], position=[0.25, 0.5]) - ) + world.set_joint_state(ctx, JointState(name=["joint1", "joint2"], position=[0.25, 0.5])) - pose = world.get_ee_pose(ctx, robot_id) + pose = world.get_ee_pose(ctx) assert pose.position.x == pytest.approx(0.75) - assert world.get_jacobian(ctx, robot_id).shape == (6, 2) + assert world.get_jacobian(ctx).shape == (6, 2) with pytest.raises(NotImplementedError, match="get_min_distance"): - world.get_min_distance(ctx, robot_id) + world.get_min_distance(ctx) def test_group_fk_and_jacobian_use_group_tip_and_local_joint_order( @@ -1276,23 +1194,22 @@ def fake_jacobian( monkeypatch.setattr(FakeScene, "forwardKinematics", fake_fk) monkeypatch.setattr(FakeScene, "computeFrameJacobian", fake_jacobian) - world, robot_id = _make_world(fake_roboplan, config) + world = _make_world(fake_roboplan, config) 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") + pose = world.get_group_ee_pose(ctx, "wrist") + jacobian = world.get_group_jacobian(ctx, "wrist") assert fk_frames == ["tcp"] assert pose.position.x == pytest.approx(6.0) np.testing.assert_allclose(jacobian, np.arange(18, dtype=np.float64).reshape(6, 3)[:, [2, 0]]) -def test_group_kinematics_reject_missing_tip_or_missing_context( +def test_group_kinematics_rejects_missing_tip( fake_roboplan: None, robot_config: RobotModelConfig ) -> None: no_tip_config = robot_config.model_copy( @@ -1304,23 +1221,12 @@ def test_group_kinematics_reject_missing_tip_or_missing_context( ] } ) - world, robot_id = _make_world(fake_roboplan, no_tip_config) + world = _make_world(fake_roboplan, no_tip_config) with pytest.raises(ValueError, match="no tip link"): - world.get_group_ee_pose(world.get_live_context(), "arm/joint_only") + world.get_group_ee_pose(world.get_live_context(), "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_ctx = jacobian_world.get_live_context() - del jacobian_ctx.q_by_robot[jacobian_robot_id] - with pytest.raises(RuntimeError, match="Missing authoritative state"): - jacobian_world.get_group_jacobian(jacobian_ctx, "arm/manipulator") + world.get_group_jacobian(world.get_live_context(), "joint_only") def test_group_jacobian_validates_projection_shape( @@ -1328,9 +1234,9 @@ def test_group_jacobian_validates_projection_shape( robot_config: RobotModelConfig, monkeypatch: pytest.MonkeyPatch, ) -> None: - world, robot_id = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) ctx = world.get_live_context() - world.set_joint_state(ctx, robot_id, JointState(name=["joint1", "joint2"], position=[0.0, 0.0])) + world.set_joint_state(ctx, JointState(name=["joint1", "joint2"], position=[0.0, 0.0])) monkeypatch.setattr( FakeScene, @@ -1338,7 +1244,7 @@ def test_group_jacobian_validates_projection_shape( 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") + world.get_group_jacobian(ctx, "manipulator") monkeypatch.setattr( FakeScene, @@ -1346,10 +1252,10 @@ def test_group_jacobian_validates_projection_shape( lambda self, q, frame_name, local=True: np.ones((6, 4)), ) with pytest.raises(ValueError, match="cannot project"): - world.get_group_jacobian(ctx, "arm/manipulator") + world.get_group_jacobian(ctx, "manipulator") -def test_legacy_kinematics_wrappers_require_unique_pose_group( +def test_kinematics_convenience_methods_require_unique_pose_group( fake_roboplan: None, robot_config: RobotModelConfig ) -> None: no_pose_config = robot_config.model_copy( @@ -1359,11 +1265,11 @@ def test_legacy_kinematics_wrappers_require_unique_pose_group( ] } ) - no_pose_world, no_pose_id = _make_world(fake_roboplan, no_pose_config) + no_pose_world = _make_world(fake_roboplan, no_pose_config) with pytest.raises(ValueError, match="no pose-targetable"): - no_pose_world.get_ee_pose(no_pose_world.get_live_context(), no_pose_id) + no_pose_world.get_ee_pose(no_pose_world.get_live_context()) with pytest.raises(ValueError, match="no pose-targetable"): - no_pose_world.get_jacobian(no_pose_world.get_live_context(), no_pose_id) + no_pose_world.get_jacobian(no_pose_world.get_live_context()) ambiguous_config = robot_config.model_copy( update={ @@ -1377,26 +1283,26 @@ def test_legacy_kinematics_wrappers_require_unique_pose_group( ] } ) - ambiguous_world, ambiguous_id = _make_world(fake_roboplan, ambiguous_config) + ambiguous_world = _make_world(fake_roboplan, ambiguous_config) with pytest.raises(ValueError, match="pose-targetable planning groups"): - ambiguous_world.get_jacobian(ambiguous_world.get_live_context(), ambiguous_id) + ambiguous_world.get_jacobian(ambiguous_world.get_live_context()) def test_group_lookup_rejects_unknown_group_id( fake_roboplan: None, robot_config: RobotModelConfig ) -> None: - world, _ = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) with pytest.raises(KeyError, match="Unknown planning group ID"): - world.get_group_ee_pose(world.get_live_context(), "other/missing") + world.get_group_ee_pose(world.get_live_context(), "missing") def test_native_planner_converts_path(fake_roboplan: None, robot_config: RobotModelConfig) -> None: - world, robot_id = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) start = JointState(name=["joint1", "joint2"], position=[0.0, 0.0]) goal = JointState(name=["joint1", "joint2"], position=[0.4, 0.2]) - result = _planner_for(world).plan_joint_path(world, robot_id, start, goal, timeout=1.0) + result = _planner_for(world).plan_joint_path(world, start, goal, timeout=1.0) assert result.status == PlanningStatus.SUCCESS assert [state.position for state in result.path] == [[0.0, 0.0], [0.2, 0.1], [0.4, 0.2]] @@ -1415,7 +1321,7 @@ def test_native_planner_shortcuts_path_with_configured_options( max_convergence_iters=8, redundant_removal_iters=5, ) - world, robot_id = _make_world( + world = _make_world( fake_roboplan, robot_config, RoboPlanPlannerConfig(path_shortcutting=shortcutting), @@ -1433,7 +1339,6 @@ def endpoints_only(shortcutter: FakePathShortcutter, path: FakeJointPath) -> Fak result = _planner_for(world).plan_joint_path( world, - robot_id, JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), JointState(name=["joint1", "joint2"], position=[0.4, 0.2]), timeout=1.0, @@ -1441,7 +1346,7 @@ def endpoints_only(shortcutter: FakePathShortcutter, path: FakeJointPath) -> Fak assert [state.position for state in result.path] == [[0.0, 0.0], [0.4, 0.2]] options = FakePathShortcutter.instances[-1].options - assert options.group_name == "arm" + assert options.group_name == "__dimos_all_configured__" assert options.max_step_size == 0.02 assert options.max_iters == 40 assert options.seed == 7 @@ -1454,7 +1359,7 @@ def test_native_planner_can_disable_path_shortcutting( robot_config: RobotModelConfig, mocker: MockerFixture, ) -> None: - world, robot_id = _make_world( + world = _make_world( fake_roboplan, robot_config, RoboPlanPlannerConfig(path_shortcutting=RoboPlanPathShortcuttingConfig(enabled=False)), @@ -1464,7 +1369,6 @@ def test_native_planner_can_disable_path_shortcutting( result = _planner_for(world).plan_joint_path( world, - robot_id, JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), JointState(name=["joint1", "joint2"], position=[0.4, 0.2]), timeout=1.0, @@ -1479,7 +1383,7 @@ def test_native_planner_uses_raw_path_when_shortcutting_fails( robot_config: RobotModelConfig, mocker: MockerFixture, ) -> None: - world, robot_id = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) mocker.patch.object( FakePathShortcutter, "shortcut", @@ -1489,7 +1393,6 @@ def test_native_planner_uses_raw_path_when_shortcutting_fails( result = _planner_for(world).plan_joint_path( world, - robot_id, JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), JointState(name=["joint1", "joint2"], position=[0.4, 0.2]), timeout=1.0, @@ -1504,7 +1407,7 @@ def test_native_planner_surfaces_unexpected_shortcutting_error( robot_config: RobotModelConfig, mocker: MockerFixture, ) -> None: - world, robot_id = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) mocker.patch.object( FakePathShortcutter, "shortcut", @@ -1515,7 +1418,6 @@ def test_native_planner_surfaces_unexpected_shortcutting_error( with pytest.raises(TypeError, match="unexpected shortcut integration error"): _planner_for(world).plan_joint_path( world, - robot_id, JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), JointState(name=["joint1", "joint2"], position=[0.4, 0.2]), timeout=1.0, @@ -1529,7 +1431,7 @@ def test_native_planner_uses_raw_path_when_shortcutting_changes_endpoint( mocker: MockerFixture, endpoint_index: int, ) -> None: - world, robot_id = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) def changed_endpoint(shortcutter: FakePathShortcutter, path: FakeJointPath) -> FakeJointPath: positions = [*path.positions] @@ -1545,7 +1447,6 @@ def changed_endpoint(shortcutter: FakePathShortcutter, path: FakeJointPath) -> F result = _planner_for(world).plan_joint_path( world, - robot_id, JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), JointState(name=["joint1", "joint2"], position=[0.4, 0.2]), timeout=1.0, @@ -1560,7 +1461,7 @@ def test_native_planner_uses_raw_path_when_shortcutting_returns_empty_path( robot_config: RobotModelConfig, mocker: MockerFixture, ) -> None: - world, robot_id = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) mocker.patch.object( FakePathShortcutter, "shortcut", @@ -1570,7 +1471,6 @@ def test_native_planner_uses_raw_path_when_shortcutting_returns_empty_path( result = _planner_for(world).plan_joint_path( world, - robot_id, JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), JointState(name=["joint1", "joint2"], position=[0.4, 0.2]), timeout=1.0, @@ -1584,7 +1484,7 @@ def test_roboplan_planner_is_distinct_and_bound_to_world( fake_roboplan: None, robot_config: RobotModelConfig, ) -> None: - world, _ = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) planner = _planner_for(world) assert planner is not world @@ -1597,7 +1497,7 @@ def test_roboplan_planner_copies_configuration( robot_config: RobotModelConfig, ) -> None: config = RoboPlanPlannerConfig(path_shortcutting=RoboPlanPathShortcuttingConfig(enabled=False)) - world, _ = _make_world(fake_roboplan, robot_config, config) + world = _make_world(fake_roboplan, robot_config, config) config.path_shortcutting.enabled = True @@ -1609,7 +1509,7 @@ def test_native_planning_blocks_obstacle_replacement( robot_config: RobotModelConfig, monkeypatch: pytest.MonkeyPatch, ) -> None: - world, robot_id = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) obstacle = Obstacle( name="box", obstacle_type=ObstacleType.BOX, @@ -1635,9 +1535,7 @@ def blocking_plan( start = JointState(name=["joint1", "joint2"], position=[0.0, 0.0]) goal = JointState(name=["joint1", "joint2"], position=[0.2, 0.1]) planning_thread = threading.Thread( - target=lambda: _planner_for(world).plan_joint_path( - world, robot_id, start, goal, timeout=1.0 - ) + target=lambda: _planner_for(world).plan_joint_path(world, start, goal, timeout=1.0) ) planning_thread.start() assert planning_started.wait(1.0) @@ -1658,11 +1556,11 @@ def blocking_plan( def test_native_planner_names_path_from_robot_config_when_start_is_unnamed( fake_roboplan: None, robot_config: RobotModelConfig ) -> None: - world, robot_id = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) start = JointState(name=[], position=[0.0, 0.0]) goal = JointState(name=["joint1", "joint2"], position=[0.4, 0.2]) - result = _planner_for(world).plan_joint_path(world, robot_id, start, goal, timeout=1.0) + result = _planner_for(world).plan_joint_path(world, start, goal, timeout=1.0) assert result.status == PlanningStatus.SUCCESS assert [state.name for state in result.path] == [["joint1", "joint2"]] * 3 @@ -1671,47 +1569,29 @@ def test_native_planner_names_path_from_robot_config_when_start_is_unnamed( def test_native_selected_planner_returns_global_selected_joint_names( fake_roboplan: None, robot_config: RobotModelConfig ) -> None: - world, _ = _make_world(fake_roboplan, robot_config) - selection = _selection((robot_config,), "arm/manipulator") + world = _make_world(fake_roboplan, robot_config) + selection = _selection((robot_config,), "manipulator") result = _planner_for(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]), + JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), + JointState(name=["joint1", "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.name for state in result.path] == [["joint1", "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) - selection = _selection((robot_config,), "arm/manipulator") - - result = _planner_for(world).plan_selected_joint_path( - world, - selection, - JointState(name=["joint2", "joint1"], position=[0.0, 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[-1].position == [0.2, 0.4] - - def test_native_selected_planner_uses_explicit_start_after_live_state_advances( fake_roboplan: None, robot_config: RobotModelConfig, mocker: MockerFixture, ) -> None: - world, robot_id = _make_world(fake_roboplan, robot_config) - selection = _selection((robot_config,), "arm/manipulator") + world = _make_world(fake_roboplan, robot_config) + selection = _selection((robot_config,), "manipulator") observed_scene_start: list[float] = [] native_plan = FakeRRT.plan @@ -1726,7 +1606,6 @@ def capture_scene_start( mocker.patch.object(FakeRRT, "plan", autospec=True, side_effect=capture_scene_start) start = JointState(name=list(selection.joint_names), position=[0.1, -0.1]) world.sync_from_joint_state( - robot_id, JointState(name=["joint1", "joint2"], position=[0.3, 0.2]), ) @@ -1753,8 +1632,8 @@ def test_native_selected_planner_rejects_multi_group_selection( ] } ) - world, _ = _make_world(fake_roboplan, config) - selection = _selection((config,), "arm/left", "arm/right") + world = _make_world(fake_roboplan, config) + selection = _selection((config,), "left", "right") result = _planner_for(world).plan_selected_joint_path( world, @@ -1767,38 +1646,11 @@ def test_native_selected_planner_rejects_multi_group_selection( assert "no generated group" in result.message -def test_native_planner_coordinates_groups_across_two_robots( - fake_roboplan: None, robot_config: RobotModelConfig -) -> None: - world, _, _, second_config = _make_two_robot_world(fake_roboplan, robot_config) - selection = _selection( - (robot_config, second_config), - "right/manipulator", - "arm/manipulator", - ) - - result = _planner_for(world).plan_selected_joint_path( - world, - selection, - JointState(name=list(selection.joint_names), position=[0.0, 0.0, 0.0, 0.0]), - JointState(name=list(selection.joint_names), position=[0.4, 0.2, 0.1, 0.3]), - ) - - assert result.status == PlanningStatus.SUCCESS - assert result.path[-1].name == [ - "arm/joint1", - "arm/joint2", - "right/joint1", - "right/joint2", - ] - assert result.path[-1].position == [0.1, 0.3, 0.4, 0.2] - - def test_cartesian_planner_returns_timed_global_joint_states_and_options( fake_roboplan: None, robot_config: RobotModelConfig ) -> None: - world, _ = _make_world(fake_roboplan, robot_config) - selection = _selection((robot_config,), "arm/manipulator") + world = _make_world(fake_roboplan, robot_config) + selection = _selection((robot_config,), "manipulator") option_overrides = { "dt": 0.02, "max_linear_speed": 0.2, @@ -1826,7 +1678,7 @@ def test_cartesian_planner_returns_timed_global_joint_states_and_options( selection, JointState(name=list(selection.joint_names), position=[0.0, 0.0]), { - "arm/manipulator": _relative_target( + "manipulator": _relative_target( Transform( translation=Vector3(0.1, 0.0, 0.0), rotation=Quaternion.from_euler(Vector3(0.0, 0.0, np.pi / 2.0)), @@ -1867,15 +1719,15 @@ def test_cartesian_planner_returns_timed_global_joint_states_and_options( def test_cartesian_zero_rotation_preserves_start_orientation( fake_roboplan: None, robot_config: RobotModelConfig ) -> None: - world, _ = _make_world(fake_roboplan, robot_config) - selection = _selection((robot_config,), "arm/manipulator") + world = _make_world(fake_roboplan, robot_config) + selection = _selection((robot_config,), "manipulator") result = _planner_for(world).plan_cartesian_path( world, selection, JointState(name=list(selection.joint_names), position=[0.0, 0.0]), { - "arm/manipulator": _relative_target( + "manipulator": _relative_target( Transform(translation=Vector3(0.05, 0.02, 0.0)), Transform(translation=Vector3(0.1, 0.0, 0.0)), ) @@ -1890,134 +1742,13 @@ def test_cartesian_zero_rotation_preserves_start_orientation( np.testing.assert_allclose(waypoint[:3, :3], track[0][:3, :3], atol=1e-12) -def test_cartesian_supports_mixed_targets_and_shared_multi_group_timing( - fake_roboplan: None, robot_config: RobotModelConfig -) -> None: - world, _, _, second_config = _make_two_robot_world(fake_roboplan, robot_config) - selection = _selection( - (robot_config, second_config), - "right/manipulator", - "arm/manipulator", - ) - - result = _planner_for(world).plan_cartesian_path( - world, - selection, - JointState(name=list(selection.joint_names), position=[0.0] * 4), - { - "right/manipulator": _absolute_target( - PoseStamped( - frame_id="world", - position=Vector3(0.2, 0.0, 0.0), - orientation=Quaternion(), - ) - ), - "arm/manipulator": _relative_target(Transform(translation=Vector3(0.05, 0.0, 0.0))), - }, - RoboPlanCartesianPathConfig(), - ) - - assert result.status == PlanningStatus.SUCCESS - assert result.timestamps == [0.0, 0.01, 0.02] - assert result.path[-1].name == list(selection.joint_names) - planner = FakeCartesianPathPlanner.instances[-1] - assert len(planner.paths) == 1 - assert len(planner.paths[0].tforms) == 2 - assert planner.paths[0].tip_frames == ["right__tcp", "arm__tcp"] - - -def test_cartesian_allows_auxiliary_groups( - fake_roboplan: None, robot_config: RobotModelConfig -) -> None: - world, _, _, second_config = _make_two_robot_world(fake_roboplan, robot_config) - selection = _selection( - (robot_config, second_config), - "arm/manipulator", - "right/manipulator", - ) - - result = _planner_for(world).plan_cartesian_path( - world, - selection, - JointState(name=list(selection.joint_names), position=[0.0] * 4), - {"arm/manipulator": _relative_target(Transform(translation=Vector3(0.05, 0.0, 0.0)))}, - RoboPlanCartesianPathConfig(), - auxiliary_groups=("right/manipulator",), - ) - - assert result.status == PlanningStatus.SUCCESS - assert result.path[-1].name == list(selection.joint_names) - assert len(FakeCartesianPathPlanner.instances[-1].paths[0].tforms) == 1 - - -@pytest.mark.parametrize( - ("targets", "auxiliary_groups", "expected_status", "message"), - [ - ({}, ("arm/manipulator",), PlanningStatus.INVALID_GOAL, "at least one target"), - ( - {"arm/manipulator": _relative_target(Transform())}, - ("arm/manipulator",), - PlanningStatus.INVALID_GOAL, - "disjoint", - ), - ( - {"arm/manipulator": _relative_target(Transform(frame_id="tool"))}, - (), - PlanningStatus.UNSUPPORTED, - "world-frame", - ), - ( - {"arm/manipulator": (Transform.identity(),)}, - (), - PlanningStatus.INVALID_GOAL, - "at least two waypoints", - ), - ( - { - "arm/manipulator": ( - Transform.identity(), - PoseStamped(frame_id="world"), - ) - }, - (), - PlanningStatus.INVALID_GOAL, - "only PoseStamped waypoints or only Transform waypoints", - ), - ], -) -def test_cartesian_rejects_invalid_requests( - fake_roboplan: None, - robot_config: RobotModelConfig, - targets: dict[str, Any], - auxiliary_groups: tuple[str, ...], - expected_status: PlanningStatus, - message: str, -) -> None: - world, _ = _make_world(fake_roboplan, robot_config) - selection = _selection((robot_config,), "arm/manipulator") - - result = _planner_for(world).plan_cartesian_path( - world, - selection, - JointState(name=list(selection.joint_names), position=[0.0, 0.0]), - targets, - RoboPlanCartesianPathConfig(), - auxiliary_groups=auxiliary_groups, - ) - - assert result.status == expected_status - assert message in result.message - assert result.path == [] - - def test_cartesian_uses_explicit_start_after_live_state_advances( fake_roboplan: None, robot_config: RobotModelConfig ) -> None: - world, robot_id = _make_world(fake_roboplan, robot_config) - selection = _selection((robot_config,), "arm/manipulator") + world = _make_world(fake_roboplan, robot_config) + selection = _selection((robot_config,), "manipulator") start = JointState(name=list(selection.joint_names), position=[0.1, 0.0]) world.sync_from_joint_state( - robot_id, JointState(name=["joint1", "joint2"], position=[0.3, 0.2]), ) @@ -2025,7 +1756,7 @@ def test_cartesian_uses_explicit_start_after_live_state_advances( world, selection, start, - {"arm/manipulator": _relative_target(Transform(translation=Vector3(0.1, 0.0, 0.0)))}, + {"manipulator": _relative_target(Transform(translation=Vector3(0.1, 0.0, 0.0)))}, RoboPlanCartesianPathConfig(), ) @@ -2040,8 +1771,8 @@ def test_cartesian_rejects_official_planner_failure( robot_config: RobotModelConfig, mocker: MockerFixture, ) -> None: - world, _ = _make_world(fake_roboplan, robot_config) - selection = _selection((robot_config,), "arm/manipulator") + world = _make_world(fake_roboplan, robot_config) + selection = _selection((robot_config,), "manipulator") mocker.patch.object( FakeCartesianPathPlanner, "plan", @@ -2053,7 +1784,7 @@ def test_cartesian_rejects_official_planner_failure( world, selection, JointState(name=list(selection.joint_names), position=[0.0, 0.0]), - {"arm/manipulator": _relative_target(Transform(translation=Vector3(0.1, 0.0, 0.0)))}, + {"manipulator": _relative_target(Transform(translation=Vector3(0.1, 0.0, 0.0)))}, RoboPlanCartesianPathConfig(), ) @@ -2061,52 +1792,13 @@ def test_cartesian_rejects_official_planner_failure( assert "tracking failed" in result.message -def test_cartesian_postvalidation_checks_combined_multi_robot_state( - fake_roboplan: None, - robot_config: RobotModelConfig, - mocker: MockerFixture, -) -> None: - world, _, _, second_config = _make_two_robot_world(fake_roboplan, robot_config) - selection = _selection( - (robot_config, second_config), - "arm/manipulator", - "right/manipulator", - ) - - def collides_only_when_both_arms_advance(scene: FakeScene, q: np.ndarray) -> bool: - by_name = dict(zip(scene.native_joint_names, q, strict=True)) - return by_name["arm__joint1"] > 0.075 and by_name["right__joint1"] > 0.075 - - mocker.patch.object( - FakeScene, - "hasCollisions", - autospec=True, - side_effect=collides_only_when_both_arms_advance, - ) - - result = _planner_for(world).plan_cartesian_path( - world, - selection, - JointState(name=list(selection.joint_names), position=[0.0] * 4), - { - "arm/manipulator": _relative_target(Transform(translation=Vector3(0.05, 0.0, 0.0))), - "right/manipulator": _relative_target(Transform(translation=Vector3(0.05, 0.0, 0.0))), - }, - RoboPlanCartesianPathConfig(), - ) - - assert result.status == PlanningStatus.NO_SOLUTION - assert result.path == [] - assert "collision post-validation" in result.message - - def test_cartesian_postvalidation_checks_between_waypoints( fake_roboplan: None, robot_config: RobotModelConfig, mocker: MockerFixture, ) -> None: - world, _ = _make_world(fake_roboplan, robot_config) - selection = _selection((robot_config,), "arm/manipulator") + world = _make_world(fake_roboplan, robot_config) + selection = _selection((robot_config,), "manipulator") def two_point_trajectory( planner: FakeCartesianPathPlanner, @@ -2144,7 +1836,7 @@ def collides_only_mid_edge(scene: FakeScene, q: np.ndarray) -> bool: world, selection, JointState(name=list(selection.joint_names), position=[0.0, 0.0]), - {"arm/manipulator": _relative_target(Transform(translation=Vector3(0.05, 0.0, 0.0)))}, + {"manipulator": _relative_target(Transform(translation=Vector3(0.05, 0.0, 0.0)))}, RoboPlanCartesianPathConfig(), ) @@ -2153,76 +1845,6 @@ def collides_only_mid_edge(scene: FakeScene, q: np.ndarray) -> bool: assert "collision post-validation" in result.message -def test_native_planner_preserves_other_robot_and_auxiliary_joint_state( - fake_roboplan: None, - robot_config: RobotModelConfig, - mocker: MockerFixture, -) -> None: - world, _, second_id, second_config = _make_two_robot_world(fake_roboplan, robot_config) - world.sync_from_joint_state( - second_id, - JointState(name=["joint1", "joint2"], position=[0.3, 0.1]), - ) - selection = _selection((robot_config, second_config), "arm/manipulator") - observed_positions: dict[str, float] = {} - native_plan = FakeRRT.plan - - def capture_scene_state( - planner: FakeRRT, - q_start: FakeJointConfiguration, - q_goal: FakeJointConfiguration, - ) -> FakeJointPath: - observed_positions.update( - zip( - planner.scene.native_joint_names, - planner.scene.current_positions, - strict=True, - ) - ) - return native_plan(planner, q_start, q_goal) - - mocker.patch.object(FakeRRT, "plan", autospec=True, side_effect=capture_scene_state) - - result = _planner_for(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.2, 0.1]), - ) - - assert result.status == PlanningStatus.SUCCESS - assert observed_positions["right__joint1"] == pytest.approx(0.3) - assert observed_positions["right__joint2"] == pytest.approx(0.1) - assert observed_positions["arm__joint3"] == pytest.approx(0.0) - assert observed_positions["right__joint3"] == pytest.approx(0.0) - - -def test_native_planner_waits_for_every_robot_state( - fake_roboplan: None, robot_config: RobotModelConfig -) -> None: - module = _import_roboplan_world(fake_roboplan) - second_config = robot_config.model_copy(update={"name": "right"}) - world = module.RoboPlanWorld() - first_id = world.add_robot(robot_config) - world.add_robot(second_config) - world.finalize() - world.sync_from_joint_state( - first_id, - JointState(name=["joint1", "joint2"], position=[0.0, 0.0]), - ) - selection = _selection((robot_config, second_config), "arm/manipulator") - - result = _planner_for(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.2, 0.1]), - ) - - assert result.status == PlanningStatus.INVALID_START - assert "authoritative state is incomplete" in result.message - - def test_native_planner_rejects_empty_path( fake_roboplan: None, robot_config: RobotModelConfig, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -2233,11 +1855,11 @@ def plan( return FakeJointPath(["joint1", "joint2"], []) monkeypatch.setattr(sys.modules["roboplan.rrt"], "RRT", EmptyPathRRT) - world, robot_id = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) start = JointState(name=["joint1", "joint2"], position=[0.0, 0.0]) goal = JointState(name=["joint1", "joint2"], position=[0.4, 0.2]) - result = _planner_for(world).plan_joint_path(world, robot_id, start, goal, timeout=1.0) + result = _planner_for(world).plan_joint_path(world, start, goal, timeout=1.0) assert result.status == PlanningStatus.NO_SOLUTION assert result.path == [] @@ -2251,7 +1873,7 @@ def test_collision_exclusion_pairs_are_written_to_generated_srdf( ("base", "link2"), ("other_base", "other_tip"), ] - world, _ = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) srdf = world._scene.constructor_kwargs["srdf"] assert 'disable_collisions link1="base" link2="link2"' in srdf @@ -2264,7 +1886,7 @@ def test_collision_exclusion_with_one_unknown_link_is_rejected( robot_config.collision_exclusion_pairs = [("base", "missing")] module = _import_roboplan_world(fake_roboplan) world = module.RoboPlanWorld() - world.add_robot(robot_config) + world.load_model(robot_config) with pytest.raises( ValueError, @@ -2276,7 +1898,7 @@ def test_collision_exclusion_with_one_unknown_link_is_rejected( def test_scene_receives_generated_model_contents_inline( fake_roboplan: None, robot_config: RobotModelConfig ) -> None: - world, _ = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) urdf = ET.fromstring(world._scene.constructor_kwargs["urdf"]) srdf = ET.fromstring(world._scene.constructor_kwargs["srdf"]) @@ -2294,7 +1916,7 @@ def test_composed_model_fills_only_missing_acceleration_limits( authored.set("acceleration", "3.5") tree.write(robot_config.model_path) - world, _ = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) urdf = ET.fromstring(world._scene.constructor_kwargs["urdf"]) acceleration_by_joint = { @@ -2315,7 +1937,7 @@ def test_base_pose_is_written_to_composed_model( robot_config.base_pose = PoseStamped( # type: ignore[call-arg] position=Vector3(1, 0, 0), orientation=Quaternion() ) - world, _ = _make_world(fake_roboplan, robot_config) + world = _make_world(fake_roboplan, robot_config) urdf_root = ET.fromstring(world._scene.constructor_kwargs["urdf"]) origin = urdf_root.find("./joint[@name='dimos_world_joint']/origin") diff --git a/dimos/manipulation/test_roboplan_integration.py b/dimos/manipulation/test_roboplan_integration.py index f9d233f502..0bc0c1f603 100644 --- a/dimos/manipulation/test_roboplan_integration.py +++ b/dimos/manipulation/test_roboplan_integration.py @@ -35,7 +35,10 @@ from dimos.msgs.geometry_msgs.Transform import Transform from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.robot.manipulators.xarm.config import make_xarm6_model_config +from dimos.robot.manipulators.xarm.config import ( + make_dual_xarm6_model_config, + make_xarm6_model_config, +) from dimos.utils.transform_utils import pose_to_matrix pytestmark = pytest.mark.self_hosted @@ -51,11 +54,9 @@ def roboplan_types() -> tuple[type[Any], type[Any]]: def _sync_zero_state( world: Any, - robot_id: str, joint_names: list[str], ) -> None: world.sync_from_joint_state( - robot_id, JointState(name=joint_names, position=[0.0] * len(joint_names)), ) @@ -63,18 +64,17 @@ def _sync_zero_state( def test_real_roboplan_plans_fixed_orientation_cartesian_path( roboplan_types: tuple[type[Any], type[Any]], ) -> None: - config = make_xarm6_model_config(name="arm") + config = make_xarm6_model_config() if not Path(config.model_path).exists(): pytest.skip(f"xArm model is unavailable: {config.model_path}") world_type, planner_type = roboplan_types world = world_type() - robot_id = world.add_robot(config) + world.load_model(config) world.finalize() planner = planner_type(world, RoboPlanPlannerConfig()) - _sync_zero_state(world, robot_id, config.joint_names) - group_id = world._planning_groups.primary_pose_group_id_for_robot("arm") - assert group_id == "arm/manipulator" + _sync_zero_state(world, config.joint_names) + group_id = "manipulator" selection = world._planning_groups.select((group_id,)) start = JointState( name=list(selection.joint_names), position=[0.0] * len(selection.joint_names) @@ -116,23 +116,18 @@ def test_real_roboplan_plans_fixed_orientation_cartesian_path( def test_real_roboplan_synchronizes_different_length_dual_arm_targets( roboplan_types: tuple[type[Any], type[Any]], ) -> None: - left_config = make_xarm6_model_config(name="left_arm", y_offset=0.3) - right_config = make_xarm6_model_config(name="right_arm", y_offset=-0.3) - if not Path(left_config.model_path).exists(): - pytest.skip(f"xArm model is unavailable: {left_config.model_path}") + config = make_dual_xarm6_model_config() + if not Path(config.model_path).exists(): + pytest.skip(f"xArm model is unavailable: {config.model_path}") world_type, planner_type = roboplan_types world = world_type() - left_id = world.add_robot(left_config) - right_id = world.add_robot(right_config) + world.load_model(config) world.finalize() planner = planner_type(world, RoboPlanPlannerConfig()) - _sync_zero_state(world, left_id, left_config.joint_names) - _sync_zero_state(world, right_id, right_config.joint_names) - left_group_id = world._planning_groups.primary_pose_group_id_for_robot("left_arm") - right_group_id = world._planning_groups.primary_pose_group_id_for_robot("right_arm") - assert left_group_id == "left_arm/manipulator" - assert right_group_id == "right_arm/manipulator" + _sync_zero_state(world, config.joint_names) + left_group_id = "left_arm" + right_group_id = "right_arm" selection = world._planning_groups.select((left_group_id, right_group_id)) start = JointState( name=list(selection.joint_names), position=[0.0] * len(selection.joint_names) diff --git a/dimos/manipulation/visualization/operator.py b/dimos/manipulation/visualization/operator.py index ebb1b67070..860ae0dec5 100644 --- a/dimos/manipulation/visualization/operator.py +++ b/dimos/manipulation/visualization/operator.py @@ -23,7 +23,7 @@ from dimos.manipulation.planning.groups.models import PlanningGroup from dimos.manipulation.planning.planners.config import CartesianPathConfig -from dimos.manipulation.planning.spec.models import GeneratedPlan, PlanningGroupID, RobotName +from dimos.manipulation.planning.spec.models import GeneratedPlan, PlanningGroupID from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.JointState import JointState @@ -104,20 +104,20 @@ def set_motion_speed(self, speed_scale: float) -> bool: """Set the runtime speed reduction used for future plans.""" return self._module.set_motion_speed(speed_scale) - 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) + def get_init_joints(self) -> JointState | None: + """Return the operator-authoritative initialization state.""" + init = self._module.get_init_joints() 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.""" + """Validate and evaluate a canonical 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._invalid(request.group_ids, "Incomplete model target state") return self._evaluate_global_target(groups, JointState(request.target), complete) def evaluate_pose_target(self, request: PoseTargetRequest) -> TargetEvaluationResult: @@ -232,7 +232,7 @@ def _validate_joint_request( 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" + request.group_ids, "Joint target must use exact selected canonical joints in order" ) if any(not math.isfinite(value) for value in positions): return None, self._invalid( @@ -275,9 +275,9 @@ def _validate_pose_request( ): 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): + if seed_names != expected: return group_ids, self._invalid( - group_ids, "Seed must use exact selected global joints in order" + group_ids, "Seed must use exact selected canonical 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") @@ -296,43 +296,34 @@ def _groups_for_ids( def _complete_states( self, groups: Sequence[PlanningGroup], target: JointState - ) -> dict[RobotName, JointState] | None: + ) -> 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): + config = self._module.get_model_config() + baseline = self._world_monitor.get_current_joint_state() + if 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 name in config.joint_names: + value = values.get(name, baseline_values.get(name)) + if value is None: 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 + positions.append(value) + return JointState({"name": list(config.joint_names), "position": positions}) def _evaluate_global_target( self, groups: Sequence[PlanningGroup], target: JointState, - complete_states: Mapping[RobotName, JointState] | None = None, + complete_state: JointState | None = None, ) -> TargetEvaluationResult: - complete = complete_states or self._complete_states(groups, target) + complete = complete_state or self._complete_states(groups, target) if complete is None: return self._invalid( tuple(group.id for group in groups), "Incomplete robot target state" @@ -340,26 +331,23 @@ def _evaluate_global_target( diagnostics: dict[PlanningGroupID, str] = {} poses: dict[PlanningGroupID, PoseStamped | None] = {} valid = True + state_valid = self._world_monitor.is_state_valid(complete) 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) - ) + group_valid = state_valid valid = valid and group_valid diagnostics[group.id] = ( - "Target is collision-free for this robot" + "Target is collision-free" 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) + poses[group.id] = self._world_monitor.get_group_ee_pose(group.id, complete) 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", + message="Target is collision-free" if valid else "Target is infeasible", collision_free=valid, group_ids=tuple(group.id for group in groups), target_joints=JointState(target), diff --git a/dimos/manipulation/visualization/test_factory.py b/dimos/manipulation/visualization/test_factory.py index 9f952f0deb..fc915911e1 100644 --- a/dimos/manipulation/visualization/test_factory.py +++ b/dimos/manipulation/visualization/test_factory.py @@ -32,7 +32,6 @@ PlanningSceneInfo, VisualizationSession, VisualizationStateFrame, - WorldRobotID, ) from dimos.manipulation.planning.spec.protocols import VisualizationSpec from dimos.manipulation.planning.world.drake_world import DRAKE_AVAILABLE, DrakeWorld @@ -85,28 +84,25 @@ def clear_vis_obstacles(self) -> None: class FakeWorld: - def add_robot(self, config: RobotModelConfig) -> WorldRobotID: - return "robot-1" - - def get_robot_ids(self) -> list[WorldRobotID]: - return [] + def load_model(self, config: RobotModelConfig) -> None: + return None - def get_robot_config(self, robot_id: WorldRobotID) -> RobotModelConfig: + def get_model_config(self) -> RobotModelConfig: return RobotModelConfig( - name="fake", model_path=Path("fake.urdf"), base_pose=PoseStamped(), - joint_names=[], + joint_names=["joint1"], planning_groups=[ PlanningGroupDefinition( - name="manipulator", joint_names=(), base_link="base_link", tip_link="ee_link" + name="manipulator", + joint_names=("joint1",), + base_link="base_link", + tip_link="ee_link", ) ], ) - def get_joint_limits( - self, robot_id: WorldRobotID - ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: + def get_joint_limits(self) -> tuple[NDArray[np.float64], NDArray[np.float64]]: return (np.array([], dtype=np.float64), np.array([], dtype=np.float64)) def add_obstacle(self, obstacle: Obstacle) -> str | None: @@ -140,42 +136,39 @@ def get_live_context(self) -> object: def scratch_context(self) -> AbstractContextManager[object | None]: return nullcontext(None) - def sync_from_joint_state(self, robot_id: WorldRobotID, joint_state: JointState) -> None: + def sync_from_joint_state(self, joint_state: JointState) -> None: return None - def set_joint_state(self, ctx: object, robot_id: WorldRobotID, joint_state: JointState) -> None: + def set_joint_state(self, ctx: object, joint_state: JointState) -> None: return None - def get_joint_state(self, ctx: object, robot_id: WorldRobotID) -> JointState: + def get_joint_state(self, ctx: object) -> JointState: return JointState({}) - def is_collision_free(self, ctx: object, robot_id: WorldRobotID) -> bool: + def is_collision_free(self, ctx: object) -> bool: return True - def get_min_distance(self, ctx: object, robot_id: WorldRobotID) -> float: + def get_min_distance(self, ctx: object) -> float: return 0.0 - def check_config_collision_free(self, robot_id: WorldRobotID, joint_state: JointState) -> bool: + def check_config_collision_free(self, joint_state: JointState) -> bool: return True def check_edge_collision_free( self, - robot_id: WorldRobotID, start: JointState, end: JointState, step_size: float = 0.05, ) -> bool: return True - def get_ee_pose(self, ctx: object, robot_id: WorldRobotID) -> PoseStamped: + def get_ee_pose(self, ctx: object) -> PoseStamped: return PoseStamped() - def get_link_pose( - self, ctx: object, robot_id: WorldRobotID, link_name: str - ) -> NDArray[np.float64]: + def get_link_pose(self, ctx: object, link_name: str) -> NDArray[np.float64]: return np.eye(4, dtype=np.float64) - def get_jacobian(self, ctx: object, robot_id: WorldRobotID) -> NDArray[np.float64]: + def get_jacobian(self, ctx: object) -> NDArray[np.float64]: return np.zeros((6, 0), dtype=np.float64) def get_group_ee_pose(self, ctx: object, group_id: str) -> PoseStamped: @@ -216,8 +209,8 @@ def animate_trajectory( ) -> None: self.visualization_calls.append(("animate_trajectory", trajectory, duration)) - def cancel_preview_animation(self, robot_ids: tuple[WorldRobotID, ...] | None = None) -> None: - self.visualization_calls.append(("cancel_preview_animation", robot_ids)) + def cancel_preview_animation(self) -> None: + self.visualization_calls.append(("cancel_preview_animation",)) def close(self) -> None: self.visualization_calls.append(("close",)) @@ -239,7 +232,7 @@ def clear_vis_obstacles(self) -> None: def test_config_defaults_to_no_visualization() -> None: - config = ManipulationModuleConfig() + config = ManipulationModuleConfig(model=FakeWorld().get_model_config()) assert isinstance(config.visualization, NoManipulationVisualizationConfig) assert config.visualization.requires_world_visualization is False @@ -247,18 +240,21 @@ def test_config_defaults_to_no_visualization() -> None: def test_config_rejects_unknown_visualization_backend() -> None: with pytest.raises(ValidationError, match="visualization"): - ManipulationModuleConfig.model_validate({"visualization": {"backend": "bad"}}) + ManipulationModuleConfig.model_validate( + {"model": FakeWorld().get_model_config(), "visualization": {"backend": "bad"}} + ) def test_config_validates_viser_visualization() -> None: config = ManipulationModuleConfig.model_validate( { + "model": FakeWorld().get_model_config(), "visualization": { "backend": "viser", "visualization_host": "0.0.0.0", "visualization_port": "8096", "viser_panel_enabled": "false", - } + }, }, ) @@ -269,7 +265,12 @@ def test_config_validates_viser_visualization() -> None: def test_config_meshcat_requires_world_visualization() -> None: - config = ManipulationModuleConfig.model_validate({"visualization": {"backend": "meshcat"}}) + config = ManipulationModuleConfig.model_validate( + { + "model": FakeWorld().get_model_config(), + "visualization": {"backend": "meshcat"}, + } + ) assert isinstance(config.visualization, MeshcatVisualizationConfig) assert config.visualization.requires_world_visualization is True @@ -298,8 +299,10 @@ def test_create_visualization_meshcat_accepts_structural_world() -> None: ) assert visualization is fake_world # type: ignore[comparison-overlap] assert isinstance(visualization, VisualizationSpec) - session = VisualizationSession(PlanningSceneInfo(robots={}), operator=object()) - frame = VisualizationStateFrame(joint_states={}) + session = VisualizationSession( + PlanningSceneInfo(model=fake_world.get_model_config()), operator=object() + ) + frame = VisualizationStateFrame(joint_state=None) trajectory = JointTrajectory(joint_names=["arm/j1"], points=[]) obstacle = Obstacle( name="box", @@ -320,7 +323,7 @@ def test_create_visualization_meshcat_accepts_structural_world() -> None: ("initialize", session), ("get_visualization_url",), ("update_state", frame), - ("cancel_preview_animation", None), + ("cancel_preview_animation",), ("animate_trajectory", trajectory, 2.5), ("close",), ("add_vis_obstacle", "box", obstacle), @@ -359,8 +362,12 @@ def test_drake_meshcat_visualization_lifecycle_is_noop_without_meshcat() -> None assert visualization is world assert isinstance(visualization, VisualizationSpec) assert world.get_visualization_url() is None - world.initialize(VisualizationSession(PlanningSceneInfo(robots={}), operator=object())) - world.update_state(VisualizationStateFrame(joint_states={})) + world.initialize( + VisualizationSession( + PlanningSceneInfo(model=FakeWorld().get_model_config()), operator=object() + ) + ) + world.update_state(VisualizationStateFrame(joint_state=None)) obstacle = Obstacle( name="box", obstacle_type=ObstacleType.BOX, diff --git a/dimos/manipulation/visualization/test_operator.py b/dimos/manipulation/visualization/test_operator.py index 84f3236071..bd0dcb60db 100644 --- a/dimos/manipulation/visualization/test_operator.py +++ b/dimos/manipulation/visualization/test_operator.py @@ -12,22 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Focused tests for the manipulation visualization operator facade.""" +"""Focused tests for the single-model visualization operator.""" from pathlib import Path +from unittest.mock import MagicMock from dimos.agents.skill_result import SkillResult -from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupDefinition +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry from dimos.manipulation.planning.planners.roboplan_config import RoboPlanCartesianPathConfig 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.planning.spec.models import GeneratedPlan, IKResult from dimos.manipulation.visualization.operator import ( CartesianTargetRequest, JointTargetRequest, @@ -35,411 +31,144 @@ 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 - ] +def _config() -> RobotModelConfig: return RobotModelConfig( - name=name, - model_path=Path("/robot.urdf"), - base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), - joint_names=joints, + model_path=Path("/model.urdf"), + joint_names=["left/j1", "left/j2", "right/j1"], base_link="base", - planning_groups=definitions, + planning_groups=[ + PlanningGroupDefinition("left_arm", ("left/j1", "left/j2"), "base", "left/tool"), + PlanningGroupDefinition("right_arm", ("right/j1",), "base", "right/tool"), + ], ) -class FakeModule: - def __init__(self) -> None: - self.state = "COMPLETED" - 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.cartesian_targets: list[ - tuple[ - dict[PlanningGroupID, tuple[PoseStamped, ...]], - RoboPlanCartesianPathConfig, - 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_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 generate_cartesian_plan( - self, - targets: dict[PlanningGroupID, tuple[PoseStamped, ...]], - config: RoboPlanCartesianPathConfig, - auxiliary_groups: tuple[PlanningGroupID, ...] = (), - ) -> GeneratedPlan | None: - self.cartesian_targets.append((targets, config, 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 _plan() -> GeneratedPlan: + trajectory = JointTrajectory( + joint_names=["left/j1", "left/j2"], + points=[TrajectoryPoint(positions=[0.4, 0.5], time_from_start=1.0)], + ) + return GeneratedPlan( + group_ids=("left_arm",), + trajectory=trajectory, + path=[JointState(name=trajectory.joint_names, position=[0.4, 0.5])], + status=PlanningStatus.SUCCESS, ) -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 == "COMPLETED" - assert status.error == "" - assert status.has_plan is True - assert module.topology_calls == 0 - assert module.telemetry_calls == 0 - assert monitor.telemetry_calls == 0 +def _operator() -> tuple[ManipulationOperator, MagicMock, MagicMock]: + config = _config() + module = MagicMock() + module.get_state.return_value = "COMPLETED" + module.get_error.return_value = "" + module.has_planned_path.return_value = True + module.get_model_config.return_value = config + module.get_init_joints.return_value = None + module.inverse_kinematics.return_value = IKResult( + status=IKStatus.SUCCESS, + joint_state=JointState(name=["left/j1", "left/j2"], position=[0.4, 0.5]), + ) + module.generate_plan_to_joint_targets.return_value = _plan() + module.generate_plan_to_pose_targets.return_value = _plan() + module.generate_cartesian_plan.return_value = _plan() + module.preview_plan.return_value = True + module.execute_plan.return_value = True + module.cancel.return_value = True + module.clear_planned_path.return_value = True + module.reset.return_value = SkillResult.ok("reset") + + monitor = MagicMock() + monitor.planning_groups = PlanningGroupRegistry([config]) + monitor.get_current_joint_state.return_value = JointState( + name=config.joint_names, position=[0.0, 0.0, 0.0] + ) + monitor.is_state_valid.return_value = True + monitor.get_group_ee_pose.return_value = PoseStamped(frame_id="world") + return ManipulationOperator(module, monitor), module, monitor -def test_evaluate_joint_target_accepts_exact_global_selection_domain() -> None: - operator, _, _ = _operator() +def test_status_is_compact() -> None: + operator, module, _ = _operator() + assert operator.status().state == "COMPLETED" + module.get_model_config.assert_not_called() - 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_evaluation_overlays_selected_target_on_complete_model_state() -> None: + operator, _, monitor = _operator() + request = JointTargetRequest( + ("left_arm",), + JointState(name=["left/j1", "left/j2"], position=[0.1, 0.2]), + ) + result = operator.evaluate_joint_target(request) + assert result.success + complete = monitor.is_state_valid.call_args.args[0] + assert complete.name == ["left/j1", "left/j2", "right/j1"] + assert complete.position == [0.1, 0.2, 0.0] -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]), - ), - ] +def test_joint_evaluation_rejects_local_unknown_and_overlapping_selection() -> None: 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"), + local = JointTargetRequest(("left_arm",), JointState(name=["j1", "j2"], position=[0.1, 0.2])) + unknown = JointTargetRequest(("missing",), JointState(name=["left/j1"], position=[0.1])) + duplicate = JointTargetRequest( + ("left_arm", "left_arm"), + JointState(name=["left/j1", "left/j2"], position=[0.1, 0.2]), ) - operator, _, _ = _operator(_robot_config(groups=groups)) - request = JointTargetRequest( - ("arm/first", "arm/second"), JointState(name=["arm/j0", "arm/j0"], position=[0.1, 0.2]) + assert all( + not operator.evaluate_joint_target(request).success + for request in (local, unknown, duplicate) ) - 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: +def test_pose_evaluation_routes_group_id_and_canonical_seed() -> 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)] + pose = PoseStamped(frame_id="world") + seed = JointState(name=["left/j1", "left/j2"], position=[0.0, 0.0]) + result = operator.evaluate_pose_target(PoseTargetRequest({"left_arm": pose}, seed=seed)) + assert result.success + module.inverse_kinematics.assert_called_once_with( + pose_targets={"left_arm": pose}, auxiliary_group_ids=(), seed=seed, check_collision=True + ) -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]), - ), - ] +def test_pose_evaluation_rejects_non_world_frame_and_ambiguous_local_seed() -> None: 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" + bad_frame = PoseTargetRequest({"left_arm": PoseStamped(frame_id="camera")}) + bad_seed = PoseTargetRequest( + {"left_arm": PoseStamped(frame_id="world")}, + seed=JointState(name=["j1", "j2"], position=[0.0, 0.0]), + ) + assert not operator.evaluate_pose_target(bad_frame).success + assert not operator.evaluate_pose_target(bad_seed).success -def test_planning_methods_return_exact_generated_plan() -> None: +def test_planning_and_actions_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_cartesian_planning_prepends_current_pose_and_routes_auxiliary_groups() -> None: - groups = ( - PlanningGroup( - "arm/manipulator", - "arm", - "manipulator", - ("arm/j0",), - ("j0",), - "base", - "tool", - ), - PlanningGroup( - "arm/gripper", - "arm", - "gripper", - ("arm/j1",), - ("j1",), - "tool", - None, - ), + joint_request = JointTargetRequest( + ("left_arm",), JointState(name=["left/j1", "left/j2"], position=[0.1, 0.2]) ) - operator, module, _ = _operator(_robot_config(groups=groups)) - pose = _pose() - config = RoboPlanCartesianPathConfig(speed_mode="time_optimal") - - result = operator.plan_cartesian( - CartesianTargetRequest( - {"arm/manipulator": pose}, - config, - ("arm/gripper",), - ) + pose_request = PoseTargetRequest({"left_arm": PoseStamped(frame_id="world")}) + assert ( + operator.plan_to_joints(joint_request) is module.generate_plan_to_joint_targets.return_value ) + assert operator.plan_to_pose(pose_request) is module.generate_plan_to_pose_targets.return_value + assert operator.preview(_plan(), 0.5) + assert operator.execute(_plan()) + assert operator.cancel() + assert operator.clear_plan() + assert operator.reset() - assert result is module.plan - assert len(module.cartesian_targets) == 1 - targets, recorded_config, auxiliary_groups = module.cartesian_targets[0] - assert targets["arm/manipulator"][0].position == Vector3(1.0, 2.0, 3.0) - assert targets["arm/manipulator"][1] is pose - assert recorded_config is config - assert auxiliary_groups == ("arm/gripper",) - -def test_actions_return_typed_results_and_cancel_fallback_ownership() -> None: +def test_cartesian_planning_uses_current_group_pose() -> 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 + target = PoseStamped(frame_id="world") + config = RoboPlanCartesianPathConfig() + result = operator.plan_cartesian(CartesianTargetRequest({"left_arm": target}, config)) + assert result is module.generate_cartesian_plan.return_value + targets = module.generate_cartesian_plan.call_args.args[0] + assert targets["left_arm"] == (monitor.get_group_ee_pose.return_value, target) diff --git a/dimos/manipulation/visualization/types.py b/dimos/manipulation/visualization/types.py index 90377a4adf..a699b00687 100644 --- a/dimos/manipulation/visualization/types.py +++ b/dimos/manipulation/visualization/types.py @@ -16,7 +16,6 @@ from typing import TypedDict -from dimos.manipulation.planning.spec.models import RobotName, WorldRobotID from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.JointState import JointState @@ -33,15 +32,12 @@ class TargetEvaluation(TypedDict, total=False): orientation_error: float -class RobotInfo(TypedDict): - name: RobotName - world_robot_id: WorldRobotID +class ModelInfo(TypedDict): joint_names: list[str] end_effector_link: str | None base_link: str max_velocity: float max_acceleration: float - has_joint_name_mapping: bool home_joints: list[float] | None pre_grasp_offset: float init_joints: list[float] | None diff --git a/dimos/manipulation/visualization/viser/animation.py b/dimos/manipulation/visualization/viser/animation.py index 6754343e21..862a873081 100644 --- a/dimos/manipulation/visualization/viser/animation.py +++ b/dimos/manipulation/visualization/viser/animation.py @@ -21,28 +21,20 @@ @dataclass(frozen=True) class PreviewFrame: - """One timestamped local robot preview frame.""" + """One timestamped model preview frame.""" time_from_start: float positions: tuple[float, ...] @dataclass(frozen=True) -class PreviewTrack: - """One fixed-baseline local robot track in a group-native preview.""" +class PreviewAnimation: + """One fixed-baseline canonical model preview.""" - robot_id: str joint_names: tuple[str, ...] frames: tuple[PreviewFrame, ...] -@dataclass(frozen=True) -class GroupPreviewAnimation: - """Validated collection of robot tracks sharing one preview transaction.""" - - tracks: tuple[PreviewTrack, ...] - - 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: @@ -55,8 +47,6 @@ def scaled_frame_delays(frames: Sequence[PreviewFrame], duration: float) -> tupl ) -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}) - ) +def preview_tick_times(preview: PreviewAnimation) -> tuple[float, ...]: + """Return stored preview timestamps without synthesizing extra samples.""" + return tuple(float(frame.time_from_start) for frame in preview.frames) diff --git a/dimos/manipulation/visualization/viser/gui.py b/dimos/manipulation/visualization/viser/gui.py index 9147bdb0c1..596657149f 100644 --- a/dimos/manipulation/visualization/viser/gui.py +++ b/dimos/manipulation/visualization/viser/gui.py @@ -14,14 +14,14 @@ from __future__ import annotations -from collections.abc import Mapping, MutableMapping, Sequence +from collections.abc import Callable, Mapping, Sequence import math from typing import TypeAlias, cast from dimos.manipulation.planning.groups.models import PlanningGroup from dimos.manipulation.planning.planners.roboplan_config import RoboPlanCartesianPathConfig from dimos.manipulation.planning.spec.config import RobotModelConfig -from dimos.manipulation.planning.spec.models import PlanningGroupID, PlanningSceneInfo, RobotName +from dimos.manipulation.planning.spec.models import PlanningGroupID, PlanningSceneInfo from dimos.manipulation.visualization.operator import ( CartesianTargetRequest, JointTargetRequest, @@ -91,11 +91,7 @@ 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}" - ) + return str(group.id) def _copy_joint_state(state: JointState | None) -> JointState | None: @@ -122,17 +118,14 @@ def __init__( server: ViserServer, scene_info: PlanningSceneInfo, operator: ManipulationOperator | object, - current_states: MutableMapping[str, JointState], + current_state: Callable[[], JointState | None], config: ViserVisualizationConfig, scene: ViserManipulationScene | None = None, ) -> None: self.server = server 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._current_state = current_state self._scene_groups_by_id = {group.id: group for group in scene_info.planning_groups} self.config = config self.scene = scene @@ -179,46 +172,17 @@ 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_model_config(self) -> RobotModelConfig: + return self.scene_info.model - 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_init_joints(self) -> JointState | None: + return _copy_joint_state(self.operator.get_init_joints()) - 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_current_joint_state(self) -> JointState | None: + return _copy_joint_state(self._current_state()) def get_group_ee_pose(self, group_id: PlanningGroupID) -> PoseStamped | None: group = self._scene_groups_by_id.get(group_id) @@ -230,23 +194,23 @@ def get_group_ee_pose(self, group_id: PlanningGroupID) -> PoseStamped | 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) + current = self.get_current_joint_state() 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): + values = self._state_values_by_name(current) + if any(name not in values for name in group.joint_names): return {} return { group.id: JointState( { "name": list(group.joint_names), - "position": [values[name] for name in group.local_joint_names], + "position": [values[name] for name in group.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 is_state_stale(self, max_age: float = 1.0) -> bool: + return self.get_current_joint_state() is None def get_module_state(self) -> str: return self.operator.status().state @@ -351,15 +315,11 @@ def execute(self) -> bool: def refresh(self) -> None: if self._closed: return - robots = self.list_robots() groups = self.list_planning_groups() - self.state.backend_status = ( - BackendConnectionStatus.READY if robots else BackendConnectionStatus.WAITING_FOR_ROBOT - ) + self.state.backend_status = BackendConnectionStatus.READY 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) @@ -367,7 +327,7 @@ def refresh(self) -> None: if set(self.state.group_joint_targets) != initialized_groups: self._build_joint_sliders() self._sync_group_selector(groups) - self._refresh_selected_robot_state() + self._refresh_model_state() self._ensure_scene_controls() self._sync_target_ghost_visibility() self._sync_robot_display_dropdown() @@ -511,8 +471,6 @@ def _toggle_group_selected(self, group_id: str) -> None: 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() @@ -600,13 +558,8 @@ 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: - robot_name = self.state.selected_robot - if robot_name is None: - self.state.current_joints = None - self.state.manipulation_state = self.get_module_state() - return - current = self.get_current_joint_state(robot_name) + def _refresh_model_state(self) -> None: + current = self.get_current_joint_state() self.state.current_joints = list(current.position) if current is not None else None self.state.manipulation_state = self.get_module_state() adapter_error = self.get_error() @@ -667,13 +620,13 @@ def _build_joint_slider_handles(self, gui: GuiApi) -> None: group = self._groups_by_id().get(group_id) if group is None: continue - config = self.get_robot_config(group.robot_name) + config = self.get_model_config() 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 + group.joint_names, group.joint_names, target.position, strict=True ): index = config_indexes.get(str(local_name)) lower, upper = DEFAULT_JOINT_LIMITS @@ -711,49 +664,17 @@ def _clear_joint_sliders(self) -> None: 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 _stale_models(self, group_ids: tuple[PlanningGroupID, ...]) -> tuple[str, ...]: + """Return the configured model label when its state is unavailable.""" + return ("model",) if group_ids and self.is_state_stale() else () - def _state_values_by_local_name(self, state: JointState | None) -> dict[str, float]: + def _state_values_by_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: @@ -761,17 +682,15 @@ def _initialize_selected_group_targets(self) -> None: group = self._groups_by_id().get(group_id) if group is None: continue - if self.is_state_stale(group.robot_name): + if self.is_state_stale(): 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): + values = self._state_values_by_name(self.get_current_joint_state()) + if any(str(name) not in values for name in group.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], + "position": [float(values[str(name)]) for name in group.joint_names], } ) if group.has_pose_target and group_id not in self.state.pose_targets: @@ -814,20 +733,18 @@ def _active_pose_targets(self) -> dict[PlanningGroupID, Pose]: if group_id in self.state.pose_targets } - def _preset_values_by_local_name(self, preset: str, robot_name: str) -> dict[str, float]: + def _preset_values_by_name(self, preset: str) -> dict[str, float]: if preset == "Current": - state = self.get_current_joint_state(robot_name) + state = self.get_current_joint_state() elif preset == "Init": - state = self.get_init_joints(robot_name) + state = self.get_init_joints() else: - config = self.get_robot_config(robot_name) - if config is None: - return {} + config = self.get_model_config() 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) + return self._state_values_by_name(state) def _remove_panel_handles(self) -> None: for key, handle in list(self._handles.items()): @@ -841,15 +758,10 @@ def _sync_preset_dropdown(self) -> None: if handle is None or not self.state.selected_group_ids: return options = ["Select preset..."] - selected_robots = self._selected_robot_names() - if any(self.get_init_joints(robot_name) is not None for robot_name in selected_robots): + if self.get_init_joints() is not None: options.append("Init") options.append("Current") - 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 - ): + if self.get_model_config().home_joints is not None: options.append("Home") for attr in ("options", "values"): if hasattr(handle, attr): @@ -870,21 +782,19 @@ def _apply_preset(self, preset: str) -> None: 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}" - ) + if preset == "Current" and self.is_state_stale(): + self._set_recoverable_error("Cannot apply Current preset without fresh telemetry") 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] + values = self._preset_values_by_name(preset) + missing = [str(name) for name in group.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] + positions = [float(values[str(name)]) for name in group.joint_names] targets[group_id] = JointState({"name": list(group.joint_names), "position": positions}) - slider_values.append((group_id, group.local_joint_names, positions)) + slider_values.append((group_id, group.joint_names, positions)) self.state.group_joint_targets.update(targets) if any( (group_id, str(local_name)) not in self._joint_sliders @@ -918,7 +828,7 @@ def _target_set_from_sliders(self) -> dict[PlanningGroupID, JointState] | None: self._set_error(f"Unknown planning group: {group_id}") return None positions: list[float] = [] - for local_name in group.local_joint_names: + for local_name in group.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}") @@ -991,60 +901,41 @@ def _move_joint_target_visuals(self, targets: Mapping[PlanningGroupID, JointStat """Optimistically move target visuals before collision/feasibility returns.""" 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) + state = self._target_ghost_state(targets) + if state is not None: + config = self.get_model_config() + self.scene.set_target_joints("model", config.joint_names, state.position) - def _target_ghost_states( + def _target_ghost_state( self, targets: Mapping[PlanningGroupID, JointState] - ) -> dict[str, JointState]: + ) -> JointState | None: groups = self._groups_by_id() - merged: dict[str, dict[str, float]] = {} - configs: dict[str, tuple[str, ...]] = {} + config = self.get_model_config() + current = self.get_current_joint_state() + values = self._state_values_by_name(current) 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],) - } + values.update(self._state_values_by_name(target)) + if not all(name in values for name in config.joint_names): + return None + return JointState( + { + "name": list(config.joint_names), + "position": [values[name] for name in config.joint_names], + } + ) def _sync_target_ghost_visibility(self) -> None: if self.scene is None: return - active_robot_ids = { - str(robot_id) + active = any( + (group := self._groups_by_id().get(group_id)) is not None and group.has_pose_target 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) + ) + self.scene.set_target_active("model", active) def _handle_target_evaluation_request( self, request: TargetEvaluationRequest @@ -1097,9 +988,7 @@ def _sync_controls_from_targets(self) -> None: 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) - ) + self._set_group_slider_values(group_id, group.joint_names, list(target.position)) self._move_joint_target_visuals(self.state.group_joint_targets) def _split_target_joints_by_group(self, target_joints: JointState) -> None: @@ -1148,9 +1037,9 @@ def _update_status_text(self) -> None: f"**State:** {status_label}", f"Target: `{self.state.target_status.value}` · Plan: `{self.state.plan_state.status.value}`", ] - stale_robots = self._stale_robot_names(self.state.selected_group_ids) + stale_models = self._stale_models(self.state.selected_group_ids) if self.state.selected_group_ids: - stale_detail = "False" if not stale_robots else f"True ({', '.join(stale_robots)})" + stale_detail = "False" if not stale_models else "True" status.append(f"State stale: `{stale_detail}`") if current is not None: status.append(f"Current joints: `{[round(v, 3) for v in current]}`") @@ -1188,15 +1077,8 @@ def _update_target_visual_state(self) -> None: 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) + if selected_groups: + self.scene.set_target_robot_visual_state("model", feasible) def _can_execute(self) -> bool: return self.state.can_execute() @@ -1243,8 +1125,8 @@ def operation() -> None: return self.state.action_status = ActionStatus.RUNNING self.state.plan_state.status = PlanStatus.PLANNING - stale_robots = self._stale_robot_names(group_ids) - if stale_robots: + stale_models = self._stale_models(group_ids) + if stale_models: if not self._operation_is_current( operation_id, selection_epoch, target_sequence_id ): @@ -1253,9 +1135,7 @@ def operation() -> None: ) return self.state.plan_state.status = PlanStatus.STALE - self.state.error = "Cannot plan without fresh telemetry for: " + ", ".join( - stale_robots - ) + self.state.error = "Cannot plan without fresh telemetry" self._finish_operation( "plan=False", clear_error=False, diff --git a/dimos/manipulation/visualization/viser/scene.py b/dimos/manipulation/visualization/viser/scene.py index 092ba61a92..e737fa0331 100644 --- a/dimos/manipulation/visualization/viser/scene.py +++ b/dimos/manipulation/visualization/viser/scene.py @@ -38,9 +38,7 @@ from dimos.manipulation.planning.spec.models import DEFAULT_OBSTACLE_RGBA, Obstacle from dimos.manipulation.planning.utils.mesh_utils import prepare_urdf_for_drake from dimos.manipulation.visualization.viser.animation import ( - GroupPreviewAnimation, - PreviewFrame, - preview_tick_times, + PreviewAnimation, scaled_frame_delays, ) from dimos.manipulation.visualization.viser.runtime import ( @@ -389,16 +387,16 @@ def robot_display_mode(self, mode: RobotDisplayMode | str) -> None: except ValueError as error: raise ValueError(f"Unsupported robot display mode: {mode!r}") from error self._robot_display_mode = normalized_mode - for robot_id in self._configs_by_id: - self._apply_robot_display_mode(robot_id) + for model_key in self._configs_by_id: + self._apply_robot_display_mode(model_key) @property def collision_geometry_available(self) -> bool: """Return whether any primary robot has loaded collision geometry.""" return any( - self._model_has_collision_geometry(self._models_by_id[robot_id]) - for robot_id in self._configs_by_id - if f"{robot_id}:current" in self._urdfs + self._model_has_collision_geometry(self._models_by_id[model_key]) + for model_key in self._configs_by_id + if f"{model_key}:current" in self._urdfs ) @staticmethod @@ -424,22 +422,26 @@ def set_reference_grid_visible(self, visible: bool) -> None: self._grid_visible = visible self._set_handle_visibility(self._grid_handle, visible) - 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) - if config.model_path and robot_id not in self._models_by_id: - self._models_by_id[robot_id] = self._load_robot_model(config) - self._ensure_robot_urdfs(robot_id, config) - - def set_target_active(self, robot_id: str, active: bool) -> None: + def _register_model_key(self, model_key: str, config: RobotModelConfig) -> None: + self._configs_by_id[model_key] = config + self._preview_visible.setdefault(model_key, False) + self._animation_generations.setdefault(model_key, 0) + self._target_active.setdefault(model_key, False) + self._target_tracks_current.setdefault(model_key, True) + if config.model_path and model_key not in self._models_by_id: + self._models_by_id[model_key] = self._load_robot_model(config) + self._ensure_robot_urdfs(model_key, config) + + def register_model(self, config: RobotModelConfig) -> None: + """Register the one configured model.""" + self._register_model_key("model", config) + + def set_target_active(self, model_key: str, active: bool) -> None: """Show the target ghost only while a pose-target group is selected.""" - self._target_active[robot_id] = active + self._target_active[model_key] = active if not active: - self._target_tracks_current[robot_id] = True - self._set_target_visibility(robot_id, active) + self._target_tracks_current[model_key] = True + self._set_target_visibility(model_key, active) def _ensure_reference_grid(self) -> None: try: @@ -471,13 +473,13 @@ def _ensure_reference_grid(self) -> None: self._grid_handle = None def ensure_target_controls( - self, robot_id: str, on_update: Callable[[TransformControlsHandle], None] + self, model_key: str, on_update: Callable[[TransformControlsHandle], None] ) -> TransformControlsHandle | None: - handle_key = f"{robot_id}:ee_control" + handle_key = f"{model_key}:ee_control" if handle_key in self._handles: return self._handles[handle_key] handle = self.server.scene.add_transform_controls( - f"/targets/{robot_id}/ee_control", scale=0.25 + f"/targets/{model_key}/ee_control", scale=0.25 ) def dispatch(event: TransformControlsEvent) -> None: @@ -490,136 +492,100 @@ def dispatch(event: TransformControlsEvent) -> None: 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: + def _update_current_model_key(self, model_key: str, joint_state: JointState | None) -> None: with self._scene_lock: - config = self._configs_by_id.get(robot_id) + config = self._configs_by_id.get(model_key) 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._ensure_robot_urdfs(model_key, config) + current = self._urdfs.get(f"{model_key}: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)) + if self._target_tracks_current.get(model_key, True): + self._set_target_joints(model_key, config.joint_names, joint_state.position) + self._set_target_visibility(model_key, self._target_active.get(model_key, False)) self.set_urdf_joints( - self._collision_fallback_urdfs.get(robot_id), + self._collision_fallback_urdfs.get(model_key), config.joint_names, joint_state.position, ) - def cancel_preview_animation(self, robot_ids: Sequence[str] | None = None) -> None: + def update_current_model(self, joint_state: JointState | None) -> None: + """Update the one configured model's canonical joint state.""" + self._update_current_model_key("model", joint_state) + + def cancel_preview_animation(self) -> 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 + for model_key in set(self._preview_visible): + self._animation_generations[model_key] = ( + self._animation_generations.get(model_key, 0) + 1 ) - if robot_id not in self._preview_visible: + if model_key not in self._preview_visible: continue - self._preview_visible[robot_id] = False - self._set_preview_visibility(robot_id, False) + self._preview_visible[model_key] = False + self._set_preview_visibility(model_key, False) - def animate_preview(self, preview: GroupPreviewAnimation, duration: float) -> bool: - """Play every robot from one normalized tick clock. + def animate_preview(self, preview: PreviewAnimation, duration: float) -> bool: + """Play the model preview from one normalized tick clock. Inputs are fully validated before ghosts become visible; a generation replacement, clear, or close stops mutation before the next tick. """ - 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 - tick_times = preview_tick_times(preview) - if not tick_times: + model_key = "model" + if not preview.frames or model_key not in self._configs_by_id: 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) + generation = self._animation_generations.get(model_key, 0) + 1 + self._animation_generations[model_key] = generation + self._preview_visible[model_key] = True + self._set_preview_visibility(model_key, True) try: - 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): + delays = scaled_frame_delays(preview.frames, duration) + for index, frame in enumerate(preview.frames): 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: + if self._animation_generations.get(model_key) != generation: 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]) + self._set_preview_ghost_joints(model_key, preview.joint_names, frame.positions) + if index < len(delays): + time.sleep(delays[index]) return True finally: 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) + if self._animation_generations.get(model_key) == generation: + self._preview_visible[model_key] = False + self._set_preview_visibility(model_key, False) def set_target_joints( - self, robot_id: str, joint_names: Sequence[str], joints: Sequence[float] + self, model_key: str, joint_names: Sequence[str], joints: Sequence[float] ) -> bool: - target = self._urdfs.get(f"{robot_id}:target") + target = self._urdfs.get(f"{model_key}:target") if target is None: return False - self._target_tracks_current[robot_id] = False - self._set_target_joints(robot_id, joint_names, joints) - self._set_target_visibility(robot_id, True) + self._target_tracks_current[model_key] = False + self._set_target_joints(model_key, joint_names, joints) + self._set_target_visibility(model_key, True) return True - def clear_target(self, robot_id: str) -> None: + def clear_target(self, model_key: str) -> None: """Return the persistent target ghost to current-state tracking.""" - self._target_tracks_current[robot_id] = True + self._target_tracks_current[model_key] = True def _set_target_joints( - self, robot_id: str, joint_names: Sequence[str], joints: Sequence[float] + self, model_key: str, joint_names: Sequence[str], joints: Sequence[float] ) -> None: - target = self._urdfs.get(f"{robot_id}:target") + target = self._urdfs.get(f"{model_key}:target") self.set_urdf_joints(target, joint_names, joints) def _set_preview_ghost_joints( - self, robot_id: str, joint_names: Sequence[str], joints: Sequence[float] + self, model_key: str, joint_names: Sequence[str], joints: Sequence[float] ) -> None: - ghost = self._urdfs.get(f"{robot_id}:preview") + ghost = self._urdfs.get(f"{model_key}:preview") self.set_urdf_joints(ghost, joint_names, joints) - def set_target_pose(self, robot_id: str, pose: Pose | None) -> None: - handle = self._handles.get(f"{robot_id}:ee_control") + def set_target_pose(self, model_key: str, pose: Pose | None) -> None: + handle = self._handles.get(f"{model_key}:ee_control") if handle is None or pose is None: return handle.position = ( @@ -634,10 +600,10 @@ def set_target_pose(self, robot_id: str, pose: Pose | None) -> None: float(pose.orientation.z), ) - 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_visual_state(self, model_key: str, feasible: bool) -> None: + """Set the matching model and control target visual state.""" + self.set_target_control_visual_state(model_key, feasible) + self.set_target_robot_visual_state(model_key, feasible) def set_target_control_visual_state(self, control_id: str, feasible: bool) -> None: """Set feasibility color for one planning-group keyed target control.""" @@ -646,11 +612,11 @@ def set_target_control_visual_state(self, control_id: str, feasible: bool) -> No if handle is not None: cast("_ColorHandle", handle).color = color - def set_target_robot_visual_state(self, robot_id: str, feasible: bool) -> None: + def set_target_robot_visual_state(self, model_key: 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") + target = self._urdfs.get(f"{model_key}:target") self._set_urdf_mesh_material(target, mesh_color, mesh_opacity) def close(self) -> None: @@ -690,17 +656,17 @@ def close(self) -> None: self._target_tracks_current.clear() self._robot_display_mode = RobotDisplayMode.VISUAL - def _ensure_robot_urdfs(self, robot_id: str, config: RobotModelConfig) -> None: + def _ensure_robot_urdfs(self, model_key: str, config: RobotModelConfig) -> None: if not config.model_path: return - model = self._models_by_id.get(robot_id) + model = self._models_by_id.get(model_key) if model is None: return for kind in ("current", "target", "preview"): - key = f"{robot_id}:{kind}" + key = f"{model_key}:{kind}" if key in self._urdfs: continue - root_node_name = self._urdf_root_node_name(robot_id, kind, config) + root_node_name = self._urdf_root_node_name(model_key, kind, config) mesh_color_override = { "current": None, "target": GOAL_ROBOT_MESH_COLOR, @@ -709,7 +675,7 @@ def _ensure_robot_urdfs(self, robot_id: str, config: RobotModelConfig) -> None: if kind == "current": # Keep both representations resident so changing the diagnostic # view does not reload or replace the primary robot. - old_fallback = self._collision_fallback_urdfs.pop(robot_id, None) + old_fallback = self._collision_fallback_urdfs.pop(model_key, None) if old_fallback is not None: self._remove_scene_handle(old_fallback) urdf = self.viser_urdf( @@ -740,7 +706,7 @@ def _ensure_robot_urdfs(self, robot_id: str, config: RobotModelConfig) -> None: fallback = self.viser_urdf( self.server, urdf_or_path=model, - root_node_name=f"/robots/{robot_id}/collision_fallback", + root_node_name=f"/robots/{model_key}/collision_fallback", mesh_color_override=( *COLLISION_MESH_COLOR, COLLISION_MESH_OPACITY, @@ -748,34 +714,34 @@ def _ensure_robot_urdfs(self, robot_id: str, config: RobotModelConfig) -> None: load_meshes=True, load_collision_meshes=False, ) - self._collision_fallback_urdfs[robot_id] = fallback + self._collision_fallback_urdfs[model_key] = fallback self._joint_names_by_urdf[id(fallback)] = tuple( str(name) for name in model.actuated_joint_names ) self._set_urdf_mesh_material( fallback, COLLISION_MESH_COLOR, COLLISION_MESH_OPACITY ) - self._apply_robot_display_mode(robot_id) + self._apply_robot_display_mode(model_key) if kind == "target": self._set_urdf_mesh_material( self._urdfs[key], GOAL_ROBOT_FEASIBLE_COLOR, GOAL_ROBOT_FEASIBLE_OPACITY ) self._set_handle_visibility( - self._urdfs[key], self._target_active.get(robot_id, False) + self._urdfs[key], self._target_active.get(model_key, False) ) elif kind == "preview": self._set_urdf_mesh_material( self._urdfs[key], PREVIEW_ROBOT_COLOR, PREVIEW_ROBOT_OPACITY ) self._set_handle_visibility( - self._urdfs[key], self._preview_visible.get(robot_id, False) + self._urdfs[key], self._preview_visible.get(model_key, False) ) - def _apply_robot_display_mode(self, robot_id: str) -> None: - current = self._urdfs.get(f"{robot_id}:current") + def _apply_robot_display_mode(self, model_key: str) -> None: + current = self._urdfs.get(f"{model_key}:current") if current is None: return - model = self._models_by_id.get(robot_id) + model = self._models_by_id.get(model_key) if model is None: return has_collision = self._model_has_collision_geometry(model) @@ -788,7 +754,7 @@ def _apply_robot_display_mode(self, robot_id: str) -> None: RobotDisplayMode.COLLISION, RobotDisplayMode.BOTH, } - fallback = self._collision_fallback_urdfs.get(robot_id) + fallback = self._collision_fallback_urdfs.get(model_key) if fallback is not None: fallback.show_visual = mode in { RobotDisplayMode.COLLISION, @@ -887,26 +853,26 @@ def _assert_base_link_is_urdf_root(config: RobotModelConfig, prepared_path: Path 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: + def _urdf_root_node_name(self, model_key: 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", + "current": f"/robots/{model_key}/current", + "target": f"/targets/{model_key}/target", + "preview": f"/previews/{model_key}/ghost", }[kind] if not self._has_non_identity_base_pose(config): return root_node_name - self._ensure_base_pose_frame(robot_id, kind, config) + self._ensure_base_pose_frame(model_key, 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" + def _ensure_base_pose_frame(self, model_key: str, kind: str, config: RobotModelConfig) -> None: + key = f"{model_key}:{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", + "current": f"/robots/{model_key}/current/base_pose", + "target": f"/targets/{model_key}/target/base_pose", + "preview": f"/previews/{model_key}/ghost/base_pose", }[kind] self._root_frames[key] = self.server.scene.add_frame( frame_name, @@ -972,11 +938,11 @@ def viser_joint_configuration( def viser_actuated_joint_names(self, urdf: ViserUrdf) -> tuple[str, ...]: return self._joint_names_by_urdf.get(id(urdf), ()) - def _set_preview_visibility(self, robot_id: str, visible: bool) -> None: - self._set_handle_visibility(self._urdfs.get(f"{robot_id}:preview"), visible) + def _set_preview_visibility(self, model_key: str, visible: bool) -> None: + self._set_handle_visibility(self._urdfs.get(f"{model_key}:preview"), visible) - def _set_target_visibility(self, robot_id: str, visible: bool) -> None: - self._set_handle_visibility(self._urdfs.get(f"{robot_id}:target"), visible) + def _set_target_visibility(self, model_key: str, visible: bool) -> None: + self._set_handle_visibility(self._urdfs.get(f"{model_key}:target"), visible) def _set_handle_visibility(self, handle: SceneHandle | None, visible: bool) -> None: if handle is None: diff --git a/dimos/manipulation/visualization/viser/state.py b/dimos/manipulation/visualization/viser/state.py index 38c0503b70..7c5bf40fe5 100644 --- a/dimos/manipulation/visualization/viser/state.py +++ b/dimos/manipulation/visualization/viser/state.py @@ -98,7 +98,6 @@ class FeasibilityState: @dataclass class PanelPlanState: status: PlanStatus = PlanStatus.NONE - robot: str | None = None group_ids: tuple[PlanningGroupID, ...] = () target_sequence_id: int = 0 plan: GeneratedPlan | None = None @@ -106,7 +105,6 @@ class PanelPlanState: @dataclass class PanelState: - selected_robot: str | None = None selected_group_ids: tuple[PlanningGroupID, ...] = () planning_mode: PlanningMode = PlanningMode.JOINT_SPACE selection_epoch: int = 0 @@ -208,7 +206,6 @@ def module_state(self) -> str: class TargetEvaluationRequest: sequence_id: int source: PreviewSource - robot_name: str | None = None selection_epoch: int = 0 group_ids: tuple[PlanningGroupID, ...] = () pose: Pose | None = None diff --git a/dimos/manipulation/visualization/viser/test_gui.py b/dimos/manipulation/visualization/viser/test_gui.py index 9d8be1e168..59a9825415 100644 --- a/dimos/manipulation/visualization/viser/test_gui.py +++ b/dimos/manipulation/visualization/viser/test_gui.py @@ -16,12 +16,14 @@ from collections.abc import Callable from dataclasses import dataclass +from pathlib import Path 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.config import RobotModelConfig 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 @@ -60,7 +62,7 @@ def __init__(self, module: FakeOperatorBackend | None = None) -> None: def status(self) -> OperatorStatus: return OperatorStatus(state="IDLE", error="", has_plan=False) - def get_init_joints(self, robot_name: str) -> None: + def get_init_joints(self) -> None: return None def cancel(self) -> bool: @@ -149,12 +151,9 @@ 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: +def planning_group(name: str, joints: tuple[str, ...]) -> PlanningGroup: return PlanningGroup( - f"{robot}/{name}", - robot, name, - tuple(f"{robot}/{joint}" for joint in joints), joints, "base", None, @@ -165,9 +164,11 @@ def make_gui(module: FakeOperatorBackend | None = None) -> ViserPanelGui: module = module or FakeOperatorBackend() return ViserPanelGui( EmptyServer(), - PlanningSceneInfo(robots={}), + PlanningSceneInfo( + model=RobotModelConfig(model_path=Path("/tmp/model.urdf"), joint_names=[]) + ), FakeOperator(module), - {}, + lambda: None, ViserVisualizationConfig(), ) @@ -235,11 +236,11 @@ def test_gui_feasibility_status_uses_exact_status_mapping( 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.selected_group_ids = ("manipulator", "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._stale_models = lambda _group_ids: ("model",) # type: ignore[method-assign] gui._set_handle_value = values.__setitem__ # type: ignore[method-assign] gui._update_status_text() @@ -247,7 +248,7 @@ def test_group_status_composes_shared_panel_state_without_robot_dropdown() -> No 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: `feasible` · Plan: `fresh`\n\nState stale: `True`", "target_summary": "Feasibility: `unknown`", } @@ -291,7 +292,7 @@ def test_gui_preview_enters_previewing_before_worker_runs( gui.state.backend_status = BackendConnectionStatus.READY gui.state.target_status = TargetStatus.FEASIBLE gui.state.manipulation_state = "COMPLETED" - gui.state.selected_group_ids = ("arm/manipulator",) + gui.state.selected_group_ids = ("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 @@ -319,8 +320,8 @@ def test_gui_selection_change_clears_invalidated_preview( monkeypatch.setattr(gui, "_operation_worker", FakeOperationSubmitWorker(submissions)) monkeypatch.setattr(gui, "refresh", lambda: None) groups = [ - planning_group("arm", "manipulator", ("j1",)), - planning_group("arm", "gripper", ("j2",)), + planning_group("manipulator", ("j1",)), + planning_group("gripper", ("j2",)), ] monkeypatch.setattr(gui, "list_planning_groups", lambda: groups) monkeypatch.setattr(gui, "_build_joint_sliders", lambda: None) @@ -355,8 +356,8 @@ def test_gui_selection_change_ignores_invalidated_preview_error( monkeypatch.setattr(gui, "_operation_worker", FakeOperationErrorWorker(errors)) monkeypatch.setattr(gui, "refresh", lambda: None) groups = [ - planning_group("arm", "manipulator", ("j1",)), - planning_group("arm", "gripper", ("j2",)), + planning_group("manipulator", ("j1",)), + planning_group("gripper", ("j2",)), ] monkeypatch.setattr(gui, "list_planning_groups", lambda: groups) monkeypatch.setattr(gui, "_build_joint_sliders", lambda: None) @@ -450,7 +451,6 @@ def test_gui_guard_errors_keep_action_idle( 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)() diff --git a/dimos/manipulation/visualization/viser/test_state.py b/dimos/manipulation/visualization/viser/test_state.py index a9ec709b93..56622f6738 100644 --- a/dimos/manipulation/visualization/viser/test_state.py +++ b/dimos/manipulation/visualization/viser/test_state.py @@ -32,8 +32,7 @@ def test_panel_cannot_plan_from_fault_without_explicit_reset() -> None: state = PanelState( - selected_robot="arm", - selected_group_ids=(PlanningGroupID("arm/manipulator"),), + selected_group_ids=(PlanningGroupID("manipulator"),), runtime=PanelRuntime.RUNNING, backend_status=BackendConnectionStatus.READY, target_status=TargetStatus.FEASIBLE, @@ -65,7 +64,7 @@ def test_sequence_change_marks_a_fresh_plan_stale() -> None: def test_selection_epoch_change_resets_plan_and_invalidates_sequence() -> None: - state = PanelState(selected_group_ids=(PlanningGroupID("arm/manipulator"),)) + state = PanelState(selected_group_ids=(PlanningGroupID("manipulator"),)) state.plan_state.status = PlanStatus.FRESH assert state.advance_selection_epoch() == 1 diff --git a/dimos/manipulation/visualization/viser/test_viser_visualization.py b/dimos/manipulation/visualization/viser/test_viser_visualization.py index 485fe8b5a3..31eb123f14 100644 --- a/dimos/manipulation/visualization/viser/test_viser_visualization.py +++ b/dimos/manipulation/visualization/viser/test_viser_visualization.py @@ -12,1860 +12,95 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Hermetic contract tests for the group-aware Viser manipulation panel.""" +"""Hermetic tests for singular Viser preview data.""" -from __future__ import annotations - -from collections.abc import Callable, Iterator, Sequence -from dataclasses import dataclass from pathlib import Path -import threading -from types import SimpleNamespace +from unittest.mock import MagicMock import pytest pytest.importorskip("viser", reason="Viser optional dependency is not installed") -from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupSelection -from dimos.manipulation.planning.spec.enums import ObstacleType, PlanningStatus +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition +from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.models import ( - GeneratedPlan, - Obstacle, PlanningSceneInfo, VisualizationSession, -) -from dimos.manipulation.visualization.operator import TargetEvaluationResult -from dimos.manipulation.visualization.viser import ( - scene as scene_module, - visualizer as visualizer_module, + VisualizationStateFrame, ) from dimos.manipulation.visualization.viser.animation import ( - GroupPreviewAnimation, + PreviewAnimation, PreviewFrame, - PreviewTrack, + preview_tick_times, scaled_frame_delays, ) -from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig -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, - RobotDisplayMode, - ViserManipulationScene, -) -from dimos.manipulation.visualization.viser.state import ( - ActionStatus, - PanelPlanState, - PlanningMode, - PlanStatus, - TargetEvaluationRequest, - TargetStatus, -) -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.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 -@dataclass -class Handle: - label: str = "" - value: object = None - options: list[str] | None = None - disabled: bool = False - 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_update(self, callback: Callable[[object], None]) -> None: - self.callback = callback - - def on_click(self, callback: Callable[[object], None]) -> None: - self.callback = callback - - def remove(self) -> None: - self.removed = True - - -class Folder(Handle): - def __init__(self, label: str, **kwargs: bool) -> None: - super().__init__(label=label) - self.kwargs = kwargs - - def __enter__(self) -> Folder: - return self - - def __exit__(self, *_: object) -> bool: - return False - - -class Gui: - def __init__(self) -> None: - 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 - - 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 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 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 - - def configure_theme(self, **kwargs: object) -> None: - self.theme_kwargs = kwargs - - -class Server: - def __init__(self) -> None: - self.gui = Gui() - self.scene = SimpleNamespace() - - -@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 - pre_grasp_offset: float = 0.0 - - 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.cartesian_plans: list[tuple[dict[str, PoseStamped], object, tuple[str, ...]]] = [] - self.cartesian_plan_success = True - self.error = "" - self.executions = 0 - self.cancelled = 0 - self.cleared = 0 - self.last_plan: GeneratedPlan | None = None - self.motion_speed = 1.0 - self.motion_speed_updates: list[float] = [] - - 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 - return plan - - def list_robots(self) -> list[str]: - return list(self.configs) - - def list_planning_groups(self) -> list[PlanningGroup]: - return self.groups - - 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}" - - 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]}) - - def get_state(self) -> str: - return "IDLE" - - def get_error(self) -> str: - return self.error - - def get_motion_speed(self) -> float: - return self.motion_speed - - def set_motion_speed(self, speed_scale: float) -> bool: - self.motion_speed = float(speed_scale) - self.motion_speed_updates.append(float(speed_scale)) - return True - - def reset(self) -> SimpleNamespace: - return SimpleNamespace(is_success=lambda: 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 - - def execute(self) -> bool: - self.executions += 1 - return True - - def cancel(self) -> bool: - self.cancelled += 1 - return True - - def clear_planned_path(self) -> bool: - self.cleared += 1 - return True - - -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 get_current_joint_state(self, robot_id: str) -> JointState: - return JointState(self.module.states[robot_id.removeprefix("id-")]) - - 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 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 Operator: - def __init__(self, module: Module, monitor: Monitor) -> None: - self.module = module - self.monitor = monitor - - def status(self) -> SimpleNamespace: - return SimpleNamespace( - state=self.module.get_state(), - error=self.module.get_error(), - has_plan=True, - ) - - def get_motion_speed(self) -> float: - return self.module.get_motion_speed() - - def set_motion_speed(self, speed_scale: float) -> bool: - return self.module.set_motion_speed(speed_scale) - - def get_init_joints(self, robot_name: str) -> JointState | None: - return self.module.get_init_joints(robot_name) - - 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 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 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 plan_to_pose(self, request: object) -> GeneratedPlan: - return self.module.make_plan(tuple(request.pose_targets)) # type: ignore[attr-defined] - - def plan_cartesian(self, request: object) -> GeneratedPlan | None: - self.module.cartesian_plans.append( - ( - dict(request.pose_targets), # type: ignore[attr-defined] - request.config, # type: ignore[attr-defined] - tuple(request.auxiliary_group_ids), # type: ignore[attr-defined] - ) - ) - group_ids = tuple( - ( - *request.pose_targets.keys(), # type: ignore[attr-defined] - *request.auxiliary_group_ids, # type: ignore[attr-defined] - ) - ) - return self.module.make_plan(group_ids) if self.module.cartesian_plan_success else None - - def preview(self, plan: GeneratedPlan, duration: float | None = None) -> bool: - return self.module.preview_plan() - - def execute(self, plan: GeneratedPlan) -> bool: - return self.module.execute() - - 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 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 panel() -> Iterator[ - Callable[[list[PlanningGroup], dict[str, JointState]], tuple[ViserPanelGui, Module, Server]] -]: - panels: list[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, scene_info, operator, current, ViserVisualizationConfig(panel_enabled=True) - ) - gui.start() - panels.append(gui) - return gui, module, server - - yield make - for gui in panels: - gui.close() - - -def states(*robots: str) -> dict[str, JointState]: - return {robot: JointState({"name": ["j1", "j2"], "position": [0.1, 0.2]}) for robot in robots} - - -def obstacle( - obstacle_type: ObstacleType, - dimensions: tuple[float, ...] = (), - *, - color: tuple[float, float, float, float] = (0.2, 0.4, 0.6, 0.75), - mesh_path: str | None = None, -) -> Obstacle: - return Obstacle( - "test obstacle", - obstacle_type, - PoseStamped( - ts=1.0, - frame_id="world", - position=[1.0, 2.0, 3.0], - orientation=[0.1, 0.2, 0.3, 0.4], - ), - dimensions, - color, - mesh_path, - ) - - -def test_panel_contract_group_order_defaults_and_controls( - 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 [(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." - ) - 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 server.gui.dropdowns[1].options == ["Joint space", "Cartesian space"] - assert [ - (slider.label, slider.min, slider.max, slider.value) - for slider in server.gui.sliders - if slider.label != "Next plan speed" - ] == [("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 and slider.label != "Next plan speed" - ] == [ - "arm/manipulator/j1", - "arm/gripper/j2", - ] - - -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: - first, second = ( - group("arm", "manipulator", ("j1",), pose=True), - group("arm", "gripper", ("j2",)), - ) - 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 - - -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 - gui._operation_worker.stop() - 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, - ) - gui.state.target_status = TargetStatus.FEASIBLE - gui._submit_execute() - assert module.executions == 1 - - -def test_cartesian_space_mode_requests_sparse_time_optimal_trajectory( - panel: Callable[..., tuple[ViserPanelGui, Module, Server]], - monkeypatch: pytest.MonkeyPatch, -) -> None: - pose_group = group("arm", "manipulator", ("j1",), pose=True) - auxiliary_group = group("arm", "gripper", ("j2",)) - gui, module, server = panel([pose_group, auxiliary_group], states("arm")) - gui._toggle_group_selected(auxiliary_group.id) - gui.state.target_status = TargetStatus.FEASIBLE - gui._operation_worker.stop() - monkeypatch.setattr( - gui, - "_operation_worker", - SimpleNamespace(submit=lambda operation, **_: operation(), stop=lambda **_: None), - ) - - server.gui.dropdowns[1].value = "Cartesian space" - server.gui.dropdowns[1].callback(SimpleNamespace(target=server.gui.dropdowns[1])) - gui._submit_plan() - - assert gui.state.planning_mode == PlanningMode.CARTESIAN_SPACE - assert len(module.cartesian_plans) == 1 - targets, config, auxiliary_ids = module.cartesian_plans[0] - assert tuple(targets) == (pose_group.id,) - assert targets[pose_group.id].frame_id == "world" - assert config.speed_mode == "time_optimal" # type: ignore[attr-defined] - assert config.dt == 0.05 # type: ignore[attr-defined] - assert auxiliary_ids == (auxiliary_group.id,) - assert gui.state.plan_state.status == PlanStatus.FRESH - assert gui.state.last_result == "plan_cartesian_space=True" - - -def test_changing_planning_mode_marks_existing_plan_stale( - panel: Callable[..., tuple[ViserPanelGui, Module, Server]], -) -> None: - selected = group("arm", "manipulator", ("j1",), pose=True) - gui, _module, server = panel([selected], states("arm")) - gui.state.plan_state.status = PlanStatus.FRESH - - server.gui.dropdowns[1].value = "Cartesian space" - server.gui.dropdowns[1].callback(SimpleNamespace(target=server.gui.dropdowns[1])) - - assert gui.state.planning_mode == PlanningMode.CARTESIAN_SPACE - assert gui.state.plan_state.status == PlanStatus.STALE - - -def test_cartesian_failure_surfaces_backend_error_without_fallback( - panel: Callable[..., tuple[ViserPanelGui, Module, Server]], - monkeypatch: pytest.MonkeyPatch, -) -> None: - selected = group("arm", "manipulator", ("j1",), pose=True) - gui, module, server = panel([selected], states("arm")) - module.cartesian_plan_success = False - module.error = "Cartesian planning failed: UNSUPPORTED" - gui.state.target_status = TargetStatus.FEASIBLE - gui._operation_worker.stop() - monkeypatch.setattr( - gui, - "_operation_worker", - SimpleNamespace(submit=lambda operation, **_: operation(), stop=lambda **_: None), - ) - server.gui.dropdowns[1].value = "Cartesian space" - server.gui.dropdowns[1].callback(SimpleNamespace(target=server.gui.dropdowns[1])) - - gui._submit_plan() - - assert len(module.cartesian_plans) == 1 - assert module.plans == [] - assert gui.state.plan_state.status == PlanStatus.FAILED - assert gui.state.error == "Cartesian planning failed: UNSUPPORTED" - - -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]})} - ) - - 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 - - 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_incomplete_preset_preserves_existing_group_targets( - panel: Callable[..., tuple[ViserPanelGui, Module, Server]], -) -> None: - 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]}) - - gui._apply_preset("Init") - - assert gui.state.group_joint_targets[selected.id] == before - assert "missing joints" in gui.state.error - - -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]})} - ) - - assert gui.state.group_joint_targets == {} - assert [ - slider.label for slider in server.gui.sliders if slider.label != "Next plan speed" - ] == [] - - 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 and slider.label != "Next plan speed" - ] == [ - "arm/manipulator/j1", - "arm/manipulator/j2", - ] - - -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), - ) - 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 []} - ) - - gui._apply_preset("Init") - - assert gui.state.group_joint_targets == before - assert "missing joints" in gui.state.error - - -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, - ), - ) - 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) - ) - assert gui.state.target_status is status_before_callback - - -class Mesh: - def __init__(self) -> None: - self.visible = False - self.color: tuple[int, int, int] | None = None - self.opacity: float | None = None - - -class Urdf: - def __init__(self, *_: object, **__: object) -> None: - self._urdf = SimpleNamespace(actuated_joint_names=("j1", "j2")) - self._meshes = [Mesh()] - self._collision_meshes = [Mesh()] - self.show_visual = True - self.show_collision = False - self.cfg: list[float] | None = None - - def update_cfg(self, cfg: Sequence[float]) -> None: - self.cfg = list(cfg) - - def remove(self) -> None: - pass - - -@pytest.fixture(autouse=True) -def fake_yourdfpy_loader(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr( - scene_module.URDF, - "load", - lambda *_args, **_kwargs: SimpleNamespace( - actuated_joint_names=("j1", "j2"), - collision_scene=SimpleNamespace(geometry={"collision": object()}), - ), - ) - - -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)), - ), - ), - ) - ) - assert scene.animate_preview(preview, 1.0) is True - assert updates[-1] == ("id-arm", [1.0, 0.2]) - - -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 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 - - -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_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 - if slider.label != "Next plan speed" - ] == [ - ("arm/manipulator/j1", -1.0, 1.0, 0.001, 0.1), - ("arm/manipulator/j2", -2.0, 2.0, 0.001, 0.2), - ] - - -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] - - gui._apply_preset("Init") - assert gui.state.group_joint_targets[selected.id].position == [-0.5, -1.0] - - gui._apply_preset("Home") - assert gui.state.group_joint_targets[selected.id].position == [0.9, 0.8] - - -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]} - ) - monitor.poses[second.id] = Pose( - {"position": [2.0, 0.0, 0.0], "orientation": [0.0, 0.0, 0.0, 1.0]} - ) - server = Server() - scene_info = PlanningSceneInfo( - robots={"id-arm": module.configs["arm"]}, planning_groups=tuple(module.groups) - ) - current = {"id-arm": JointState(module.states["arm"])} - gui = ViserPanelGui( - server, - scene_info, - Operator(module, monitor), - current, - ViserVisualizationConfig(panel_enabled=True), - ) - try: - gui.start() - gui._toggle_group_selected(second.id) - - 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_panel_action_controls_are_present_in_source_order( - panel: Callable[..., tuple[ViserPanelGui, Module, Server]], -) -> None: - 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_next_plan_speed_slider_updates_future_speed_without_staling_plan( - panel: Callable[..., tuple[ViserPanelGui, Module, Server]], -) -> None: - selected = group("arm", "manipulator", ("j1",), pose=True) - gui, module, server = panel([selected], states("arm")) - accepted = module.make_plan((selected.id,)) - gui.state.plan_state = PanelPlanState(status=PlanStatus.FRESH, plan=accepted) - speed_slider = next( - slider for slider in server.gui.sliders if slider.label == "Next plan speed" - ) - speed_slider.value = 0.5 - assert speed_slider.callback is not None - - speed_slider.callback(SimpleNamespace(target=speed_slider)) - - assert module.motion_speed_updates == [0.5] - assert module.last_plan is accepted - assert gui.state.plan_state.plan is accepted - assert gui.state.plan_state.status == PlanStatus.FRESH - - -def test_next_plan_speed_slider_is_disabled_during_panel_operation( - panel: Callable[..., tuple[ViserPanelGui, Module, Server]], -) -> None: - selected = group("arm", "manipulator", ("j1",), pose=True) - gui, _module, server = panel([selected], states("arm")) - speed_slider = next( - slider for slider in server.gui.sliders if slider.label == "Next plan speed" - ) - - gui.state.action_status = ActionStatus.RUNNING - gui.refresh() - - assert speed_slider.disabled is True - - -def test_target_callbacks_require_current_target_identity( - panel: Callable[..., tuple[ViserPanelGui, Module, Server]], -) -> None: - first, second = ( - group("arm", "manipulator", ("j1",), pose=True), - group("arm", "gripper", ("j2",)), - ) - 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,), - ) - - gui._apply_target_evaluation_result(request, TargetEvaluationResult(True, "FEASIBLE", "", True)) - - assert gui.state.target_status == TargetStatus.CHECKING - - -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), +def _model() -> RobotModelConfig: + return RobotModelConfig( + model_path=Path("/model.urdf"), + joint_names=["left/j1", "right/j1"], + planning_groups=[ + PlanningGroupDefinition("left_arm", ("left/j1",), "base", "left/tool"), + PlanningGroupDefinition("right_arm", ("right/j1",), "base", "right/tool"), + ], ) - 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, - ) - - gui._apply_target_evaluation_result(request, TargetEvaluationResult(True, "FEASIBLE", "", True)) - - 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) - ) - 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 - 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._apply_target_evaluation_result( - request, TargetEvaluationResult(True, "COLLISION", "", 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_INFEASIBLE_COLOR - for item in (arm_primary, arm_secondary, other) - ) - 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_scene_shared_clock_uses_stored_unequal_robot_frames( - monkeypatch: pytest.MonkeyPatch, -) -> 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))), - ) - 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, - ) - assert updates == [ - ("left", [0.0]), - ("right", [10.0]), - ("left", [0.0]), - ("right", [11.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 scaled_frame_delays(frames, 2.0) == (0.5, 1.5) - - -def test_panel_disables_plan_preview_and_execute_until_a_feasible_target( - panel: Callable[..., tuple[ViserPanelGui, Module, Server]], -) -> None: - selected = group("arm", "manipulator", ("j1",), pose=True) - _gui, _module, server = panel([selected], states("arm")) - - assert [button.disabled for button in server.gui.buttons[1:4]] == [True, True, True] - - -def test_panel_status_reports_target_and_plan_defaults( - panel: Callable[..., tuple[ViserPanelGui, Module, Server]], -) -> None: - selected = group("arm", "manipulator", ("j1",), pose=True) - gui, _module, server = panel([selected], states("arm")) - - assert server.gui.markdown[0].value == "### Status\n**State:** Ready" - assert gui.state.error == "" - - -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) - - 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 - - -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) - - 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 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) - - scene.cancel_preview_animation() - - assert scene._preview_visible == {"id-arm": False} - assert scene._animation_generation == 1 - - -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"), - ) - - with pytest.raises(ValueError, match="base_link 'base'.*URDF root 'world'"): - scene._assert_base_link_is_urdf_root(SimpleNamespace(base_link="base"), "robot.urdf") - - -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]}) - ) - - assert ViserManipulationScene._has_non_identity_base_pose(identity) is False - assert ViserManipulationScene._has_non_identity_base_pose(translated) is True - - -@pytest.mark.parametrize("mode", [RobotDisplayMode.COLLISION, RobotDisplayMode.BOTH]) -def test_scene_display_mode_survives_primary_recreation_and_keeps_ghosts_unchanged( - mode: RobotDisplayMode, -) -> None: - server = Server() - 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) - current = scene._urdfs["id-arm:current"] - target = scene._urdfs["id-arm:target"] - preview = scene._urdfs["id-arm:preview"] - target_state = (target._meshes[0].visible, target._meshes[0].color, target._meshes[0].opacity) - preview_state = ( - preview._meshes[0].visible, - preview._meshes[0].color, - preview._meshes[0].opacity, - ) - - scene.robot_display_mode = mode - assert (current.show_visual, current.show_collision) == ( - mode is RobotDisplayMode.BOTH, - True, - ) - assert ( - target._meshes[0].visible, - target._meshes[0].color, - target._meshes[0].opacity, - ) == target_state - assert ( - preview._meshes[0].visible, - preview._meshes[0].color, - preview._meshes[0].opacity, - ) == preview_state - - scene._urdfs.pop("id-arm:current") - scene.register_robot("id-arm", config) - recreated = scene._urdfs["id-arm:current"] - scene.update_current_robot("id-arm", JointState({"name": ["j1", "j2"], "position": [0.7, 0.2]})) - assert recreated is not current - assert scene.robot_display_mode is mode - assert (recreated.show_visual, recreated.show_collision) == ( - mode is RobotDisplayMode.BOTH, - True, - ) - assert recreated.cfg == [0.7, 0.2] - - -def test_panel_robot_display_selector_and_collision_warning_use_session_scene( - panel: Callable[..., tuple[ViserPanelGui, Module, Server]], -) -> None: - selected = group("arm", "manipulator", ("j1",), pose=True) - gui, _module, server = panel([selected], states("arm")) - # The panel fixture intentionally uses a session without a scene; attach the - # already-created scene controls through the normal panel API contract. - scene = ViserManipulationScene(server, Urdf) - scene.prepared_urdf_path = lambda _config: "robot.urdf" # type: ignore[method-assign] - scene.register_robot( - "id-arm", Config("arm", ["j1", "j2"], [-1.0, -2.0], [1.0, 2.0], [0.0, 0.0]) - ) - gui.scene = scene - gui._build_scene_controls(server.gui) - display = gui._handles["robot_display"] - assert display.options == ["Visual", "Collision", "Both"] - assert display.value == "Visual" - warning = gui._handles["robot_display_warning"] - assert warning.visible is False - display.callback(SimpleNamespace(target=SimpleNamespace(value="Collision"))) - assert scene.robot_display_mode is RobotDisplayMode.COLLISION - assert warning.visible is False - - -@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) - - monkeypatch.setattr( - "dimos.manipulation.visualization.viser.scene.time.sleep", block_after_first_tick - ) - old = GroupPreviewAnimation( +def test_preview_timing_uses_one_model_track() -> None: + preview = PreviewAnimation( + ("left/j1", "right/j1"), ( - PreviewTrack( - "id-arm", - ("j1",), - ( - PreviewFrame(0.0, (0.0,)), - PreviewFrame(1.0, (1.0,)), - ), - ), - ) - ) - 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() - worker.join(timeout=2.0) - - assert not worker.is_alive() - assert updates == stable_updates - - -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_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]}) - - gui._apply_target_evaluation_result( - request, - TargetEvaluationResult( - True, "FEASIBLE", "", True, group_poses={selected.id: computed_pose} + PreviewFrame(0.0, (0.0, 0.0)), + PreviewFrame(1.0, (1.0, 0.5)), + PreviewFrame(3.0, (2.0, 1.0)), ), ) + assert preview_tick_times(preview) == (0.0, 1.0, 3.0) + assert scaled_frame_delays(preview.frames, 6.0) == (2.0, 4.0) - assert control.position == (0.7, 0.8, 0.9) - assert control.wxyz == (0.4, 0.1, 0.2, 0.3) - gui.close() - - -@pytest.mark.parametrize( - ("obstacle_type", "dimensions", "method", "shape"), - [ - (ObstacleType.BOX, (1.0, 2.0, 3.0), "add_box", (1.0, 2.0, 3.0)), - (ObstacleType.SPHERE, (0.4,), "add_icosphere", 0.4), - (ObstacleType.CYLINDER, (0.5, 1.5), "add_cylinder", (0.5, 1.5)), - ], -) -def test_scene_renders_obstacle_geometry_with_pose_color_and_visibility( - obstacle_type: ObstacleType, - dimensions: tuple[float, ...], - method: str, - shape: object, -) -> None: - server = Server() - server.scene.add_grid = lambda *_args, **_kwargs: Handle() - calls: list[tuple[str, str, dict[str, object]]] = [] - - def add_shape(path: str, **kwargs: object) -> Handle: - calls.append((method, path, kwargs)) - return Handle(visible=bool(kwargs["visible"])) - setattr(server.scene, method, add_shape) - scene = ViserManipulationScene(server, Urdf) - - scene.add_vis_obstacle("shape", obstacle(obstacle_type, dimensions)) - - assert calls == [ - ( - method, - "/manipulation/obstacles/shape", - { - "dimensions" if method == "add_box" else "radius": shape - if method != "add_cylinder" - else shape[0], - **({} if method != "add_cylinder" else {"height": shape[1]}), - "color": (51, 102, 153), - "opacity": 0.75, - "position": (1.0, 2.0, 3.0), - "wxyz": (0.4, 0.1, 0.2, 0.3), - "visible": True, - }, - ) - ] - - -@pytest.mark.parametrize( - "bad_color", - [ - (0.1, 0.2, 0.3), - (float("nan"), 0.2, 0.3, 0.4), - (0.1, 0.2, 0.3, 1.5), - ], -) -def test_scene_uses_fallback_appearance_and_proxy_for_invalid_obstacles( - bad_color: tuple[float, ...], -) -> None: - server = Server() - server.scene.add_grid = lambda *_args, **_kwargs: Handle() - calls: list[tuple[str, dict[str, object]]] = [] - - def add_box(path: str, **kwargs: object) -> Handle: - calls.append((path, kwargs)) - return Handle(visible=bool(kwargs["visible"])) - - server.scene.add_box = add_box - server.scene.add_label = ( - lambda path, text, **kwargs: calls.append((path, {"text": text, **kwargs})) or Handle() - ) - scene = ViserManipulationScene(server, Urdf) - - scene.add_vis_obstacle("bad", obstacle(ObstacleType.BOX, (1.0, 2.0, 3.0), color=bad_color)) - - assert calls[0][0].endswith("/bad") - assert calls[0][1]["color"] == (55, 190, 210) - assert calls[0][1]["opacity"] == 0.55 - - -def test_scene_replaces_invalid_box_geometry_with_a_visible_proxy() -> None: - server = Server() - server.scene.add_grid = lambda *_args, **_kwargs: Handle() - calls: list[tuple[str, dict[str, object]]] = [] - server.scene.add_box = lambda path, **kwargs: calls.append((path, kwargs)) or Handle() - server.scene.add_label = ( - lambda path, text, **kwargs: calls.append((path, {"text": text, **kwargs})) or Handle() +def test_visualizer_builds_full_model_preview_from_selected_canonical_joints() -> None: + visualizer = ViserManipulationVisualizer() + visualizer._model_config = _model() + visualizer._current_state = JointState(name=["left/j1", "right/j1"], position=[0.1, 0.2]) + trajectory = JointTrajectory( + joint_names=["right/j1"], + points=[TrajectoryPoint(positions=[0.8], time_from_start=1.0)], ) - scene = ViserManipulationScene(server, Urdf) + preview = visualizer._raw_preview_animation(trajectory) + assert preview == PreviewAnimation(("left/j1", "right/j1"), (PreviewFrame(1.0, (0.1, 0.8)),)) - scene.add_vis_obstacle("invalid", obstacle(ObstacleType.BOX, (1.0,))) - assert calls[0][0].endswith("mesh-failure-proxy") - assert calls[0][1]["visible"] is True - assert "box dimensions" in str(calls[1][1]["text"]) - - -def test_scene_mesh_rendering_accepts_scene_meshes_and_falls_back_on_load_failure( - monkeypatch: pytest.MonkeyPatch, -) -> None: - server = Server() - server.scene.add_grid = lambda *_args, **_kwargs: Handle() - mesh_calls: list[tuple[str, object, object, dict[str, object]]] = [] - server.scene.add_mesh_simple = lambda path, vertices, faces, **kwargs: ( - mesh_calls.append((path, vertices, faces, kwargs)) or Handle() - ) - scene = ViserManipulationScene(server, Urdf) - - mesh = SimpleNamespace( - dump=lambda concatenate: SimpleNamespace( - vertices=[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], faces=[[0, 1, 2]] +def test_visualizer_rejects_unknown_or_duplicate_trajectory_joints() -> None: + visualizer = ViserManipulationVisualizer() + visualizer._model_config = _model() + visualizer._current_state = JointState(name=["left/j1", "right/j1"], position=[0.1, 0.2]) + for names in (["unknown"], ["left/j1", "left/j1"]): + trajectory = JointTrajectory( + joint_names=names, + points=[TrajectoryPoint(positions=[0.0] * len(names), time_from_start=1.0)], ) - ) - monkeypatch.setattr(scene_module.trimesh, "load_mesh", lambda *_args, **_kwargs: mesh) - scene.add_vis_obstacle("mesh", obstacle(ObstacleType.MESH, mesh_path="triangle.obj")) - assert mesh_calls[0][0] == "/manipulation/obstacles/mesh" - assert mesh_calls[0][1].shape == (3, 3) - assert mesh_calls[0][2].shape == (1, 3) - - fallback_paths: list[str] = [] - server.scene.add_box = lambda path, **kwargs: fallback_paths.append(path) or Handle( - visible=bool(kwargs["visible"]) - ) - server.scene.add_label = lambda path, *_args, **_kwargs: fallback_paths.append(path) or Handle() - monkeypatch.setattr( - scene_module.trimesh, - "load_mesh", - lambda *_args, **_kwargs: (_ for _ in ()).throw(OSError("missing")), - ) - scene.add_vis_obstacle("missing", obstacle(ObstacleType.MESH, mesh_path="missing.obj")) - assert fallback_paths == [ - "/manipulation/obstacles/missing/mesh-failure-proxy", - "/manipulation/obstacles/missing/mesh-failure-label", - ] - - -def test_scene_obstacle_visibility_replacement_cleanup_and_closed_state() -> None: - server = Server() - server.scene.add_grid = lambda *_args, **_kwargs: Handle() - handles: list[Handle] = [] - server.scene.add_box = ( - lambda _path, **kwargs: handles.append(Handle(visible=bool(kwargs["visible"]))) - or handles[-1] - ) - scene = ViserManipulationScene(server, Urdf) - item = obstacle(ObstacleType.BOX, (1.0, 1.0, 1.0)) - - scene.add_vis_obstacle("box", item) - checkbox = server.gui.add_checkbox("unused", initial_value=True) - checkbox.on_update(lambda event: scene.set_obstacles_visible(event.target.value)) - checkbox.callback(SimpleNamespace(target=SimpleNamespace(value=False))) - assert handles[0].visible is False - scene.add_vis_obstacle("box", item) - assert handles[0].removed is True - scene.clear_vis_obstacles() - assert handles[-1].removed is True - scene.close() - count = len(handles) - scene.add_vis_obstacle("closed", item) - scene.set_obstacles_visible(True) - assert len(handles) == count - - -def test_scene_complete_and_pose_updates_preserve_atomic_obstacle_state() -> None: - server = Server() - server.scene.add_grid = lambda *_args, **_kwargs: Handle() - calls: list[tuple[str, dict[str, object], Handle]] = [] - - def add_box(path: str, **kwargs: object) -> Handle: - handle = Handle(visible=bool(kwargs["visible"])) - calls.append((path, kwargs, handle)) - return handle - - server.scene.add_box = add_box - scene = ViserManipulationScene(server, Urdf) - item = obstacle(ObstacleType.BOX, (1.0, 2.0, 3.0)) - scene.add_vis_obstacle("shape", item) - moved_pose = PoseStamped(position=[7.0, 8.0, 9.0], orientation=[0.0, 0.0, 0.0, 1.0]) - - scene.update_vis_obstacle_pose("shape", moved_pose) - - assert len(calls) == 2 - assert calls[0][2].removed is True - assert calls[1][1]["dimensions"] == (1.0, 2.0, 3.0) - assert calls[1][1]["position"] == (7.0, 8.0, 9.0) - assert scene._obstacles["shape"].dimensions == (1.0, 2.0, 3.0) - - -def test_scene_updates_handle_unknown_ids_proxy_failures_and_warning_lifecycle( - monkeypatch: pytest.MonkeyPatch, -) -> None: - server = Server() - server.scene.add_grid = lambda *_args, **_kwargs: Handle() - server.scene.add_box = lambda *_args, **_kwargs: Handle() - server.scene.add_icosphere = lambda *_args, **_kwargs: Handle() - server.scene.add_label = lambda *_args, **_kwargs: Handle() - scene = ViserManipulationScene(server, Urdf) - - scene.update_vis_obstacle_pose("missing", PoseStamped()) - monkeypatch.setattr( - scene_module.trimesh, - "load_mesh", - lambda *_args, **_kwargs: (_ for _ in ()).throw(OSError("missing")), - ) - with pytest.raises(RuntimeError, match="renderer used a proxy"): - scene.update_vis_obstacle(obstacle(ObstacleType.MESH, mesh_path="missing.obj")) - - scene.show_obstacle_warning("first warning") - warning = scene._obstacle_warning_handle - assert warning is not None - scene.show_obstacle_warning("updated warning") - assert warning.content == "updated warning" - monkeypatch.setattr( - server.gui, - "add_markdown", - lambda _message: (_ for _ in ()).throw(RuntimeError("GUI unavailable")), - ) - scene._obstacle_warning_handle = None - scene.show_obstacle_warning("ignored warning") - - -def test_visualizer_forwards_obstacle_operations_and_ignores_them_after_close( - monkeypatch: pytest.MonkeyPatch, -) -> None: - calls: list[tuple[object, ...]] = [] - - class FakeRuntime: - url = "http://localhost:8095" - - def __init__(self, _config: ViserVisualizationConfig) -> None: - pass - - def start(self) -> Server: - calls.append(("start",)) - return Server() - - def close(self) -> None: - calls.append(("runtime-close",)) - - class FakeScene: - def __init__(self, _server: Server, _viser_urdf: object) -> None: - calls.append(("scene-create",)) - - def add_vis_obstacle(self, obstacle_id: str, value: Obstacle) -> None: - calls.append(("add", obstacle_id, value)) - - def update_vis_obstacle(self, value: Obstacle) -> None: - calls.append(("update", value)) - - def update_vis_obstacle_pose(self, obstacle_id: str, pose: PoseStamped) -> None: - calls.append(("update-pose", obstacle_id, pose)) - - def show_obstacle_warning(self, message: str) -> None: - calls.append(("warning", message)) - - def remove_vis_obstacle(self, obstacle_id: str) -> None: - calls.append(("remove", obstacle_id)) - - def clear_vis_obstacles(self) -> None: - calls.append(("clear",)) - - def close(self) -> None: - calls.append(("scene-close",)) - - monkeypatch.setattr(visualizer_module, "ViserRuntime", FakeRuntime) - monkeypatch.setattr(visualizer_module, "ViserManipulationScene", FakeScene) - visualizer = ViserManipulationVisualizer(config=ViserVisualizationConfig(panel_enabled=False)) - item = obstacle(ObstacleType.SPHERE, (0.5,)) - - visualizer.initialize(VisualizationSession(PlanningSceneInfo(robots={}))) - visualizer.add_vis_obstacle("sphere", item) - visualizer.update_vis_obstacle(item) - visualizer.update_vis_obstacle_pose("sphere", item.pose) - visualizer.remove_vis_obstacle("sphere") - visualizer.clear_vis_obstacles() - assert calls == [ - ("start",), - ("scene-create",), - ("add", "sphere", item), - ("update", item), - ("update-pose", "sphere", item.pose), - ("remove", "sphere"), - ("clear",), - ] - - visualizer.close() - visualizer.add_vis_obstacle("ignored", item) - visualizer.update_vis_obstacle(item) - visualizer.update_vis_obstacle_pose("ignored", item.pose) - visualizer.remove_vis_obstacle("ignored") - visualizer.clear_vis_obstacles() - assert calls == [ - ("start",), - ("scene-create",), - ("add", "sphere", item), - ("update", item), - ("update-pose", "sphere", item.pose), - ("remove", "sphere"), - ("clear",), - ("scene-close",), - ("runtime-close",), - ] - - -def test_visualizer_handles_update_failure_once_and_exposes_warning( - monkeypatch: pytest.MonkeyPatch, -) -> None: - calls: list[tuple[object, ...]] = [] - - class FakeRuntime: - url = "http://localhost:8095" - - def __init__(self, _config: ViserVisualizationConfig) -> None: - pass - - def start(self) -> Server: - return Server() - - def close(self) -> None: - return None - - class FakeScene: - def __init__(self, _server: Server, _viser_urdf: object) -> None: - pass - - def update_vis_obstacle(self, value: Obstacle) -> None: - calls.append(("update", value.name)) - raise RuntimeError("renderer failed") - - def update_vis_obstacle_pose(self, obstacle_id: str, _pose: PoseStamped) -> None: - calls.append(("update-pose", obstacle_id)) - raise RuntimeError("pose renderer failed") - - def show_obstacle_warning(self, message: str) -> None: - calls.append(("warning", message)) - - def close(self) -> None: - return None - - monkeypatch.setattr(visualizer_module, "ViserRuntime", FakeRuntime) - monkeypatch.setattr(visualizer_module, "ViserManipulationScene", FakeScene) - visualizer = ViserManipulationVisualizer(config=ViserVisualizationConfig(panel_enabled=False)) - item = obstacle(ObstacleType.SPHERE, (0.5,)) - visualizer.initialize(VisualizationSession(PlanningSceneInfo(robots={}))) + assert visualizer._raw_preview_animation(trajectory) is None - visualizer.update_vis_obstacle(item) - visualizer.update_vis_obstacle_pose(item.name, item.pose) - assert calls[0] == ("update", item.name) - assert calls[1][0] == "warning" - assert item.name in str(calls[1][1]) - assert calls[2] == ("update-pose", item.name) - assert calls[3][0] == "warning" - assert item.name in str(calls[3][1]) - assert len(calls) == 4 +def test_visualizer_initializes_and_updates_one_scene_model() -> None: + visualizer = ViserManipulationVisualizer() + scene = MagicMock() + visualizer._scene = scene + visualizer._runtime = MagicMock() + visualizer._initialize_scene(PlanningSceneInfo(model=_model())) + scene.register_model.assert_called_once() + state = JointState(name=["left/j1", "right/j1"], position=[0.1, 0.2]) + visualizer.update_state(VisualizationStateFrame(joint_state=state)) + scene.update_current_model.assert_called_once_with(state) -def test_visualizer_update_is_a_noop_when_start_produces_no_scene( - monkeypatch: pytest.MonkeyPatch, -) -> None: - visualizer = ViserManipulationVisualizer(config=ViserVisualizationConfig(panel_enabled=False)) - monkeypatch.setattr(visualizer, "_ensure_started", lambda: None) - item = obstacle(ObstacleType.SPHERE, (0.5,)) - visualizer.update_vis_obstacle(item) - visualizer.update_vis_obstacle_pose(item.name, item.pose) +def test_visualization_session_contains_singular_model() -> None: + session = VisualizationSession(PlanningSceneInfo(model=_model()), operator=object()) + assert session.scene.model.joint_names == ["left/j1", "right/j1"] diff --git a/dimos/manipulation/visualization/viser/test_visualizer_lifecycle.py b/dimos/manipulation/visualization/viser/test_visualizer_lifecycle.py index a6f7011d91..09d377269b 100644 --- a/dimos/manipulation/visualization/viser/test_visualizer_lifecycle.py +++ b/dimos/manipulation/visualization/viser/test_visualizer_lifecycle.py @@ -83,7 +83,6 @@ def stop(self) -> None: def fake_robot_config(name: str) -> RobotModelConfig: return RobotModelConfig( - name=name, model_path=Path(f"{name}.urdf"), base_pose=PoseStamped(), joint_names=[], @@ -135,8 +134,8 @@ def __init__( ) -> None: calls.append(("create", "scene")) - def register_robot(self, robot_id: str, config: RobotModelConfig) -> None: - calls.append((robot_id, config.name)) + def register_model(self, config: RobotModelConfig) -> None: + calls.append(("model", config.model_path.stem)) def close(self) -> None: calls.append(("close", "scene")) @@ -147,7 +146,7 @@ def __init__( server: FakeServer, scene_info: PlanningSceneInfo, operator: object, - current_states: dict[str, JointState], + current_state: object, config: ViserVisualizationConfig, scene: FakeScene, ) -> None: @@ -169,12 +168,7 @@ def close(self) -> None: visualizer = ViserManipulationVisualizer( config=ViserVisualizationConfig(panel_enabled=True), ) - scene = PlanningSceneInfo( - robots={ - "robot-1": fake_robot_config("arm1"), - "robot-2": fake_robot_config("arm2"), - } - ) + scene = PlanningSceneInfo(model=fake_robot_config("arm")) visualizer.initialize(VisualizationSession(scene, operator=FakeDependency())) @@ -183,8 +177,7 @@ def close(self) -> None: ("create", "scene"), ("create", "gui"), ("start", "gui"), - ("robot-1", "arm1"), - ("robot-2", "arm2"), + ("model", "arm"), ("refresh", "gui"), ] @@ -214,6 +207,9 @@ def __init__( ) -> None: pass + def register_model(self, config: RobotModelConfig) -> None: + pass + def close(self) -> None: closed.append("scene") @@ -223,7 +219,7 @@ def __init__( server: FakeServer, scene_info: PlanningSceneInfo, operator: object, - current_states: dict[str, JointState], + current_state: object, config: ViserVisualizationConfig, scene: FakeScene, ) -> None: @@ -245,7 +241,9 @@ def close(self) -> None: with pytest.raises(RuntimeError, match="gui failed"): visualizer.initialize( - VisualizationSession(PlanningSceneInfo(robots={}), operator=FakeDependency()) + VisualizationSession( + PlanningSceneInfo(model=fake_robot_config("model")), operator=FakeDependency() + ) ) assert closed == ["gui", "scene", "runtime"] @@ -286,7 +284,9 @@ def __init__( with pytest.raises(RuntimeError, match="scene failed"): visualizer.initialize( - VisualizationSession(PlanningSceneInfo(robots={}), operator=FakeDependency()) + VisualizationSession( + PlanningSceneInfo(model=fake_robot_config("model")), operator=FakeDependency() + ) ) assert closed == ["runtime"] @@ -318,6 +318,9 @@ def __init__( ) -> None: pass + def register_model(self, config: RobotModelConfig) -> None: + pass + def close(self) -> None: closed.append("scene") @@ -327,7 +330,7 @@ def __init__( server: FakeServer, scene_info: PlanningSceneInfo, operator: object, - current_states: dict[str, JointState], + current_state: object, config: ViserVisualizationConfig, scene: FakeScene, ) -> None: @@ -351,7 +354,9 @@ def close(self) -> None: config=ViserVisualizationConfig(panel_enabled=True), ) visualizer.initialize( - VisualizationSession(PlanningSceneInfo(robots={}), operator=FakeDependency()) + VisualizationSession( + PlanningSceneInfo(model=fake_robot_config("model")), operator=FakeDependency() + ) ) with pytest.raises(RuntimeError, match="gui close failed"): @@ -417,12 +422,12 @@ def __init__( ) -> None: calls.append(("scene", "create")) - def update_current_robot(self, robot_id: str, joint_state: JointState | None) -> None: + def update_current_model(self, joint_state: JointState | None) -> None: assert joint_state == current - calls.append(("update", robot_id)) + calls.append(("update", "model")) - def register_robot(self, robot_id: str, config: RobotModelConfig) -> None: - calls.append(("register", robot_id)) + def register_model(self, config: RobotModelConfig) -> None: + calls.append(("register", "model")) def cancel_preview_animation(self) -> None: calls.append(("cancel", "preview")) @@ -442,24 +447,21 @@ def close(self) -> None: ) assert hasattr(ViserManipulationVisualizer, "cancel_preview_animation") - visualizer.initialize( - VisualizationSession(PlanningSceneInfo({"robot-1": fake_robot_config("arm")})) - ) + visualizer.initialize(VisualizationSession(PlanningSceneInfo(fake_robot_config("arm")))) visualizer.cancel_preview_animation() - visualizer.update_state(VisualizationStateFrame({"robot-1": current})) + visualizer.update_state(VisualizationStateFrame(current)) visualizer.cancel_preview_animation() visualizer.animate_trajectory(JointTrajectory(joint_names=["arm/joint1"]), duration=1.5) visualizer.close() - visualizer.update_state(VisualizationStateFrame({"robot-1": current})) + visualizer.update_state(VisualizationStateFrame(current)) assert calls == [ ("runtime", "start"), ("scene", "create"), - ("register", "robot-1"), + ("register", "model"), ("cancel", "preview"), - ("update", "robot-1"), + ("update", "model"), ("cancel", "preview"), - ("animate", "groups"), ("scene", "close"), ("runtime", "close"), ] @@ -534,12 +536,12 @@ def parse_prepared_model(path: Path) -> SimpleNamespace: scene = ViserManipulationScene(Server(), Urdf) monkeypatch.setattr(scene, "_model_has_collision_geometry", lambda _model: True) - scene.register_robot("robot-1", config) + scene.register_model(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", + "/robots/model/current/base_pose/urdf", + "/targets/model/target/base_pose/urdf", + "/previews/model/ghost/base_pose/urdf", ] assert prepared == [{"package_paths": {}, "xacro_args": {}, "convert_meshes": False}] assert all(path == fixed_world_root for path, _ in created) @@ -567,14 +569,14 @@ def test_selected_display_mode_survives_primary_recreation_and_joint_updates( config = fake_robot_config("arm") config.joint_names = ["joint1"] - scene.register_robot("robot-1", config) + scene.register_model(config) scene.robot_display_mode = mode - old_current = scene._urdfs["robot-1:current"] - scene._urdfs.pop("robot-1:current") + old_current = scene._urdfs["model:current"] + scene._urdfs.pop("model:current") - scene.register_robot("robot-1", config) - current = scene._urdfs["robot-1:current"] - scene.update_current_robot("robot-1", JointState({"name": ["joint1"], "position": [0.75]})) + scene.register_model(config) + current = scene._urdfs["model:current"] + scene.update_current_model(JointState({"name": ["joint1"], "position": [0.75]})) assert current is not old_current assert scene.robot_display_mode == mode diff --git a/dimos/manipulation/visualization/viser/visualizer.py b/dimos/manipulation/visualization/viser/visualizer.py index 1b64601128..e4b099510c 100644 --- a/dimos/manipulation/visualization/viser/visualizer.py +++ b/dimos/manipulation/visualization/viser/visualizer.py @@ -14,14 +14,12 @@ from __future__ import annotations -from collections.abc import Sequence from contextlib import suppress from typing import TYPE_CHECKING from dimos.manipulation.visualization.viser.animation import ( - GroupPreviewAnimation, + PreviewAnimation, PreviewFrame, - PreviewTrack, ) from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig from dimos.manipulation.visualization.viser.gui import ViserPanelGui @@ -75,10 +73,8 @@ def __init__( 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._current_state: JointState | None = None + self._model_config: RobotModelConfig | None = None self._closed = False def _ensure_started(self) -> None: @@ -96,7 +92,7 @@ def _ensure_started(self) -> None: server, self._session_scene, self._operator, - self._current_states, + lambda: self._current_state, self.config, scene, ) @@ -133,13 +129,7 @@ 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._model_config = session.scene.model self._initialize_scene(session.scene) def _initialize_scene(self, scene: PlanningSceneInfo) -> None: @@ -150,8 +140,7 @@ def _initialize_scene(self, scene: PlanningSceneInfo) -> None: if self._scene is None: return try: - for robot_id, config in scene.robots.items(): - self._scene.register_robot(str(robot_id), config) + self._scene.register_model(scene.model) if self._gui is not None: self._gui.refresh() except Exception: @@ -234,10 +223,9 @@ def update_state(self, frame: VisualizationStateFrame) -> None: self._ensure_started() if self._scene is None: return - 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 frame.joint_state is not None: + self._current_state = JointState(frame.joint_state) + self._scene.update_current_model(frame.joint_state) if self._gui is not None: self._gui.refresh() @@ -255,7 +243,7 @@ def animate_trajectory( 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: + def cancel_preview_animation(self) -> None: """Cancel preview playback without starting a renderer or waiting for it. The world monitor deliberately invokes this outside its visualization @@ -267,42 +255,24 @@ def cancel_preview_animation(self, robot_ids: Sequence[str] | None = None) -> No """ scene = self._scene if scene is not None: - if robot_ids is None: - scene.cancel_preview_animation() - else: - scene.cancel_preview_animation(robot_ids) + scene.cancel_preview_animation() - 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 + def _raw_preview_animation(self, trajectory: JointTrajectory) -> PreviewAnimation | None: + config = self._model_config + if config is None or len(trajectory.joint_names) != len(set(trajectory.joint_names)): + return None + indices = {str(name): index for index, name in enumerate(trajectory.joint_names)} + if not set(indices).issubset(config.joint_names): + return None + baseline = self._baseline_values(config, self._current_state) + if baseline is None: + return None + frames: list[PreviewFrame] = [] + for point in trajectory.points: + selected = {name: float(point.positions[index]) for name, index in indices.items()} + positions = [float(selected.get(name, baseline[name])) for name in config.joint_names] + frames.append(PreviewFrame(float(point.time_from_start), tuple(positions))) + return PreviewAnimation(tuple(config.joint_names), tuple(frames)) if frames else None @staticmethod def _baseline_values( diff --git a/dimos/robot/manipulators/_modeling.py b/dimos/robot/manipulators/_modeling.py index d9b528d0d5..ae43ab2044 100644 --- a/dimos/robot/manipulators/_modeling.py +++ b/dimos/robot/manipulators/_modeling.py @@ -19,17 +19,13 @@ import math from typing import TypeAlias -from dimos.manipulation.planning.spec.models import RobotName from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.Quaternion import Quaternion from dimos.msgs.geometry_msgs.Vector3 import Vector3 DegreesOfFreedom: TypeAlias = int JointPrefix: TypeAlias = str -UrdfJointPrefix: TypeAlias = str UrdfJointName: TypeAlias = str -CoordinatorJointName: TypeAlias = str -JointNameMapping: TypeAlias = dict[CoordinatorJointName, UrdfJointName] def base_pose( @@ -51,16 +47,3 @@ def joint_names( prefix: JointPrefix = "joint", ) -> list[UrdfJointName]: return [f"{prefix}{i}" for i in range(1, dof + 1)] - - -def coordinator_joint_mapping( - name: RobotName, - dof: DegreesOfFreedom, - *, - joint_prefix: JointPrefix | None = None, - urdf_joint_prefix: UrdfJointPrefix = "", -) -> JointNameMapping: - prefix = f"{name}/" if joint_prefix is None else joint_prefix - if not prefix: - return {} - return {f"{prefix}joint{i}": f"{urdf_joint_prefix}joint{i}" for i in range(1, dof + 1)} diff --git a/dimos/robot/manipulators/a1z/blueprints/basic.py b/dimos/robot/manipulators/a1z/blueprints/basic.py index 10f2e63afb..7a7fe9f993 100644 --- a/dimos/robot/manipulators/a1z/blueprints/basic.py +++ b/dimos/robot/manipulators/a1z/blueprints/basic.py @@ -28,7 +28,7 @@ _a1z_planner_hw = a1z_hardware("arm") a1z_planner_coordinator = autoconnect( - planner(robots=[make_a1z_model_config(name="arm")]), + planner(model=make_a1z_model_config()), coordinator( hardware=[_a1z_planner_hw], tasks=[trajectory_task(_a1z_planner_hw)], diff --git a/dimos/robot/manipulators/a1z/blueprints/teleop.py b/dimos/robot/manipulators/a1z/blueprints/teleop.py index 4cc2a1a8da..4be7df5113 100644 --- a/dimos/robot/manipulators/a1z/blueprints/teleop.py +++ b/dimos/robot/manipulators/a1z/blueprints/teleop.py @@ -53,7 +53,7 @@ ], ), ManipulationModule.blueprint( - robots=[_a1z_model], + model=_a1z_model, visualization={"backend": "viser"}, ), ) @@ -83,7 +83,7 @@ ], ), ManipulationModule.blueprint( - robots=[_a1z_quest_model], + model=_a1z_quest_model, visualization={"backend": "viser"}, ), ) diff --git a/dimos/robot/manipulators/a1z/config.py b/dimos/robot/manipulators/a1z/config.py index 0c2b67f1c7..bbbf8418b8 100644 --- a/dimos/robot/manipulators/a1z/config.py +++ b/dimos/robot/manipulators/a1z/config.py @@ -20,7 +20,7 @@ import attrs -from dimos.control.components import HardwareComponent, HardwareType, make_joints +from dimos.control.components import HardwareComponent, HardwareType from dimos.core.global_config import global_config from dimos.hardware.manipulators.galaxea_a1z.config import ( A1ZConfig, @@ -30,7 +30,6 @@ from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.robot.manipulators._modeling import ( base_pose, - coordinator_joint_mapping, joint_names, ) from dimos.utils.data import LfsPath @@ -79,7 +78,7 @@ def a1z_hardware( return HardwareComponent( hardware_id=hw_id, hardware_type=HardwareType.MANIPULATOR, - joints=make_joints(hw_id, A1Z_DOF), + joints=joint_names(A1Z_DOF, prefix="arm_joint"), adapter_type=adapter_type, address=address, auto_enable=True, @@ -91,23 +90,20 @@ def a1z_hardware( def make_a1z_model_config( - name: str = "arm", *, has_gripper: bool = True, - joint_prefix: str | None = None, home_joints: list[float] | None = None, ) -> RobotModelConfig: - local_joint_names = joint_names(A1Z_DOF, prefix="arm_joint") + model_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=local_joint_names, + joint_names=model_joint_names, base_link="base_link", planning_groups=[ PlanningGroupDefinition( name="manipulator", - joint_names=tuple(local_joint_names), + joint_names=tuple(model_joint_names), base_link="base_link", tip_link=("gripper_eef_link" if has_gripper else "arm_link6"), ) @@ -115,12 +111,6 @@ def make_a1z_model_config( package_paths=A1Z_PACKAGE_PATHS, auto_convert_meshes=True, collision_exclusion_pairs=A1Z_COLLISION_EXCLUSIONS, - joint_name_mapping=coordinator_joint_mapping( - name, - A1Z_DOF, - joint_prefix=joint_prefix, - urdf_joint_prefix="arm_", - ), - gripper_hardware_id=name if has_gripper else None, + gripper_hardware_id="arm" if has_gripper else None, home_joints=home_joints or [0.0] * A1Z_DOF, ) diff --git a/dimos/robot/manipulators/a750/blueprints/teleop.py b/dimos/robot/manipulators/a750/blueprints/teleop.py index 90435cb267..c63683f4ea 100644 --- a/dimos/robot/manipulators/a750/blueprints/teleop.py +++ b/dimos/robot/manipulators/a750/blueprints/teleop.py @@ -44,7 +44,7 @@ ], ), ManipulationModule.blueprint( - robots=[_a750_model], + model=_a750_model, visualization={"backend": "meshcat"}, ), ) diff --git a/dimos/robot/manipulators/a750/config.py b/dimos/robot/manipulators/a750/config.py index b3705ddccc..7e793f4a27 100644 --- a/dimos/robot/manipulators/a750/config.py +++ b/dimos/robot/manipulators/a750/config.py @@ -19,13 +19,12 @@ import math from pathlib import Path -from dimos.control.components import HardwareComponent, HardwareType, make_joints +from dimos.control.components import HardwareComponent, HardwareType 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, - coordinator_joint_mapping, joint_names, ) from dimos.utils.data import LfsPath @@ -70,7 +69,7 @@ def make_a750_hardware( auto_enable: bool = True, home_joints: list[float] | None = None, ) -> HardwareComponent: - joints = make_joints(hw_id, 6) + joints = joint_names(6) return HardwareComponent( hardware_id=hw_id, hardware_type=HardwareType.MANIPULATOR, @@ -96,23 +95,18 @@ def a750_hardware(hw_id: str = "arm", *, mock_without_address: bool = False) -> ) -def make_a750_model_config( - name: str = "arm", - *, - joint_prefix: str | None = None, -) -> RobotModelConfig: +def make_a750_model_config() -> RobotModelConfig: dof = 6 - local_joint_names = joint_names(dof) + model_joint_names = joint_names(dof) return RobotModelConfig( - name=name, model_path=A750_MODEL_PATH, base_pose=base_pose(), - joint_names=local_joint_names, + joint_names=model_joint_names, base_link="base_link", planning_groups=[ PlanningGroupDefinition( name="manipulator", - joint_names=tuple(local_joint_names), + joint_names=tuple(model_joint_names), base_link="base_link", tip_link="gripper_base", ) @@ -120,11 +114,6 @@ def make_a750_model_config( package_paths=A750_PACKAGE_PATHS, auto_convert_meshes=True, collision_exclusion_pairs=A750_GRIPPER_COLLISION_EXCLUSIONS, - joint_name_mapping=coordinator_joint_mapping( - name, - dof, - joint_prefix=joint_prefix, - ), - gripper_hardware_id=name, + gripper_hardware_id="arm", home_joints=A750_HOME_JOINTS, ) diff --git a/dimos/robot/manipulators/common/blueprints.py b/dimos/robot/manipulators/common/blueprints.py index 725b1255ff..b74ab87045 100644 --- a/dimos/robot/manipulators/common/blueprints.py +++ b/dimos/robot/manipulators/common/blueprints.py @@ -88,9 +88,8 @@ def _resolve_control_ik( robot_model: RobotModelConfig, control_ik: PinkControlIKOverrides | None, ) -> dict[str, Any]: - coordinator_joints = robot_model.get_coordinator_joint_names() - if hardware.joints != coordinator_joints: - raise ValueError("hardware joints must match RobotModelConfig coordinator joints") + if list(hardware.joints) != list(robot_model.joint_names): + raise ValueError("hardware and RobotModelConfig joints must match exactly") payload = dict(control_ik or {}) payload["robot_model"] = robot_model return payload @@ -222,13 +221,13 @@ def coordinator( def planner( *, - robots: Sequence[RobotModelConfig], + model: RobotModelConfig, planning_timeout: float = 10.0, visualization: dict[str, Any] | None = None, **kwargs: Any, ) -> Blueprint: module_kwargs: dict[str, Any] = { - "robots": list(robots), + "model": model, "planning_timeout": planning_timeout, **kwargs, } diff --git a/dimos/robot/manipulators/common/mixed.py b/dimos/robot/manipulators/common/mixed.py index fcd0615c2a..4eb846c638 100644 --- a/dimos/robot/manipulators/common/mixed.py +++ b/dimos/robot/manipulators/common/mixed.py @@ -59,6 +59,7 @@ adapter_type="xarm", address=global_config.xarm6_ip, gripper=True, + canonical_joint_names=[f"xarm_arm/joint{i}" for i in range(1, 7)], ) _piper_teleop_hw = make_piper_hardware( "piper_arm", @@ -66,8 +67,8 @@ address=global_config.can_port or "can0", gripper=True, ) -_xarm6_teleop_model = make_xarm6_model_config(name="xarm_arm", add_gripper=False) -_piper_teleop_model = make_piper_model_config(name="piper_arm") +_xarm6_teleop_model = make_xarm6_model_config(add_gripper=False, prefix="xarm_arm/") +_piper_teleop_model = make_piper_model_config() coordinator_teleop_dual = ControlCoordinator.blueprint( hardware=[_xarm6_teleop_hw, _piper_teleop_hw], diff --git a/dimos/robot/manipulators/openarm/blueprints/planner.py b/dimos/robot/manipulators/openarm/blueprints/planner.py index 6872b15157..75631a28c2 100644 --- a/dimos/robot/manipulators/openarm/blueprints/planner.py +++ b/dimos/robot/manipulators/openarm/blueprints/planner.py @@ -24,15 +24,10 @@ mock_right, right_hw, ) -from dimos.robot.manipulators.openarm.config import openarm_model_config +from dimos.robot.manipulators.openarm.config import openarm_dual_model_config openarm_mock_planner_coordinator = autoconnect( - planner( - robots=[ - openarm_model_config("left"), - openarm_model_config("right"), - ], - ), + planner(model=openarm_dual_model_config()), coordinator( hardware=[mock_left, mock_right], tasks=[trajectory_task(mock_left, mock_right)], @@ -40,12 +35,7 @@ ) openarm_planner_coordinator = autoconnect( - planner( - robots=[ - openarm_model_config("left"), - openarm_model_config("right"), - ], - ), + planner(model=openarm_dual_model_config()), coordinator( hardware=[left_hw, right_hw], tasks=[trajectory_task(left_hw, right_hw)], diff --git a/dimos/robot/manipulators/openarm/blueprints/teleop.py b/dimos/robot/manipulators/openarm/blueprints/teleop.py index eb6baae066..418bec5d51 100644 --- a/dimos/robot/manipulators/openarm/blueprints/teleop.py +++ b/dimos/robot/manipulators/openarm/blueprints/teleop.py @@ -42,7 +42,7 @@ ], ), ManipulationModule.blueprint( - robots=[_openarm_model], + model=_openarm_model, visualization={"backend": "meshcat"}, ), ) @@ -61,7 +61,7 @@ ], ), ManipulationModule.blueprint( - robots=[_openarm_model], + model=_openarm_model, visualization={"backend": "meshcat"}, ), ) diff --git a/dimos/robot/manipulators/openarm/config.py b/dimos/robot/manipulators/openarm/config.py index 307450d054..26816be6e9 100644 --- a/dimos/robot/manipulators/openarm/config.py +++ b/dimos/robot/manipulators/openarm/config.py @@ -34,6 +34,7 @@ OPENARM_LEFT_MODEL = OPENARM_PKG / "urdf/robot/openarm_v10_left.urdf" OPENARM_RIGHT_MODEL = OPENARM_PKG / "urdf/robot/openarm_v10_right.urdf" OPENARM_V10_FK_MODEL = OPENARM_PKG / "urdf/robot/openarm_v10_single.urdf" +OPENARM_V10_BIMANUAL_MODEL = OPENARM_PKG / "urdf/robot/openarm_v10_bimanual.urdf" OPENARM_PACKAGE_PATHS: dict[str, Path] = {"openarm_description": OPENARM_PKG} # Linux assigns can0/can1 in USB enumeration order, which is not guaranteed stable. @@ -79,20 +80,18 @@ def openarm_hardware( ) -def openarm_model_config(side: str, name: str | None = None) -> RobotModelConfig: +def openarm_model_config(side: str) -> RobotModelConfig: validate_side(side) - resolved_name = name or f"{side}_arm" - local_joint_names = openarm_joints(side) + model_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=local_joint_names, + joint_names=model_joint_names, base_link="openarm_body_link0", planning_groups=[ PlanningGroupDefinition( name="manipulator", - joint_names=tuple(local_joint_names), + joint_names=tuple(model_joint_names), base_link="openarm_body_link0", tip_link=f"openarm_{side}_link7", ) @@ -113,24 +112,22 @@ def openarm_single_hardware( ) -> HardwareComponent: return openarm_hardware( "left", - name="arm", adapter_type=adapter_type, address=address, ) def openarm_single_model_config() -> RobotModelConfig: - local_joint_names = openarm_joints("left") + model_joint_names = openarm_joints("left") return RobotModelConfig( - name="arm", model_path=OPENARM_V10_FK_MODEL, base_pose=base_pose(), - joint_names=local_joint_names, + joint_names=model_joint_names, base_link="openarm_body_link0", planning_groups=[ PlanningGroupDefinition( name="manipulator", - joint_names=tuple(local_joint_names), + joint_names=tuple(model_joint_names), base_link="openarm_body_link0", tip_link="openarm_left_link7", ) @@ -141,3 +138,40 @@ def openarm_single_model_config() -> RobotModelConfig: max_acceleration=1.0, home_joints=[0.0] * 7, ) + + +def openarm_dual_model_config() -> RobotModelConfig: + """Return the canonical two-arm OpenArm model and its planning groups.""" + left_joints = openarm_joints("left") + right_joints = openarm_joints("right") + all_joints = [*left_joints, *right_joints] + return RobotModelConfig( + model_path=OPENARM_V10_BIMANUAL_MODEL, + base_pose=base_pose(), + joint_names=all_joints, + base_link="openarm_body_link0", + planning_groups=[ + PlanningGroupDefinition( + name="left_arm", + joint_names=tuple(left_joints), + base_link="openarm_body_link0", + tip_link="openarm_left_link7", + ), + PlanningGroupDefinition( + name="right_arm", + joint_names=tuple(right_joints), + base_link="openarm_body_link0", + tip_link="openarm_right_link7", + ), + PlanningGroupDefinition( + name="both_arms", + joint_names=tuple(all_joints), + base_link="openarm_body_link0", + ), + ], + package_paths=OPENARM_PACKAGE_PATHS, + auto_convert_meshes=True, + max_velocity=0.5, + max_acceleration=1.0, + home_joints=[0.0] * 14, + ) diff --git a/dimos/robot/manipulators/openarm/test_openarm.py b/dimos/robot/manipulators/openarm/test_openarm.py new file mode 100644 index 0000000000..80a6260294 --- /dev/null +++ b/dimos/robot/manipulators/openarm/test_openarm.py @@ -0,0 +1,29 @@ +# 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. + +from dimos.manipulation.planning.spec.validation import validate_robot_model_config +from dimos.robot.manipulators.openarm.config import openarm_dual_model_config + + +def test_openarm_dual_model_contains_every_canonical_joint_and_group_frame() -> None: + config = openarm_dual_model_config() + + description = validate_robot_model_config(config) + + assert description.actuated_joint_names == config.joint_names + assert [group.name for group in config.planning_groups] == [ + "left_arm", + "right_arm", + "both_arms", + ] diff --git a/dimos/robot/manipulators/openyam/blueprints/basic.py b/dimos/robot/manipulators/openyam/blueprints/basic.py index cd9fbbcbb5..7bd997e20d 100644 --- a/dimos/robot/manipulators/openyam/blueprints/basic.py +++ b/dimos/robot/manipulators/openyam/blueprints/basic.py @@ -27,7 +27,7 @@ _openyam_planner_hw = make_openyam_hardware("arm") openyam_planner_coordinator = autoconnect( - planner(robots=[make_openyam_model_config(name="arm")]), + planner(model=make_openyam_model_config()), coordinator( hardware=[_openyam_planner_hw], tasks=[trajectory_task(_openyam_planner_hw)], diff --git a/dimos/robot/manipulators/openyam/blueprints/teleop.py b/dimos/robot/manipulators/openyam/blueprints/teleop.py index 16a5fdffff..bb25b89d42 100644 --- a/dimos/robot/manipulators/openyam/blueprints/teleop.py +++ b/dimos/robot/manipulators/openyam/blueprints/teleop.py @@ -27,7 +27,7 @@ from dimos.teleop.keyboard.keyboard_teleop_module import KeyboardTeleopModule _openyam_keyboard_hw = make_openyam_hardware("arm") -_openyam_model = make_openyam_model_config(name="arm") +_openyam_model = make_openyam_model_config() keyboard_teleop_openyam = autoconnect( KeyboardTeleopModule.blueprint(), @@ -41,7 +41,7 @@ ], ), ManipulationModule.blueprint( - robots=[_openyam_model], + model=_openyam_model, visualization={"backend": "viser"}, ), ) diff --git a/dimos/robot/manipulators/openyam/config.py b/dimos/robot/manipulators/openyam/config.py index a2972eacdb..d1788f41be 100644 --- a/dimos/robot/manipulators/openyam/config.py +++ b/dimos/robot/manipulators/openyam/config.py @@ -18,12 +18,11 @@ from pathlib import Path -from dimos.control.components import HardwareComponent, HardwareType, make_joints +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, - coordinator_joint_mapping, joint_names, ) from dimos.utils.data import LfsPath @@ -47,7 +46,7 @@ def make_openyam_hardware( return HardwareComponent( hardware_id=hw_id, hardware_type=HardwareType.MANIPULATOR, - joints=make_joints(hw_id, OPENYAM_DOF), + joints=joint_names(OPENYAM_DOF, prefix="yam_joint"), adapter_type="mock", address=None, auto_enable=auto_enable, @@ -66,23 +65,20 @@ def openyam_hardware( def make_openyam_model_config( - name: str = "arm", *, - joint_prefix: str | None = None, home_joints: list[float] | None = None, ) -> RobotModelConfig: """Build a planning config for the gripper-equipped OpenYAM.""" - local_joint_names = joint_names(OPENYAM_DOF, prefix="yam_joint") + model_joint_names = joint_names(OPENYAM_DOF, prefix="yam_joint") return RobotModelConfig( - name=name, model_path=OPENYAM_MODEL_PATH, base_pose=base_pose(), - joint_names=local_joint_names, + joint_names=model_joint_names, base_link="yam_base_link", planning_groups=[ PlanningGroupDefinition( name="manipulator", - joint_names=tuple(local_joint_names), + joint_names=tuple(model_joint_names), base_link="yam_base_link", tip_link="yam_hand_tcp", ) @@ -90,12 +86,6 @@ def make_openyam_model_config( package_paths=OPENYAM_PACKAGE_PATHS, auto_convert_meshes=True, collision_exclusion_pairs=[], - joint_name_mapping=coordinator_joint_mapping( - name, - OPENYAM_DOF, - joint_prefix=joint_prefix, - urdf_joint_prefix="yam_", - ), - gripper_hardware_id=name, + gripper_hardware_id="arm", home_joints=home_joints or [0.0] * OPENYAM_DOF, ) diff --git a/dimos/robot/manipulators/openyam/test_openyam.py b/dimos/robot/manipulators/openyam/test_openyam.py index 4f501a2b73..875ffda131 100644 --- a/dimos/robot/manipulators/openyam/test_openyam.py +++ b/dimos/robot/manipulators/openyam/test_openyam.py @@ -38,13 +38,10 @@ def _coordinator_kwargs(blueprint: Blueprint) -> dict[str, Any]: return _module_kwargs(blueprint, ControlCoordinator) -def test_openyam_model_config_has_expected_links_and_mapping() -> None: - config = make_openyam_model_config(name="arm") +def test_openyam_model_config_has_expected_links_and_canonical_joints() -> None: + config = make_openyam_model_config() assert config.joint_names == [f"yam_joint{i}" for i in range(1, OPENYAM_DOF + 1)] - assert config.joint_name_mapping == { - f"arm/joint{i}": f"yam_joint{i}" for i in range(1, OPENYAM_DOF + 1) - } assert config.base_link == "yam_base_link" assert config.planning_groups[0].tip_link == "yam_hand_tcp" assert list(config.package_paths) == list(OPENYAM_PACKAGE_PATHS) @@ -55,7 +52,7 @@ def test_openyam_mock_hardware_has_gripper() -> None: hardware = make_openyam_hardware("arm") assert hardware.adapter_type == "mock" - assert hardware.joints == [f"arm/joint{i}" for i in range(1, OPENYAM_DOF + 1)] + assert hardware.joints == [f"yam_joint{i}" for i in range(1, OPENYAM_DOF + 1)] assert hardware.gripper_joints == ["arm/gripper"] @@ -74,15 +71,14 @@ def test_openyam_mock_adapter_set_get_behavior() -> None: def test_openyam_planner_blueprint_preserves_model_config() -> None: blueprint = openyam_planner_coordinator kwargs = _module_kwargs(blueprint, ManipulationModule) - config = ManipulationModuleConfig(**kwargs).robots[0] + config = ManipulationModuleConfig(**kwargs).model - assert config.name == "arm" assert config.joint_names == [f"yam_joint{i}" for i in range(1, OPENYAM_DOF + 1)] assert config.planning_groups[0].tip_link == "yam_hand_tcp" assert config.gripper_hardware_id == "arm" task = _coordinator_kwargs(blueprint)["tasks"][0] assert task.type == "trajectory" - assert task.joint_names == [f"arm/joint{i}" for i in range(1, OPENYAM_DOF + 1)] + assert task.joint_names == config.joint_names def test_openyam_coordinator_blueprint_uses_six_arm_joints() -> None: diff --git a/dimos/robot/manipulators/piper/blueprints/teleop.py b/dimos/robot/manipulators/piper/blueprints/teleop.py index 8933aa14b0..8f65c8d400 100644 --- a/dimos/robot/manipulators/piper/blueprints/teleop.py +++ b/dimos/robot/manipulators/piper/blueprints/teleop.py @@ -68,7 +68,7 @@ ], ), ManipulationModule.blueprint( - robots=[_piper_model], + model=_piper_model, visualization={"backend": "viser"}, ), ) @@ -111,7 +111,7 @@ class _PiperTeleopCoordinator(ControlCoordinator): ], ), ManipulationModule.blueprint( - robots=[_piper_model], + model=_piper_model, visualization={"backend": "viser"}, ), *mujoco_if_sim(PIPER_SIM_PATH, len(_piper_teleop_hw.joints)), diff --git a/dimos/robot/manipulators/piper/config.py b/dimos/robot/manipulators/piper/config.py index 1d390aa553..56c37a25f9 100644 --- a/dimos/robot/manipulators/piper/config.py +++ b/dimos/robot/manipulators/piper/config.py @@ -18,13 +18,12 @@ from pathlib import Path -from dimos.control.components import HardwareComponent, HardwareType, make_joints +from dimos.control.components import HardwareComponent, HardwareType 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, - coordinator_joint_mapping, joint_names, ) from dimos.utils.data import LfsPath @@ -69,6 +68,7 @@ def make_piper_hardware( auto_enable: bool = True, adapter_kwargs: dict[str, object] | None = None, home_joints: list[float] | None = None, + canonical_joint_names: list[str] | None = None, ) -> HardwareComponent: kwargs = _adapter_kwargs(home_joints) if adapter_kwargs: @@ -76,7 +76,7 @@ def make_piper_hardware( return HardwareComponent( hardware_id=hw_id, hardware_type=HardwareType.MANIPULATOR, - joints=make_joints(hw_id, 6), + joints=canonical_joint_names or joint_names(6), adapter_type=adapter_type, address=address, auto_enable=auto_enable, @@ -95,6 +95,7 @@ def piper_hardware( gripper_closed_position: float | None = None, mock_without_address: bool = True, home_joints: list[float] | None = None, + canonical_joint_names: list[str] | None = None, ) -> HardwareComponent: if global_config.simulation: return make_piper_hardware( @@ -105,6 +106,7 @@ def piper_hardware( gripper_open_position=gripper_open_position, gripper_closed_position=gripper_closed_position, home_joints=home_joints, + canonical_joint_names=canonical_joint_names, ) address = global_config.can_port or "can0" if mock_without_address and not global_config.can_port: @@ -114,6 +116,7 @@ def piper_hardware( gripper_open_position=gripper_open_position, gripper_closed_position=gripper_closed_position, home_joints=home_joints, + canonical_joint_names=canonical_joint_names, ) return make_piper_hardware( hw_id, @@ -123,28 +126,26 @@ def piper_hardware( gripper_open_position=gripper_open_position, gripper_closed_position=gripper_closed_position, home_joints=home_joints, + canonical_joint_names=canonical_joint_names, ) def make_piper_model_config( - name: str = "arm", *, - joint_prefix: str | None = None, home_joints: list[float] | None = None, ) -> RobotModelConfig: dof = 6 - local_joint_names = joint_names(dof) + model_joint_names = joint_names(dof) model_home_joints = list(home_joints) if home_joints is not None else list(PIPER_HOME_JOINTS) return RobotModelConfig( - name=name, model_path=PIPER_MODEL_PATH, base_pose=base_pose(), - joint_names=local_joint_names, + joint_names=model_joint_names, base_link="base_link", planning_groups=[ PlanningGroupDefinition( name="manipulator", - joint_names=tuple(local_joint_names), + joint_names=tuple(model_joint_names), base_link="base_link", tip_link="gripper_base", ) @@ -152,11 +153,6 @@ def make_piper_model_config( package_paths=PIPER_PACKAGE_PATHS, auto_convert_meshes=True, collision_exclusion_pairs=PIPER_GRIPPER_COLLISION_EXCLUSIONS, - joint_name_mapping=coordinator_joint_mapping( - name, - dof, - joint_prefix=joint_prefix, - ), - gripper_hardware_id=name, + gripper_hardware_id="arm", home_joints=model_home_joints, ) diff --git a/dimos/robot/manipulators/xarm/blueprints/basic.py b/dimos/robot/manipulators/xarm/blueprints/basic.py index 31bf7bb022..0668ae051e 100644 --- a/dimos/robot/manipulators/xarm/blueprints/basic.py +++ b/dimos/robot/manipulators/xarm/blueprints/basic.py @@ -23,22 +23,24 @@ from dimos.robot.manipulators.xarm.config import ( XARM6_SIM_PATH, XARM7_SIM_PATH, - make_xarm6_model_config, + make_dual_xarm6_model_config, make_xarm7_model_config, make_xarm_hardware, xarm6_hardware, xarm7_hardware, ) -_mock_left_xarm6_hw = make_xarm_hardware("left_arm", 6) -_mock_right_xarm6_hw = make_xarm_hardware("right_arm", 6) +_dual_xarm6_model = make_dual_xarm6_model_config() +_mock_left_xarm6_hw = make_xarm_hardware( + "left_arm", 6, canonical_joint_names=list(_dual_xarm6_model.planning_groups[0].joint_names) +) +_mock_right_xarm6_hw = make_xarm_hardware( + "right_arm", 6, canonical_joint_names=list(_dual_xarm6_model.planning_groups[1].joint_names) +) 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), - ], + model=_dual_xarm6_model, visualization={"backend": "viser"}, ), coordinator( @@ -50,7 +52,7 @@ _xarm7_hw = xarm7_hardware("arm", gripper=True, mock_without_address=True) xarm7_planner_coordinator = autoconnect( - planner(robots=[make_xarm7_model_config(name="arm", add_gripper=True)]), + planner(model=make_xarm7_model_config(add_gripper=True)), coordinator( hardware=[_xarm7_hw], tasks=[trajectory_task(_xarm7_hw)], @@ -77,8 +79,12 @@ *mujoco_if_sim(XARM6_SIM_PATH, len(_coordinator_xarm6_hw.joints)), ) -_xarm7_left = xarm7_hardware("left_arm") -_xarm6_right = xarm6_hardware("right_arm") +_xarm7_left = xarm7_hardware( + "left_arm", canonical_joint_names=[f"left_arm/joint{i}" for i in range(1, 8)] +) +_xarm6_right = xarm6_hardware( + "right_arm", canonical_joint_names=[f"right_arm/joint{i}" for i in range(1, 7)] +) coordinator_dual_xarm = ControlCoordinator.blueprint( hardware=[_xarm7_left, _xarm6_right], diff --git a/dimos/robot/manipulators/xarm/blueprints/perception.py b/dimos/robot/manipulators/xarm/blueprints/perception.py index f187e7ab50..d72ef3c397 100644 --- a/dimos/robot/manipulators/xarm/blueprints/perception.py +++ b/dimos/robot/manipulators/xarm/blueprints/perception.py @@ -34,14 +34,11 @@ xarm_perception = autoconnect( PickAndPlaceModule.blueprint( - robots=[ - make_xarm7_model_config( - name="arm", - add_gripper=True, - pitch=math.radians(45), - tf_extra_links=["link7"], - ) - ], + model=make_xarm7_model_config( + add_gripper=True, + pitch=math.radians(45), + tf_extra_links=["link7"], + ), planning_timeout=10.0, visualization={"backend": "meshcat"}, floor_z=-0.02, diff --git a/dimos/robot/manipulators/xarm/blueprints/simulation.py b/dimos/robot/manipulators/xarm/blueprints/simulation.py index fb6e21f09a..de82961f2f 100644 --- a/dimos/robot/manipulators/xarm/blueprints/simulation.py +++ b/dimos/robot/manipulators/xarm/blueprints/simulation.py @@ -33,7 +33,7 @@ xarm_perception_sim = autoconnect( PickAndPlaceModule.blueprint( - robots=[make_xarm7_sim_robot_config()], + model=make_xarm7_sim_robot_config(), planning_timeout=10.0, visualization={"backend": "meshcat"}, ), diff --git a/dimos/robot/manipulators/xarm/blueprints/teleop.py b/dimos/robot/manipulators/xarm/blueprints/teleop.py index ef28d2274f..fc03305c1c 100644 --- a/dimos/robot/manipulators/xarm/blueprints/teleop.py +++ b/dimos/robot/manipulators/xarm/blueprints/teleop.py @@ -66,7 +66,7 @@ ], ), ManipulationModule.blueprint( - robots=[make_xarm6_model_config(add_gripper=True)], + model=make_xarm6_model_config(add_gripper=True), visualization={"backend": "viser"}, ), ) @@ -88,7 +88,7 @@ ], ), ManipulationModule.blueprint( - robots=[make_xarm7_model_config(add_gripper=True)], + model=make_xarm7_model_config(add_gripper=True), visualization={"backend": "viser"}, ), ) @@ -195,7 +195,7 @@ class _XArm7TeleopCoordinator(ControlCoordinator): ], ), ManipulationModule.blueprint( - robots=[_xarm7_teleop_model], + model=_xarm7_teleop_model, visualization={"backend": "viser"}, ), *mujoco_if_sim(XARM7_SIM_PATH, len(_xarm7_teleop_hw.joints)), @@ -224,7 +224,7 @@ class _XArm7TeleopCoordinator(ControlCoordinator): ], ), ManipulationModule.blueprint( - robots=[_xarm6_teleop_model], + model=_xarm6_teleop_model, visualization={"backend": "viser"}, ), *mujoco_if_sim(XARM6_SIM_PATH, len(_xarm6_teleop_hw.joints)), diff --git a/dimos/robot/manipulators/xarm/config.py b/dimos/robot/manipulators/xarm/config.py index 410a1bf87c..f297a12792 100644 --- a/dimos/robot/manipulators/xarm/config.py +++ b/dimos/robot/manipulators/xarm/config.py @@ -23,14 +23,12 @@ HardwareComponent, HardwareType, make_gripper_joints, - 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, - coordinator_joint_mapping, joint_names, ) from dimos.utils.data import LfsPath @@ -70,7 +68,6 @@ def make_xarm7_sim_robot_config() -> RobotModelConfig: return make_xarm7_model_config( - name="arm", add_gripper=True, tf_extra_links=["link7"], home_joints=XARM7_SIM_HOME, @@ -78,13 +75,12 @@ def make_xarm7_sim_robot_config() -> RobotModelConfig: ) -def make_dual_xarm6_model_config(name: str = "robot") -> RobotModelConfig: +def make_dual_xarm6_model_config() -> RobotModelConfig: """Return one statically authored model containing two canonical xArm6 chains.""" left_joints = joint_names(6, prefix="left/joint") right_joints = joint_names(6, prefix="right/joint") canonical_joints = [*left_joints, *right_joints] return RobotModelConfig( - name=name, model_path=XARM_DUAL_MODEL_PATH, srdf_path=XARM_DUAL_SRDF_PATH, joint_names=canonical_joints, @@ -161,6 +157,7 @@ def make_xarm_hardware( auto_enable: bool = True, adapter_kwargs: dict[str, object] | None = None, home_joints: list[float] | None = None, + canonical_joint_names: list[str] | None = None, ) -> HardwareComponent: kwargs = _adapter_kwargs(home_joints) if adapter_kwargs: @@ -168,7 +165,7 @@ def make_xarm_hardware( return HardwareComponent( hardware_id=hw_id, hardware_type=HardwareType.MANIPULATOR, - joints=make_joints(hw_id, dof), + joints=canonical_joint_names or joint_names(dof), adapter_type=adapter_type, address=address, auto_enable=auto_enable, @@ -187,6 +184,7 @@ def xarm7_hardware( gripper_closed_position: float | None = None, mock_without_address: bool = False, home_joints: list[float] | None = None, + canonical_joint_names: list[str] | None = None, ) -> HardwareComponent: if global_config.simulation: return make_xarm_hardware( @@ -198,6 +196,7 @@ def xarm7_hardware( gripper_open_position=gripper_open_position, gripper_closed_position=gripper_closed_position, home_joints=home_joints, + canonical_joint_names=canonical_joint_names, ) address = global_config.xarm7_ip if mock_without_address and not address: @@ -208,6 +207,7 @@ def xarm7_hardware( gripper_open_position=gripper_open_position, gripper_closed_position=gripper_closed_position, home_joints=home_joints, + canonical_joint_names=canonical_joint_names, ) return make_xarm_hardware( hw_id, @@ -218,6 +218,7 @@ def xarm7_hardware( gripper_open_position=gripper_open_position, gripper_closed_position=gripper_closed_position, home_joints=home_joints, + canonical_joint_names=canonical_joint_names, ) @@ -229,6 +230,7 @@ def xarm6_hardware( gripper_closed_position: float | None = None, mock_without_address: bool = False, home_joints: list[float] | None = None, + canonical_joint_names: list[str] | None = None, ) -> HardwareComponent: if global_config.simulation: return make_xarm_hardware( @@ -240,6 +242,7 @@ def xarm6_hardware( gripper_open_position=gripper_open_position, gripper_closed_position=gripper_closed_position, home_joints=home_joints, + canonical_joint_names=canonical_joint_names, ) address = global_config.xarm6_ip if mock_without_address and not address: @@ -250,6 +253,7 @@ def xarm6_hardware( gripper_open_position=gripper_open_position, gripper_closed_position=gripper_closed_position, home_joints=home_joints, + canonical_joint_names=canonical_joint_names, ) return make_xarm_hardware( hw_id, @@ -260,25 +264,26 @@ def xarm6_hardware( gripper_open_position=gripper_open_position, gripper_closed_position=gripper_closed_position, home_joints=home_joints, + canonical_joint_names=canonical_joint_names, ) def make_xarm_model_config( - name: str, dof: int, *, + prefix: str = "", add_gripper: bool = True, x_offset: float = 0.0, y_offset: float = 0.0, z_offset: float = 0.0, pitch: float = 0.0, - joint_prefix: str | None = None, tf_extra_links: list[str] | None = None, home_joints: list[float] | None = None, pre_grasp_offset: float = 0.10, ) -> RobotModelConfig: xacro_args = { "dof": str(dof), + "prefix": prefix, "limited": "true", "attach_xyz": "0 0 0", "attach_rpy": "0 0 0", @@ -286,47 +291,42 @@ def make_xarm_model_config( 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}" + model_joint_names = joint_names(dof, prefix=f"{prefix}joint") + tip_link = f"{prefix}link_tcp" if add_gripper else f"{prefix}link{dof}" + collision_exclusions = [ + (f"{prefix}{left}", f"{prefix}{right}") for left, right in XARM_GRIPPER_COLLISION_EXCLUSIONS + ] return RobotModelConfig( - name=name, model_path=XARM_MODEL_PATH, base_pose=base_pose(x_offset, y_offset, z_offset, pitch), - joint_names=local_joint_names, - base_link="link_base", + joint_names=model_joint_names, + base_link=f"{prefix}link_base", planning_groups=[ PlanningGroupDefinition( name="manipulator", - joint_names=tuple(local_joint_names), - base_link="link_base", + joint_names=tuple(model_joint_names), + base_link=f"{prefix}link_base", tip_link=tip_link, ) ], package_paths=XARM_PACKAGE_PATHS, xacro_args=xacro_args, auto_convert_meshes=True, - collision_exclusion_pairs=(XARM_GRIPPER_COLLISION_EXCLUSIONS if add_gripper else []), - joint_name_mapping=coordinator_joint_mapping( - name, - dof, - joint_prefix=joint_prefix, - ), - gripper_hardware_id=name if add_gripper else None, - tf_extra_links=tf_extra_links or [], + collision_exclusion_pairs=collision_exclusions if add_gripper else [], + gripper_hardware_id="arm" if add_gripper else None, + tf_extra_links=[f"{prefix}{link}" for link in (tf_extra_links or [])], home_joints=home_joints or [0.0] * dof, pre_grasp_offset=pre_grasp_offset, ) def make_xarm6_model_config( - name: str = "arm", **kwargs: Any, ) -> RobotModelConfig: - return make_xarm_model_config(name, 6, **kwargs) + return make_xarm_model_config(6, **kwargs) def make_xarm7_model_config( - name: str = "arm", **kwargs: Any, ) -> RobotModelConfig: - return make_xarm_model_config(name, 7, **kwargs) + return make_xarm_model_config(7, **kwargs) diff --git a/dimos/robot/manipulators/xarm/test_model_config.py b/dimos/robot/manipulators/xarm/test_model_config.py index 6d6167669a..90605b044e 100644 --- a/dimos/robot/manipulators/xarm/test_model_config.py +++ b/dimos/robot/manipulators/xarm/test_model_config.py @@ -15,7 +15,10 @@ """xArm prepared-model configuration tests.""" from dimos.manipulation.planning.spec.validation import validate_robot_model_config -from dimos.robot.manipulators.xarm.config import make_dual_xarm6_model_config +from dimos.robot.manipulators.xarm.config import ( + make_dual_xarm6_model_config, + make_xarm6_model_config, +) def test_dual_xarm6_is_one_prepared_model_with_canonical_groups() -> None: @@ -32,3 +35,13 @@ def test_dual_xarm6_is_one_prepared_model_with_canonical_groups() -> None: ] assert config.planning_groups[0].joint_names == tuple(config.joint_names[:6]) assert config.planning_groups[1].joint_names == tuple(config.joint_names[6:]) + + +def test_prefixed_xarm_model_uses_coordinator_facing_names_in_asset() -> None: + config = make_xarm6_model_config(add_gripper=False, prefix="xarm_arm/") + + model = validate_robot_model_config(config) + + assert model.actuated_joint_names == config.joint_names + assert config.joint_names == [f"xarm_arm/joint{i}" for i in range(1, 7)] + assert config.planning_groups[0].tip_link == "xarm_arm/link6" diff --git a/docs/capabilities/manipulation/adding_a_custom_arm.md b/docs/capabilities/manipulation/adding_a_custom_arm.md index 3754ce4545..7b7e3e7202 100644 --- a/docs/capabilities/manipulation/adding_a_custom_arm.md +++ b/docs/capabilities/manipulation/adding_a_custom_arm.md @@ -610,16 +610,15 @@ yourarm_planner = manipulation_module( | Field | Description | |-------|-------------| | `model_path` | Path to `.urdf` or `.xacro` file | -| `joint_names` | Ordered controllable local model joint set (must match URDF); not itself a planning group | +| `joint_names` | Ordered canonical model joint set (must match the URDF and coordinator); 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) | | `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. +Coordinator-facing joint states and trajectories use the model's canonical +joint names unchanged. Keep hardware-native name translation inside the +hardware adapter. Planning-group `base_link`/`tip_link` values define kinematic chains and pose target frames. `base_link` is only the robot-scoped link placed by diff --git a/docs/capabilities/manipulation/openarm_integration.md b/docs/capabilities/manipulation/openarm_integration.md index 4dd5b587dd..37df1ed900 100644 --- a/docs/capabilities/manipulation/openarm_integration.md +++ b/docs/capabilities/manipulation/openarm_integration.md @@ -143,32 +143,33 @@ This gives you an interactive Python prompt with these functions: | Function | Purpose | |---|---| -| `robots()` | List configured robots (here: `["left_arm", "right_arm"]`) | -| `joints(robot_name)` | Read current joint positions (7 floats) | -| `ee(robot_name)` | Read current end-effector pose | +| `info()` | Inspect the configured model and its canonical joints | +| `groups()` | List planning groups (`left_arm`, `right_arm`, `both_arms`) | +| `joints()` | Read the complete 14-joint canonical model state | +| `ee(group_id)` | Read a planning group's end-effector pose | | `state()` | Module state: `IDLE`, `PLANNING`, `EXECUTING`, `FAULT`, etc. | -| `plan([q1..q7], robot_name)` | Plan a collision-free trajectory to a joint configuration | -| `plan_pose(x, y, z, robot_name=...)` | Plan to a Cartesian EE pose (preserves current orientation) | -| `preview(robot_name)` | Animate the planned path in Meshcat without executing | +| `plan_group(group_id, [q1..q7])` | Plan a collision-free trajectory for one arm | +| `plan_group_pose(group_id, x, y, z)` | Plan one arm to a Cartesian EE pose | +| `preview()` | Animate the planned path in Meshcat without executing | | `execute()` | Send the complete planned trajectory to the coordinator | -| `home(robot_name)` | Plan + execute to home joints | +| `home(group_id)` | Plan + execute one group to home joints | | `commands()` | Print all available functions | #### Example session — simple joint moves ```python skip ->>> robots() -['left_arm', 'right_arm'] +>>> [group.id for group in groups()] +['left_arm', 'right_arm', 'both_arms'] ->>> joints(robot_name="left_arm") -[0.02, -0.01, -0.13, 0.15, 0.17, -0.07, 0.10] +>>> joints() +[0.02, -0.01, -0.13, 0.15, 0.17, -0.07, 0.10, 0.01, 0.0, -0.1, 0.1, 0.2, 0.0, 0.0] >>> # One-liner: plan → preview in Meshcat → execute on hardware ->>> plan([0.3, 0, 0, 0, 0, 0, 0], robot_name="left_arm") and preview(robot_name="left_arm") and execute() +>>> plan_group("left_arm", [0.3, 0, 0, 0, 0, 0, 0]) and preview() and execute() True ->>> joints(robot_name="left_arm") -[0.30, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00] # arm is now at the commanded pose +>>> joints()[:7] +[0.30, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00] ``` `plan()` returns `True` on success, `False` if planning failed (check the coordinator terminal for `COLLISION_AT_GOAL`, `INVALID_START`, `NO_SOLUTION`, etc). The `and` chaining is an idiom — if any step fails, the next one is short-circuited. @@ -184,9 +185,9 @@ If you ever get stuck in a `FAULT` state (e.g. an invalid plan was sent), reset ```python skip >>> # Move both arms to mirrored poses ->>> plan([0.5, 0, 0, 0, 0, 0, 0], robot_name="left_arm") and execute() +>>> plan_group("left_arm", [0.5, 0, 0, 0, 0, 0, 0]) and execute() True ->>> plan([-0.5, 0, 0, 0, 0, 0, 0], robot_name="right_arm") and execute() +>>> plan_group("right_arm", [-0.5, 0, 0, 0, 0, 0, 0]) and execute() True ``` @@ -195,8 +196,8 @@ Each arm plans and executes independently — the coordinator runs both trajecto #### Example session — Cartesian target ```python skip ->>> ee(robot_name="left_arm") # see where the EE currently is ->>> plan_pose(0.1, 0.3, 0.5, robot_name="left_arm") and preview(robot_name="left_arm") +>>> ee("left_arm") # see where the EE currently is +>>> plan_group_pose("left_arm", 0.1, 0.3, 0.5) and preview() True >>> execute() True @@ -209,7 +210,7 @@ If you don't know which Cartesian targets are reachable, check first with the wo ```python skip >>> add_box("table", 0.4, 0.0, 0.1, w=0.6, h=0.4, d=0.05) # rectangular obstacle >>> add_sphere("ball", 0.3, 0.2, 0.4, radius=0.05) ->>> plan_pose(0.4, 0.0, 0.3, robot_name="left_arm") # now plans around it +>>> plan_group_pose("left_arm", 0.4, 0.0, 0.3) # now plans around it >>> remove("table") # id returned by add_* ``` diff --git a/docs/capabilities/manipulation/planning_groups.md b/docs/capabilities/manipulation/planning_groups.md index 67b611474f..0bdc7d2a88 100644 --- a/docs/capabilities/manipulation/planning_groups.md +++ b/docs/capabilities/manipulation/planning_groups.md @@ -8,18 +8,15 @@ torso, without confusing that group with the robot's hardware identity. | 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. | +| Model | The single configured `RobotModelConfig`. | +| Planning group | A named subset of the model's controllable joints. | +| Planning group ID | Stable declared group name, such as `left_arm`. | +| Canonical joint name | Joint name used unchanged by the model, planner, coordinator, and visualization. | +| Generated plan | Planning artifact containing selected group IDs, geometric waypoints, and one synchronized canonical 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`. +The configured model owns one canonical joint namespace. Planning groups select +subsets of that namespace; they do not rename or prefix joints. ## Discovering planning groups @@ -33,7 +30,7 @@ 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 +3. Fallback generation of one `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. @@ -63,7 +60,7 @@ unique serial target frame. 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 +This field is the model's ordered canonical joint set, not an implicit planning group. Fallback removes terminal prismatic leaves first, including branched finger @@ -84,16 +81,14 @@ 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. +unnamed in the group's joint order or named with the group's canonical joints. ```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], + "left_arm": JointState( + name=["left/joint1", "left/joint2", "left/joint3"], + position=[0.0, -0.4, 0.2], ) } ) @@ -105,8 +100,8 @@ 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"], + {"left_arm": target_pose}, + auxiliary_groups=["torso"], ) ``` @@ -125,49 +120,37 @@ manip.preview_plan(plan) manip.execute_plan(plan) ``` -A generated plan is the execution boundary: execution never filters a -multi-robot plan. To execute one robot, first plan only that robot's planning -group. - -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. +A generated plan is the execution boundary. The canonical trajectory is sent +unchanged to the coordinator. To execute a subset, first plan only that subset's +planning group. ## 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 +- a geometric path of `JointState` waypoints keyed by canonical joint names; +- one materialized synchronized `JointTrajectory` over the same selected canonical 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 -translates selected joint names at the coordinator boundary and invokes the -coordinator's sole trajectory task once without filling omitted joints in the -RPC trajectory. The task remains planning-group agnostic, claims its full -configured joint set, and holds omitted joints while executing the active -planned subset. A newly accepted trajectory replaces the task's current -trajectory. +parameterize the geometric path. Preview and execution forward the canonical +trajectory without renaming, splitting, or merging it. The coordinator's +trajectory task remains planning-group agnostic and holds omitted joints while +executing the selected subset. ## 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. +stripping. This is model 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. +Planning-group `base_link` and `tip_link` values are the source for chain bases +and pose target frames. Convenience pose APIs require an explicit group when +the model has multiple pose-targetable groups. 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. +depending on the blueprint. `joint_names` describes the ordered controllable +canonical model joint set. diff --git a/openspec/changes/refactor-control-coordinator-connections/tasks.md b/openspec/changes/refactor-control-coordinator-connections/tasks.md index 6b5c3a2d5d..c4a60692b7 100644 --- a/openspec/changes/refactor-control-coordinator-connections/tasks.md +++ b/openspec/changes/refactor-control-coordinator-connections/tasks.md @@ -27,18 +27,18 @@ The numbered sections are review and merge units. PR4a–PR4c may proceed in par ## 2. PR2 — Single-Robot Manipulation Cutover -- [ ] 2.1 Change manipulation configuration from a robot list/registry to exactly one prepared model configuration. -- [ ] 2.2 Remove `RobotName`, `WorldRobotID`, public robot selectors, and manipulation robot lookup/list APIs. -- [ ] 2.3 Remove local/global joint-name types, robot-prefix parse/make helpers, and mapping inversion from planning and execution. -- [ ] 2.4 Change world state, collision, FK, and model access APIs to operate on the one logical model without a robot-ID parameter. -- [ ] 2.5 Change planning-group selection to select canonical joint subsets and group-specific frames without robot membership. -- [ ] 2.6 Simplify Pink and other multi-target solvers to solve several groups within one model rather than grouping by robot. -- [ ] 2.7 Replace per-robot initialization state, trajectory splitting, and execution targets with one canonical joint-state and trajectory flow. -- [ ] 2.8 Replace visualization and world-monitor maps keyed by robot ID with one logical model state. -- [ ] 2.9 Remove robot selectors from semantic manipulation RPCs, DimOS `Spec` Protocols, skills, MCP schemas, and callers. -- [ ] 2.10 Convert the dual-xArm mock and authored bimanual blueprints to the static prepared model and planning groups. -- [ ] 2.11 Delete obsolete multi-robot code paths and compatibility properties exposed by the migrated manipulation surfaces. -- [ ] 2.12 Run focused manipulation, planning, collision, FK, execution, visualization, skill-schema, and blueprint-build tests. +- [x] 2.1 Change manipulation configuration from a robot list/registry to exactly one prepared model configuration. +- [x] 2.2 Remove `RobotName`, `WorldRobotID`, public robot selectors, and manipulation robot lookup/list APIs. +- [x] 2.3 Remove local/global joint-name types, robot-prefix parse/make helpers, and mapping inversion from planning and execution. +- [x] 2.4 Change world state, collision, FK, and model access APIs to operate on the one logical model without a robot-ID parameter. +- [x] 2.5 Change planning-group selection to select canonical joint subsets and group-specific frames without robot membership. +- [x] 2.6 Simplify Pink and other multi-target solvers to solve several groups within one model rather than grouping by robot. +- [x] 2.7 Replace per-robot initialization state, trajectory splitting, and execution targets with one canonical joint-state and trajectory flow. +- [x] 2.8 Replace visualization and world-monitor maps keyed by robot ID with one logical model state. +- [x] 2.9 Remove robot selectors from semantic manipulation RPCs, DimOS `Spec` Protocols, skills, MCP schemas, and callers. +- [x] 2.10 Convert the dual-xArm mock and authored bimanual blueprints to the static prepared model and planning groups. +- [x] 2.11 Delete obsolete multi-robot code paths and compatibility properties exposed by the migrated manipulation surfaces. +- [x] 2.12 Run focused manipulation, planning, collision, FK, execution, visualization, skill-schema, and blueprint-build tests. ## 3. PR3 — Scalar Control and Lifecycle Contracts