From c7eb5d93497aa9364f48dc0c86c74b0409b4f827 Mon Sep 17 00:00:00 2001 From: cc Date: Wed, 1 Jul 2026 16:47:31 -0700 Subject: [PATCH 01/14] spec: archive 1 --- .../planning-groups-foundation/.openspec.yaml | 2 -- .../planning-groups-foundation/README.md | 3 -- .../planning-groups-foundation/design.md | 29 ------------------- .../planning-groups-foundation/proposal.md | 28 ------------------ .../planning-groups-foundation/tasks.md | 18 ------------ .../manipulation-planning-config/spec.md | 4 ++- .../specs/planning-group-foundation/spec.md | 4 ++- 7 files changed, 6 insertions(+), 82 deletions(-) delete mode 100644 openspec/changes/planning-groups-foundation/.openspec.yaml delete mode 100644 openspec/changes/planning-groups-foundation/README.md delete mode 100644 openspec/changes/planning-groups-foundation/design.md delete mode 100644 openspec/changes/planning-groups-foundation/proposal.md delete mode 100644 openspec/changes/planning-groups-foundation/tasks.md rename openspec/{changes/planning-groups-foundation => }/specs/manipulation-planning-config/spec.md (90%) rename openspec/{changes/planning-groups-foundation => }/specs/planning-group-foundation/spec.md (95%) diff --git a/openspec/changes/planning-groups-foundation/.openspec.yaml b/openspec/changes/planning-groups-foundation/.openspec.yaml deleted file mode 100644 index 34f9314d22..0000000000 --- a/openspec/changes/planning-groups-foundation/.openspec.yaml +++ /dev/null @@ -1,2 +0,0 @@ -schema: spec-driven -created: 2026-06-29 diff --git a/openspec/changes/planning-groups-foundation/README.md b/openspec/changes/planning-groups-foundation/README.md deleted file mode 100644 index 820f444c27..0000000000 --- a/openspec/changes/planning-groups-foundation/README.md +++ /dev/null @@ -1,3 +0,0 @@ -# planning-groups-foundation - -Extract PR 1 from reference branch cc/spec/movegroup: introduce planning-group foundations, RobotConfig group generation, manipulator group declarations, and foundation docs/tests. diff --git a/openspec/changes/planning-groups-foundation/design.md b/openspec/changes/planning-groups-foundation/design.md deleted file mode 100644 index 287006f709..0000000000 --- a/openspec/changes/planning-groups-foundation/design.md +++ /dev/null @@ -1,29 +0,0 @@ -## Context - -The current branch introduces planning groups across the stack, but the reviewable foundation is the abstraction itself: stable group IDs, local/global joint mapping, group definitions, RobotConfig discovery/generation, and removal of robot-scoped end-effector metadata. - -## Goals / Non-Goals - -**Goals:** -- Define planning groups as the only source for chain base links and pose target tip links. -- Preserve robot-scoped `base_link` as placement/weld/strip metadata, not a planning-chain selector. -- Let robot configs create explicit or discovered planning groups. -- Provide small, direct tests that prove group IDs, joint-name mapping, discovery, and strict config validation. - -**Non-Goals:** -- Do not migrate world backends or planners in this PR. -- Do not migrate `ManipulationModule` public APIs in this PR. -- Do not include visualization or control changes. - -## Decisions - -- Use `PlanningGroupDefinition` for config-time group declarations and `PlanningGroup` for resolved runtime group IDs. -- Use globally scoped group IDs such as `/` at API boundaries. -- Keep `RobotModelConfig.base_link` for robot placement semantics only. -- Reject legacy robot-scoped `end_effector_link`; pose target frames come from group `tip_link` or SRDF discovery. -- Keep this PR compatible enough for later stacked PRs; if a later change is needed for tests, prefer a small compatibility wrapper over pulling later files forward. - -## Risks / Trade-offs - -- Earlier PRs may need temporary compatibility paths because the full branch removes legacy behavior across several layers. -- Discovery behavior must be clear: high-level `RobotConfig.to_robot_model_config()` can discover/fallback groups, while direct `RobotModelConfig(...)` requires explicit groups. diff --git a/openspec/changes/planning-groups-foundation/proposal.md b/openspec/changes/planning-groups-foundation/proposal.md deleted file mode 100644 index c6c7261bc1..0000000000 --- a/openspec/changes/planning-groups-foundation/proposal.md +++ /dev/null @@ -1,28 +0,0 @@ -## Why - -The planning-group refactor is currently bundled in one large branch. Reviewers need a small first PR that establishes the core planning-group abstraction and robot configuration plumbing without also reviewing planner algorithms, module APIs, visualization, or control integration. - -This change extracts PR 1 from the reference branch/worktree `cc/spec/movegroup`. Treat the existing branch as a working reference implementation; do not redesign the feature unless required to keep this slice independently valid. - -## What Changes - -- Add planning-group identifiers, definitions, runtime models, registry helpers, and joint-name conversion utilities. -- Update planning spec/config models so robot-scoped `end_effector_link` is removed and pose target frames are represented by planning-group `tip_link` values. -- Add planning-group conversion/discovery support through manipulation planning config models. -- Update manipulator robot configs to declare or discover explicit planning groups. -- Add foundation-level unit tests and minimal concept docs. - -## Capabilities - -### New Capabilities -- `planning-group-foundation`: Core planning-group data model, discovery, robot config generation, and validation contracts. - -### Modified Capabilities -- `manipulation-planning-config`: Robot model configuration uses planning groups as the source of chain and pose-target metadata. - -## Impact - -- Base branch: `main`. -- Reference implementation: branch/worktree `cc/spec/movegroup`; stabilize to a commit before assigning agents. -- Primary files: `dimos/manipulation/planning/groups/*`, `dimos/manipulation/planning/spec/*`, `dimos/robot/manipulators/*/config.py`, `dimos/manipulation/planning/groups/test_planning_groups.py`, foundation tests, and planning-group docs. -- Out of scope: world backends, IK/RRT algorithms, `ManipulationModule`, Viser UI, and control changes. diff --git a/openspec/changes/planning-groups-foundation/tasks.md b/openspec/changes/planning-groups-foundation/tasks.md deleted file mode 100644 index 0b625f1e79..0000000000 --- a/openspec/changes/planning-groups-foundation/tasks.md +++ /dev/null @@ -1,18 +0,0 @@ -## 1. Extract foundation files - -- [x] 1.1 Start from `main` and use `cc/spec/movegroup` as the reference implementation. -- [x] 1.2 Bring over `dimos/manipulation/planning/groups/*`. -- [x] 1.3 Bring over required planning spec/config changes only. -- [x] 1.4 Bring over manipulator config group declarations without `dimos/robot/config.py`. - -## 2. Tests and docs - -- [x] 2.1 Bring over `dimos/manipulation/planning/groups/test_planning_groups.py`. -- [x] 2.2 Remove stale `dimos/robot/test_config.py` scope from this PR. -- [x] 2.3 Bring over only foundation docs needed to explain planning groups and custom-arm config. - -## 3. Validation - -- [x] 3.1 Run targeted planning-group tests under `dimos/manipulation/planning/groups/test_planning_groups.py`. -- [x] 3.2 Run targeted mypy on changed production files. -- [x] 3.3 Confirm no files from world/IK/RRT/module/Viser/control were pulled in accidentally. diff --git a/openspec/changes/planning-groups-foundation/specs/manipulation-planning-config/spec.md b/openspec/specs/manipulation-planning-config/spec.md similarity index 90% rename from openspec/changes/planning-groups-foundation/specs/manipulation-planning-config/spec.md rename to openspec/specs/manipulation-planning-config/spec.md index ecb76b89b5..bd13e49196 100644 --- a/openspec/changes/planning-groups-foundation/specs/manipulation-planning-config/spec.md +++ b/openspec/specs/manipulation-planning-config/spec.md @@ -1,4 +1,6 @@ -## MODIFIED Requirements +# manipulation-planning-config Specification + +## Requirements ### Requirement: Robot model configuration must separate placement from planning chains Robot model configuration MUST treat robot-scoped `base_link` as the link placed by `base_pose` and used by backend weld/strip behavior. Planning-chain base links and pose target frames MUST come from planning groups. diff --git a/openspec/changes/planning-groups-foundation/specs/planning-group-foundation/spec.md b/openspec/specs/planning-group-foundation/spec.md similarity index 95% rename from openspec/changes/planning-groups-foundation/specs/planning-group-foundation/spec.md rename to openspec/specs/planning-group-foundation/spec.md index 2da6cd9e9e..6124a3bf5c 100644 --- a/openspec/changes/planning-groups-foundation/specs/planning-group-foundation/spec.md +++ b/openspec/specs/planning-group-foundation/spec.md @@ -1,4 +1,6 @@ -## ADDED Requirements +# planning-group-foundation Specification + +## Requirements ### Requirement: Planning group definitions must identify kinematic chains The planning configuration MUST represent kinematic chain metadata with planning-group definitions containing group-local joint names, a chain base link, and an optional pose-target tip link. From 94cd3ce72d4b347aa5fc1caa13ff9b24ec38c52c Mon Sep 17 00:00:00 2001 From: cc Date: Wed, 1 Jul 2026 19:31:20 -0700 Subject: [PATCH 02/14] feat: add planning group world queries --- .../planning/monitor/test_world_monitor.py | 150 +++++++++++++++- .../planning/monitor/world_monitor.py | 166 +++++++++++++++++- dimos/manipulation/planning/spec/protocols.py | 9 + .../planning/world/drake_world.py | 145 ++++++++++++--- .../planning/world/roboplan_world.py | 98 +++++++++-- .../world/test_drake_world_planning_groups.py | 153 ++++++++++++++++ dimos/manipulation/test_roboplan_world.py | 90 ++++++++++ .../planning-world-monitor-groups/tasks.md | 21 +-- 8 files changed, 780 insertions(+), 52 deletions(-) create mode 100644 dimos/manipulation/planning/world/test_drake_world_planning_groups.py diff --git a/dimos/manipulation/planning/monitor/test_world_monitor.py b/dimos/manipulation/planning/monitor/test_world_monitor.py index 355c1e6712..aed6d6a9a5 100644 --- a/dimos/manipulation/planning/monitor/test_world_monitor.py +++ b/dimos/manipulation/planning/monitor/test_world_monitor.py @@ -17,6 +17,9 @@ from pathlib import Path from typing import Any +import numpy as np +import pytest + from dimos.manipulation.planning import factory as planning_factory from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.monitor import world_monitor as world_monitor_module @@ -25,15 +28,48 @@ 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 + + +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: + _ = max_age + return self._stale + + +class _ScratchContext: + def __enter__(self) -> str: + return "scratch" + + def __exit__(self, exc_type: object, exc: object, traceback: object) -> bool: + return False class FakeWorld: def __init__(self) -> None: - self.calls: list[tuple[str, Any]] = [] + self.calls: list[tuple[Any, ...]] = [] + self.configs: dict[str, RobotModelConfig] = {} def add_robot(self, config): self.calls.append(("add_robot", config)) - return "robot-1" + robot_id = f"robot-{len(self.configs) + 1}" + self.configs[robot_id] = config + return robot_id def get_robot_ids(self): return [] @@ -70,12 +106,14 @@ def get_live_context(self): return None def scratch_context(self): - return self + self.calls.append(("scratch_context", None)) + return _ScratchContext() def sync_from_joint_state(self, robot_id, joint_state) -> None: return None def set_joint_state(self, ctx, robot_id, joint_state) -> None: + self.calls.append(("set_joint_state", ctx, robot_id, joint_state)) return None def get_joint_state(self, ctx, robot_id): @@ -96,12 +134,20 @@ def check_edge_collision_free(self, robot_id, start, end, step_size: float = 0.0 def get_ee_pose(self, ctx, robot_id): return None + def get_group_ee_pose(self, ctx, group_id): + self.calls.append(("get_group_ee_pose", ctx, group_id)) + return PoseStamped(position=Vector3(1, 2, 3), orientation=Quaternion([0, 0, 0, 1])) + def get_link_pose(self, ctx, robot_id, link_name): return [] def get_jacobian(self, ctx, robot_id): return [] + def get_group_jacobian(self, ctx, group_id): + self.calls.append(("get_group_jacobian", ctx, group_id)) + return np.ones((6, 2)) + def get_visualization_url(self): return None @@ -165,6 +211,26 @@ def _robot_config() -> RobotModelConfig: ) +def _robot_config_with_groups(groups: list[PlanningGroupDefinition]) -> RobotModelConfig: + return _robot_config().model_copy(update={"planning_groups": groups}) + + +def _three_joint_reordered_group_config() -> RobotModelConfig: + return _robot_config().model_copy( + update={ + "joint_names": ["j1", "j2", "j3"], + "planning_groups": [ + PlanningGroupDefinition( + name="manipulator", + joint_names=("j2", "j1"), + base_link="base", + tip_link="ee", + ) + ], + } + ) + + def test_world_monitor_add_robot_records_scene_without_visualization_probe() -> None: fake_world = FakeWorld() fake_viz = FakeViz() @@ -209,3 +275,81 @@ def test_create_planning_specs_wraps_existing_world(monkeypatch) -> None: assert planning_specs.world_monitor.visualization is None assert planning_specs.kinematics is fake_kinematics assert planning_specs.planner is fake_planner + + +def test_world_monitor_exposes_planning_groups_and_duplicate_names_do_not_mutate() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + monitor.add_robot(_robot_config()) + + assert [group.id for group in monitor.planning_groups.list()] == ["arm/manipulator"] + with pytest.raises(ValueError, match="already registered"): + monitor.add_robot(_robot_config()) + assert [call[0] for call in fake_world.calls].count("add_robot") == 1 + + +def test_current_group_joint_state_uses_public_names_in_group_order() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + robot_id = monitor.add_robot(_three_joint_reordered_group_config()) + monitor._state_monitors[robot_id] = _FakeStateMonitor([0.1, 0.2, 0.3]) # type: ignore[attr-defined] + + state = monitor.get_current_group_joint_state("arm/manipulator") + + assert state.name == ["arm/j2", "arm/j1"] + assert state.position == [0.2, 0.1] + + +def test_group_kinematics_route_to_backend_and_merge_group_state() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + robot_id = monitor.add_robot(_three_joint_reordered_group_config()) + monitor._state_monitors[robot_id] = _FakeStateMonitor([0.1, 0.2, 0.3]) # type: ignore[attr-defined] + + pose = monitor.get_group_ee_pose( + "arm/manipulator", JointState({"name": ["arm/j2", "arm/j1"], "position": [0.8, 0.9]}) + ) + jacobian = monitor.get_group_jacobian( + "arm/manipulator", JointState({"name": ["arm/j2", "arm/j1"], "position": [0.3, 0.4]}) + ) + + set_calls = [call for call in fake_world.calls if call[0] == "set_joint_state"] + assert set_calls[0][3].name == ["j1", "j2", "j3"] + assert set_calls[0][3].position == [0.9, 0.8, 0.3] + assert set_calls[1][3].name == ["j1", "j2", "j3"] + assert set_calls[1][3].position == [0.4, 0.3, 0.3] + assert pose.position.x == 1 + assert jacobian.shape == (6, 2) + assert ("get_group_ee_pose", "scratch", "arm/manipulator") in fake_world.calls + assert ("get_group_jacobian", "scratch", "arm/manipulator") in fake_world.calls + + +def test_legacy_wrappers_fail_for_no_pose_and_ambiguous_pose_groups() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + no_pose_id = monitor.add_robot( + _robot_config_with_groups( + [PlanningGroupDefinition(name="base", joint_names=("j1",), base_link="base")] + ) + ) + with pytest.raises(ValueError, match="no pose-targetable"): + monitor.get_ee_pose(no_pose_id, JointState({"name": ["j1", "j2"], "position": [0.0, 0.0]})) + + fake_world2 = FakeWorld() + monitor2 = world_monitor_module.WorldMonitor(world=fake_world2) # type: ignore[arg-type] + ambiguous_id = monitor2.add_robot( + _robot_config_with_groups( + [ + PlanningGroupDefinition( + name="a", joint_names=("j1",), base_link="base", tip_link="ee1" + ), + PlanningGroupDefinition( + name="b", joint_names=("j2",), base_link="base", tip_link="ee2" + ), + ] + ) + ) + with pytest.raises(ValueError, match="pose-targetable planning groups"): + monitor2.get_jacobian( + ambiguous_id, JointState({"name": ["j1", "j2"], "position": [0.0, 0.0]}) + ) diff --git a/dimos/manipulation/planning/monitor/world_monitor.py b/dimos/manipulation/planning/monitor/world_monitor.py index 5e12568874..3cd54d05c4 100644 --- a/dimos/manipulation/planning/monitor/world_monitor.py +++ b/dimos/manipulation/planning/monitor/world_monitor.py @@ -21,6 +21,8 @@ from typing import TYPE_CHECKING, Any from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT +from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry +from dimos.manipulation.planning.groups.utils import filter_joint_state_to_selected_joints from dimos.manipulation.planning.monitor.robot_state_monitor import RobotStateMonitor from dimos.manipulation.planning.monitor.world_obstacle_monitor import WorldObstacleMonitor from dimos.manipulation.planning.spec.models import PlanningSceneInfo @@ -40,6 +42,8 @@ CollisionObjectMessage, JointPath, Obstacle, + PlanningGroupID, + RobotName, WorldRobotID, ) from dimos.msgs.vision_msgs.Detection3D import Detection3D @@ -61,6 +65,8 @@ def __init__( self._lock = threading.RLock() self._robot_joints: dict[WorldRobotID, list[str]] = {} self._robot_configs: dict[WorldRobotID, RobotModelConfig] = {} + self._robot_ids_by_name: dict[RobotName, WorldRobotID] = {} + self._planning_groups = PlanningGroupRegistry() self._state_monitors: dict[WorldRobotID, RobotStateMonitor] = {} self._obstacle_monitor: WorldObstacleMonitor | None = None self._viz_thread: threading.Thread | None = None @@ -72,12 +78,23 @@ def __init__( def add_robot(self, config: RobotModelConfig) -> WorldRobotID: """Add a robot. Returns robot_id.""" with self._lock: + if config.name in self._robot_ids_by_name: + raise ValueError(f"Robot name '{config.name}' is already registered") + new_groups = PlanningGroupRegistry() + new_groups.add_robot(config) robot_id = self._world.add_robot(config) self._robot_joints[robot_id] = config.joint_names self._robot_configs[robot_id] = config + self._robot_ids_by_name[config.name] = robot_id + self._planning_groups.add_robot(config) logger.info(f"Added robot '{config.name}' as '{robot_id}'") return robot_id + @property + def planning_groups(self) -> PlanningGroupRegistry: + """Registered public planning groups.""" + return self._planning_groups + def planning_scene_info(self) -> PlanningSceneInfo: """Return a stable metadata snapshot of the initialized planning scene.""" with self._lock: @@ -294,6 +311,116 @@ def get_current_joint_state(self, robot_id: WorldRobotID) -> JointState | None: ctx = self._world.get_live_context() return self._world.get_joint_state(ctx, robot_id) + def current_global_joint_state(self, max_age: float = 1.0) -> JointState: + """Return current state for all fresh robots with public global joint names.""" + names: list[str] = [] + positions: list[float] = [] + for robot_name, robot_id in self._robot_ids_by_name.items(): + if robot_id in self._state_monitors and self.is_state_stale(robot_id, max_age): + continue + state = self.get_current_joint_state(robot_id) + if state is None: + continue + pos_by_name = dict(zip(state.name, state.position, strict=True)) + for local_name in self._robot_joints.get(robot_id, []): + if local_name in pos_by_name: + names.append(f"{robot_name}/{local_name}") + positions.append(float(pos_by_name[local_name])) + return JointState({"name": names, "position": positions}) + + def get_current_group_joint_state( + self, group_id: PlanningGroupID, max_age: float = 1.0 + ) -> JointState: + """Return current joint state scoped and ordered for one planning group.""" + group = self._planning_groups.get(group_id) + robot_id = self._robot_ids_by_name[group.robot_name] + if robot_id in self._state_monitors and self.is_state_stale(robot_id, max_age): + raise ValueError(f"Current state for robot '{group.robot_name}' is stale") + state = self.get_current_joint_state(robot_id) + if state is None: + raise ValueError(f"Current state for robot '{group.robot_name}' is unavailable") + return filter_joint_state_to_selected_joints( + state, group.joint_names, group.local_joint_names + ) + + def _current_robot_joint_state_for_group( + self, group_id: PlanningGroupID, max_age: float = 1.0 + ) -> JointState: + """Return current full robot state for a planning group.""" + group = self._planning_groups.get(group_id) + robot_id = self._robot_ids_by_name[group.robot_name] + if robot_id in self._state_monitors and self.is_state_stale(robot_id, max_age): + raise ValueError(f"Current state for robot '{group.robot_name}' is stale") + state = self.get_current_joint_state(robot_id) + if state is None: + raise ValueError(f"Current state for robot '{group.robot_name}' is unavailable") + return state + + def _full_robot_joint_state_from_input( + self, robot_id: WorldRobotID, joint_state: JointState + ) -> JointState | None: + """Return a full robot-local state if input contains all robot joints.""" + config = self._robot_configs[robot_id] + if not joint_state.name: + if len(joint_state.position) != len(config.joint_names): + return None + return JointState( + {"name": list(config.joint_names), "position": list(joint_state.position)} + ) + if len(joint_state.name) != len(joint_state.position): + raise ValueError("JointState name and position lengths must match") + resolved_positions: dict[str, float] = {} + global_prefix = f"{config.name}/" + for name, position in zip(joint_state.name, joint_state.position, strict=True): + if name in config.joint_names: + resolved_name = name + elif name in config.joint_name_mapping: + resolved_name = config.joint_name_mapping[name] + elif name.startswith(global_prefix): + resolved_name = name[len(global_prefix) :] + else: + resolved_name = config.get_urdf_joint_name(name) + if resolved_name not in config.joint_names: + return None + if resolved_name in resolved_positions: + raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") + resolved_positions[resolved_name] = float(position) + if set(resolved_positions) != set(config.joint_names): + return None + return JointState( + { + "name": list(config.joint_names), + "position": [resolved_positions[name] for name in config.joint_names], + } + ) + + def _joint_state_for_group_query( + self, group_id: PlanningGroupID, joint_state: JointState | None + ) -> JointState: + """Normalize full-robot or group-scoped input to full robot-local state.""" + group = self._planning_groups.get(group_id) + robot_id = self._robot_ids_by_name[group.robot_name] + if joint_state is None: + return self._current_robot_joint_state_for_group(group_id) + full_state = self._full_robot_joint_state_from_input(robot_id, joint_state) + if full_state is not None: + return full_state + current_state = self._current_robot_joint_state_for_group(group_id) + group_state = filter_joint_state_to_selected_joints( + joint_state, group.joint_names, group.local_joint_names + ) + positions_by_name = dict(zip(current_state.name, current_state.position, strict=True)) + for local_name, position in zip(group.local_joint_names, group_state.position, strict=True): + if local_name not in positions_by_name: + raise ValueError(f"Current state is missing group joint '{local_name}'") + positions_by_name[local_name] = float(position) + return JointState( + { + "name": list(current_state.name), + "position": [positions_by_name[name] for name in current_state.name], + } + ) + def get_current_velocities(self, robot_id: WorldRobotID) -> JointState | None: """Get current joint velocities as JointState. Returns None if not available.""" if robot_id in self._state_monitors: @@ -367,15 +494,23 @@ def get_ee_pose( self, robot_id: WorldRobotID, joint_state: JointState | None = None ) -> PoseStamped: """Get end-effector pose. Uses current state if joint_state is None.""" + robot_name = self._robot_configs[robot_id].name + group_id = self._planning_groups.primary_pose_group_id_for_robot(robot_name) + if group_id is None: + raise ValueError(f"Robot '{robot_name}' has no pose-targetable planning group") + return self.get_group_ee_pose(group_id, joint_state) + + def get_group_ee_pose( + self, group_id: PlanningGroupID, joint_state: JointState | None = None + ) -> PoseStamped: + """Get planning-group tip pose. Uses current group state if joint_state is None.""" + group = self._planning_groups.get(group_id) + robot_id = self._robot_ids_by_name[group.robot_name] with self._world.scratch_context() as ctx: - # If no state provided, fetch current from state monitor - if joint_state is None: - joint_state = self.get_current_joint_state(robot_id) - - if joint_state is not None: - self._world.set_joint_state(ctx, robot_id, joint_state) + normalized_state = self._joint_state_for_group_query(group_id, joint_state) + self._world.set_joint_state(ctx, robot_id, normalized_state) - return self._world.get_ee_pose(ctx, robot_id) + return self._world.get_group_ee_pose(ctx, group_id) def get_link_pose( self, robot_id: WorldRobotID, link_name: str, joint_state: JointState | None = None @@ -411,9 +546,22 @@ def get_link_pose( def get_jacobian(self, robot_id: WorldRobotID, joint_state: JointState) -> NDArray[np.float64]: """Get 6xN Jacobian matrix.""" + robot_name = self._robot_configs[robot_id].name + group_id = self._planning_groups.primary_pose_group_id_for_robot(robot_name) + if group_id is None: + raise ValueError(f"Robot '{robot_name}' has no pose-targetable planning group") + return self.get_group_jacobian(group_id, joint_state) + + def get_group_jacobian( + self, group_id: PlanningGroupID, joint_state: JointState + ) -> NDArray[np.float64]: + """Get 6xN planning-group Jacobian matrix.""" + group = self._planning_groups.get(group_id) + robot_id = self._robot_ids_by_name[group.robot_name] with self._world.scratch_context() as ctx: - self._world.set_joint_state(ctx, robot_id, joint_state) - return self._world.get_jacobian(ctx, robot_id) + normalized_state = self._joint_state_for_group_query(group_id, joint_state) + self._world.set_joint_state(ctx, robot_id, normalized_state) + return self._world.get_group_jacobian(ctx, group_id) # Lifecycle diff --git a/dimos/manipulation/planning/spec/protocols.py b/dimos/manipulation/planning/spec/protocols.py index c7ee95ee0a..b5fea13a36 100644 --- a/dimos/manipulation/planning/spec/protocols.py +++ b/dimos/manipulation/planning/spec/protocols.py @@ -33,6 +33,7 @@ IKResult, JointPath, Obstacle, + PlanningGroupID, PlanningResult, PlanningSceneInfo, WorldRobotID, @@ -170,6 +171,14 @@ def get_jacobian(self, ctx: Any, robot_id: WorldRobotID) -> NDArray[np.float64]: """Get end-effector Jacobian (6 x n_joints).""" ... + def get_group_ee_pose(self, ctx: Any, group_id: PlanningGroupID) -> PoseStamped: + """Get planning-group tip pose.""" + ... + + def get_group_jacobian(self, ctx: Any, group_id: PlanningGroupID) -> NDArray[np.float64]: + """Get planning-group Jacobian (6 x n_group_joints).""" + ... + @runtime_checkable class VisualizationSpec(Protocol): diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index ca426ba340..12067f647c 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -25,11 +25,14 @@ import numpy as np +from dimos.manipulation.planning.groups.identifiers import is_global_joint_name +from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import ObstacleType from dimos.manipulation.planning.spec.models import ( JointPath, Obstacle, + PlanningGroupID, PlanningSceneInfo, WorldRobotID, ) @@ -184,6 +187,7 @@ def __init__(self, time_step: float = 0.0, enable_viz: bool = False) -> None: # Tracking data self._robots: dict[WorldRobotID, _RobotData] = {} + self._planning_groups = PlanningGroupRegistry() self._obstacles: dict[str, _ObstacleData] = {} self._robot_counter = 0 self._obstacle_counter = 0 @@ -207,6 +211,11 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: raise RuntimeError("Cannot add robot after world is finalized") with self._lock: + if self._planning_groups.groups_for_robot(config.name): + raise ValueError(f"Robot name '{config.name}' is already registered") + new_groups = PlanningGroupRegistry() + new_groups.add_robot(config) + self._robot_counter += 1 robot_id = f"robot_{self._robot_counter}" @@ -215,9 +224,16 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: self._validate_joints(config, model_instance) - ee_frame = self._plant.GetBodyByName( - config.end_effector_link, model_instance - ).body_frame() + ee_link = config.base_link + try: + primary_group_id = new_groups.primary_pose_group_id_for_robot(config.name) + except ValueError: + primary_group_id = None + if primary_group_id is not None: + primary_group = new_groups.get(primary_group_id) + if primary_group.tip_link is not None: + ee_link = primary_group.tip_link + ee_frame = self._plant.GetBodyByName(ee_link, model_instance).body_frame() base_frame = self._plant.GetBodyByName(config.base_link, model_instance).body_frame() # Preview (yellow ghost) — always a separate instance per robot @@ -235,6 +251,7 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: base_frame=base_frame, preview_model_instance=preview_model_instance, ) + self._planning_groups.add_robot(config) logger.info(f"Added robot '{robot_id}' ({config.name})") return robot_id @@ -317,6 +334,11 @@ def get_robot_config(self, robot_id: WorldRobotID) -> RobotModelConfig: raise KeyError(f"Robot '{robot_id}' not found") return self._robots[robot_id].config + @property + def planning_groups(self) -> PlanningGroupRegistry: + """Registered planning groups for this world.""" + return self._planning_groups + def get_joint_limits( self, robot_id: WorldRobotID ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: @@ -768,8 +790,7 @@ def sync_from_joint_state(self, robot_id: WorldRobotID, joint_state: JointState) if not self._finalized or self._plant_context is None: return # Silently ignore before finalization - # Extract positions as numpy array for internal use - positions = np.array(joint_state.position, dtype=np.float64) + positions = self._joint_state_to_q(robot_id, joint_state) with self._lock: self._set_positions_internal(self._plant_context, robot_id, positions) @@ -787,8 +808,7 @@ def set_joint_state( if not self._finalized: raise RuntimeError("World must be finalized first") - # Extract positions as numpy array for internal use - positions = np.array(joint_state.position, dtype=np.float64) + positions = self._joint_state_to_q(robot_id, joint_state) # Get plant context from diagram context plant_ctx = self._diagram.GetMutableSubsystemContext(self._plant, ctx) @@ -809,6 +829,51 @@ 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] + if not joint_state.name: + if len(joint_state.position) != len(robot_data.config.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") + name_to_pos: dict[str, float] = {} + for name, position in zip(joint_state.name, joint_state.position, strict=True): + if name in robot_data.config.joint_names: + resolved_name = name + elif name in robot_data.config.joint_name_mapping: + resolved_name = robot_data.config.joint_name_mapping[name] + elif is_global_joint_name(name): + prefix = f"{robot_data.config.name}/" + if not name.startswith(prefix): + continue + resolved_name = name[len(prefix) :] + if resolved_name not in robot_data.config.joint_names: + raise ValueError(f"Unknown global joint name for DrakeWorld: {name}") + else: + resolved_name = robot_data.config.get_urdf_joint_name(name) + if resolved_name in name_to_pos: + raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") + name_to_pos[resolved_name] = float(position) + missing = [name for name in robot_data.config.joint_names if name not in name_to_pos] + if missing: + raise ValueError(f"JointState missing joints for DrakeWorld: {missing}") + return np.asarray([name_to_pos[name] for name in robot_data.config.joint_names]) + + def _robot_id_for_group(self, group_id: PlanningGroupID) -> WorldRobotID: + group = self._planning_groups.get(group_id) + matches = [ + rid for rid, data in self._robots.items() if data.config.name == group.robot_name + ] + if not matches: + raise KeyError(f"No robot registered for planning group '{group_id}'") + return matches[0] + def get_joint_state(self, ctx: Context, robot_id: WorldRobotID) -> JointState: """Get robot joint state from given context.""" if not self._finalized: @@ -905,16 +970,28 @@ def check_edge_collision_free( def get_ee_pose(self, ctx: Context, robot_id: WorldRobotID) -> PoseStamped: """Get end-effector pose.""" - if not self._finalized: - raise RuntimeError("World must be finalized first") - if robot_id not in self._robots: raise KeyError(f"Robot '{robot_id}' not found") - robot_data = self._robots[robot_id] + group_id = self._planning_groups.primary_pose_group_id_for_robot(robot_data.config.name) + if group_id is None: + raise ValueError( + f"Robot '{robot_data.config.name}' has no pose-targetable planning group" + ) + return self.get_group_ee_pose(ctx, group_id) + + def get_group_ee_pose(self, ctx: Context, group_id: PlanningGroupID) -> PoseStamped: + """Get planning-group tip pose.""" + if not self._finalized: + raise RuntimeError("World must be finalized first") + + group = self._planning_groups.get(group_id) + if group.tip_link is None: + raise ValueError(f"Planning group '{group_id}' has no tip link") + robot_data = self._robots[self._robot_id_for_group(group_id)] plant_ctx = self._diagram.GetSubsystemContext(self._plant, ctx) - ee_body = robot_data.ee_frame.body() + ee_body = self._plant.GetBodyByName(group.tip_link, robot_data.model_instance) X_WE = self._plant.EvalBodyPoseInWorld(plant_ctx, ee_body) # Extract position and quaternion from Drake transform @@ -955,30 +1032,58 @@ def get_jacobian(self, ctx: Context, robot_id: WorldRobotID) -> NDArray[np.float Rows: [vx, vy, vz, wx, wy, wz] (linear, then angular) """ - if not self._finalized: - raise RuntimeError("World must be finalized first") - if robot_id not in self._robots: raise KeyError(f"Robot '{robot_id}' not found") - robot_data = self._robots[robot_id] + group_id = self._planning_groups.primary_pose_group_id_for_robot(robot_data.config.name) + if group_id is None: + raise ValueError( + f"Robot '{robot_data.config.name}' has no pose-targetable planning group" + ) + return self.get_group_jacobian(ctx, group_id) + + def get_group_jacobian(self, ctx: Context, group_id: PlanningGroupID) -> NDArray[np.float64]: + """Get geometric Jacobian (6 x group joints) in group-local order.""" + if not self._finalized: + raise RuntimeError("World must be finalized first") + + group = self._planning_groups.get(group_id) + if group.tip_link is None: + raise ValueError(f"Planning group '{group_id}' has no tip link") + robot_data = self._robots[self._robot_id_for_group(group_id)] plant_ctx = self._diagram.GetSubsystemContext(self._plant, ctx) + tip_frame = self._plant.GetBodyByName( + group.tip_link, robot_data.model_instance + ).body_frame() # Compute full Jacobian J_full = self._plant.CalcJacobianSpatialVelocity( plant_ctx, JacobianWrtVariable.kQDot, - robot_data.ee_frame, + tip_frame, np.array([0.0, 0.0, 0.0]), # type: ignore[arg-type] # Point on end-effector self._plant.world_frame(), self._plant.world_frame(), ) - # Extract columns for this robot's joints - n_joints = len(robot_data.joint_indices) + # Extract columns for configured controllable joints only. + joint_indices_by_name = dict( + zip(robot_data.config.joint_names, robot_data.joint_indices, strict=True) + ) + missing = [ + joint_name + for joint_name in group.local_joint_names + if joint_name not in joint_indices_by_name + ] + if missing: + raise ValueError( + f"Planning group '{group_id}' references non-controllable joints: {missing}" + ) + group_joint_indices = [joint_indices_by_name[name] for name in group.local_joint_names] + n_joints = len(group_joint_indices) J_robot = np.zeros((6, n_joints)) - for i, joint_idx in enumerate(robot_data.joint_indices): + for i, joint_idx in enumerate(group_joint_indices): J_robot[:, i] = J_full[:, joint_idx] # Reorder rows: Drake uses [angular, linear], we want [linear, angular] diff --git a/dimos/manipulation/planning/world/roboplan_world.py b/dimos/manipulation/planning/world/roboplan_world.py index d252d3fdc9..e224ff583a 100644 --- a/dimos/manipulation/planning/world/roboplan_world.py +++ b/dimos/manipulation/planning/world/roboplan_world.py @@ -41,9 +41,16 @@ "Install the manipulation extra before selecting the roboplan backend." ) from exc +from dimos.manipulation.planning.groups.identifiers import is_global_joint_name +from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import ObstacleType, PlanningStatus -from dimos.manipulation.planning.spec.models import Obstacle, PlanningResult, WorldRobotID +from dimos.manipulation.planning.spec.models import ( + Obstacle, + PlanningGroupID, + PlanningResult, + WorldRobotID, +) from dimos.manipulation.planning.utils.mesh_utils import prepare_urdf_for_drake from dimos.manipulation.planning.utils.path_utils import compute_path_length from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped @@ -88,6 +95,7 @@ def __init__(self, enable_viz: bool = False, **_: object) -> None: logger.warning("RoboPlanWorld does not currently provide manipulation visualization") self._robots: dict[WorldRobotID, _RoboPlanRobotData] = {} + self._planning_groups = PlanningGroupRegistry() self._obstacles: dict[str, Obstacle] = {} self._robot_counter = 0 self._finalized = False @@ -104,6 +112,8 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: raise ValueError("RoboPlanWorld currently supports one robot per Scene") if not Path(config.model_path).exists(): raise FileNotFoundError(f"Robot model not found: {Path(config.model_path).resolve()}") + if self._planning_groups.groups_for_robot(config.name): + raise ValueError(f"Robot name '{config.name}' is already registered") self._validate_robot_config(config) self._robot_counter += 1 @@ -119,6 +129,7 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: self._live_context.q_by_robot[robot_id] = np.zeros( len(config.joint_names), dtype=np.float64 ) + self._planning_groups.add_robot(config) logger.info(f"Added RoboPlan robot '{robot_id}' ({config.name})") return robot_id @@ -130,6 +141,11 @@ def get_robot_config(self, robot_id: WorldRobotID) -> RobotModelConfig: """Get robot configuration by ID.""" return self._get_robot(robot_id).config + @property + def planning_groups(self) -> PlanningGroupRegistry: + """Registered planning groups for this world.""" + return self._planning_groups + def get_joint_limits( self, robot_id: WorldRobotID ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: @@ -272,7 +288,17 @@ def check_edge_collision_free( def get_ee_pose(self, ctx: RoboPlanContext, robot_id: WorldRobotID) -> PoseStamped: """Get end-effector pose if RoboPlan exposes FK.""" robot = self._get_robot(robot_id) - mat = self.get_link_pose(ctx, robot_id, robot.config.end_effector_link) + group_id = self._planning_groups.primary_pose_group_id_for_robot(robot.config.name) + if group_id is None: + raise ValueError(f"Robot '{robot.config.name}' has no pose-targetable planning group") + return self.get_group_ee_pose(ctx, group_id) + + def get_group_ee_pose(self, ctx: RoboPlanContext, group_id: PlanningGroupID) -> PoseStamped: + """Get planning-group tip pose if RoboPlan exposes FK.""" + group = self._planning_groups.get(group_id) + if group.tip_link is None: + raise ValueError(f"Planning group '{group_id}' has no tip link") + mat = self.get_link_pose(ctx, self._robot_id_for_group(group_id), group.tip_link) pose = matrix_to_pose(mat) return PoseStamped( frame_id="world", @@ -299,17 +325,46 @@ def get_link_pose( def get_jacobian(self, ctx: RoboPlanContext, robot_id: WorldRobotID) -> NDArray[np.float64]: """Get end-effector Jacobian if RoboPlan exposes a compatible API.""" robot = self._get_robot(robot_id) + group_id = self._planning_groups.primary_pose_group_id_for_robot(robot.config.name) + if group_id is None: + raise ValueError(f"Robot '{robot.config.name}' has no pose-targetable planning group") + return self.get_group_jacobian(ctx, group_id) + + def get_group_jacobian( + self, ctx: RoboPlanContext, group_id: PlanningGroupID + ) -> NDArray[np.float64]: + """Get planning-group Jacobian projected to group-local joint order.""" + group = self._planning_groups.get(group_id) + if group.tip_link is None: + raise ValueError(f"Planning group '{group_id}' has no tip link") + robot_id = self._robot_id_for_group(group_id) + robot = self._get_robot(robot_id) q = ctx.q_by_robot.get(robot_id) if q is None: raise KeyError(f"Robot '{robot_id}' not found in context") scene = self._require_scene() - result = scene.computeFrameJacobian( - self._to_scene_q(robot_id, q), robot.config.end_effector_link, True - ) + result = scene.computeFrameJacobian(self._to_scene_q(robot_id, q), group.tip_link, True) arr = np.asarray(result, dtype=np.float64) if arr.shape[0] != 6: raise ValueError(f"Unexpected RoboPlan Jacobian shape: {arr.shape}; expected 6 x n") - return arr + scene_joint_order = self._query_scene_joint_order(scene, robot.config) + if scene_joint_order is not None and arr.shape[1] == len(scene_joint_order): + missing = [name for name in group.local_joint_names if name not in scene_joint_order] + if missing: + raise ValueError(f"Unknown joints for planning group '{group_id}': {missing}") + indices = [scene_joint_order.index(name) for name in group.local_joint_names] + return arr[:, indices] + if arr.shape[1] == len(robot.config.joint_names): + missing = [ + name for name in group.local_joint_names if name not in robot.config.joint_names + ] + if missing: + raise ValueError(f"Unknown joints for planning group '{group_id}': {missing}") + indices = [robot.config.joint_names.index(name) for name in group.local_joint_names] + return arr[:, indices] + raise ValueError( + f"Unexpected RoboPlan Jacobian shape: {arr.shape}; cannot project group '{group_id}'" + ) # PlannerSpec for native RoboPlan planning @@ -507,6 +562,15 @@ def _get_robot(self, robot_id: WorldRobotID) -> _RoboPlanRobotData: raise KeyError(f"Robot '{robot_id}' not found") return self._robots[robot_id] + def _robot_id_for_group(self, group_id: PlanningGroupID) -> WorldRobotID: + group = self._planning_groups.get(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]: @@ -515,10 +579,24 @@ def _joint_state_to_q( raise ValueError("JointState position length must match configured joint count") if not joint_state.name: return np.asarray(joint_state.position, dtype=np.float64) - name_to_pos = { - robot.config.get_urdf_joint_name(name): position - for name, position in zip(joint_state.name, joint_state.position, strict=True) - } + name_to_pos: dict[str, float] = {} + for name, position in zip(joint_state.name, joint_state.position, strict=True): + if name in robot.config.joint_names: + resolved_name = name + elif name in robot.config.joint_name_mapping: + resolved_name = robot.config.joint_name_mapping[name] + elif is_global_joint_name(name): + prefix = f"{robot.config.name}/" + if not name.startswith(prefix): + continue + resolved_name = name[len(prefix) :] + if resolved_name not in robot.config.joint_names: + raise ValueError(f"Unknown global joint name for RoboPlanWorld: {name}") + else: + resolved_name = robot.config.get_urdf_joint_name(name) + if resolved_name in name_to_pos: + raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") + name_to_pos[resolved_name] = float(position) missing = [name for name in robot.config.joint_names if name not in name_to_pos] if missing: raise ValueError(f"JointState missing joints for RoboPlanWorld: {missing}") diff --git a/dimos/manipulation/planning/world/test_drake_world_planning_groups.py b/dimos/manipulation/planning/world/test_drake_world_planning_groups.py new file mode 100644 index 0000000000..841ed27f99 --- /dev/null +++ b/dimos/manipulation/planning/world/test_drake_world_planning_groups.py @@ -0,0 +1,153 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + +pydrake = pytest.importorskip("pydrake") +_ = pydrake + +from dimos.manipulation.planning.groups.models import PlanningGroupDefinition +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.world.drake_world import DrakeWorld +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.sensor_msgs.JointState import JointState + + +def _write_urdf(path: Path) -> None: + path.write_text( + """ + + + + + + + + + + + + + + + +""" + ) + + +def _config( + path: Path, groups: list[PlanningGroupDefinition], joints: list[str] | None = None +) -> RobotModelConfig: + return RobotModelConfig( + name="arm", + model_path=path, + base_pose=PoseStamped(position=[0, 0, 0], orientation=[0, 0, 0, 1]), + joint_names=joints or ["joint1", "joint2"], + base_link="base_link", + planning_groups=groups, + ) + + +def _arm_group( + *joint_names: str, tip_link: str | None = "tool0", name: str = "arm" +) -> PlanningGroupDefinition: + return PlanningGroupDefinition( + name=name, joint_names=joint_names, base_link="base_link", tip_link=tip_link + ) + + +def test_drake_group_fk_uses_tip_link_and_legacy_unique_pose_group(tmp_path: Path) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + world = DrakeWorld() + robot_id = world.add_robot(_config(urdf, [_arm_group("joint1", "joint2")])) + world.finalize() + ctx = world.get_live_context() + world.set_joint_state( + ctx, robot_id, JointState({"name": ["joint1", "joint2"], "position": [0.0, 0.0]}) + ) + + group_pose = world.get_group_ee_pose(ctx, "arm/arm") + legacy_pose = world.get_ee_pose(ctx, robot_id) + + assert group_pose.position.x == pytest.approx(2.0) + assert legacy_pose.position.x == pytest.approx(group_pose.position.x) + assert world.get_jacobian(ctx, robot_id).shape == (6, 2) + + +def test_drake_group_jacobian_shape_and_group_local_order(tmp_path: Path) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + world = DrakeWorld() + robot_id = world.add_robot( + _config( + urdf, + [ + _arm_group("joint1", "joint2", name="wrist_forward"), + _arm_group("joint2", "joint1", name="wrist_reverse"), + ], + ) + ) + world.finalize() + ctx = world.get_live_context() + world.set_joint_state( + ctx, robot_id, JointState({"name": ["joint1", "joint2"], "position": [0.0, 0.0]}) + ) + + forward_jacobian = world.get_group_jacobian(ctx, "arm/wrist_forward") + reverse_jacobian = world.get_group_jacobian(ctx, "arm/wrist_reverse") + + assert reverse_jacobian.shape == (6, 2) + np.testing.assert_allclose(reverse_jacobian[:, 0], forward_jacobian[:, 1]) + np.testing.assert_allclose(reverse_jacobian[:, 1], forward_jacobian[:, 0]) + + +def test_drake_legacy_wrappers_fail_at_call_time_for_no_or_ambiguous_pose(tmp_path: Path) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + no_pose = DrakeWorld() + no_pose_id = no_pose.add_robot(_config(urdf, [_arm_group("joint1", tip_link=None)])) + no_pose.finalize() + with pytest.raises(ValueError, match="no pose-targetable"): + no_pose.get_ee_pose(no_pose.get_live_context(), no_pose_id) + + ambiguous = DrakeWorld() + ambiguous_id = ambiguous.add_robot( + _config( + urdf, + [ + _arm_group("joint1", tip_link="link1", name="a"), + _arm_group("joint2", tip_link="tool0", name="b"), + ], + ) + ) + ambiguous.finalize() + with pytest.raises(ValueError, match="multiple pose-targetable"): + ambiguous.get_jacobian(ambiguous.get_live_context(), ambiguous_id) + + +def test_drake_group_jacobian_rejects_non_controllable_group_joints(tmp_path: Path) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + world = DrakeWorld() + world.add_robot(_config(urdf, [_arm_group("joint1", "joint2")], joints=["joint1"])) + world.finalize() + + with pytest.raises(ValueError, match="non-controllable"): + world.get_group_jacobian(world.get_live_context(), "arm/arm") diff --git a/dimos/manipulation/test_roboplan_world.py b/dimos/manipulation/test_roboplan_world.py index 65e4000a03..f3879297d5 100644 --- a/dimos/manipulation/test_roboplan_world.py +++ b/dimos/manipulation/test_roboplan_world.py @@ -387,6 +387,37 @@ def test_joint_name_mapping_is_applied_to_input_states( assert live_round_trip.position == [0.2, 0.3] +def test_global_joint_names_are_mapped_without_regressing_coordinator_names( + fake_roboplan: None, robot_config: RobotModelConfig +) -> None: + robot_config.joint_name_mapping = {"arm/j1": "joint1", "arm/j2": "joint2"} + world, robot_id = _make_world(fake_roboplan, robot_config) + world.finalize() + + world.sync_from_joint_state( + robot_id, JointState(name=["arm/j1", "arm/j2"], position=[0.4, 0.5]) + ) + assert world.get_joint_state(world.get_live_context(), robot_id).position == [0.4, 0.5] + + world.sync_from_joint_state( + robot_id, JointState(name=["arm/joint1", "arm/joint2"], position=[0.2, 0.3]) + ) + assert world.get_joint_state(world.get_live_context(), robot_id).position == [0.2, 0.3] + + +def test_duplicate_resolved_joint_names_fail_clearly( + fake_roboplan: None, robot_config: RobotModelConfig +) -> None: + robot_config.joint_name_mapping = {"alias": "joint1"} + world, robot_id = _make_world(fake_roboplan, robot_config) + world.finalize() + + with pytest.raises(ValueError, match="duplicate joint 'joint1'"): + world.sync_from_joint_state( + robot_id, JointState(name=["joint1", "alias"], position=[0.1, 0.2]) + ) + + def test_obstacle_mutation_updates_scene_and_stored_pose( fake_roboplan: None, robot_config: RobotModelConfig ) -> None: @@ -473,6 +504,65 @@ def test_fk_jacobian_and_explicit_min_distance_unsupported( world.get_min_distance(ctx, robot_id) +def test_group_fk_and_jacobian_use_group_tip_and_local_joint_order( + fake_roboplan: None, robot_config: RobotModelConfig, monkeypatch: pytest.MonkeyPatch +) -> None: + config = robot_config.model_copy( + update={ + "joint_names": ["joint1", "joint2", "joint3"], + "planning_groups": [ + PlanningGroupDefinition( + name="wrist", + joint_names=("joint3", "joint1"), + base_link="base", + tip_link="wrist_tip", + ) + ], + "joint_limits_lower": [-1.0, -2.0, -3.0], + "joint_limits_upper": [1.0, 2.0, 3.0], + } + ) + monkeypatch.setattr(FakeScene, "joint_group_joint_names", ["joint2", "joint1", "joint3"]) + monkeypatch.setattr(FakeScene, "position_limits_lower", [-2.0, -1.0, -3.0]) + monkeypatch.setattr(FakeScene, "position_limits_upper", [2.0, 1.0, 3.0]) + fk_frames: list[str] = [] + + def fake_fk( + self: FakeScene, q: np.ndarray, frame_name: str, base_frame: str = "" + ) -> np.ndarray: + _ = (self, base_frame) + fk_frames.append(frame_name) + mat = np.eye(4) + mat[0, 3] = float(np.sum(q)) + return mat + + def fake_jacobian( + self: FakeScene, q: np.ndarray, frame_name: str, local: bool = True + ) -> np.ndarray: + _ = (self, q) + assert frame_name == "wrist_tip" + assert local is True + return np.arange(18, dtype=np.float64).reshape(6, 3) + + monkeypatch.setattr(FakeScene, "forwardKinematics", fake_fk) + monkeypatch.setattr(FakeScene, "computeFrameJacobian", fake_jacobian) + world, robot_id = _make_world(fake_roboplan, config) + world.finalize() + ctx = world.get_live_context() + world.set_joint_state( + ctx, + robot_id, + JointState({"name": ["joint1", "joint2", "joint3"], "position": [1.0, 2.0, 3.0]}), + ) + + pose = world.get_group_ee_pose(ctx, "arm/wrist") + jacobian = world.get_group_jacobian(ctx, "arm/wrist") + + assert fk_frames == ["wrist_tip"] + assert pose.position.x == pytest.approx(6.0) + np.testing.assert_allclose(jacobian, np.arange(18, dtype=np.float64).reshape(6, 3)[:, [2, 1]]) + + def test_native_planner_converts_path(fake_roboplan: None, robot_config: RobotModelConfig) -> None: world, robot_id = _make_world(fake_roboplan, robot_config) world.finalize() diff --git a/openspec/changes/planning-world-monitor-groups/tasks.md b/openspec/changes/planning-world-monitor-groups/tasks.md index acb647d94a..8e31566a05 100644 --- a/openspec/changes/planning-world-monitor-groups/tasks.md +++ b/openspec/changes/planning-world-monitor-groups/tasks.md @@ -1,18 +1,19 @@ ## 1. Extract backend group support -- [ ] 1.1 Start from PR 1 and use `cc/spec/movegroup` as reference. -- [ ] 1.2 Extract Drake world group FK/Jacobian and base-link handling changes. -- [ ] 1.3 Extract RoboPlan world group FK/Jacobian, joint-name normalization, and URDF strip handling changes. -- [ ] 1.4 Extract world monitor and robot state monitor group-aware query changes. +- [x] 1.1 Start from PR 1 and use `cc/spec/movegroup` as reference. +- [x] 1.2 Extract Drake world group FK/Jacobian and base-link handling changes. +- [x] 1.3 Extract RoboPlan world group FK/Jacobian, joint-name normalization, and URDF strip handling changes. +- [x] 1.4 Extract world monitor and robot state monitor group-aware query changes. ## 2. Tests -- [ ] 2.1 Bring over Drake group world tests. -- [ ] 2.2 Bring over RoboPlan world tests, including planning split file if needed. -- [ ] 2.3 Bring over WorldMonitor tests for group state and ambiguity behavior. +- [x] 2.1 Bring over Drake group world tests. +- [x] 2.2 Bring over RoboPlan world tests, including planning split file if needed. +- [x] 2.3 Bring over WorldMonitor tests for group state and ambiguity behavior. ## 3. Validation -- [ ] 3.1 Run `uv run pytest dimos/manipulation/planning/world/test_drake_world_planning_groups.py dimos/manipulation/test_roboplan_world.py dimos/manipulation/test_roboplan_world_planning.py dimos/manipulation/planning/monitor/test_world_monitor.py -q`. -- [ ] 3.2 Run targeted mypy on changed backend/monitor files. -- [ ] 3.3 Optional manual smoke: load a robot config, print group list, FK pose, Jacobian shape, and collision status. +- [x] 3.1 Run `uv run pytest dimos/manipulation/planning/world/test_drake_world_planning_groups.py dimos/manipulation/test_roboplan_world.py dimos/manipulation/test_roboplan_world_planning.py dimos/manipulation/planning/monitor/test_world_monitor.py -q`. + - Note: `dimos/manipulation/test_roboplan_world_planning.py` does not exist in this branch; ran the same targeted command without that optional split file. +- [x] 3.2 Run targeted mypy on changed backend/monitor files. +- [x] 3.3 Optional manual smoke: load a robot config, print group list, FK pose, Jacobian shape, and collision status. From e59b9ef66eb9f07d4a8bd77b3927d95dd457d11a Mon Sep 17 00:00:00 2001 From: cc Date: Wed, 1 Jul 2026 22:09:34 -0700 Subject: [PATCH 03/14] refactor: centralize planning group registry --- dimos/manipulation/planning/groups/utils.py | 149 +++++++++++++++++- .../planning/monitor/test_world_monitor.py | 58 ++++++- .../planning/monitor/world_monitor.py | 93 +++-------- .../planning/world/drake_world.py | 39 ++--- .../planning/world/roboplan_world.py | 32 ++-- 5 files changed, 264 insertions(+), 107 deletions(-) diff --git a/dimos/manipulation/planning/groups/utils.py b/dimos/manipulation/planning/groups/utils.py index d128d18e3a..b642251471 100644 --- a/dimos/manipulation/planning/groups/utils.py +++ b/dimos/manipulation/planning/groups/utils.py @@ -20,12 +20,17 @@ assert_global_joint_names, assert_local_joint_names, is_global_joint_name, + make_global_joint_names, + make_planning_group_id, + parse_planning_group_id, ) -from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupDefinition +from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.models import ( GlobalJointName, LocalModelJointName, PlanningGroupID, + RobotName, ) from dimos.msgs.sensor_msgs.JointState import JointState from dimos.utils.logging_config import setup_logger @@ -55,6 +60,144 @@ def matching_global_joint_name( return None +def planning_group_from_configs( + group_id: PlanningGroupID, configs: Sequence[RobotModelConfig] +) -> PlanningGroup: + """Resolve a public planning-group ID from stored robot configs.""" + robot_name, group_name = parse_planning_group_id(group_id) + config = _config_for_robot_name(robot_name, configs) + definition = _definition_for_group_name(group_name, config.planning_groups, group_id) + return planning_group_from_definition(config, definition) + + +def planning_group_from_definition( + config: RobotModelConfig, definition: PlanningGroupDefinition +) -> PlanningGroup: + """Build the public planning-group model for one config definition.""" + return PlanningGroup( + id=make_planning_group_id(config.name, definition.name), + robot_name=config.name, + group_name=definition.name, + joint_names=tuple(make_global_joint_names(config.name, definition.joint_names)), + local_joint_names=definition.joint_names, + base_link=definition.base_link, + tip_link=definition.tip_link, + source=definition.source, + ) + + +def validate_planning_group_config(config: RobotModelConfig) -> None: + """Validate config-derived planning groups without constructing a registry.""" + seen_group_ids: set[PlanningGroupID] = set() + seen_group_names: set[str] = set() + for definition in config.planning_groups: + group_id = make_planning_group_id(config.name, definition.name) + if definition.name in seen_group_names or group_id in seen_group_ids: + raise ValueError(f"Planning group '{group_id}' is already registered") + planning_group_from_definition(config, definition) + seen_group_names.add(definition.name) + seen_group_ids.add(group_id) + + +def primary_pose_group_id_for_config(config: RobotModelConfig) -> PlanningGroupID | None: + """Return the unique pose-targetable group ID for robot-scoped wrappers.""" + pose_groups = [group for group in config.planning_groups if group.has_pose_target] + if not pose_groups: + return None + if len(pose_groups) > 1: + raise ValueError( + f"Robot '{config.name}' has {len(pose_groups)} pose-targetable planning groups; " + "use an explicit planning group ID" + ) + return make_planning_group_id(config.name, pose_groups[0].name) + + +def full_robot_joint_state_from_input( + config: RobotModelConfig, joint_state: JointState +) -> JointState | None: + """Return a full robot-local state if input contains all robot joints.""" + if not joint_state.name: + if len(joint_state.position) != len(config.joint_names): + return None + return JointState( + {"name": list(config.joint_names), "position": list(joint_state.position)} + ) + if len(joint_state.name) != len(joint_state.position): + raise ValueError("JointState name and position lengths must match") + resolved_positions: dict[str, float] = {} + global_prefix = f"{config.name}/" + for name, position in zip(joint_state.name, joint_state.position, strict=True): + if name in config.joint_names: + resolved_name = name + elif name in config.joint_name_mapping: + resolved_name = config.joint_name_mapping[name] + elif name.startswith(global_prefix): + resolved_name = name[len(global_prefix) :] + else: + resolved_name = config.get_urdf_joint_name(name) + if resolved_name not in config.joint_names: + return None + if resolved_name in resolved_positions: + raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") + resolved_positions[resolved_name] = float(position) + if set(resolved_positions) != set(config.joint_names): + return None + return JointState( + { + "name": list(config.joint_names), + "position": [resolved_positions[name] for name in config.joint_names], + } + ) + + +def joint_state_for_group_query( + config: RobotModelConfig, + group: PlanningGroup, + current_state: JointState, + joint_state: JointState | None, +) -> JointState: + """Normalize full-robot or group-scoped input to full robot-local state.""" + if joint_state is None: + return current_state + full_state = full_robot_joint_state_from_input(config, joint_state) + if full_state is not None: + return full_state + group_state = filter_joint_state_to_selected_joints( + joint_state, group.joint_names, group.local_joint_names + ) + positions_by_name = dict(zip(current_state.name, current_state.position, strict=True)) + for local_name, position in zip(group.local_joint_names, group_state.position, strict=True): + if local_name not in positions_by_name: + raise ValueError(f"Current state is missing group joint '{local_name}'") + positions_by_name[local_name] = float(position) + return JointState( + { + "name": list(current_state.name), + "position": [positions_by_name[name] for name in current_state.name], + } + ) + + +def _config_for_robot_name( + robot_name: RobotName, configs: Sequence[RobotModelConfig] +) -> RobotModelConfig: + for config in configs: + if config.name == robot_name: + return config + raise KeyError(f"No robot registered for planning group robot '{robot_name}'") + + +def _definition_for_group_name( + group_name: str, + definitions: Sequence[PlanningGroupDefinition], + group_id: PlanningGroupID, +) -> PlanningGroupDefinition: + for definition in definitions: + if definition.name == group_name: + return definition + raise KeyError(f"Unknown planning group ID: {group_id}") + + def filter_joint_state_to_selected_joints( joint_state: JointState, global_joint_names: Sequence[GlobalJointName], @@ -104,7 +247,7 @@ def joint_target_to_global_names( f"Target for '{group.id}' has {len(target.position)} positions, " f"expected {len(group.joint_names)}" ) - return JointState(name=list(group.joint_names), position=list(target.position)) + return JointState({"name": list(group.joint_names), "position": list(target.position)}) if len(target.name) != len(target.position): raise ValueError( @@ -140,4 +283,4 @@ def joint_target_to_global_names( extra = set(target_names) - set(expected_names) if extra: raise ValueError(f"Target for '{group.id}' has extra joints: {sorted(extra)}") - return JointState(name=list(group.joint_names), position=global_positions) + return JointState({"name": list(group.joint_names), "position": global_positions}) diff --git a/dimos/manipulation/planning/monitor/test_world_monitor.py b/dimos/manipulation/planning/monitor/test_world_monitor.py index aed6d6a9a5..b7d9016922 100644 --- a/dimos/manipulation/planning/monitor/test_world_monitor.py +++ b/dimos/manipulation/planning/monitor/test_world_monitor.py @@ -288,18 +288,74 @@ def test_world_monitor_exposes_planning_groups_and_duplicate_names_do_not_mutate assert [call[0] for call in fake_world.calls].count("add_robot") == 1 +def test_world_monitor_invalid_duplicate_group_config_does_not_mutate_backend() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + invalid_config = _robot_config_with_groups( + [ + PlanningGroupDefinition( + name="manipulator", joint_names=("j1",), base_link="base", tip_link="ee" + ), + PlanningGroupDefinition( + name="manipulator", joint_names=("j2",), base_link="base", tip_link="ee" + ), + ] + ) + + with pytest.raises(ValueError, match="already registered"): + monitor.add_robot(invalid_config) + + assert [call[0] for call in fake_world.calls].count("add_robot") == 0 + + +def test_world_monitor_invalid_group_joint_name_does_not_mutate_backend() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + invalid_config = _robot_config_with_groups( + [ + PlanningGroupDefinition( + name="manipulator", + joint_names=("j1", "bad/joint"), + base_link="base", + tip_link="ee", + ) + ] + ) + + with pytest.raises(ValueError, match="Invalid local joint name"): + monitor.add_robot(invalid_config) + + assert [call[0] for call in fake_world.calls].count("add_robot") == 0 + + def test_current_group_joint_state_uses_public_names_in_group_order() -> None: fake_world = FakeWorld() monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] robot_id = monitor.add_robot(_three_joint_reordered_group_config()) monitor._state_monitors[robot_id] = _FakeStateMonitor([0.1, 0.2, 0.3]) # type: ignore[attr-defined] - state = monitor.get_current_group_joint_state("arm/manipulator") + state = monitor.current_group_joint_state("arm/manipulator") assert state.name == ["arm/j2", "arm/j1"] assert state.position == [0.2, 0.1] +def test_group_kinematics_with_full_state_does_not_require_current_state() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + monitor.add_robot(_three_joint_reordered_group_config()) + + pose = monitor.get_group_ee_pose( + "arm/manipulator", + JointState({"name": ["j1", "j2", "j3"], "position": [0.1, 0.2, 0.3]}), + ) + + set_calls = [call for call in fake_world.calls if call[0] == "set_joint_state"] + assert set_calls[0][3].name == ["j1", "j2", "j3"] + assert set_calls[0][3].position == [0.1, 0.2, 0.3] + assert pose.position.x == 1 + + def test_group_kinematics_route_to_backend_and_merge_group_state() -> None: fake_world = FakeWorld() monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] diff --git a/dimos/manipulation/planning/monitor/world_monitor.py b/dimos/manipulation/planning/monitor/world_monitor.py index 3cd54d05c4..13ff2d2fc3 100644 --- a/dimos/manipulation/planning/monitor/world_monitor.py +++ b/dimos/manipulation/planning/monitor/world_monitor.py @@ -22,7 +22,12 @@ from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT 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.groups.utils import ( + filter_joint_state_to_selected_joints, + full_robot_joint_state_from_input, + joint_state_for_group_query, + validate_planning_group_config, +) from dimos.manipulation.planning.monitor.robot_state_monitor import RobotStateMonitor from dimos.manipulation.planning.monitor.world_obstacle_monitor import WorldObstacleMonitor from dimos.manipulation.planning.spec.models import PlanningSceneInfo @@ -80,8 +85,7 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: with self._lock: if config.name in self._robot_ids_by_name: raise ValueError(f"Robot name '{config.name}' is already registered") - new_groups = PlanningGroupRegistry() - new_groups.add_robot(config) + 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 @@ -301,9 +305,11 @@ def get_current_joint_state(self, robot_id: WorldRobotID) -> JointState | None: if positions is not None: joint_names = self._robot_joints.get(robot_id, []) return JointState( - name=joint_names, - position=positions.tolist(), - velocity=velocities.tolist() if velocities is not None else [], + { + "name": joint_names, + "position": positions.tolist(), + "velocity": velocities.tolist() if velocities is not None else [], + } ) # Fall back to world's live context @@ -328,17 +334,12 @@ def current_global_joint_state(self, max_age: float = 1.0) -> JointState: positions.append(float(pos_by_name[local_name])) return JointState({"name": names, "position": positions}) - def get_current_group_joint_state( + def current_group_joint_state( self, group_id: PlanningGroupID, max_age: float = 1.0 ) -> JointState: """Return current joint state scoped and ordered for one planning group.""" group = self._planning_groups.get(group_id) - robot_id = self._robot_ids_by_name[group.robot_name] - if robot_id in self._state_monitors and self.is_state_stale(robot_id, max_age): - raise ValueError(f"Current state for robot '{group.robot_name}' is stale") - state = self.get_current_joint_state(robot_id) - if state is None: - raise ValueError(f"Current state for robot '{group.robot_name}' is unavailable") + state = self._current_robot_joint_state_for_group(group_id, max_age) return filter_joint_state_to_selected_joints( state, group.joint_names, group.local_joint_names ) @@ -356,69 +357,21 @@ def _current_robot_joint_state_for_group( raise ValueError(f"Current state for robot '{group.robot_name}' is unavailable") return state - def _full_robot_joint_state_from_input( - self, robot_id: WorldRobotID, joint_state: JointState - ) -> JointState | None: - """Return a full robot-local state if input contains all robot joints.""" - config = self._robot_configs[robot_id] - if not joint_state.name: - if len(joint_state.position) != len(config.joint_names): - return None - return JointState( - {"name": list(config.joint_names), "position": list(joint_state.position)} - ) - if len(joint_state.name) != len(joint_state.position): - raise ValueError("JointState name and position lengths must match") - resolved_positions: dict[str, float] = {} - global_prefix = f"{config.name}/" - for name, position in zip(joint_state.name, joint_state.position, strict=True): - if name in config.joint_names: - resolved_name = name - elif name in config.joint_name_mapping: - resolved_name = config.joint_name_mapping[name] - elif name.startswith(global_prefix): - resolved_name = name[len(global_prefix) :] - else: - resolved_name = config.get_urdf_joint_name(name) - if resolved_name not in config.joint_names: - return None - if resolved_name in resolved_positions: - raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") - resolved_positions[resolved_name] = float(position) - if set(resolved_positions) != set(config.joint_names): - return None - return JointState( - { - "name": list(config.joint_names), - "position": [resolved_positions[name] for name in config.joint_names], - } - ) - def _joint_state_for_group_query( self, group_id: PlanningGroupID, joint_state: JointState | None ) -> JointState: """Normalize full-robot or group-scoped input to full robot-local state.""" group = self._planning_groups.get(group_id) robot_id = self._robot_ids_by_name[group.robot_name] - if joint_state is None: - return self._current_robot_joint_state_for_group(group_id) - full_state = self._full_robot_joint_state_from_input(robot_id, joint_state) - if full_state is not None: - return full_state + if joint_state is not None: + full_state = full_robot_joint_state_from_input( + self._robot_configs[robot_id], joint_state + ) + if full_state is not None: + return full_state current_state = self._current_robot_joint_state_for_group(group_id) - group_state = filter_joint_state_to_selected_joints( - joint_state, group.joint_names, group.local_joint_names - ) - positions_by_name = dict(zip(current_state.name, current_state.position, strict=True)) - for local_name, position in zip(group.local_joint_names, group_state.position, strict=True): - if local_name not in positions_by_name: - raise ValueError(f"Current state is missing group joint '{local_name}'") - positions_by_name[local_name] = float(position) - return JointState( - { - "name": list(current_state.name), - "position": [positions_by_name[name] for name in current_state.name], - } + return joint_state_for_group_query( + self._robot_configs[robot_id], group, current_state, joint_state ) def get_current_velocities(self, robot_id: WorldRobotID) -> JointState | None: @@ -427,7 +380,7 @@ def get_current_velocities(self, robot_id: WorldRobotID) -> JointState | None: velocities = self._state_monitors[robot_id].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": joint_names, "velocity": velocities.tolist()}) return None def wait_for_state(self, robot_id: WorldRobotID, timeout: float = 1.0) -> bool: diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index 12067f647c..b1aa4da024 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -26,7 +26,11 @@ import numpy as np from dimos.manipulation.planning.groups.identifiers import is_global_joint_name -from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry +from dimos.manipulation.planning.groups.utils import ( + planning_group_from_configs, + primary_pose_group_id_for_config, + validate_planning_group_config, +) from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import ObstacleType from dimos.manipulation.planning.spec.models import ( @@ -187,7 +191,6 @@ def __init__(self, time_step: float = 0.0, enable_viz: bool = False) -> None: # Tracking data self._robots: dict[WorldRobotID, _RobotData] = {} - self._planning_groups = PlanningGroupRegistry() self._obstacles: dict[str, _ObstacleData] = {} self._robot_counter = 0 self._obstacle_counter = 0 @@ -211,10 +214,9 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: raise RuntimeError("Cannot add robot after world is finalized") with self._lock: - if self._planning_groups.groups_for_robot(config.name): + if any(data.config.name == config.name for data in self._robots.values()): raise ValueError(f"Robot name '{config.name}' is already registered") - new_groups = PlanningGroupRegistry() - new_groups.add_robot(config) + validate_planning_group_config(config) self._robot_counter += 1 robot_id = f"robot_{self._robot_counter}" @@ -226,11 +228,11 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: ee_link = config.base_link try: - primary_group_id = new_groups.primary_pose_group_id_for_robot(config.name) + primary_group_id = primary_pose_group_id_for_config(config) except ValueError: primary_group_id = None if primary_group_id is not None: - primary_group = new_groups.get(primary_group_id) + primary_group = planning_group_from_configs(primary_group_id, [config]) if primary_group.tip_link is not None: ee_link = primary_group.tip_link ee_frame = self._plant.GetBodyByName(ee_link, model_instance).body_frame() @@ -251,8 +253,6 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: base_frame=base_frame, preview_model_instance=preview_model_instance, ) - self._planning_groups.add_robot(config) - logger.info(f"Added robot '{robot_id}' ({config.name})") return robot_id @@ -334,11 +334,6 @@ def get_robot_config(self, robot_id: WorldRobotID) -> RobotModelConfig: raise KeyError(f"Robot '{robot_id}' not found") return self._robots[robot_id].config - @property - def planning_groups(self) -> PlanningGroupRegistry: - """Registered planning groups for this world.""" - return self._planning_groups - def get_joint_limits( self, robot_id: WorldRobotID ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: @@ -866,7 +861,9 @@ def _joint_state_to_q( return np.asarray([name_to_pos[name] for name in robot_data.config.joint_names]) def _robot_id_for_group(self, group_id: PlanningGroupID) -> WorldRobotID: - group = self._planning_groups.get(group_id) + group = planning_group_from_configs( + group_id, [data.config for data in self._robots.values()] + ) matches = [ rid for rid, data in self._robots.items() if data.config.name == group.robot_name ] @@ -973,7 +970,7 @@ def get_ee_pose(self, ctx: Context, robot_id: WorldRobotID) -> PoseStamped: if robot_id not in self._robots: raise KeyError(f"Robot '{robot_id}' not found") robot_data = self._robots[robot_id] - group_id = self._planning_groups.primary_pose_group_id_for_robot(robot_data.config.name) + group_id = 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" @@ -985,7 +982,9 @@ def get_group_ee_pose(self, ctx: Context, group_id: PlanningGroupID) -> PoseStam if not self._finalized: raise RuntimeError("World must be finalized first") - group = self._planning_groups.get(group_id) + group = planning_group_from_configs( + group_id, [data.config for data in self._robots.values()] + ) 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)] @@ -1035,7 +1034,7 @@ def get_jacobian(self, ctx: Context, robot_id: WorldRobotID) -> NDArray[np.float if robot_id not in self._robots: raise KeyError(f"Robot '{robot_id}' not found") robot_data = self._robots[robot_id] - group_id = self._planning_groups.primary_pose_group_id_for_robot(robot_data.config.name) + group_id = 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" @@ -1047,7 +1046,9 @@ def get_group_jacobian(self, ctx: Context, group_id: PlanningGroupID) -> NDArray if not self._finalized: raise RuntimeError("World must be finalized first") - group = self._planning_groups.get(group_id) + group = planning_group_from_configs( + group_id, [data.config for data in self._robots.values()] + ) 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)] diff --git a/dimos/manipulation/planning/world/roboplan_world.py b/dimos/manipulation/planning/world/roboplan_world.py index e224ff583a..11a19c883c 100644 --- a/dimos/manipulation/planning/world/roboplan_world.py +++ b/dimos/manipulation/planning/world/roboplan_world.py @@ -42,7 +42,11 @@ ) from exc from dimos.manipulation.planning.groups.identifiers import is_global_joint_name -from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry +from dimos.manipulation.planning.groups.utils import ( + planning_group_from_configs, + primary_pose_group_id_for_config, + validate_planning_group_config, +) from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import ObstacleType, PlanningStatus from dimos.manipulation.planning.spec.models import ( @@ -95,7 +99,6 @@ def __init__(self, enable_viz: bool = False, **_: object) -> None: logger.warning("RoboPlanWorld does not currently provide manipulation visualization") self._robots: dict[WorldRobotID, _RoboPlanRobotData] = {} - self._planning_groups = PlanningGroupRegistry() self._obstacles: dict[str, Obstacle] = {} self._robot_counter = 0 self._finalized = False @@ -112,8 +115,9 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: raise ValueError("RoboPlanWorld currently supports one robot per Scene") if not Path(config.model_path).exists(): raise FileNotFoundError(f"Robot model not found: {Path(config.model_path).resolve()}") - if self._planning_groups.groups_for_robot(config.name): + if any(data.config.name == config.name for data in self._robots.values()): raise ValueError(f"Robot name '{config.name}' is already registered") + validate_planning_group_config(config) self._validate_robot_config(config) self._robot_counter += 1 @@ -129,7 +133,6 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: self._live_context.q_by_robot[robot_id] = np.zeros( len(config.joint_names), dtype=np.float64 ) - self._planning_groups.add_robot(config) logger.info(f"Added RoboPlan robot '{robot_id}' ({config.name})") return robot_id @@ -141,11 +144,6 @@ def get_robot_config(self, robot_id: WorldRobotID) -> RobotModelConfig: """Get robot configuration by ID.""" return self._get_robot(robot_id).config - @property - def planning_groups(self) -> PlanningGroupRegistry: - """Registered planning groups for this world.""" - return self._planning_groups - def get_joint_limits( self, robot_id: WorldRobotID ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: @@ -288,14 +286,16 @@ def check_edge_collision_free( def get_ee_pose(self, ctx: RoboPlanContext, robot_id: WorldRobotID) -> PoseStamped: """Get end-effector pose if RoboPlan exposes FK.""" robot = self._get_robot(robot_id) - group_id = self._planning_groups.primary_pose_group_id_for_robot(robot.config.name) + group_id = primary_pose_group_id_for_config(robot.config) if group_id is None: raise ValueError(f"Robot '{robot.config.name}' has no pose-targetable planning group") return self.get_group_ee_pose(ctx, group_id) def get_group_ee_pose(self, ctx: RoboPlanContext, group_id: PlanningGroupID) -> PoseStamped: """Get planning-group tip pose if RoboPlan exposes FK.""" - group = self._planning_groups.get(group_id) + group = planning_group_from_configs( + group_id, [data.config for data in self._robots.values()] + ) 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) @@ -325,7 +325,7 @@ def get_link_pose( def get_jacobian(self, ctx: RoboPlanContext, robot_id: WorldRobotID) -> NDArray[np.float64]: """Get end-effector Jacobian if RoboPlan exposes a compatible API.""" robot = self._get_robot(robot_id) - group_id = self._planning_groups.primary_pose_group_id_for_robot(robot.config.name) + group_id = primary_pose_group_id_for_config(robot.config) if group_id is None: raise ValueError(f"Robot '{robot.config.name}' has no pose-targetable planning group") return self.get_group_jacobian(ctx, group_id) @@ -334,7 +334,9 @@ def get_group_jacobian( self, ctx: RoboPlanContext, group_id: PlanningGroupID ) -> NDArray[np.float64]: """Get planning-group Jacobian projected to group-local joint order.""" - group = self._planning_groups.get(group_id) + group = planning_group_from_configs( + group_id, [data.config for data in self._robots.values()] + ) 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) @@ -563,7 +565,9 @@ def _get_robot(self, robot_id: WorldRobotID) -> _RoboPlanRobotData: return self._robots[robot_id] def _robot_id_for_group(self, group_id: PlanningGroupID) -> WorldRobotID: - group = self._planning_groups.get(group_id) + group = planning_group_from_configs( + group_id, [data.config for data in self._robots.values()] + ) matches = [ rid for rid, data in self._robots.items() if data.config.name == group.robot_name ] From db1b72902b1a9a70a24abdee516fcd202672d32e Mon Sep 17 00:00:00 2001 From: cc Date: Wed, 1 Jul 2026 22:45:00 -0700 Subject: [PATCH 04/14] refactor: align planning group utilities --- .../planning/groups/test_planning_groups.py | 18 -- dimos/manipulation/planning/groups/utils.py | 158 +----------------- .../planning/monitor/world_monitor.py | 82 +++++++-- .../planning/world/drake_world.py | 75 ++++++--- .../planning/world/roboplan_world.py | 78 +++++++-- .../world/test_drake_world_planning_groups.py | 2 +- 6 files changed, 192 insertions(+), 221 deletions(-) diff --git a/dimos/manipulation/planning/groups/test_planning_groups.py b/dimos/manipulation/planning/groups/test_planning_groups.py index ae5776fd74..e716366202 100644 --- a/dimos/manipulation/planning/groups/test_planning_groups.py +++ b/dimos/manipulation/planning/groups/test_planning_groups.py @@ -20,7 +20,6 @@ import pytest -from dimos.manipulation.planning.groups import utils as planning_group_utils from dimos.manipulation.planning.groups.discovery import ( FALLBACK_PLANNING_GROUP_NAME, PlanningGroupDiscoveryError, @@ -474,23 +473,6 @@ def test_matching_global_joint_name_requires_unique_suffix_match() -> None: assert matching_global_joint_name({"left/j1": 1.0}, "j2") is None -def test_matching_global_joint_name_warns_when_suffix_match_is_not_unique( - monkeypatch: pytest.MonkeyPatch, -) -> None: - messages: list[str] = [] - - def capture_warning(message: str) -> None: - messages.append(message) - - monkeypatch.setattr(planning_group_utils.logger, "warning", capture_warning) - - assert matching_global_joint_name({"left/j1": 1.0, "right/j1": 2.0}, "j1") is None - assert messages == [ - "Expected exactly one global joint ending with local joint name 'j1', " - "found 2 matches: ['left/j1', 'right/j1']" - ] - - def test_planning_group_id_from_selector_accepts_id_or_group() -> None: group = _make_group() diff --git a/dimos/manipulation/planning/groups/utils.py b/dimos/manipulation/planning/groups/utils.py index b642251471..e55f997a33 100644 --- a/dimos/manipulation/planning/groups/utils.py +++ b/dimos/manipulation/planning/groups/utils.py @@ -20,22 +20,14 @@ assert_global_joint_names, assert_local_joint_names, is_global_joint_name, - make_global_joint_names, - make_planning_group_id, - parse_planning_group_id, ) -from dimos.manipulation.planning.groups.models import PlanningGroup, PlanningGroupDefinition -from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.groups.models import PlanningGroup from dimos.manipulation.planning.spec.models import ( GlobalJointName, LocalModelJointName, PlanningGroupID, - RobotName, ) from dimos.msgs.sensor_msgs.JointState import JointState -from dimos.utils.logging_config import setup_logger - -logger = setup_logger() def planning_group_id_from_selector(selector: PlanningGroupID | PlanningGroup) -> PlanningGroupID: @@ -53,151 +45,9 @@ def matching_global_joint_name( matches = [name for name in positions_by_name if name.endswith(suffix)] if len(matches) == 1: return matches[0] - logger.warning( - f"Expected exactly one global joint ending with local joint name " - f"'{local_joint_name}', found {len(matches)} matches: {matches}" - ) return None -def planning_group_from_configs( - group_id: PlanningGroupID, configs: Sequence[RobotModelConfig] -) -> PlanningGroup: - """Resolve a public planning-group ID from stored robot configs.""" - robot_name, group_name = parse_planning_group_id(group_id) - config = _config_for_robot_name(robot_name, configs) - definition = _definition_for_group_name(group_name, config.planning_groups, group_id) - return planning_group_from_definition(config, definition) - - -def planning_group_from_definition( - config: RobotModelConfig, definition: PlanningGroupDefinition -) -> PlanningGroup: - """Build the public planning-group model for one config definition.""" - return PlanningGroup( - id=make_planning_group_id(config.name, definition.name), - robot_name=config.name, - group_name=definition.name, - joint_names=tuple(make_global_joint_names(config.name, definition.joint_names)), - local_joint_names=definition.joint_names, - base_link=definition.base_link, - tip_link=definition.tip_link, - source=definition.source, - ) - - -def validate_planning_group_config(config: RobotModelConfig) -> None: - """Validate config-derived planning groups without constructing a registry.""" - seen_group_ids: set[PlanningGroupID] = set() - seen_group_names: set[str] = set() - for definition in config.planning_groups: - group_id = make_planning_group_id(config.name, definition.name) - if definition.name in seen_group_names or group_id in seen_group_ids: - raise ValueError(f"Planning group '{group_id}' is already registered") - planning_group_from_definition(config, definition) - seen_group_names.add(definition.name) - seen_group_ids.add(group_id) - - -def primary_pose_group_id_for_config(config: RobotModelConfig) -> PlanningGroupID | None: - """Return the unique pose-targetable group ID for robot-scoped wrappers.""" - pose_groups = [group for group in config.planning_groups if group.has_pose_target] - if not pose_groups: - return None - if len(pose_groups) > 1: - raise ValueError( - f"Robot '{config.name}' has {len(pose_groups)} pose-targetable planning groups; " - "use an explicit planning group ID" - ) - return make_planning_group_id(config.name, pose_groups[0].name) - - -def full_robot_joint_state_from_input( - config: RobotModelConfig, joint_state: JointState -) -> JointState | None: - """Return a full robot-local state if input contains all robot joints.""" - if not joint_state.name: - if len(joint_state.position) != len(config.joint_names): - return None - return JointState( - {"name": list(config.joint_names), "position": list(joint_state.position)} - ) - if len(joint_state.name) != len(joint_state.position): - raise ValueError("JointState name and position lengths must match") - resolved_positions: dict[str, float] = {} - global_prefix = f"{config.name}/" - for name, position in zip(joint_state.name, joint_state.position, strict=True): - if name in config.joint_names: - resolved_name = name - elif name in config.joint_name_mapping: - resolved_name = config.joint_name_mapping[name] - elif name.startswith(global_prefix): - resolved_name = name[len(global_prefix) :] - else: - resolved_name = config.get_urdf_joint_name(name) - if resolved_name not in config.joint_names: - return None - if resolved_name in resolved_positions: - raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") - resolved_positions[resolved_name] = float(position) - if set(resolved_positions) != set(config.joint_names): - return None - return JointState( - { - "name": list(config.joint_names), - "position": [resolved_positions[name] for name in config.joint_names], - } - ) - - -def joint_state_for_group_query( - config: RobotModelConfig, - group: PlanningGroup, - current_state: JointState, - joint_state: JointState | None, -) -> JointState: - """Normalize full-robot or group-scoped input to full robot-local state.""" - if joint_state is None: - return current_state - full_state = full_robot_joint_state_from_input(config, joint_state) - if full_state is not None: - return full_state - group_state = filter_joint_state_to_selected_joints( - joint_state, group.joint_names, group.local_joint_names - ) - positions_by_name = dict(zip(current_state.name, current_state.position, strict=True)) - for local_name, position in zip(group.local_joint_names, group_state.position, strict=True): - if local_name not in positions_by_name: - raise ValueError(f"Current state is missing group joint '{local_name}'") - positions_by_name[local_name] = float(position) - return JointState( - { - "name": list(current_state.name), - "position": [positions_by_name[name] for name in current_state.name], - } - ) - - -def _config_for_robot_name( - robot_name: RobotName, configs: Sequence[RobotModelConfig] -) -> RobotModelConfig: - for config in configs: - if config.name == robot_name: - return config - raise KeyError(f"No robot registered for planning group robot '{robot_name}'") - - -def _definition_for_group_name( - group_name: str, - definitions: Sequence[PlanningGroupDefinition], - group_id: PlanningGroupID, -) -> PlanningGroupDefinition: - for definition in definitions: - if definition.name == group_name: - return definition - raise KeyError(f"Unknown planning group ID: {group_id}") - - def filter_joint_state_to_selected_joints( joint_state: JointState, global_joint_names: Sequence[GlobalJointName], @@ -226,7 +76,7 @@ def filter_joint_state_to_selected_joints( missing.append(global_name) if missing: - raise ValueError(f"Joint state is missing selected joints: {missing}") + raise ValueError(f"IK result is missing selected joints: {missing}") return JointState({"name": list(global_joint_names), "position": selected_positions}) @@ -247,7 +97,7 @@ def joint_target_to_global_names( f"Target for '{group.id}' has {len(target.position)} positions, " f"expected {len(group.joint_names)}" ) - return JointState({"name": list(group.joint_names), "position": list(target.position)}) + return JointState(name=list(group.joint_names), position=list(target.position)) if len(target.name) != len(target.position): raise ValueError( @@ -283,4 +133,4 @@ def joint_target_to_global_names( extra = set(target_names) - set(expected_names) if extra: raise ValueError(f"Target for '{group.id}' has extra joints: {sorted(extra)}") - return JointState({"name": list(group.joint_names), "position": global_positions}) + return JointState(name=list(group.joint_names), position=global_positions) diff --git a/dimos/manipulation/planning/monitor/world_monitor.py b/dimos/manipulation/planning/monitor/world_monitor.py index 13ff2d2fc3..cf38697678 100644 --- a/dimos/manipulation/planning/monitor/world_monitor.py +++ b/dimos/manipulation/planning/monitor/world_monitor.py @@ -21,13 +21,12 @@ from typing import TYPE_CHECKING, Any from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT -from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry -from dimos.manipulation.planning.groups.utils import ( - filter_joint_state_to_selected_joints, - full_robot_joint_state_from_input, - joint_state_for_group_query, - validate_planning_group_config, +from dimos.manipulation.planning.groups.identifiers import ( + make_global_joint_names, + make_planning_group_id, ) +from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry +from dimos.manipulation.planning.groups.utils import filter_joint_state_to_selected_joints from dimos.manipulation.planning.monitor.robot_state_monitor import RobotStateMonitor from dimos.manipulation.planning.monitor.world_obstacle_monitor import WorldObstacleMonitor from dimos.manipulation.planning.spec.models import PlanningSceneInfo @@ -85,7 +84,7 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: with self._lock: if config.name in self._robot_ids_by_name: raise ValueError(f"Robot name '{config.name}' is already registered") - validate_planning_group_config(config) + 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 @@ -364,15 +363,74 @@ def _joint_state_for_group_query( group = self._planning_groups.get(group_id) robot_id = self._robot_ids_by_name[group.robot_name] if joint_state is not None: - full_state = full_robot_joint_state_from_input( - self._robot_configs[robot_id], joint_state - ) + full_state = self._full_robot_joint_state_from_input(robot_id, joint_state) if full_state is not None: return full_state current_state = self._current_robot_joint_state_for_group(group_id) - return joint_state_for_group_query( - self._robot_configs[robot_id], group, current_state, joint_state + if joint_state is None: + return current_state + group_state = filter_joint_state_to_selected_joints( + joint_state, group.joint_names, group.local_joint_names ) + positions_by_name = dict(zip(current_state.name, current_state.position, strict=True)) + for local_name, position in zip(group.local_joint_names, group_state.position, strict=True): + if local_name not in positions_by_name: + raise ValueError(f"Current state is missing group joint '{local_name}'") + positions_by_name[local_name] = float(position) + return JointState( + { + "name": list(current_state.name), + "position": [positions_by_name[name] for name in current_state.name], + } + ) + + def _full_robot_joint_state_from_input( + self, robot_id: WorldRobotID, joint_state: JointState + ) -> JointState | None: + """Return a full robot-local state if input contains all robot joints.""" + config = self._robot_configs[robot_id] + if not joint_state.name: + if len(joint_state.position) != len(config.joint_names): + return None + return JointState( + {"name": list(config.joint_names), "position": list(joint_state.position)} + ) + if len(joint_state.name) != len(joint_state.position): + raise ValueError("JointState name and position lengths must match") + resolved_positions: dict[str, float] = {} + global_prefix = f"{config.name}/" + for name, position in zip(joint_state.name, joint_state.position, strict=True): + if name in config.joint_names: + resolved_name = name + elif name in config.joint_name_mapping: + resolved_name = config.joint_name_mapping[name] + elif name.startswith(global_prefix): + resolved_name = name[len(global_prefix) :] + else: + resolved_name = config.get_urdf_joint_name(name) + if resolved_name not in config.joint_names: + return None + if resolved_name in resolved_positions: + raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") + resolved_positions[resolved_name] = float(position) + if set(resolved_positions) != set(config.joint_names): + return None + return JointState( + { + "name": list(config.joint_names), + "position": [resolved_positions[name] for name in config.joint_names], + } + ) + + def _validate_planning_group_config(self, config: RobotModelConfig) -> None: + """Validate planning groups before mutating world/backend state.""" + seen_group_names: set[str] = set() + for definition in config.planning_groups: + group_id = make_planning_group_id(config.name, definition.name) + if definition.name in seen_group_names: + raise ValueError(f"Planning group '{group_id}' is already registered") + make_global_joint_names(config.name, definition.joint_names) + seen_group_names.add(definition.name) def get_current_velocities(self, robot_id: WorldRobotID) -> JointState | None: """Get current joint velocities as JointState. Returns None if not available.""" diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index b1aa4da024..b84c735556 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -25,12 +25,12 @@ import numpy as np -from dimos.manipulation.planning.groups.identifiers import is_global_joint_name -from dimos.manipulation.planning.groups.utils import ( - planning_group_from_configs, - primary_pose_group_id_for_config, - validate_planning_group_config, +from dimos.manipulation.planning.groups.identifiers import ( + is_global_joint_name, + make_global_joint_names, + make_planning_group_id, ) +from dimos.manipulation.planning.groups.models import PlanningGroup from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import ObstacleType from dimos.manipulation.planning.spec.models import ( @@ -216,7 +216,7 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: 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") - validate_planning_group_config(config) + self._validate_planning_group_config(config) self._robot_counter += 1 robot_id = f"robot_{self._robot_counter}" @@ -228,11 +228,11 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: ee_link = config.base_link try: - primary_group_id = primary_pose_group_id_for_config(config) + primary_group_id = self._primary_pose_group_id_for_config(config) except ValueError: primary_group_id = None if primary_group_id is not None: - primary_group = planning_group_from_configs(primary_group_id, [config]) + primary_group = self._planning_group_from_config(config, primary_group_id) if primary_group.tip_link is not None: ee_link = primary_group.tip_link ee_frame = self._plant.GetBodyByName(ee_link, model_instance).body_frame() @@ -334,6 +334,49 @@ def get_robot_config(self, robot_id: WorldRobotID) -> RobotModelConfig: raise KeyError(f"Robot '{robot_id}' not found") return self._robots[robot_id].config + def _validate_planning_group_config(self, config: RobotModelConfig) -> None: + seen_group_names: set[str] = set() + for definition in config.planning_groups: + make_planning_group_id(config.name, definition.name) + if definition.name in seen_group_names: + raise ValueError(f"Planning group '{definition.name}' is already registered") + make_global_joint_names(config.name, definition.joint_names) + seen_group_names.add(definition.name) + + def _planning_group_from_config( + self, config: RobotModelConfig, group_id: PlanningGroupID + ) -> PlanningGroup: + for definition in config.planning_groups: + if make_planning_group_id(config.name, definition.name) == group_id: + joint_names = tuple(make_global_joint_names(config.name, definition.joint_names)) + return PlanningGroup( + group_id, + config.name, + definition.name, + joint_names, + definition.joint_names, + definition.base_link, + definition.tip_link, + definition.source, + ) + raise KeyError(f"Unknown planning group ID: {group_id}") + + def _planning_group_from_id(self, group_id: PlanningGroupID) -> PlanningGroup: + for robot_data in self._robots.values(): + try: + return self._planning_group_from_config(robot_data.config, group_id) + except KeyError: + continue + raise KeyError(f"Unknown planning group ID: {group_id}") + + def _primary_pose_group_id_for_config(self, config: RobotModelConfig) -> PlanningGroupID | None: + pose_groups = [group for group in config.planning_groups if group.has_pose_target] + if not pose_groups: + return None + if len(pose_groups) > 1: + raise ValueError(f"Robot '{config.name}' has multiple pose groups") + return make_planning_group_id(config.name, pose_groups[0].name) + def get_joint_limits( self, robot_id: WorldRobotID ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: @@ -861,9 +904,7 @@ def _joint_state_to_q( return np.asarray([name_to_pos[name] for name in robot_data.config.joint_names]) def _robot_id_for_group(self, group_id: PlanningGroupID) -> WorldRobotID: - group = planning_group_from_configs( - group_id, [data.config for data in self._robots.values()] - ) + group = self._planning_group_from_id(group_id) matches = [ rid for rid, data in self._robots.items() if data.config.name == group.robot_name ] @@ -970,7 +1011,7 @@ def get_ee_pose(self, ctx: Context, robot_id: WorldRobotID) -> PoseStamped: if robot_id not in self._robots: raise KeyError(f"Robot '{robot_id}' not found") robot_data = self._robots[robot_id] - group_id = primary_pose_group_id_for_config(robot_data.config) + 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" @@ -982,9 +1023,7 @@ def get_group_ee_pose(self, ctx: Context, group_id: PlanningGroupID) -> PoseStam if not self._finalized: raise RuntimeError("World must be finalized first") - group = planning_group_from_configs( - group_id, [data.config for data in self._robots.values()] - ) + 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)] @@ -1034,7 +1073,7 @@ def get_jacobian(self, ctx: Context, robot_id: WorldRobotID) -> NDArray[np.float if robot_id not in self._robots: raise KeyError(f"Robot '{robot_id}' not found") robot_data = self._robots[robot_id] - group_id = primary_pose_group_id_for_config(robot_data.config) + 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" @@ -1046,9 +1085,7 @@ def get_group_jacobian(self, ctx: Context, group_id: PlanningGroupID) -> NDArray if not self._finalized: raise RuntimeError("World must be finalized first") - group = planning_group_from_configs( - group_id, [data.config for data in self._robots.values()] - ) + 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)] diff --git a/dimos/manipulation/planning/world/roboplan_world.py b/dimos/manipulation/planning/world/roboplan_world.py index 11a19c883c..0a3c69d30c 100644 --- a/dimos/manipulation/planning/world/roboplan_world.py +++ b/dimos/manipulation/planning/world/roboplan_world.py @@ -41,12 +41,12 @@ "Install the manipulation extra before selecting the roboplan backend." ) from exc -from dimos.manipulation.planning.groups.identifiers import is_global_joint_name -from dimos.manipulation.planning.groups.utils import ( - planning_group_from_configs, - primary_pose_group_id_for_config, - validate_planning_group_config, +from dimos.manipulation.planning.groups.identifiers import ( + is_global_joint_name, + make_global_joint_names, + make_planning_group_id, ) +from dimos.manipulation.planning.groups.models import PlanningGroup from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import ObstacleType, PlanningStatus from dimos.manipulation.planning.spec.models import ( @@ -117,7 +117,7 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: raise FileNotFoundError(f"Robot model not found: {Path(config.model_path).resolve()}") if any(data.config.name == config.name for data in self._robots.values()): raise ValueError(f"Robot name '{config.name}' is already registered") - validate_planning_group_config(config) + self._validate_planning_group_config(config) self._validate_robot_config(config) self._robot_counter += 1 @@ -286,16 +286,14 @@ def check_edge_collision_free( def get_ee_pose(self, ctx: RoboPlanContext, robot_id: WorldRobotID) -> PoseStamped: """Get end-effector pose if RoboPlan exposes FK.""" robot = self._get_robot(robot_id) - group_id = primary_pose_group_id_for_config(robot.config) + group_id = self._primary_pose_group_id_for_config(robot.config) if group_id is None: raise ValueError(f"Robot '{robot.config.name}' has no pose-targetable planning group") return self.get_group_ee_pose(ctx, group_id) def get_group_ee_pose(self, ctx: RoboPlanContext, group_id: PlanningGroupID) -> PoseStamped: """Get planning-group tip pose if RoboPlan exposes FK.""" - group = planning_group_from_configs( - group_id, [data.config for data in self._robots.values()] - ) + 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) @@ -325,7 +323,7 @@ def get_link_pose( def get_jacobian(self, ctx: RoboPlanContext, robot_id: WorldRobotID) -> NDArray[np.float64]: """Get end-effector Jacobian if RoboPlan exposes a compatible API.""" robot = self._get_robot(robot_id) - group_id = primary_pose_group_id_for_config(robot.config) + group_id = self._primary_pose_group_id_for_config(robot.config) if group_id is None: raise ValueError(f"Robot '{robot.config.name}' has no pose-targetable planning group") return self.get_group_jacobian(ctx, group_id) @@ -334,9 +332,7 @@ def get_group_jacobian( self, ctx: RoboPlanContext, group_id: PlanningGroupID ) -> NDArray[np.float64]: """Get planning-group Jacobian projected to group-local joint order.""" - group = planning_group_from_configs( - group_id, [data.config for data in self._robots.values()] - ) + 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) @@ -559,15 +555,63 @@ def _query_scene_joint_order( return None return list(group_info.joint_names) + def _validate_planning_group_config(self, config: RobotModelConfig) -> None: + """Validate planning groups before mutating backend state.""" + seen_group_names: set[str] = set() + for definition in config.planning_groups: + group_id = make_planning_group_id(config.name, definition.name) + if definition.name in seen_group_names: + raise ValueError(f"Planning group '{group_id}' is already registered") + make_global_joint_names(config.name, definition.joint_names) + seen_group_names.add(definition.name) + + def _planning_group_from_config( + self, config: RobotModelConfig, group_id: PlanningGroupID + ) -> PlanningGroup: + for definition in config.planning_groups: + if make_planning_group_id(config.name, definition.name) == group_id: + return PlanningGroup( + id=group_id, + robot_name=config.name, + group_name=definition.name, + joint_names=tuple(make_global_joint_names(config.name, definition.joint_names)), + local_joint_names=definition.joint_names, + base_link=definition.base_link, + tip_link=definition.tip_link, + source=definition.source, + ) + raise KeyError(f"Unknown planning group ID: {group_id}") + + def _planning_group_from_id(self, group_id: PlanningGroupID) -> PlanningGroup: + for robot in self._robots.values(): + try: + return self._planning_group_from_config(robot.config, group_id) + except KeyError: + continue + raise KeyError(f"Unknown planning group ID: {group_id}") + + def _primary_pose_group_id_for_config(self, config: RobotModelConfig) -> PlanningGroupID | None: + pose_group_ids = [ + make_planning_group_id(config.name, group.name) + for group in config.planning_groups + if group.has_pose_target + ] + if not pose_group_ids: + return None + if len(pose_group_ids) > 1: + raise ValueError( + f"Robot '{config.name}' has {len(pose_group_ids)} pose-targetable " + "planning groups; use an explicit planning group ID" + ) + return pose_group_ids[0] + def _get_robot(self, robot_id: WorldRobotID) -> _RoboPlanRobotData: if robot_id not in self._robots: raise KeyError(f"Robot '{robot_id}' not found") return self._robots[robot_id] def _robot_id_for_group(self, group_id: PlanningGroupID) -> WorldRobotID: - group = planning_group_from_configs( - group_id, [data.config for data in self._robots.values()] - ) + group = self._planning_group_from_id(group_id) matches = [ rid for rid, data in self._robots.items() if data.config.name == group.robot_name ] 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 841ed27f99..1331e447b7 100644 --- a/dimos/manipulation/planning/world/test_drake_world_planning_groups.py +++ b/dimos/manipulation/planning/world/test_drake_world_planning_groups.py @@ -138,7 +138,7 @@ 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-targetable"): + with pytest.raises(ValueError, match="multiple pose"): ambiguous.get_jacobian(ambiguous.get_live_context(), ambiguous_id) From f9b269c848c53c8519a15a93b8e9d3f1c757ca5c Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 2 Jul 2026 10:53:25 -0700 Subject: [PATCH 05/14] chore: revert kwarg call --- .../planning/monitor/world_monitor.py | 28 +++++++------------ 1 file changed, 10 insertions(+), 18 deletions(-) diff --git a/dimos/manipulation/planning/monitor/world_monitor.py b/dimos/manipulation/planning/monitor/world_monitor.py index cf38697678..5d5f1f4707 100644 --- a/dimos/manipulation/planning/monitor/world_monitor.py +++ b/dimos/manipulation/planning/monitor/world_monitor.py @@ -304,11 +304,9 @@ def get_current_joint_state(self, robot_id: WorldRobotID) -> JointState | None: if positions is not None: joint_names = self._robot_joints.get(robot_id, []) return JointState( - { - "name": joint_names, - "position": positions.tolist(), - "velocity": velocities.tolist() if velocities is not None else [], - } + name=joint_names, + position=positions.tolist(), + velocity=velocities.tolist() if velocities is not None else [], ) # Fall back to world's live context @@ -331,7 +329,7 @@ def current_global_joint_state(self, max_age: float = 1.0) -> JointState: if local_name in pos_by_name: names.append(f"{robot_name}/{local_name}") positions.append(float(pos_by_name[local_name])) - return JointState({"name": names, "position": positions}) + return JointState(name=names, position=positions) def current_group_joint_state( self, group_id: PlanningGroupID, max_age: float = 1.0 @@ -378,10 +376,8 @@ def _joint_state_for_group_query( raise ValueError(f"Current state is missing group joint '{local_name}'") positions_by_name[local_name] = float(position) return JointState( - { - "name": list(current_state.name), - "position": [positions_by_name[name] for name in current_state.name], - } + name=list(current_state.name), + position=[positions_by_name[name] for name in current_state.name], ) def _full_robot_joint_state_from_input( @@ -392,9 +388,7 @@ def _full_robot_joint_state_from_input( if not joint_state.name: if len(joint_state.position) != len(config.joint_names): return None - return JointState( - {"name": list(config.joint_names), "position": list(joint_state.position)} - ) + return JointState(name=list(config.joint_names), position=list(joint_state.position)) if len(joint_state.name) != len(joint_state.position): raise ValueError("JointState name and position lengths must match") resolved_positions: dict[str, float] = {} @@ -416,10 +410,8 @@ def _full_robot_joint_state_from_input( if set(resolved_positions) != set(config.joint_names): return None return JointState( - { - "name": list(config.joint_names), - "position": [resolved_positions[name] for name in config.joint_names], - } + name=list(config.joint_names), + position=[resolved_positions[name] for name in config.joint_names], ) def _validate_planning_group_config(self, config: RobotModelConfig) -> None: @@ -438,7 +430,7 @@ def get_current_velocities(self, robot_id: WorldRobotID) -> JointState | None: velocities = self._state_monitors[robot_id].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=joint_names, velocity=velocities.tolist()) return None def wait_for_state(self, robot_id: WorldRobotID, timeout: float = 1.0) -> bool: From 8d0527112967c833addaf771737f12d0563744ca Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 2 Jul 2026 11:34:25 -0700 Subject: [PATCH 06/14] refactor: simplify group state queries --- .../planning/monitor/test_world_monitor.py | 19 ++--- .../planning/monitor/world_monitor.py | 84 +++---------------- 2 files changed, 22 insertions(+), 81 deletions(-) diff --git a/dimos/manipulation/planning/monitor/test_world_monitor.py b/dimos/manipulation/planning/monitor/test_world_monitor.py index b7d9016922..84a17160d1 100644 --- a/dimos/manipulation/planning/monitor/test_world_monitor.py +++ b/dimos/manipulation/planning/monitor/test_world_monitor.py @@ -347,7 +347,7 @@ def test_group_kinematics_with_full_state_does_not_require_current_state() -> No pose = monitor.get_group_ee_pose( "arm/manipulator", - JointState({"name": ["j1", "j2", "j3"], "position": [0.1, 0.2, 0.3]}), + 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"] @@ -356,17 +356,18 @@ def test_group_kinematics_with_full_state_does_not_require_current_state() -> No assert pose.position.x == 1 -def test_group_kinematics_route_to_backend_and_merge_group_state() -> None: +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] - 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.add_robot(_three_joint_reordered_group_config()) pose = monitor.get_group_ee_pose( - "arm/manipulator", JointState({"name": ["arm/j2", "arm/j1"], "position": [0.8, 0.9]}) + "arm/manipulator", + JointState(name=["j1", "j2", "j3"], position=[0.9, 0.8, 0.3]), ) jacobian = monitor.get_group_jacobian( - "arm/manipulator", JointState({"name": ["arm/j2", "arm/j1"], "position": [0.3, 0.4]}) + "arm/manipulator", + JointState(name=["j1", "j2", "j3"], position=[0.4, 0.3, 0.3]), ) set_calls = [call for call in fake_world.calls if call[0] == "set_joint_state"] @@ -389,7 +390,7 @@ def test_legacy_wrappers_fail_for_no_pose_and_ambiguous_pose_groups() -> None: ) ) with pytest.raises(ValueError, match="no pose-targetable"): - monitor.get_ee_pose(no_pose_id, JointState({"name": ["j1", "j2"], "position": [0.0, 0.0]})) + monitor.get_ee_pose(no_pose_id, JointState(name=["j1", "j2"], position=[0.0, 0.0])) fake_world2 = FakeWorld() monitor2 = world_monitor_module.WorldMonitor(world=fake_world2) # type: ignore[arg-type] @@ -406,6 +407,4 @@ 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(ambiguous_id, JointState(name=["j1", "j2"], position=[0.0, 0.0])) diff --git a/dimos/manipulation/planning/monitor/world_monitor.py b/dimos/manipulation/planning/monitor/world_monitor.py index 5d5f1f4707..9a0241d6a2 100644 --- a/dimos/manipulation/planning/monitor/world_monitor.py +++ b/dimos/manipulation/planning/monitor/world_monitor.py @@ -335,16 +335,12 @@ 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) - state = self._current_robot_joint_state_for_group(group_id, max_age) - return filter_joint_state_to_selected_joints( - state, group.joint_names, group.local_joint_names - ) + return self._current_robot_joint_state_for_group(group_id, max_age) def _current_robot_joint_state_for_group( self, group_id: PlanningGroupID, max_age: float = 1.0 ) -> JointState: - """Return current full robot state for a planning group.""" + """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): @@ -352,66 +348,8 @@ def _current_robot_joint_state_for_group( state = self.get_current_joint_state(robot_id) if state is None: raise ValueError(f"Current state for robot '{group.robot_name}' is unavailable") - return state - - def _joint_state_for_group_query( - self, group_id: PlanningGroupID, joint_state: JointState | None - ) -> JointState: - """Normalize full-robot or group-scoped input to full robot-local state.""" - group = self._planning_groups.get(group_id) - robot_id = self._robot_ids_by_name[group.robot_name] - if joint_state is not None: - full_state = self._full_robot_joint_state_from_input(robot_id, joint_state) - if full_state is not None: - return full_state - current_state = self._current_robot_joint_state_for_group(group_id) - if joint_state is None: - return current_state - group_state = filter_joint_state_to_selected_joints( - joint_state, group.joint_names, group.local_joint_names - ) - positions_by_name = dict(zip(current_state.name, current_state.position, strict=True)) - for local_name, position in zip(group.local_joint_names, group_state.position, strict=True): - if local_name not in positions_by_name: - raise ValueError(f"Current state is missing group joint '{local_name}'") - positions_by_name[local_name] = float(position) - return JointState( - name=list(current_state.name), - position=[positions_by_name[name] for name in current_state.name], - ) - - def _full_robot_joint_state_from_input( - self, robot_id: WorldRobotID, joint_state: JointState - ) -> JointState | None: - """Return a full robot-local state if input contains all robot joints.""" - config = self._robot_configs[robot_id] - if not joint_state.name: - if len(joint_state.position) != len(config.joint_names): - return None - return JointState(name=list(config.joint_names), position=list(joint_state.position)) - if len(joint_state.name) != len(joint_state.position): - raise ValueError("JointState name and position lengths must match") - resolved_positions: dict[str, float] = {} - global_prefix = f"{config.name}/" - for name, position in zip(joint_state.name, joint_state.position, strict=True): - if name in config.joint_names: - resolved_name = name - elif name in config.joint_name_mapping: - resolved_name = config.joint_name_mapping[name] - elif name.startswith(global_prefix): - resolved_name = name[len(global_prefix) :] - else: - resolved_name = config.get_urdf_joint_name(name) - if resolved_name not in config.joint_names: - return None - if resolved_name in resolved_positions: - raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") - resolved_positions[resolved_name] = float(position) - if set(resolved_positions) != set(config.joint_names): - return None - return JointState( - name=list(config.joint_names), - position=[resolved_positions[name] for name in config.joint_names], + return filter_joint_state_to_selected_joints( + state, group.joint_names, group.local_joint_names ) def _validate_planning_group_config(self, config: RobotModelConfig) -> None: @@ -506,12 +444,17 @@ def get_ee_pose( def get_group_ee_pose( self, group_id: PlanningGroupID, joint_state: JointState | None = None ) -> PoseStamped: - """Get planning-group tip pose. Uses current group state if joint_state is None.""" + """Get planning-group tip pose. Uses current robot state if joint_state is None.""" group = self._planning_groups.get(group_id) robot_id = self._robot_ids_by_name[group.robot_name] with self._world.scratch_context() as ctx: - normalized_state = self._joint_state_for_group_query(group_id, joint_state) - self._world.set_joint_state(ctx, robot_id, normalized_state) + if joint_state is None: + if robot_id in self._state_monitors and self.is_state_stale(robot_id): + raise ValueError(f"Current state for robot '{group.robot_name}' is stale") + joint_state = self.get_current_joint_state(robot_id) + if joint_state is None: + raise ValueError(f"Current state for robot '{group.robot_name}' is unavailable") + self._world.set_joint_state(ctx, robot_id, joint_state) return self._world.get_group_ee_pose(ctx, group_id) @@ -562,8 +505,7 @@ def get_group_jacobian( group = self._planning_groups.get(group_id) robot_id = self._robot_ids_by_name[group.robot_name] with self._world.scratch_context() as ctx: - normalized_state = self._joint_state_for_group_query(group_id, joint_state) - self._world.set_joint_state(ctx, robot_id, normalized_state) + self._world.set_joint_state(ctx, robot_id, joint_state) return self._world.get_group_jacobian(ctx, group_id) # Lifecycle From 340aef7b906dbb2fc66f2a8a271027fe6610cefe Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 2 Jul 2026 12:04:50 -0700 Subject: [PATCH 07/14] refactor: simplify current group state --- .../manipulation/planning/monitor/world_monitor.py | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/dimos/manipulation/planning/monitor/world_monitor.py b/dimos/manipulation/planning/monitor/world_monitor.py index 9a0241d6a2..9b54332de8 100644 --- a/dimos/manipulation/planning/monitor/world_monitor.py +++ b/dimos/manipulation/planning/monitor/world_monitor.py @@ -324,21 +324,13 @@ def current_global_joint_state(self, max_age: float = 1.0) -> JointState: state = self.get_current_joint_state(robot_id) if state is None: continue - pos_by_name = dict(zip(state.name, state.position, strict=True)) - for local_name in self._robot_joints.get(robot_id, []): - if local_name in pos_by_name: - names.append(f"{robot_name}/{local_name}") - positions.append(float(pos_by_name[local_name])) + for name, position in zip(state.name, state.position, strict=True): + names.append(f"{robot_name}/{name}") + positions.append(float(position)) return JointState(name=names, position=positions) def current_group_joint_state( self, group_id: PlanningGroupID, max_age: float = 1.0 - ) -> JointState: - """Return current joint state scoped and ordered for one planning group.""" - return self._current_robot_joint_state_for_group(group_id, max_age) - - def _current_robot_joint_state_for_group( - 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) From a33f8ec0e147e45c571570915ba127e9ccbc70b9 Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 2 Jul 2026 12:26:16 -0700 Subject: [PATCH 08/14] refactor: share joint state ordering --- dimos/manipulation/planning/groups/utils.py | 47 +++++++++++++++++++ .../planning/world/drake_world.py | 37 ++++----------- .../planning/world/roboplan_world.py | 35 +++----------- 3 files changed, 62 insertions(+), 57 deletions(-) diff --git a/dimos/manipulation/planning/groups/utils.py b/dimos/manipulation/planning/groups/utils.py index e55f997a33..ae3e865c8f 100644 --- a/dimos/manipulation/planning/groups/utils.py +++ b/dimos/manipulation/planning/groups/utils.py @@ -16,6 +16,9 @@ from collections.abc import Mapping, Sequence +import numpy as np +from numpy.typing import NDArray + from dimos.manipulation.planning.groups.identifiers import ( assert_global_joint_names, assert_local_joint_names, @@ -134,3 +137,47 @@ def joint_target_to_global_names( if extra: raise ValueError(f"Target for '{group.id}' has extra joints: {sorted(extra)}") return JointState(name=list(group.joint_names), position=global_positions) + + +def joint_state_to_ordered_positions( + joint_state: JointState, + *, + robot_name: str, + joint_names: Sequence[str], + joint_name_mapping: Mapping[str, str], + context: str, +) -> NDArray[np.float64]: + """Convert a JointState to an array ordered by local robot joint names.""" + if not joint_state.name: + if len(joint_state.position) != len(joint_names): + raise ValueError("JointState position length must match configured joint count") + return np.asarray(joint_state.position, dtype=np.float64) + + if len(joint_state.name) != len(joint_state.position): + raise ValueError("JointState name and position lengths must match") + + joint_name_set = set(joint_names) + name_to_pos: dict[str, float] = {} + prefix = f"{robot_name}/" + 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): + if not name.startswith(prefix): + continue + resolved_name = name[len(prefix) :] + if resolved_name not in joint_name_set: + raise ValueError(f"Unknown global joint name for {context}: {name}") + else: + resolved_name = joint_name_mapping.get(name, name) + + if resolved_name in name_to_pos: + raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") + name_to_pos[resolved_name] = float(position) + + missing = [name for name in joint_names if name not in name_to_pos] + if missing: + raise ValueError(f"JointState missing joints for {context}: {missing}") + return np.asarray([name_to_pos[name] for name in joint_names], dtype=np.float64) diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index b84c735556..e5cb3cb72a 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -26,11 +26,11 @@ import numpy as np from dimos.manipulation.planning.groups.identifiers import ( - is_global_joint_name, make_global_joint_names, make_planning_group_id, ) from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.manipulation.planning.groups.utils import joint_state_to_ordered_positions from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import ObstacleType from dimos.manipulation.planning.spec.models import ( @@ -874,34 +874,13 @@ def _joint_state_to_q( if robot_id not in self._robots: raise KeyError(f"Robot '{robot_id}' not found") robot_data = self._robots[robot_id] - if not joint_state.name: - if len(joint_state.position) != len(robot_data.config.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") - name_to_pos: dict[str, float] = {} - for name, position in zip(joint_state.name, joint_state.position, strict=True): - if name in robot_data.config.joint_names: - resolved_name = name - elif name in robot_data.config.joint_name_mapping: - resolved_name = robot_data.config.joint_name_mapping[name] - elif is_global_joint_name(name): - prefix = f"{robot_data.config.name}/" - if not name.startswith(prefix): - continue - resolved_name = name[len(prefix) :] - if resolved_name not in robot_data.config.joint_names: - raise ValueError(f"Unknown global joint name for DrakeWorld: {name}") - else: - resolved_name = robot_data.config.get_urdf_joint_name(name) - if resolved_name in name_to_pos: - raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") - name_to_pos[resolved_name] = float(position) - missing = [name for name in robot_data.config.joint_names if name not in name_to_pos] - if missing: - raise ValueError(f"JointState missing joints for DrakeWorld: {missing}") - return np.asarray([name_to_pos[name] for name in robot_data.config.joint_names]) + return joint_state_to_ordered_positions( + joint_state, + robot_name=robot_data.config.name, + joint_names=robot_data.config.joint_names, + joint_name_mapping=robot_data.config.joint_name_mapping, + context="DrakeWorld", + ) def _robot_id_for_group(self, group_id: PlanningGroupID) -> WorldRobotID: group = self._planning_group_from_id(group_id) diff --git a/dimos/manipulation/planning/world/roboplan_world.py b/dimos/manipulation/planning/world/roboplan_world.py index 0a3c69d30c..5038fdeebe 100644 --- a/dimos/manipulation/planning/world/roboplan_world.py +++ b/dimos/manipulation/planning/world/roboplan_world.py @@ -42,11 +42,11 @@ ) from exc from dimos.manipulation.planning.groups.identifiers import ( - is_global_joint_name, make_global_joint_names, make_planning_group_id, ) from dimos.manipulation.planning.groups.models import PlanningGroup +from dimos.manipulation.planning.groups.utils import joint_state_to_ordered_positions from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import ObstacleType, PlanningStatus from dimos.manipulation.planning.spec.models import ( @@ -623,33 +623,12 @@ def _joint_state_to_q( self, robot_id: WorldRobotID, joint_state: JointState ) -> NDArray[np.float64]: robot = self._get_robot(robot_id) - if len(joint_state.position) != len(robot.config.joint_names): - raise ValueError("JointState position length must match configured joint count") - if not joint_state.name: - return np.asarray(joint_state.position, dtype=np.float64) - name_to_pos: dict[str, float] = {} - for name, position in zip(joint_state.name, joint_state.position, strict=True): - if name in robot.config.joint_names: - resolved_name = name - elif name in robot.config.joint_name_mapping: - resolved_name = robot.config.joint_name_mapping[name] - elif is_global_joint_name(name): - prefix = f"{robot.config.name}/" - if not name.startswith(prefix): - continue - resolved_name = name[len(prefix) :] - if resolved_name not in robot.config.joint_names: - raise ValueError(f"Unknown global joint name for RoboPlanWorld: {name}") - else: - resolved_name = robot.config.get_urdf_joint_name(name) - if resolved_name in name_to_pos: - raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") - name_to_pos[resolved_name] = float(position) - missing = [name for name in robot.config.joint_names if name not in name_to_pos] - if missing: - raise ValueError(f"JointState missing joints for RoboPlanWorld: {missing}") - return np.asarray( - [name_to_pos[name] for name in robot.config.joint_names], dtype=np.float64 + return joint_state_to_ordered_positions( + joint_state, + robot_name=robot.config.name, + joint_names=robot.config.joint_names, + joint_name_mapping=robot.config.joint_name_mapping, + context="RoboPlanWorld", ) def _require_finalized(self) -> None: From df05bcbfa13149e99841804ed0989162cfbfabe2 Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 2 Jul 2026 12:32:30 -0700 Subject: [PATCH 09/14] refactor: simplify joint state ordering util --- dimos/manipulation/planning/groups/utils.py | 11 +++-------- dimos/manipulation/planning/world/drake_world.py | 2 -- dimos/manipulation/planning/world/roboplan_world.py | 2 -- 3 files changed, 3 insertions(+), 12 deletions(-) diff --git a/dimos/manipulation/planning/groups/utils.py b/dimos/manipulation/planning/groups/utils.py index ae3e865c8f..66d33ea359 100644 --- a/dimos/manipulation/planning/groups/utils.py +++ b/dimos/manipulation/planning/groups/utils.py @@ -142,10 +142,8 @@ def joint_target_to_global_names( def joint_state_to_ordered_positions( joint_state: JointState, *, - robot_name: str, joint_names: Sequence[str], joint_name_mapping: Mapping[str, str], - context: str, ) -> NDArray[np.float64]: """Convert a JointState to an array ordered by local robot joint names.""" if not joint_state.name: @@ -158,18 +156,15 @@ def joint_state_to_ordered_positions( joint_name_set = set(joint_names) name_to_pos: dict[str, float] = {} - prefix = f"{robot_name}/" 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): - if not name.startswith(prefix): - continue - resolved_name = name[len(prefix) :] + resolved_name = name.split("/", maxsplit=1)[1] if resolved_name not in joint_name_set: - raise ValueError(f"Unknown global joint name for {context}: {name}") + raise ValueError(f"Unknown global joint name: {name}") else: resolved_name = joint_name_mapping.get(name, name) @@ -179,5 +174,5 @@ def joint_state_to_ordered_positions( missing = [name for name in joint_names if name not in name_to_pos] if missing: - raise ValueError(f"JointState missing joints for {context}: {missing}") + raise ValueError(f"JointState missing joints: {missing}") return np.asarray([name_to_pos[name] for name in joint_names], dtype=np.float64) diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index e5cb3cb72a..172e7f40a1 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -876,10 +876,8 @@ def _joint_state_to_q( robot_data = self._robots[robot_id] return joint_state_to_ordered_positions( joint_state, - robot_name=robot_data.config.name, joint_names=robot_data.config.joint_names, joint_name_mapping=robot_data.config.joint_name_mapping, - context="DrakeWorld", ) def _robot_id_for_group(self, group_id: PlanningGroupID) -> WorldRobotID: diff --git a/dimos/manipulation/planning/world/roboplan_world.py b/dimos/manipulation/planning/world/roboplan_world.py index 5038fdeebe..b5ac50d828 100644 --- a/dimos/manipulation/planning/world/roboplan_world.py +++ b/dimos/manipulation/planning/world/roboplan_world.py @@ -625,10 +625,8 @@ def _joint_state_to_q( robot = self._get_robot(robot_id) return joint_state_to_ordered_positions( joint_state, - robot_name=robot.config.name, joint_names=robot.config.joint_names, joint_name_mapping=robot.config.joint_name_mapping, - context="RoboPlanWorld", ) def _require_finalized(self) -> None: From 3189191b78983b933c25d3a6ff0aa44220a0f56b Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 2 Jul 2026 13:09:13 -0700 Subject: [PATCH 10/14] spec: archive 2 --- .../.openspec.yaml | 2 -- .../planning-world-monitor-groups/README.md | 3 --- .../planning-world-monitor-groups/design.md | 27 ------------------- .../planning-world-monitor-groups/proposal.md | 26 ------------------ .../planning-world-monitor-groups/tasks.md | 19 ------------- .../planning-world-monitor-groups/spec.md | 6 ++++- 6 files changed, 5 insertions(+), 78 deletions(-) delete mode 100644 openspec/changes/planning-world-monitor-groups/.openspec.yaml delete mode 100644 openspec/changes/planning-world-monitor-groups/README.md delete mode 100644 openspec/changes/planning-world-monitor-groups/design.md delete mode 100644 openspec/changes/planning-world-monitor-groups/proposal.md delete mode 100644 openspec/changes/planning-world-monitor-groups/tasks.md rename openspec/{changes/planning-world-monitor-groups => }/specs/planning-world-monitor-groups/spec.md (92%) diff --git a/openspec/changes/planning-world-monitor-groups/.openspec.yaml b/openspec/changes/planning-world-monitor-groups/.openspec.yaml deleted file mode 100644 index 34f9314d22..0000000000 --- a/openspec/changes/planning-world-monitor-groups/.openspec.yaml +++ /dev/null @@ -1,2 +0,0 @@ -schema: spec-driven -created: 2026-06-29 diff --git a/openspec/changes/planning-world-monitor-groups/README.md b/openspec/changes/planning-world-monitor-groups/README.md deleted file mode 100644 index c5b27282f4..0000000000 --- a/openspec/changes/planning-world-monitor-groups/README.md +++ /dev/null @@ -1,3 +0,0 @@ -# planning-world-monitor-groups - -Extract PR 2 from reference branch cc/spec/movegroup: teach world backends and world monitors to answer planning-group FK, Jacobian, state, and collision queries. diff --git a/openspec/changes/planning-world-monitor-groups/design.md b/openspec/changes/planning-world-monitor-groups/design.md deleted file mode 100644 index eb1e36a5e6..0000000000 --- a/openspec/changes/planning-world-monitor-groups/design.md +++ /dev/null @@ -1,27 +0,0 @@ -## Context - -Planning groups become useful when world backends can answer queries for a specific group. This layer should prove that configured groups map correctly onto backend model instances, joint order, FK target frames, and Jacobian columns. - -## Goals / Non-Goals - -**Goals:** -- Implement group FK/Jacobian support in Drake and RoboPlan worlds. -- Make WorldMonitor expose group-scoped state and query helpers. -- Validate group-local joint ordering and global/local joint-state mapping. -- Keep base pose/base link semantics safe by rejecting unsupported backend cases rather than silently producing wrong transforms. - -**Non-Goals:** -- Do not migrate planner or IK algorithms yet. -- Do not expose new public module APIs yet. -- Do not include Viser UI changes. - -## Decisions - -- Robot-scoped FK/Jacobian compatibility wrappers may remain temporarily but must resolve through exactly one pose-targetable group or fail clearly. -- RoboPlan Jacobians must be projected into group-local joint order; returning a full backend Jacobian is not sufficient. -- Backends should reject unsupported base-pose/base-link combinations explicitly. - -## Risks / Trade-offs - -- Backend capabilities differ. Tests should pin the supported subset and error messages. -- Some existing tests may need compatibility shims until PR 3 migrates algorithms and module APIs. diff --git a/openspec/changes/planning-world-monitor-groups/proposal.md b/openspec/changes/planning-world-monitor-groups/proposal.md deleted file mode 100644 index 2d8d1681d1..0000000000 --- a/openspec/changes/planning-world-monitor-groups/proposal.md +++ /dev/null @@ -1,26 +0,0 @@ -## Why - -Once planning groups exist, the backend query layer must answer group-scoped FK, Jacobian, joint-state, and collision questions. This is the second reviewable PR in the stack: it consumes the foundation without exposing the public `ManipulationModule` API or UI rewrite. - -## What Changes - -- Teach Drake and RoboPlan world backends to resolve planning groups. -- Add group FK/Jacobian APIs and group-local joint ordering behavior. -- Update world monitor and robot state monitor logic for group-aware state access. -- Preserve legacy robot-scoped wrappers only where needed for compatibility, resolving through a unique pose-targetable group. -- Add backend and monitor tests. - -## Capabilities - -### New Capabilities -- `planning-world-monitor-groups`: World and monitor support for planning-group queries. - -### Modified Capabilities -- `manipulation-world-backends`: Drake/RoboPlan backends understand group-local chains and target frames. - -## Impact - -- Base branch: PR 1 `planning-groups-foundation`. -- Reference implementation: `cc/spec/movegroup`. -- Primary files: `dimos/manipulation/planning/world/*`, `dimos/manipulation/planning/monitor/*`, RoboPlan world tests, Drake planning-group tests. -- Out of scope: IK/RRT algorithm migration, `ManipulationModule`, Viser UI, and control integration. diff --git a/openspec/changes/planning-world-monitor-groups/tasks.md b/openspec/changes/planning-world-monitor-groups/tasks.md deleted file mode 100644 index 8e31566a05..0000000000 --- a/openspec/changes/planning-world-monitor-groups/tasks.md +++ /dev/null @@ -1,19 +0,0 @@ -## 1. Extract backend group support - -- [x] 1.1 Start from PR 1 and use `cc/spec/movegroup` as reference. -- [x] 1.2 Extract Drake world group FK/Jacobian and base-link handling changes. -- [x] 1.3 Extract RoboPlan world group FK/Jacobian, joint-name normalization, and URDF strip handling changes. -- [x] 1.4 Extract world monitor and robot state monitor group-aware query changes. - -## 2. Tests - -- [x] 2.1 Bring over Drake group world tests. -- [x] 2.2 Bring over RoboPlan world tests, including planning split file if needed. -- [x] 2.3 Bring over WorldMonitor tests for group state and ambiguity behavior. - -## 3. Validation - -- [x] 3.1 Run `uv run pytest dimos/manipulation/planning/world/test_drake_world_planning_groups.py dimos/manipulation/test_roboplan_world.py dimos/manipulation/test_roboplan_world_planning.py dimos/manipulation/planning/monitor/test_world_monitor.py -q`. - - Note: `dimos/manipulation/test_roboplan_world_planning.py` does not exist in this branch; ran the same targeted command without that optional split file. -- [x] 3.2 Run targeted mypy on changed backend/monitor files. -- [x] 3.3 Optional manual smoke: load a robot config, print group list, FK pose, Jacobian shape, and collision status. diff --git a/openspec/changes/planning-world-monitor-groups/specs/planning-world-monitor-groups/spec.md b/openspec/specs/planning-world-monitor-groups/spec.md similarity index 92% rename from openspec/changes/planning-world-monitor-groups/specs/planning-world-monitor-groups/spec.md rename to openspec/specs/planning-world-monitor-groups/spec.md index 3bff6f8bd0..5a737525bc 100644 --- a/openspec/changes/planning-world-monitor-groups/specs/planning-world-monitor-groups/spec.md +++ b/openspec/specs/planning-world-monitor-groups/spec.md @@ -1,4 +1,8 @@ -## ADDED Requirements +## Purpose + +Specify world and monitor support for planning-group queries. + +## Requirements ### Requirement: World backends must answer group pose queries World backends MUST provide end-effector pose queries for pose-targetable planning groups using the group's configured target frame. From 8a4ddbdaf5b13c1e3b4d365fa472d62fe4514116 Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 2 Jul 2026 13:43:39 -0700 Subject: [PATCH 11/14] test: improve planning group coverage --- .../planning/groups/test_planning_groups.py | 105 +++++++++++++ .../planning/monitor/test_world_monitor.py | 81 ++++++++++ .../planning/world/drake_world.py | 9 +- .../world/test_drake_world_planning_groups.py | 50 +++++- dimos/manipulation/test_roboplan_world.py | 142 ++++++++++++++++++ 5 files changed, 380 insertions(+), 7 deletions(-) diff --git a/dimos/manipulation/planning/groups/test_planning_groups.py b/dimos/manipulation/planning/groups/test_planning_groups.py index e716366202..0d5f90a2c9 100644 --- a/dimos/manipulation/planning/groups/test_planning_groups.py +++ b/dimos/manipulation/planning/groups/test_planning_groups.py @@ -32,6 +32,7 @@ from dimos.manipulation.planning.groups.registry import PlanningGroupRegistry from dimos.manipulation.planning.groups.utils import ( filter_joint_state_to_selected_joints, + joint_state_to_ordered_positions, joint_target_to_global_names, matching_global_joint_name, planning_group_id_from_selector, @@ -473,6 +474,37 @@ def test_matching_global_joint_name_requires_unique_suffix_match() -> 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() @@ -551,3 +583,76 @@ def test_robot_model_config_end_effector_link_rejects_ambiguous_pose_groups() -> with pytest.raises(ValueError, match="multiple pose-target planning groups"): _ = config.end_effector_link + + +def test_joint_state_to_ordered_positions_accepts_all_supported_name_forms() -> None: + joint_names = ["joint1", "joint2", "joint3"] + mapping = {"hw1": "joint1", "hw2": "joint2", "hw3": "joint3"} + + unnamed = joint_state_to_ordered_positions( + JointState(name=[], position=[1.0, 2.0, 3.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + local = joint_state_to_ordered_positions( + JointState(name=["joint3", "joint1", "joint2"], position=[30.0, 10.0, 20.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + coordinator = joint_state_to_ordered_positions( + JointState(name=["hw2", "hw3", "hw1"], position=[200.0, 300.0, 100.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + global_names = joint_state_to_ordered_positions( + JointState(name=["arm/joint2", "arm/joint1", "arm/joint3"], position=[2.0, 1.0, 3.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + + assert unnamed.tolist() == [1.0, 2.0, 3.0] + assert local.tolist() == [10.0, 20.0, 30.0] + assert coordinator.tolist() == [100.0, 200.0, 300.0] + assert global_names.tolist() == [1.0, 2.0, 3.0] + + +def test_joint_state_to_ordered_positions_rejects_invalid_inputs() -> None: + joint_names = ["joint1", "joint2"] + mapping = {"hw1": "joint1"} + + with pytest.raises(ValueError, match="position length"): + joint_state_to_ordered_positions( + JointState(name=[], position=[1.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + with pytest.raises(ValueError, match="name and position"): + joint_state_to_ordered_positions( + JointState(name=["joint1", "joint2"], position=[1.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + with pytest.raises(ValueError, match="duplicate"): + joint_state_to_ordered_positions( + JointState(name=["joint1", "hw1"], position=[1.0, 2.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + with pytest.raises(ValueError, match="Unknown global"): + joint_state_to_ordered_positions( + JointState(name=["arm/joint3", "joint2"], position=[1.0, 2.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + with pytest.raises(ValueError, match="missing joints"): + joint_state_to_ordered_positions( + JointState(name=["joint1"], position=[1.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) + with pytest.raises(ValueError, match="missing joints"): + joint_state_to_ordered_positions( + JointState(name=["mystery", "joint2"], position=[1.0, 2.0]), + joint_names=joint_names, + joint_name_mapping=mapping, + ) diff --git a/dimos/manipulation/planning/monitor/test_world_monitor.py b/dimos/manipulation/planning/monitor/test_world_monitor.py index 84a17160d1..a78957858c 100644 --- a/dimos/manipulation/planning/monitor/test_world_monitor.py +++ b/dimos/manipulation/planning/monitor/test_world_monitor.py @@ -340,6 +340,87 @@ def test_current_group_joint_state_uses_public_names_in_group_order() -> None: assert state.position == [0.2, 0.1] +def test_current_global_joint_state_skips_stale_robots_and_preserves_state_order() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + fresh_id = monitor.add_robot(_three_joint_reordered_group_config()) + stale_id = monitor.add_robot( + RobotModelConfig( + name="arm2", + model_path=Path("/tmp/arm2.urdf"), + joint_names=["a", "b"], + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", joint_names=("a", "b"), base_link="base", tip_link="ee" + ) + ], + ) + ) + monitor._state_monitors[fresh_id] = _FakeStateMonitor([0.1, 0.2, 0.3]) # type: ignore[attr-defined] + monitor._state_monitors[stale_id] = _FakeStateMonitor([1.0, 2.0], stale=True) # type: ignore[attr-defined] + monitor.add_robot( + RobotModelConfig( + name="arm3", + model_path=Path("/tmp/arm3.urdf"), + joint_names=["x"], + planning_groups=[ + PlanningGroupDefinition( + name="manipulator", joint_names=("x",), base_link="base", tip_link="ee" + ) + ], + ) + ) + + state = monitor.current_global_joint_state(max_age=0.5) + + assert state.name == ["arm/j1", "arm/j2", "arm/j3"] + assert state.position == [0.1, 0.2, 0.3] + + +def test_current_group_joint_state_rejects_stale_or_unavailable_state() -> None: + stale_world = FakeWorld() + stale_monitor = world_monitor_module.WorldMonitor(world=stale_world) # type: ignore[arg-type] + stale_id = stale_monitor.add_robot(_three_joint_reordered_group_config()) + stale_monitor._state_monitors[stale_id] = _FakeStateMonitor([0.1, 0.2, 0.3], stale=True) # type: ignore[attr-defined] + + with pytest.raises(ValueError, match="stale"): + stale_monitor.current_group_joint_state("arm/manipulator") + + unavailable_monitor = world_monitor_module.WorldMonitor(world=FakeWorld()) # type: ignore[arg-type] + unavailable_monitor.add_robot(_three_joint_reordered_group_config()) + with pytest.raises(ValueError, match="unavailable"): + unavailable_monitor.current_group_joint_state("arm/manipulator") + + +def test_group_ee_pose_uses_current_state_when_no_joint_state_is_provided() -> None: + fake_world = FakeWorld() + monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] + robot_id = monitor.add_robot(_three_joint_reordered_group_config()) + monitor._state_monitors[robot_id] = _FakeStateMonitor([0.1, 0.2, 0.3]) # type: ignore[attr-defined] + + pose = monitor.get_group_ee_pose("arm/manipulator") + + set_calls = [call for call in fake_world.calls if call[0] == "set_joint_state"] + assert set_calls[0][3].name == ["j1", "j2", "j3"] + assert set_calls[0][3].position == [0.1, 0.2, 0.3] + assert pose.position.x == 1 + + +def test_group_ee_pose_without_joint_state_rejects_stale_or_unavailable_state() -> None: + stale_world = FakeWorld() + stale_monitor = world_monitor_module.WorldMonitor(world=stale_world) # type: ignore[arg-type] + stale_id = stale_monitor.add_robot(_three_joint_reordered_group_config()) + stale_monitor._state_monitors[stale_id] = _FakeStateMonitor([0.1, 0.2, 0.3], stale=True) # type: ignore[attr-defined] + + with pytest.raises(ValueError, match="stale"): + stale_monitor.get_group_ee_pose("arm/manipulator") + + unavailable_monitor = world_monitor_module.WorldMonitor(world=FakeWorld()) # type: ignore[arg-type] + unavailable_monitor.add_robot(_three_joint_reordered_group_config()) + with pytest.raises(ValueError, match="unavailable"): + unavailable_monitor.get_group_ee_pose("arm/manipulator") + + def test_group_kinematics_with_full_state_does_not_require_current_state() -> None: fake_world = FakeWorld() monitor = world_monitor_module.WorldMonitor(world=fake_world) # type: ignore[arg-type] diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index 172e7f40a1..dae0d3f73c 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -334,7 +334,8 @@ def get_robot_config(self, robot_id: WorldRobotID) -> RobotModelConfig: raise KeyError(f"Robot '{robot_id}' not found") return self._robots[robot_id].config - def _validate_planning_group_config(self, config: RobotModelConfig) -> None: + @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) @@ -343,8 +344,9 @@ def _validate_planning_group_config(self, config: RobotModelConfig) -> None: make_global_joint_names(config.name, definition.joint_names) seen_group_names.add(definition.name) + @staticmethod def _planning_group_from_config( - self, config: RobotModelConfig, group_id: PlanningGroupID + config: RobotModelConfig, group_id: PlanningGroupID ) -> PlanningGroup: for definition in config.planning_groups: if make_planning_group_id(config.name, definition.name) == group_id: @@ -369,7 +371,8 @@ def _planning_group_from_id(self, group_id: PlanningGroupID) -> PlanningGroup: continue raise KeyError(f"Unknown planning group ID: {group_id}") - def _primary_pose_group_id_for_config(self, config: RobotModelConfig) -> PlanningGroupID | None: + @staticmethod + def _primary_pose_group_id_for_config(config: RobotModelConfig) -> PlanningGroupID | None: pose_groups = [group for group in config.planning_groups if group.has_pose_target] if not pose_groups: return None 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 1331e447b7..c981c76896 100644 --- a/dimos/manipulation/planning/world/test_drake_world_planning_groups.py +++ b/dimos/manipulation/planning/world/test_drake_world_planning_groups.py @@ -19,15 +19,17 @@ import numpy as np import pytest -pydrake = pytest.importorskip("pydrake") -_ = pydrake - from dimos.manipulation.planning.groups.models import PlanningGroupDefinition from dimos.manipulation.planning.spec.config import RobotModelConfig -from dimos.manipulation.planning.world.drake_world import DrakeWorld +from dimos.manipulation.planning.world.drake_world import DRAKE_AVAILABLE, DrakeWorld from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.JointState import JointState +requires_drake = pytest.mark.skipif( + not DRAKE_AVAILABLE, + reason="Drake planning-group tests require the manipulation extra", +) + def _write_urdf(path: Path) -> None: path.write_text( @@ -72,6 +74,43 @@ def _arm_group( ) +def test_drake_config_group_helpers_resolve_groups_without_drake_runtime(tmp_path: Path) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + config = _config(urdf, [_arm_group("joint2", "joint1", name="wrist")]) + + group = DrakeWorld._planning_group_from_config(config, "arm/wrist") + + assert DrakeWorld._primary_pose_group_id_for_config(config) == "arm/wrist" + assert group.id == "arm/wrist" + assert group.joint_names == ("arm/joint2", "arm/joint1") + assert group.local_joint_names == ("joint2", "joint1") + assert group.tip_link == "tool0" + + +def test_drake_config_group_helpers_validate_duplicate_and_ambiguous_groups( + tmp_path: Path, +) -> None: + urdf = tmp_path / "robot.urdf" + _write_urdf(urdf) + duplicate = _config( + urdf, + [_arm_group("joint1", name="same"), _arm_group("joint2", name="same")], + ) + ambiguous = _config( + urdf, + [_arm_group("joint1", name="a"), _arm_group("joint2", name="b")], + ) + + with pytest.raises(ValueError, match="already registered"): + DrakeWorld._validate_planning_group_config(duplicate) + with pytest.raises(ValueError, match="multiple pose"): + DrakeWorld._primary_pose_group_id_for_config(ambiguous) + with pytest.raises(KeyError, match="Unknown planning group ID"): + DrakeWorld._planning_group_from_config(ambiguous, "arm/missing") + + +@requires_drake def test_drake_group_fk_uses_tip_link_and_legacy_unique_pose_group(tmp_path: Path) -> None: urdf = tmp_path / "robot.urdf" _write_urdf(urdf) @@ -91,6 +130,7 @@ def test_drake_group_fk_uses_tip_link_and_legacy_unique_pose_group(tmp_path: Pat assert world.get_jacobian(ctx, robot_id).shape == (6, 2) +@requires_drake def test_drake_group_jacobian_shape_and_group_local_order(tmp_path: Path) -> None: urdf = tmp_path / "robot.urdf" _write_urdf(urdf) @@ -118,6 +158,7 @@ def test_drake_group_jacobian_shape_and_group_local_order(tmp_path: Path) -> Non np.testing.assert_allclose(reverse_jacobian[:, 1], forward_jacobian[:, 0]) +@requires_drake def test_drake_legacy_wrappers_fail_at_call_time_for_no_or_ambiguous_pose(tmp_path: Path) -> None: urdf = tmp_path / "robot.urdf" _write_urdf(urdf) @@ -142,6 +183,7 @@ def test_drake_legacy_wrappers_fail_at_call_time_for_no_or_ambiguous_pose(tmp_pa ambiguous.get_jacobian(ambiguous.get_live_context(), ambiguous_id) +@requires_drake def test_drake_group_jacobian_rejects_non_controllable_group_joints(tmp_path: Path) -> None: urdf = tmp_path / "robot.urdf" _write_urdf(urdf) diff --git a/dimos/manipulation/test_roboplan_world.py b/dimos/manipulation/test_roboplan_world.py index f3879297d5..32eaa97bfa 100644 --- a/dimos/manipulation/test_roboplan_world.py +++ b/dimos/manipulation/test_roboplan_world.py @@ -563,6 +563,148 @@ def fake_jacobian( np.testing.assert_allclose(jacobian, np.arange(18, dtype=np.float64).reshape(6, 3)[:, [2, 1]]) +def test_group_kinematics_reject_missing_tip_or_missing_context( + fake_roboplan: None, robot_config: RobotModelConfig +) -> None: + no_tip_config = robot_config.model_copy( + update={ + "planning_groups": [ + PlanningGroupDefinition( + name="joint_only", joint_names=("joint1", "joint2"), base_link="base" + ) + ] + } + ) + world, robot_id = _make_world(fake_roboplan, no_tip_config) + world.finalize() + + with pytest.raises(ValueError, match="no tip link"): + world.get_group_ee_pose(world.get_live_context(), "arm/joint_only") + with pytest.raises(ValueError, match="no tip link"): + world.get_group_jacobian(world.get_live_context(), "arm/joint_only") + + ctx = world.get_live_context() + del ctx.q_by_robot[robot_id] + with pytest.raises(KeyError, match=robot_id): + world.get_link_pose(ctx, robot_id, "tcp") + + jacobian_world, jacobian_robot_id = _make_world(fake_roboplan, robot_config) + jacobian_world.finalize() + jacobian_ctx = jacobian_world.get_live_context() + del jacobian_ctx.q_by_robot[jacobian_robot_id] + with pytest.raises(KeyError, match=jacobian_robot_id): + jacobian_world.get_group_jacobian(jacobian_ctx, "arm/manipulator") + + +def test_group_jacobian_validates_projection_shape_and_joint_names( + fake_roboplan: None, + robot_config: RobotModelConfig, + monkeypatch: pytest.MonkeyPatch, +) -> None: + world, robot_id = _make_world(fake_roboplan, robot_config) + world.finalize() + ctx = world.get_live_context() + world.set_joint_state(ctx, robot_id, JointState(name=["joint1", "joint2"], position=[0.0, 0.0])) + + monkeypatch.setattr(FakeScene, "joint_group_joint_names", ["joint1", "other"]) + monkeypatch.setattr( + FakeScene, + "computeFrameJacobian", + lambda self, q, frame_name, local=True: np.ones((6, 2)), + ) + with pytest.raises(ValueError, match="Unknown joints"): + world.get_group_jacobian(ctx, "arm/manipulator") + + monkeypatch.setattr(FakeScene, "joint_group_joint_names", ["joint1", "joint2"]) + monkeypatch.setattr( + FakeScene, + "computeFrameJacobian", + lambda self, q, frame_name, local=True: np.ones((5, 2)), + ) + with pytest.raises(ValueError, match="Unexpected RoboPlan Jacobian shape"): + world.get_group_jacobian(ctx, "arm/manipulator") + + monkeypatch.setattr( + FakeScene, + "computeFrameJacobian", + lambda self, q, frame_name, local=True: np.ones((6, 3)), + ) + with pytest.raises(ValueError, match="cannot project"): + world.get_group_jacobian(ctx, "arm/manipulator") + + +def test_group_jacobian_falls_back_to_configured_joint_order_when_scene_order_is_unavailable( + fake_roboplan: None, + robot_config: RobotModelConfig, + monkeypatch: pytest.MonkeyPatch, +) -> None: + world, robot_id = _make_world(fake_roboplan, robot_config) + world.finalize() + ctx = world.get_live_context() + world.set_joint_state(ctx, robot_id, JointState(name=["joint1", "joint2"], position=[0.0, 0.0])) + + def missing_group_info(self: FakeScene, name: str) -> FakeJointGroupInfo: + _ = (self, name) + raise AttributeError("no joint group info") + + monkeypatch.setattr(FakeScene, "getJointGroupInfo", missing_group_info) + monkeypatch.setattr( + FakeScene, + "computeFrameJacobian", + lambda self, q, frame_name, local=True: np.arange(12, dtype=np.float64).reshape(6, 2), + ) + + np.testing.assert_allclose( + world.get_group_jacobian(ctx, "arm/manipulator"), + np.arange(12, dtype=np.float64).reshape(6, 2), + ) + + +def test_legacy_kinematics_wrappers_require_unique_pose_group( + fake_roboplan: None, robot_config: RobotModelConfig +) -> None: + no_pose_config = robot_config.model_copy( + update={ + "planning_groups": [ + PlanningGroupDefinition(name="base", joint_names=("joint1",), base_link="base") + ] + } + ) + no_pose_world, no_pose_id = _make_world(fake_roboplan, no_pose_config) + no_pose_world.finalize() + with pytest.raises(ValueError, match="no pose-targetable"): + no_pose_world.get_ee_pose(no_pose_world.get_live_context(), no_pose_id) + with pytest.raises(ValueError, match="no pose-targetable"): + no_pose_world.get_jacobian(no_pose_world.get_live_context(), no_pose_id) + + ambiguous_config = robot_config.model_copy( + update={ + "planning_groups": [ + PlanningGroupDefinition( + name="a", joint_names=("joint1",), base_link="base", tip_link="a_tip" + ), + PlanningGroupDefinition( + name="b", joint_names=("joint2",), base_link="base", tip_link="b_tip" + ), + ] + } + ) + ambiguous_world, ambiguous_id = _make_world(fake_roboplan, ambiguous_config) + ambiguous_world.finalize() + with pytest.raises(ValueError, match="pose-targetable planning groups"): + ambiguous_world.get_jacobian(ambiguous_world.get_live_context(), ambiguous_id) + + +def test_group_lookup_rejects_unknown_group_id( + fake_roboplan: None, robot_config: RobotModelConfig +) -> None: + world, _ = _make_world(fake_roboplan, robot_config) + world.finalize() + + with pytest.raises(KeyError, match="Unknown planning group ID"): + world.get_group_ee_pose(world.get_live_context(), "other/missing") + + def test_native_planner_converts_path(fake_roboplan: None, robot_config: RobotModelConfig) -> None: world, robot_id = _make_world(fake_roboplan, robot_config) world.finalize() From 5c3893567bf0d1799231b2f9244bf901208053f6 Mon Sep 17 00:00:00 2001 From: cc <55869557+TomCC7@users.noreply.github.com> Date: Thu, 2 Jul 2026 13:45:50 -0700 Subject: [PATCH 12/14] Update dimos/manipulation/planning/groups/utils.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- dimos/manipulation/planning/groups/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dimos/manipulation/planning/groups/utils.py b/dimos/manipulation/planning/groups/utils.py index 66d33ea359..31130084de 100644 --- a/dimos/manipulation/planning/groups/utils.py +++ b/dimos/manipulation/planning/groups/utils.py @@ -166,7 +166,7 @@ def joint_state_to_ordered_positions( if resolved_name not in joint_name_set: raise ValueError(f"Unknown global joint name: {name}") else: - resolved_name = joint_name_mapping.get(name, name) + raise ValueError(f"Unrecognized joint name '{name}': not a known local name, not in joint_name_mapping, and not a global name") if resolved_name in name_to_pos: raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") From 562f86ae98acfea7242d2ddaa79eb94e25e80b65 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Thu, 2 Jul 2026 20:46:27 +0000 Subject: [PATCH 13/14] [autofix.ci] apply automated fixes --- dimos/manipulation/planning/groups/utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/dimos/manipulation/planning/groups/utils.py b/dimos/manipulation/planning/groups/utils.py index 31130084de..283681adcc 100644 --- a/dimos/manipulation/planning/groups/utils.py +++ b/dimos/manipulation/planning/groups/utils.py @@ -166,7 +166,9 @@ def joint_state_to_ordered_positions( 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") + raise ValueError( + f"Unrecognized joint name '{name}': not a known local name, not in joint_name_mapping, and not a global name" + ) if resolved_name in name_to_pos: raise ValueError(f"JointState resolves duplicate joint '{resolved_name}'") From 90baf52f582de827555436ce82b8e5f138e412a3 Mon Sep 17 00:00:00 2001 From: cc Date: Thu, 2 Jul 2026 14:10:26 -0700 Subject: [PATCH 14/14] test: fix joint state ordering assertion --- dimos/manipulation/planning/groups/test_planning_groups.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dimos/manipulation/planning/groups/test_planning_groups.py b/dimos/manipulation/planning/groups/test_planning_groups.py index 0d5f90a2c9..8b399354a6 100644 --- a/dimos/manipulation/planning/groups/test_planning_groups.py +++ b/dimos/manipulation/planning/groups/test_planning_groups.py @@ -650,7 +650,7 @@ def test_joint_state_to_ordered_positions_rejects_invalid_inputs() -> None: joint_names=joint_names, joint_name_mapping=mapping, ) - with pytest.raises(ValueError, match="missing joints"): + 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,