From 9215021fa04cc57708949f5fa01d1e33ae7d11bb Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Mon, 12 Jan 2026 08:36:01 -0800 Subject: [PATCH 01/31] added robot_description folders to data/ folder with lfs tracking --- data/.lfs/piper_description.tar.gz | 3 +++ data/.lfs/xarm_description.tar.gz | 3 +++ 2 files changed, 6 insertions(+) create mode 100644 data/.lfs/piper_description.tar.gz create mode 100644 data/.lfs/xarm_description.tar.gz diff --git a/data/.lfs/piper_description.tar.gz b/data/.lfs/piper_description.tar.gz new file mode 100644 index 0000000000..857b5d6f92 --- /dev/null +++ b/data/.lfs/piper_description.tar.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d1b9ba17e1b85b71edea6d0106dd31e39962dcfa366ca7339223e01a64d87116 +size 16725107 diff --git a/data/.lfs/xarm_description.tar.gz b/data/.lfs/xarm_description.tar.gz new file mode 100644 index 0000000000..41cb07eca2 --- /dev/null +++ b/data/.lfs/xarm_description.tar.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8cca498b7a9ab98d70451d42aa86f6daf9085061a919ac82a567d3d3af5e0991 +size 12703025 From 73e9dcef888f8a7a098472ca3309ff23f7adda49 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Mon, 12 Jan 2026 15:58:11 -0800 Subject: [PATCH 02/31] Base types: RobotModelConfig, Obstacle, Protocol specs for defining manipulation scenes --- dimos/manipulation/planning/spec.py | 715 ++++++++++++++++++++++++++++ 1 file changed, 715 insertions(+) create mode 100644 dimos/manipulation/planning/spec.py diff --git a/dimos/manipulation/planning/spec.py b/dimos/manipulation/planning/spec.py new file mode 100644 index 0000000000..c8bf854b4a --- /dev/null +++ b/dimos/manipulation/planning/spec.py @@ -0,0 +1,715 @@ +# 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. + +""" +Manipulation Planning Specifications + +Contains Protocol definitions and data classes for the manipulation planning stack. + +## Protocols + +- WorldSpec: Core backend owning physics/collision (implemented by DrakeWorld) +- KinematicsSpec: Stateless IK operations (implemented by DrakeKinematics) +- PlannerSpec: Joint-space path planning (implemented by DrakePlanner) +- VizSpec: Visualization backend (future) + +## Data Classes + +- RobotModelConfig: Robot configuration for adding to world +- Obstacle: Obstacle specification +- IKResult: Result of IK solve +- PlanningResult: Result of path planning + +## Usage + +All code should use Protocol types (not concrete classes): + +```python +from dimos.manipulation.planning.spec import WorldSpec, KinematicsSpec + +def plan_motion(world: WorldSpec, kinematics: KinematicsSpec, ...): + # Works with any conforming implementation + pass +``` + +Use factory functions to create concrete instances: + +```python +from dimos.manipulation.planning.factory import create_world, create_kinematics + +world = create_world(backend="drake") # Returns WorldSpec +kinematics = create_kinematics() # Returns KinematicsSpec +``` +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum, auto +from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable + +if TYPE_CHECKING: + from contextlib import AbstractContextManager + + import numpy as np + from numpy.typing import NDArray + + +# ============================================================================= +# Enums +# ============================================================================= + + +class ObstacleType(Enum): + """Type of obstacle geometry.""" + + BOX = auto() + SPHERE = auto() + CYLINDER = auto() + MESH = auto() + + +class IKStatus(Enum): + """Status of IK solution.""" + + SUCCESS = auto() + NO_SOLUTION = auto() + SINGULARITY = auto() + JOINT_LIMITS = auto() + COLLISION = auto() + TIMEOUT = auto() + + +class PlanningStatus(Enum): + """Status of motion planning.""" + + SUCCESS = auto() + NO_SOLUTION = auto() + TIMEOUT = auto() + INVALID_START = auto() + INVALID_GOAL = auto() + COLLISION_AT_START = auto() + COLLISION_AT_GOAL = auto() + + +# ============================================================================= +# Data Classes +# ============================================================================= + + +@dataclass +class RobotModelConfig: + """Configuration for adding a robot to the world. + + Attributes: + name: Human-readable robot name + urdf_path: Path to URDF file (can be .urdf or .xacro) + base_pose: 4x4 homogeneous transform for robot base + joint_names: Ordered list of controlled joint names (in URDF namespace) + end_effector_link: Name of the end-effector link for FK/IK + base_link: Name of the base link (default: "base_link") + package_paths: Dict mapping package names to filesystem paths + joint_limits_lower: Lower joint limits (radians) + joint_limits_upper: Upper joint limits (radians) + velocity_limits: Joint velocity limits (rad/s) + auto_convert_meshes: Auto-convert DAE/STL meshes to OBJ for Drake + xacro_args: Arguments to pass to xacro processor (for .xacro files) + collision_exclusion_pairs: List of (link1, link2) pairs to exclude from collision. + Useful for parallel linkage mechanisms like grippers where non-adjacent + links may legitimately overlap (e.g., mimic joints). + max_velocity: Maximum joint velocity for trajectory generation (rad/s) + max_acceleration: Maximum joint acceleration for trajectory generation (rad/s^2) + joint_name_mapping: Maps orchestrator joint names to URDF joint names. + Example: {"left_joint1": "joint1"} means orchestrator's "left_joint1" + corresponds to URDF's "joint1". If empty, names are assumed to match. + orchestrator_task_name: Task name for executing trajectories via orchestrator RPC. + If set, trajectories can be executed via execute_trajectory() RPC. + """ + + name: str + urdf_path: str + base_pose: NDArray[np.float64] + joint_names: list[str] + end_effector_link: str + base_link: str = "base_link" + package_paths: dict[str, str] = field(default_factory=dict) + joint_limits_lower: NDArray[np.float64] | None = None + joint_limits_upper: NDArray[np.float64] | None = None + velocity_limits: NDArray[np.float64] | None = None + auto_convert_meshes: bool = False + xacro_args: dict[str, str] = field(default_factory=dict) + collision_exclusion_pairs: list[tuple[str, str]] = field(default_factory=list) + # Motion constraints for trajectory generation + max_velocity: float = 1.0 + max_acceleration: float = 2.0 + # Orchestrator integration + joint_name_mapping: dict[str, str] = field(default_factory=dict) + orchestrator_task_name: str | None = None + + def get_urdf_joint_name(self, orchestrator_name: str) -> str: + """Translate orchestrator joint name to URDF joint name. + + Args: + orchestrator_name: Joint name from orchestrator (e.g., "left_joint1") + + Returns: + Corresponding URDF joint name (e.g., "joint1") + """ + return self.joint_name_mapping.get(orchestrator_name, orchestrator_name) + + def get_orchestrator_joint_name(self, urdf_name: str) -> str: + """Translate URDF joint name to orchestrator joint name. + + Args: + urdf_name: Joint name from URDF (e.g., "joint1") + + Returns: + Corresponding orchestrator joint name (e.g., "left_joint1") + """ + # Reverse lookup + for orch_name, u_name in self.joint_name_mapping.items(): + if u_name == urdf_name: + return orch_name + return urdf_name + + def get_orchestrator_joint_names(self) -> list[str]: + """Get joint names in orchestrator namespace. + + Returns: + List of joint names as they appear in orchestrator messages. + If no mapping, returns the URDF joint names. + """ + if not self.joint_name_mapping: + return self.joint_names + return [self.get_orchestrator_joint_name(j) for j in self.joint_names] + + +@dataclass +class Obstacle: + """Obstacle specification for collision avoidance. + + Attributes: + name: Unique name for the obstacle + obstacle_type: Type of geometry (BOX, SPHERE, CYLINDER, MESH) + pose: 4x4 homogeneous transform + dimensions: Type-specific dimensions: + - BOX: (width, height, depth) + - SPHERE: (radius,) + - CYLINDER: (radius, height) + - MESH: Not used + color: RGBA color tuple (0-1 range) + mesh_path: Path to mesh file (for MESH type) + """ + + name: str + obstacle_type: ObstacleType + pose: NDArray[np.float64] + dimensions: tuple[float, ...] = () + color: tuple[float, float, float, float] = (0.8, 0.2, 0.2, 0.8) + mesh_path: str | None = None + + +@dataclass +class IKResult: + """Result of an IK solve. + + Attributes: + status: Solution status + joint_positions: Solution joint positions (None if failed) + position_error: Cartesian position error (meters) + orientation_error: Orientation error (radians) + iterations: Number of iterations taken + message: Human-readable status message + """ + + status: IKStatus + joint_positions: NDArray[np.float64] | None = None + position_error: float = 0.0 + orientation_error: float = 0.0 + iterations: int = 0 + message: str = "" + + def is_success(self) -> bool: + """Check if IK was successful.""" + return self.status == IKStatus.SUCCESS + + +@dataclass +class PlanningResult: + """Result of motion planning. + + Attributes: + status: Planning status + path: List of joint configurations (empty if failed) + planning_time: Time taken to plan (seconds) + path_length: Total path length in joint space (radians) + iterations: Number of iterations/nodes expanded + message: Human-readable status message + """ + + status: PlanningStatus + path: list[NDArray[np.float64]] = field(default_factory=list) + planning_time: float = 0.0 + path_length: float = 0.0 + iterations: int = 0 + message: str = "" + + def is_success(self) -> bool: + """Check if planning was successful.""" + return self.status == PlanningStatus.SUCCESS + + +@dataclass +class CollisionObjectMessage: + """Message for adding/updating/removing obstacles. + + Used by monitors to handle obstacle updates from external sources. + + Attributes: + id: Unique identifier for the object + operation: "add", "update", or "remove" + primitive_type: "box", "sphere", or "cylinder" (for add/update) + pose: 4x4 transform (for add/update) + dimensions: Type-specific dimensions (for add/update) + color: RGBA color tuple + """ + + id: str + operation: str # "add", "update", "remove" + primitive_type: str | None = None + pose: NDArray[np.float64] | None = None + dimensions: tuple[float, ...] | None = None + color: tuple[float, float, float, float] = (0.8, 0.2, 0.2, 0.8) + + +@dataclass +class Detection3D: + """3D detection result from perception pipeline. + + Used by monitors to handle perception updates. + + Attributes: + id: Unique detection ID + label: Object class label + pose: 4x4 transform + dimensions: (width, height, depth) + confidence: Detection confidence (0-1) + """ + + id: str + label: str + pose: NDArray[np.float64] + dimensions: tuple[float, float, float] + confidence: float = 1.0 + + +# ============================================================================= +# Protocols +# ============================================================================= + + +@runtime_checkable +class WorldSpec(Protocol): + """Protocol for the world/scene backend. + + The world owns the physics/collision backend and provides: + - Robot/obstacle management + - Collision checking + - Forward kinematics + - Context management for thread safety + + ## Context Management + + The world maintains: + - Live context: Mirrors current robot state (synced from driver) + - Scratch contexts: Thread-safe clones for planning/IK operations + + All state-dependent operations (collision checking, FK, Jacobian) take + a context parameter to ensure thread safety. + + ## Implementations + + - DrakeWorld: Uses Drake's MultibodyPlant and SceneGraph + """ + + # ========================================================================= + # Robot Management + # ========================================================================= + + def add_robot(self, config: RobotModelConfig) -> str: + """Add a robot to the world. + + Args: + config: Robot configuration + + Returns: + Unique robot ID + """ + ... + + def get_robot_ids(self) -> list[str]: + """Get all robot IDs.""" + ... + + def get_robot_config(self, robot_id: str) -> RobotModelConfig: + """Get robot configuration.""" + ... + + def get_joint_limits(self, robot_id: str) -> tuple[NDArray[np.float64], NDArray[np.float64]]: + """Get joint limits (lower, upper) for a robot.""" + ... + + # ========================================================================= + # Obstacle Management + # ========================================================================= + + def add_obstacle(self, obstacle: Obstacle) -> str: + """Add an obstacle to the world. + + Args: + obstacle: Obstacle specification + + Returns: + Unique obstacle ID + """ + ... + + def remove_obstacle(self, obstacle_id: str) -> bool: + """Remove an obstacle. + + Returns: + True if obstacle was removed + """ + ... + + def update_obstacle_pose(self, obstacle_id: str, pose: NDArray[np.float64]) -> bool: + """Update obstacle pose. + + Returns: + True if obstacle was updated + """ + ... + + def clear_obstacles(self) -> None: + """Remove all obstacles.""" + ... + + # ========================================================================= + # Lifecycle + # ========================================================================= + + def finalize(self) -> None: + """Finalize the world for simulation/collision checking. + + Must be called after adding all robots and before any operations. + """ + ... + + @property + def is_finalized(self) -> bool: + """Check if world is finalized.""" + ... + + # ========================================================================= + # Context Management + # ========================================================================= + + def get_live_context(self) -> Any: + """Get the live context (mirrors real robot state). + + The live context is synced with real robot state via + sync_from_joint_state(). Use with caution - modifications + affect all users. + + For planning, prefer scratch_context() instead. + """ + ... + + def scratch_context(self) -> AbstractContextManager[Any]: + """Get a scratch context for planning (thread-safe clone). + + Example: + with world.scratch_context() as ctx: + world.set_positions(ctx, robot_id, q_test) + if world.is_collision_free(ctx, robot_id): + # Valid configuration + pass + + Yields: + Context that can be used for state operations + """ + ... + + def sync_from_joint_state(self, robot_id: str, positions: NDArray[np.float64]) -> None: + """Sync live context from joint state (called by driver/monitor). + + Args: + robot_id: Robot to update + positions: Joint positions from driver + """ + ... + + # ========================================================================= + # State Operations (require context) + # ========================================================================= + + def set_positions(self, ctx: Any, robot_id: str, positions: NDArray[np.float64]) -> None: + """Set robot joint positions in a context. + + Args: + ctx: Context from scratch_context() or get_live_context() + robot_id: Robot to update + positions: Joint positions (radians) + """ + ... + + def get_positions(self, ctx: Any, robot_id: str) -> NDArray[np.float64]: + """Get robot joint positions from a context. + + Args: + ctx: Context + robot_id: Robot to query + + Returns: + Joint positions (radians) + """ + ... + + # ========================================================================= + # Collision Checking (require context) + # ========================================================================= + + def is_collision_free(self, ctx: Any, robot_id: str) -> bool: + """Check if robot configuration is collision-free. + + Args: + ctx: Context with robot state set + robot_id: Robot to check + + Returns: + True if no collisions + """ + ... + + def get_min_distance(self, ctx: Any, robot_id: str) -> float: + """Get minimum distance to obstacles. + + Args: + ctx: Context with robot state set + robot_id: Robot to check + + Returns: + Minimum distance (meters), negative if in collision + """ + ... + + # ========================================================================= + # Forward Kinematics (require context) + # ========================================================================= + + def get_ee_pose(self, ctx: Any, robot_id: str) -> NDArray[np.float64]: + """Get end-effector pose. + + Args: + ctx: Context with robot state set + robot_id: Robot to query + + Returns: + 4x4 homogeneous transform + """ + ... + + def get_jacobian(self, ctx: Any, robot_id: str) -> NDArray[np.float64]: + """Get end-effector Jacobian. + + Args: + ctx: Context with robot state set + robot_id: Robot to query + + Returns: + 6 x n_joints Jacobian matrix [linear; angular] + """ + ... + + +@runtime_checkable +class KinematicsSpec(Protocol): + """Protocol for inverse kinematics solver. + + Kinematics solvers are stateless (except for configuration) and + use WorldSpec for all FK/collision operations. + + ## Methods + + - solve(): Full optimization-based IK with collision checking + - solve_iterative(): Iterative Jacobian-based IK + - solve_differential(): Single Jacobian step for velocity control + + ## Implementations + + - DrakeKinematics: Uses Drake's InverseKinematics + SNOPT/IPOPT + """ + + def solve( + self, + world: WorldSpec, + robot_id: str, + target_pose: NDArray[np.float64], + seed: NDArray[np.float64] | None = None, + position_tolerance: float = 0.001, + orientation_tolerance: float = 0.01, + check_collision: bool = True, + max_attempts: int = 10, + ) -> IKResult: + """Solve full IK with optional collision checking. + + Args: + world: World for FK/collision + robot_id: Robot to solve for + target_pose: 4x4 target end-effector transform + seed: Initial guess (uses current state if None) + position_tolerance: Position tolerance (meters) + orientation_tolerance: Orientation tolerance (radians) + check_collision: Check solution is collision-free + max_attempts: Random restarts for robustness + + Returns: + IKResult with status and solution + """ + ... + + def solve_iterative( + self, + world: WorldSpec, + robot_id: str, + target_pose: NDArray[np.float64], + seed: NDArray[np.float64], + max_iterations: int = 100, + position_tolerance: float = 0.001, + orientation_tolerance: float = 0.01, + ) -> IKResult: + """Solve IK iteratively using Jacobian method. + + Slower but more predictable behavior near singularities. + + Args: + world: World for FK/Jacobian + robot_id: Robot to solve for + target_pose: 4x4 target transform + seed: Initial joint configuration + max_iterations: Maximum iterations + position_tolerance: Convergence tolerance (meters) + orientation_tolerance: Convergence tolerance (radians) + + Returns: + IKResult with status and solution + """ + ... + + def solve_differential( + self, + world: WorldSpec, + robot_id: str, + current_joints: NDArray[np.float64], + twist: NDArray[np.float64], + dt: float, + ) -> NDArray[np.float64] | None: + """Single Jacobian step for velocity control. + + Args: + world: World for Jacobian computation + robot_id: Robot to control + current_joints: Current joint positions + twist: Desired end-effector twist [vx, vy, vz, wx, wy, wz] + dt: Time step (seconds) + + Returns: + Joint velocities (rad/s) or None if near singularity + """ + ... + + +@runtime_checkable +class PlannerSpec(Protocol): + """Protocol for motion planner. + + Planners find collision-free paths from start to goal configurations. + They use WorldSpec for collision checking and are stateless. + + ## Implementations + + - DrakePlanner: RRT-Connect planner + - DrakeRRTStarPlanner: RRT* planner (asymptotically optimal) + """ + + def plan_joint_path( + self, + world: WorldSpec, + robot_id: str, + q_start: NDArray[np.float64], + q_goal: NDArray[np.float64], + timeout: float = 10.0, + ) -> PlanningResult: + """Plan a collision-free joint-space path. + + Args: + world: World for collision checking + robot_id: Robot to plan for + q_start: Start configuration + q_goal: Goal configuration + timeout: Planning timeout (seconds) + + Returns: + PlanningResult with status and path + """ + ... + + def get_name(self) -> str: + """Get planner name.""" + ... + + +@runtime_checkable +class VizSpec(Protocol): + """Protocol for visualization backend. + + Provides methods to update robot/obstacle visualization. + + Note: For Drake, visualization is typically integrated into DrakeWorld + via enable_viz=True. This protocol is for advanced use cases. + + ## Implementations + + - DrakeWorld (integrated): Use create_world(enable_viz=True) + """ + + def set_robot_state(self, robot_id: str, positions: NDArray[np.float64]) -> None: + """Update robot visualization state.""" + ... + + def add_obstacle(self, obstacle: Obstacle) -> str: + """Add obstacle to visualization.""" + ... + + def remove_obstacle(self, obstacle_id: str) -> None: + """Remove obstacle from visualization.""" + ... + + def get_url(self) -> str | None: + """Get visualization URL (e.g., Meshcat URL).""" + ... + + def publish(self) -> None: + """Force publish current state to visualization.""" + ... From 265f5b4b28eaa9841990176873c6aa212656a3db Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Mon, 12 Jan 2026 15:58:55 -0800 Subject: [PATCH 03/31] Mesh conversion utils for converting any urdf to drake compatible mesh files --- dimos/manipulation/planning/mesh_utils.py | 279 ++++++++++++++++++++++ 1 file changed, 279 insertions(+) create mode 100644 dimos/manipulation/planning/mesh_utils.py diff --git a/dimos/manipulation/planning/mesh_utils.py b/dimos/manipulation/planning/mesh_utils.py new file mode 100644 index 0000000000..d7645f359e --- /dev/null +++ b/dimos/manipulation/planning/mesh_utils.py @@ -0,0 +1,279 @@ +# 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. + +""" +Mesh Utilities for Drake + +Provides utilities for preparing URDF files for use with Drake: +- Xacro processing +- Mesh format conversion (DAE/STL to OBJ) +- Package path resolution + +Example: + urdf_path = prepare_urdf_for_drake( + urdf_path="/path/to/robot.xacro", + package_paths={"robot_description": "/path/to/robot_description"}, + xacro_args={"use_sim": "true"}, + convert_meshes=True, + ) +""" + +from __future__ import annotations + +import hashlib +import os +from pathlib import Path +import re +import shutil +import tempfile +from typing import TYPE_CHECKING + +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +# Cache directory for processed URDFs +_CACHE_DIR = Path(tempfile.gettempdir()) / "dimos_urdf_cache" + + +def prepare_urdf_for_drake( + urdf_path: Path | str, + package_paths: dict[str, str] | None = None, + xacro_args: dict[str, str] | None = None, + convert_meshes: bool = False, +) -> str: + """Prepare a URDF/xacro file for use with Drake. + + This function: + 1. Processes xacro files if needed + 2. Resolves package:// URIs in mesh paths + 3. Optionally converts DAE/STL meshes to OBJ format + + Args: + urdf_path: Path to URDF or xacro file + package_paths: Dict mapping package names to filesystem paths + xacro_args: Arguments to pass to xacro processor + convert_meshes: Convert DAE/STL meshes to OBJ for Drake compatibility + + Returns: + Path to the prepared URDF file (may be cached) + """ + urdf_path = Path(urdf_path) + package_paths = package_paths or {} + xacro_args = xacro_args or {} + + # Generate cache key + cache_key = _generate_cache_key(urdf_path, package_paths, xacro_args, convert_meshes) + cache_path = _CACHE_DIR / cache_key / urdf_path.stem + cache_path.mkdir(parents=True, exist_ok=True) + cached_urdf = cache_path / f"{urdf_path.stem}.urdf" + + # Check cache + if cached_urdf.exists(): + logger.debug(f"Using cached URDF: {cached_urdf}") + return str(cached_urdf) + + # Process xacro if needed + if urdf_path.suffix in (".xacro", ".urdf.xacro"): + urdf_content = _process_xacro(urdf_path, package_paths, xacro_args) + else: + urdf_content = urdf_path.read_text() + + # Strip transmission blocks (Drake doesn't need them, and they can cause issues) + urdf_content = _strip_transmission_blocks(urdf_content) + + # Resolve package:// URIs + urdf_content = _resolve_package_uris(urdf_content, package_paths, cache_path) + + # Convert meshes if requested + if convert_meshes: + urdf_content = _convert_meshes(urdf_content, cache_path) + + # Write processed URDF + cached_urdf.write_text(urdf_content) + logger.info(f"Prepared URDF cached at: {cached_urdf}") + + return str(cached_urdf) + + +def _generate_cache_key( + urdf_path: Path, + package_paths: dict[str, str], + xacro_args: dict[str, str], + convert_meshes: bool, +) -> str: + """Generate a cache key for the URDF configuration. + + Includes a version number to invalidate cache when processing logic changes. + """ + # Include file modification time + mtime = urdf_path.stat().st_mtime if urdf_path.exists() else 0 + + # Version number to invalidate cache when processing logic changes + # Increment this when adding new processing steps (e.g., stripping transmission blocks) + processing_version = "v2" + + key_data = f"{processing_version}:{urdf_path}:{mtime}:{sorted(package_paths.items())}:{sorted(xacro_args.items())}:{convert_meshes}" + return hashlib.md5(key_data.encode()).hexdigest()[:16] + + +def _process_xacro( + xacro_path: Path, + package_paths: dict[str, str], + xacro_args: dict[str, str], +) -> str: + """Process xacro file to URDF.""" + try: + import xacro # type: ignore[import-untyped] + except ImportError: + raise ImportError( + "xacro is required for processing .xacro files. Install with: pip install xacro" + ) + + # Create a custom substitution_args_context that resolves $(find pkg) to our paths + # This avoids requiring ROS package discovery + from xacro import substitution_args + + # Store original function + original_find = substitution_args._find + + def custom_find(resolved: str, a: str, args: list[str], context: dict[str, str]) -> str: + """Custom $(find pkg) handler that uses our package_paths.""" + pkg_name = args[0] if args else "" + if pkg_name in package_paths: + pkg_path = str(Path(package_paths[pkg_name]).resolve()) + return resolved.replace(f"$({a})", pkg_path) + # Fall back to original behavior + return str(original_find(resolved, a, args, context)) + + # Monkey-patch the find function temporarily + substitution_args._find = custom_find + + try: + # Process xacro with our mappings + doc = xacro.process_file( + str(xacro_path), + mappings=xacro_args, + ) + return str(doc.toprettyxml(indent=" ")) + finally: + # Restore original function + substitution_args._find = original_find + + +def _strip_transmission_blocks(urdf_content: str) -> str: + """Remove transmission blocks from URDF content. + + Drake doesn't need transmission blocks (they're for Gazebo/ROS control), + and they can cause parsing errors if they contain malformed actuator names. + + Args: + urdf_content: URDF XML content as string + + Returns: + URDF content with transmission blocks removed + """ + # Pattern to match ... blocks (including nested content) + # Uses non-greedy matching and handles nested tags + pattern = r"]*>.*?" + + # Remove transmission blocks (with flags for multiline and dotall) + result = re.sub(pattern, "", urdf_content, flags=re.DOTALL | re.MULTILINE) + + # Also remove any standalone blocks that might reference transmissions + # (some URDFs have gazebo plugins that reference transmissions) + gazebo_pattern = r".*?]*gazebo_ros_control[^>]*>.*?.*?" + result = re.sub(gazebo_pattern, "", result, flags=re.DOTALL | re.MULTILINE) + + return result + + +def _resolve_package_uris( + urdf_content: str, + package_paths: dict[str, str], + output_dir: Path, +) -> str: + """Resolve package:// URIs to filesystem paths.""" + # Pattern for package:// URIs + pattern = r'package://([^/]+)/(.+?)(["\s<>])' + + def replace_uri(match: re.Match[str]) -> str: + pkg_name = match.group(1) + rel_path = match.group(2) + suffix = match.group(3) + + if pkg_name in package_paths: + # Ensure absolute path for proper resolution + pkg_path = Path(package_paths[pkg_name]).resolve() + full_path = pkg_path / rel_path + if full_path.exists(): + return f"{full_path}{suffix}" + else: + logger.warning(f"File not found: {full_path}") + + # Return original if not found + return match.group(0) + + return re.sub(pattern, replace_uri, urdf_content) + + +def _convert_meshes(urdf_content: str, output_dir: Path) -> str: + """Convert DAE/STL meshes to OBJ format for Drake compatibility.""" + try: + import trimesh + except ImportError: + logger.warning("trimesh not installed, skipping mesh conversion") + return urdf_content + + mesh_dir = output_dir / "meshes" + mesh_dir.mkdir(exist_ok=True) + + # Find mesh file references + pattern = r'filename="([^"]+\.(dae|stl|DAE|STL))"' + + converted: dict[str, str] = {} + + def convert_mesh(match: re.Match[str]) -> str: + original_path = match.group(1) + + if original_path in converted: + return f'filename="{converted[original_path]}"' + + try: + # Load mesh + mesh = trimesh.load(original_path, force="mesh") + + # Generate output path + mesh_name = Path(original_path).stem + obj_path = mesh_dir / f"{mesh_name}.obj" + + # Export as OBJ (trimesh.export returns None, ignore) + mesh.export(str(obj_path), file_type="obj") # type: ignore[no-untyped-call] + logger.debug(f"Converted mesh: {original_path} -> {obj_path}") + + converted[original_path] = str(obj_path) + return f'filename="{obj_path}"' + + except Exception as e: + logger.warning(f"Failed to convert mesh {original_path}: {e}") + return match.group(0) + + return re.sub(pattern, convert_mesh, urdf_content) + + +def clear_cache() -> None: + """Clear the URDF cache directory.""" + if _CACHE_DIR.exists(): + shutil.rmtree(_CACHE_DIR) + logger.info(f"Cleared URDF cache: {_CACHE_DIR}") From d0dfbfa7aa1699fa9441e7539c1f5e3d234a9ea4 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Mon, 12 Jan 2026 16:01:38 -0800 Subject: [PATCH 04/31] DrakeWorld implementation that maintains all objects, robots in the scene and owns the scene graph --- dimos/manipulation/planning/world/__init__.py | 27 + .../planning/world/drake_world.py | 948 ++++++++++++++++++ 2 files changed, 975 insertions(+) create mode 100644 dimos/manipulation/planning/world/__init__.py create mode 100644 dimos/manipulation/planning/world/drake_world.py diff --git a/dimos/manipulation/planning/world/__init__.py b/dimos/manipulation/planning/world/__init__.py new file mode 100644 index 0000000000..8ddef7fdff --- /dev/null +++ b/dimos/manipulation/planning/world/__init__.py @@ -0,0 +1,27 @@ +# Copyright 2025 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. + +""" +World Module + +Contains world implementations that own the physics/collision backend. + +## Implementations + +- DrakeWorld: Uses Drake MultibodyPlant + SceneGraph +""" + +from dimos.manipulation.planning.world.drake_world import DrakeWorld + +__all__ = ["DrakeWorld"] diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py new file mode 100644 index 0000000000..5d5705a8b0 --- /dev/null +++ b/dimos/manipulation/planning/world/drake_world.py @@ -0,0 +1,948 @@ +# 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. + +""" +Drake World Implementation + +Implements WorldSpec using Drake's MultibodyPlant and SceneGraph. + +## Architecture + +Uses Drake's DiagramBuilder pattern: +``` +DiagramBuilder +├── MultibodyPlant (kinematics, dynamics) +├── SceneGraph (collision geometry) +└── MeshcatVisualizer (optional) +``` + +## Context Management + +- Live context: Mirrors current robot state (synced from driver) +- Scratch contexts: Thread-safe clones for planning/IK operations + +Example: + world = DrakeWorld(enable_viz=True) + robot_id = world.add_robot(config) + world.add_obstacle(table) + world.finalize() + + # Sync live state from driver + world.sync_from_joint_state(robot_id, joint_positions) + + # Planning uses scratch contexts + with world.scratch_context() as ctx: + world.set_positions(ctx, robot_id, q_test) + if world.is_collision_free(ctx, robot_id): + ee_pose = world.get_ee_pose(ctx, robot_id) +""" + +from __future__ import annotations + +from contextlib import contextmanager +import copy +from dataclasses import dataclass +from pathlib import Path +from threading import RLock +from typing import TYPE_CHECKING, Any + +import numpy as np + +from dimos.manipulation.planning.mesh_utils import prepare_urdf_for_drake +from dimos.manipulation.planning.spec import ( + Obstacle, + ObstacleType, + RobotModelConfig, +) +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from collections.abc import Generator + + from numpy.typing import NDArray + +try: + from pydrake.geometry import ( + AddContactMaterial, + Box, + CollisionFilterDeclaration, + Cylinder, + GeometryInstance, + GeometrySet, + MakePhongIllustrationProperties, + Meshcat, + MeshcatVisualizer, + MeshcatVisualizerParams, + ProximityProperties, + QueryObject, + Rgba, + Role, + SceneGraph, + Sphere, + ) + from pydrake.math import RigidTransform + from pydrake.multibody.parsing import Parser + from pydrake.multibody.plant import ( + AddMultibodyPlantSceneGraph, + CoulombFriction, + MultibodyPlant, + ) + from pydrake.multibody.tree import JacobianWrtVariable + from pydrake.systems.framework import Context, DiagramBuilder + + DRAKE_AVAILABLE = True +except ImportError: + DRAKE_AVAILABLE = False + +logger = setup_logger() + + +@dataclass +class _RobotData: + """Internal data for tracking a robot in the world.""" + + robot_id: str + config: RobotModelConfig + model_instance: Any # ModelInstanceIndex + joint_indices: list[int] # Indices into plant's position vector + ee_frame: Any # BodyFrame for end-effector + base_frame: Any # BodyFrame for base + + +@dataclass +class _ObstacleData: + """Internal data for tracking an obstacle in the world.""" + + obstacle_id: str + obstacle: Obstacle + geometry_id: Any # GeometryId + source_id: Any # SourceId + + +class DrakeWorld: + """Drake implementation of WorldSpec. + + Owns a single Drake Diagram containing: + - MultibodyPlant (kinematics, dynamics) + - SceneGraph (collision geometry) + - Optional MeshcatVisualizer + + ## Context Management + + - _live_context: Mirrors current robot state (synced from driver) + - scratch_context(): Returns cloned context for planning/IK + + ## Thread Safety + + - Live context writes are protected by RLock + - scratch_context() returns independent clones + - All public methods that modify state are thread-safe + """ + + def __init__( + self, + time_step: float = 0.0, + enable_viz: bool = False, + ): + """Create a Drake world. + + Args: + time_step: Simulation time step (0 for kinematics-only) + enable_viz: If True, enable Meshcat visualization + """ + if not DRAKE_AVAILABLE: + raise ImportError("Drake is not installed. Install with: pip install drake") + + self._time_step = time_step + self._enable_viz = enable_viz + self._lock = RLock() + + # Build Drake diagram + self._builder = DiagramBuilder() + self._plant: MultibodyPlant + self._scene_graph: SceneGraph + self._plant, self._scene_graph = AddMultibodyPlantSceneGraph( + self._builder, time_step=time_step + ) + self._parser = Parser(self._plant) + + # Visualization + self._meshcat: Meshcat | None = None + self._meshcat_visualizer: MeshcatVisualizer | None = None + if enable_viz: + self._meshcat = Meshcat() + + # Create model instance for obstacles + self._obstacles_model_instance = self._plant.AddModelInstance("obstacles") + + # Tracking data + self._robots: dict[str, _RobotData] = {} + self._obstacles: dict[str, _ObstacleData] = {} + self._robot_counter = 0 + self._obstacle_counter = 0 + + # Built diagram and contexts (created after finalize) + self._diagram: Any = None + self._live_context: Context | None = None + self._plant_context: Context | None = None + self._scene_graph_context: Context | None = None + self._finalized = False + + # Obstacle source for dynamic obstacles + self._obstacle_source_id: Any = None + + # ============= Robot Management ============= + + def add_robot(self, config: RobotModelConfig) -> str: + """Add a robot to the world. + + Args: + config: Robot configuration including URDF path and joint names + + Returns: + robot_id: Unique identifier for the robot + + Raises: + RuntimeError: If world is already finalized + """ + if self._finalized: + raise RuntimeError("Cannot add robot after world is finalized") + + with self._lock: + # Generate unique robot ID + self._robot_counter += 1 + robot_id = f"robot_{self._robot_counter}" + + # Prepare URDF: process xacro and convert STL meshes if needed + original_path = Path(config.urdf_path).resolve() + if not original_path.exists(): + raise FileNotFoundError(f"URDF/xacro not found: {original_path}") + + urdf_path = prepare_urdf_for_drake( + urdf_path=original_path, + package_paths=config.package_paths, + xacro_args=config.xacro_args, + # Always convert meshes when requested - Drake needs OBJ for collision + convert_meshes=config.auto_convert_meshes, + ) + + urdf_path_obj = Path(urdf_path) + logger.info(f"Using prepared URDF: {urdf_path_obj}") + + # Register package paths for mesh resolution + if config.package_paths: + for pkg_name, pkg_path in config.package_paths.items(): + self._parser.package_map().Add(pkg_name, Path(pkg_path)) + else: + # Fallback: register URDF directory as package + package_dir = urdf_path_obj.parent + self._parser.package_map().Add( + package_name=f"{config.name}_description", + package_path=package_dir, + ) + + # Parse the URDF + model_instances = self._parser.AddModels(urdf_path_obj) + if not model_instances: + raise ValueError(f"Failed to parse URDF: {urdf_path}") + + model_instance = model_instances[0] + + # Get base body + base_body = self._plant.GetBodyByName(config.base_link, model_instance) + + # Check if URDF already welded base to world + base_already_welded = False + try: + world_joint = self._plant.GetJointByName("world_joint", model_instance) + if ( + world_joint.parent_body().name() == "world" + and world_joint.child_body().name() == config.base_link + ): + base_already_welded = True + logger.info("URDF has 'world_joint', skipping weld") + except RuntimeError: + pass + + # Weld base to world if needed + if not base_already_welded: + world_frame = self._plant.world_frame() + base_transform = RigidTransform(config.base_pose) + self._plant.WeldFrames( + world_frame, + base_body.body_frame(), + base_transform, + ) + + # Verify joints exist + for joint_name in config.joint_names: + try: + self._plant.GetJointByName(joint_name, model_instance) + except RuntimeError: + raise ValueError(f"Joint '{joint_name}' not found in URDF") + + # Get end-effector frame + try: + ee_body = self._plant.GetBodyByName(config.end_effector_link, model_instance) + ee_frame = ee_body.body_frame() + except RuntimeError: + raise ValueError( + f"End-effector link '{config.end_effector_link}' not found in URDF" + ) + + # Store robot data (joint_indices computed after finalize) + self._robots[robot_id] = _RobotData( + robot_id=robot_id, + config=config, + model_instance=model_instance, + joint_indices=[], + ee_frame=ee_frame, + base_frame=base_body.body_frame(), + ) + + logger.info(f"Added robot '{robot_id}' ({config.name})") + return robot_id + + def get_robot_ids(self) -> list[str]: + """Get all robot IDs in the world.""" + return list(self._robots.keys()) + + def get_robot_config(self, robot_id: str) -> RobotModelConfig: + """Get robot configuration by ID.""" + if robot_id not in self._robots: + raise KeyError(f"Robot '{robot_id}' not found") + return self._robots[robot_id].config + + def get_joint_limits(self, robot_id: str) -> tuple[NDArray[np.float64], NDArray[np.float64]]: + """Get joint limits (lower, upper) in radians.""" + if robot_id not in self._robots: + raise KeyError(f"Robot '{robot_id}' not found") + + config = self._robots[robot_id].config + + if config.joint_limits_lower is not None and config.joint_limits_upper is not None: + return ( + np.array(config.joint_limits_lower), + np.array(config.joint_limits_upper), + ) + + # Default to ±π + n_joints = len(config.joint_names) + return ( + np.full(n_joints, -np.pi), + np.full(n_joints, np.pi), + ) + + # ============= Obstacle Management ============= + + def add_obstacle(self, obstacle: Obstacle) -> str: + """Add an obstacle to the world.""" + with self._lock: + # Use obstacle's name as ID (allows external ID management) + obstacle_id = obstacle.name + + # Check for duplicate in our tracking + if obstacle_id in self._obstacles: + logger.debug(f"Obstacle '{obstacle_id}' already exists, skipping") + return obstacle_id + + try: + if not self._finalized: + geometry_id = self._add_obstacle_to_plant(obstacle, obstacle_id) + self._obstacles[obstacle_id] = _ObstacleData( + obstacle_id=obstacle_id, + obstacle=obstacle, + geometry_id=geometry_id, + source_id=self._plant.get_source_id(), + ) + else: + geometry_id = self._add_obstacle_to_scene_graph(obstacle, obstacle_id) + self._obstacles[obstacle_id] = _ObstacleData( + obstacle_id=obstacle_id, + obstacle=obstacle, + geometry_id=geometry_id, + source_id=self._obstacle_source_id, + ) + + logger.debug(f"Added obstacle '{obstacle_id}': {obstacle.obstacle_type.value}") + except RuntimeError as e: + # Handle case where geometry name already exists in SceneGraph + # (can happen with concurrent access) + if "already been used" in str(e): + logger.debug(f"Obstacle '{obstacle_id}' already in SceneGraph, skipping") + else: + raise + + return obstacle_id + + def _add_obstacle_to_plant(self, obstacle: Obstacle, obstacle_id: str) -> Any: + """Add obstacle to plant (before finalization).""" + shape = self._create_shape(obstacle) + + body = self._plant.AddRigidBody( + obstacle_id, + self._obstacles_model_instance, + ) + + transform = RigidTransform(obstacle.pose) + geometry_id = self._plant.RegisterCollisionGeometry( + body, + RigidTransform(), + shape, + obstacle_id + "_collision", + ProximityProperties(), + ) + + diffuse_color = np.array(obstacle.color) + self._plant.RegisterVisualGeometry( + body, + RigidTransform(), + shape, + obstacle_id + "_visual", + diffuse_color, + ) + + self._plant.WeldFrames( + self._plant.world_frame(), + body.body_frame(), + transform, + ) + + return geometry_id + + def _add_obstacle_to_scene_graph(self, obstacle: Obstacle, obstacle_id: str) -> Any: + """Add obstacle to scene graph (after finalization).""" + if self._obstacle_source_id is None: + raise RuntimeError("Obstacle source not initialized") + + shape = self._create_shape(obstacle) + transform = RigidTransform(obstacle.pose) + # MakePhongIllustrationProperties expects numpy array, not Rgba + rgba_array = np.array(obstacle.color, dtype=np.float64) + + # Create proximity properties with contact material for collision detection + # Without these properties, the geometry is invisible to collision queries + proximity_props = ProximityProperties() + AddContactMaterial( + dissipation=0.0, + point_stiffness=1e6, + friction=CoulombFriction(static_friction=1.0, dynamic_friction=1.0), + properties=proximity_props, + ) + + geometry_instance = GeometryInstance( + X_PG=transform, + shape=shape, + name=obstacle_id, + ) + geometry_instance.set_illustration_properties(MakePhongIllustrationProperties(rgba_array)) + geometry_instance.set_proximity_properties(proximity_props) + + frame_id = self._scene_graph.world_frame_id() + geometry_id = self._scene_graph.RegisterGeometry( + self._obstacle_source_id, + frame_id, + geometry_instance, + ) + + # Also add to Meshcat directly (MeshcatVisualizer doesn't show dynamic geometries) + if self._meshcat is not None: + self._add_obstacle_to_meshcat(obstacle, obstacle_id) + + return geometry_id + + def _add_obstacle_to_meshcat(self, obstacle: Obstacle, obstacle_id: str) -> None: + """Add obstacle visualization directly to Meshcat.""" + if self._meshcat is None: + return + + # Use Drake's geometry types for Meshcat + path = f"obstacles/{obstacle_id}" + transform = RigidTransform(obstacle.pose) + rgba = Rgba(*obstacle.color) + + # Create Drake shape and add to Meshcat + shape: Box | Sphere | Cylinder + if obstacle.obstacle_type == ObstacleType.BOX: + shape = Box(*obstacle.dimensions) + elif obstacle.obstacle_type == ObstacleType.SPHERE: + shape = Sphere(obstacle.dimensions[0]) + elif obstacle.obstacle_type == ObstacleType.CYLINDER: + shape = Cylinder(obstacle.dimensions[0], obstacle.dimensions[1]) + else: + logger.warning(f"Cannot visualize obstacle type: {obstacle.obstacle_type}") + return + + # Use Drake's Meshcat.SetObject with shape and color + self._meshcat.SetObject(path, shape, rgba) + self._meshcat.SetTransform(path, transform) + + def _create_shape(self, obstacle: Obstacle) -> Any: + """Create Drake shape from obstacle specification.""" + if obstacle.obstacle_type == ObstacleType.BOX: + return Box(*obstacle.dimensions) + elif obstacle.obstacle_type == ObstacleType.SPHERE: + return Sphere(obstacle.dimensions[0]) + elif obstacle.obstacle_type == ObstacleType.CYLINDER: + return Cylinder(obstacle.dimensions[0], obstacle.dimensions[1]) + else: + raise ValueError(f"Unsupported obstacle type: {obstacle.obstacle_type}") + + def remove_obstacle(self, obstacle_id: str) -> bool: + """Remove an obstacle by ID.""" + with self._lock: + if obstacle_id not in self._obstacles: + return False + + obstacle_data = self._obstacles[obstacle_id] + + if self._finalized and self._scene_graph_context is not None: + self._scene_graph.RemoveGeometry( + obstacle_data.source_id, + obstacle_data.geometry_id, + ) + + # Also remove from Meshcat + if self._meshcat is not None: + path = f"obstacles/{obstacle_id}" + self._meshcat.Delete(path) + + del self._obstacles[obstacle_id] + logger.debug(f"Removed obstacle '{obstacle_id}'") + return True + + def update_obstacle_pose(self, obstacle_id: str, pose: NDArray[np.float64]) -> bool: + """Update obstacle pose (4x4 transform).""" + with self._lock: + if obstacle_id not in self._obstacles: + return False + + if pose.shape != (4, 4): + raise ValueError(f"Pose must be 4x4, got {pose.shape}") + + self._obstacles[obstacle_id].obstacle.pose = pose.copy() + + # Update Meshcat visualization + if self._meshcat is not None: + path = f"obstacles/{obstacle_id}" + transform = RigidTransform(pose) + self._meshcat.SetTransform(path, transform) + + # Note: SceneGraph geometry pose is fixed after registration + # Meshcat is updated for visualization, but collision checking + # uses the original pose. For dynamic obstacles, remove and re-add. + + return True + + def clear_obstacles(self) -> None: + """Remove all obstacles.""" + with self._lock: + obstacle_ids = list(self._obstacles.keys()) + for obs_id in obstacle_ids: + self.remove_obstacle(obs_id) + + # ============= Lifecycle ============= + + def finalize(self) -> None: + """Finalize world - locks robot topology, enables collision checking.""" + if self._finalized: + logger.warning("World already finalized") + return + + with self._lock: + # Finalize plant + self._plant.Finalize() + + # Compute joint indices for each robot + for robot_id, robot_data in self._robots.items(): + joint_indices: list[int] = [] + for joint_name in robot_data.config.joint_names: + joint = self._plant.GetJointByName(joint_name, robot_data.model_instance) + start_idx = joint.position_start() + num_positions = joint.num_positions() + joint_indices.extend(range(start_idx, start_idx + num_positions)) + robot_data.joint_indices = joint_indices + logger.debug(f"Robot '{robot_id}' joint indices: {joint_indices}") + + # Setup collision filters + self._setup_collision_filters() + + # Register obstacle source for dynamic obstacles + self._obstacle_source_id = self._scene_graph.RegisterSource("dynamic_obstacles") + + # Add visualization if enabled + if self._meshcat is not None: + params = MeshcatVisualizerParams() + params.role = Role.kIllustration + self._meshcat_visualizer = MeshcatVisualizer.AddToBuilder( + self._builder, + self._scene_graph, + self._meshcat, + params, + ) + + # Build diagram + self._diagram = self._builder.Build() + self._live_context = self._diagram.CreateDefaultContext() + + # Get subsystem contexts + self._plant_context = self._diagram.GetMutableSubsystemContext( + self._plant, self._live_context + ) + self._scene_graph_context = self._diagram.GetMutableSubsystemContext( + self._scene_graph, self._live_context + ) + + self._finalized = True + logger.info(f"World finalized with {len(self._robots)} robots") + + # Initial visualization publish + if self._meshcat_visualizer is not None: + self._meshcat_visualizer.ForcedPublish( + self._diagram.GetSubsystemContext(self._meshcat_visualizer, self._live_context) + ) + + @property + def is_finalized(self) -> bool: + """Check if world is finalized.""" + return self._finalized + + def _setup_collision_filters(self) -> None: + """Filter out collisions between adjacent links in kinematic chain.""" + scene_graph = self._scene_graph + + for _robot_id, robot_data in self._robots.items(): + model_instance = robot_data.model_instance + body_indices = self._plant.GetBodyIndices(model_instance) + bodies = [self._plant.get_body(bi) for bi in body_indices] + + # Build parent-child relationships + parent_map: dict[str, str] = {} + for joint_idx in self._plant.GetJointIndices(model_instance): + joint = self._plant.get_joint(joint_idx) + parent_body = joint.parent_body() + child_body = joint.child_body() + if parent_body.index() != self._plant.world_body().index(): + parent_map[child_body.name()] = parent_body.name() + + # Compute kinematic distance + def get_chain_distance(name1: str, name2: str) -> int: + def path_to_root(name: str) -> list[str]: + path = [name] + while name in parent_map: + name = parent_map[name] + path.append(name) + return path + + path1 = path_to_root(name1) + path2 = path_to_root(name2) + + set1 = set(path1) + for i, node in enumerate(path2): + if node in set1: + idx1 = path1.index(node) + return idx1 + i + + return len(path1) + len(path2) + + # Filter collisions within 2 joints + filter_distance = 2 + filtered_pairs: set[tuple[str, str]] = set() + + for i, body1 in enumerate(bodies): + for body2 in bodies[i + 1 :]: + name1, name2 = body1.name(), body2.name() + dist = get_chain_distance(name1, name2) + + if dist <= filter_distance: + pair = (min(name1, name2), max(name1, name2)) + if pair not in filtered_pairs: + filtered_pairs.add(pair) + + try: + geoms1 = self._plant.GetCollisionGeometriesForBody(body1) + geoms2 = self._plant.GetCollisionGeometriesForBody(body2) + + if geoms1 and geoms2: + scene_graph.collision_filter_manager().Apply( + CollisionFilterDeclaration().ExcludeBetween( + GeometrySet(geoms1), GeometrySet(geoms2) + ) + ) + except Exception as e: + logger.warning(f"Could not filter {name1}<->{name2}: {e}") + + # Add user-specified collision exclusions from config + # Useful for parallel linkage mechanisms (grippers) where non-adjacent + # links may legitimately overlap due to mimic joints + exclusion_pairs = robot_data.config.collision_exclusion_pairs + if exclusion_pairs: + body_name_map = {body.name(): body for body in bodies} + for name1, name2 in exclusion_pairs: + if name1 in body_name_map and name2 in body_name_map: + body1 = body_name_map[name1] + body2 = body_name_map[name2] + try: + geoms1 = self._plant.GetCollisionGeometriesForBody(body1) + geoms2 = self._plant.GetCollisionGeometriesForBody(body2) + if geoms1 and geoms2: + scene_graph.collision_filter_manager().Apply( + CollisionFilterDeclaration().ExcludeBetween( + GeometrySet(geoms1), GeometrySet(geoms2) + ) + ) + logger.debug(f"Excluded collision pair: {name1}<->{name2}") + except Exception as e: + logger.warning(f"Could not filter pair {name1}<->{name2}: {e}") + else: + missing = [n for n in (name1, name2) if n not in body_name_map] + logger.warning(f"Collision exclusion: links not found: {missing}") + + logger.info("Collision filters applied") + + # ============= Context Management ============= + + def get_live_context(self) -> Context: + """Get the live context (mirrors current robot state). + + WARNING: Not thread-safe for reads during writes. + Use scratch_context() for planning operations. + """ + if not self._finalized or self._live_context is None: + raise RuntimeError("World must be finalized first") + return self._live_context + + @contextmanager + def scratch_context(self) -> Generator[Context, None, None]: + """Get a scratch context for planning (thread-safe). + + Uses CreateDefaultContext() instead of Clone() so that dynamically + added obstacles are visible to collision queries. + + Usage: + with world.scratch_context() as ctx: + world.set_positions(ctx, robot_id, q) + if world.is_collision_free(ctx, robot_id): + ... + """ + if not self._finalized: + raise RuntimeError("World must be finalized first") + + # Create fresh context - this sees all geometries including dynamic obstacles + # (Clone() doesn't see geometries added after the original context was created) + ctx = self._diagram.CreateDefaultContext() + + yield ctx + # Context automatically cleaned up when exiting + + def sync_from_joint_state(self, robot_id: str, positions: NDArray[np.float64]) -> None: + """Sync live context from driver's joint state. + + Called by StateMonitor when new JointState arrives. + """ + if not self._finalized or self._plant_context is None: + return # Silently ignore before finalization + + with self._lock: + self._set_positions_internal(self._plant_context, robot_id, positions) + + # NOTE: ForcedPublish is intentionally NOT called here. + # Calling ForcedPublish from the LCM callback thread blocks message processing. + # Visualization can be updated via publish_to_meshcat() from non-callback contexts. + + # ============= State Operations (context-based) ============= + + def set_positions(self, ctx: Context, robot_id: str, positions: NDArray[np.float64]) -> None: + """Set robot positions in given context.""" + if not self._finalized: + raise RuntimeError("World must be finalized first") + + # Get plant context from diagram context + plant_ctx = self._diagram.GetMutableSubsystemContext(self._plant, ctx) + self._set_positions_internal(plant_ctx, robot_id, positions) + + def _set_positions_internal( + self, plant_ctx: Context, robot_id: str, positions: NDArray[np.float64] + ) -> None: + """Internal: Set positions in a plant context.""" + if robot_id not in self._robots: + raise KeyError(f"Robot '{robot_id}' not found") + + robot_data = self._robots[robot_id] + full_positions = self._plant.GetPositions(plant_ctx).copy() + + for i, joint_idx in enumerate(robot_data.joint_indices): + full_positions[joint_idx] = positions[i] + + self._plant.SetPositions(plant_ctx, full_positions) + + def get_positions(self, ctx: Context, robot_id: str) -> NDArray[np.float64]: + """Get robot positions from given context.""" + 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] + plant_ctx = self._diagram.GetSubsystemContext(self._plant, ctx) + full_positions = self._plant.GetPositions(plant_ctx) + + return np.array([full_positions[idx] for idx in robot_data.joint_indices]) + + # ============= Collision Checking (context-based) ============= + + def is_collision_free(self, ctx: Context, robot_id: str) -> bool: + """Check if current configuration in context is collision-free.""" + 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") + + scene_graph_ctx = self._diagram.GetSubsystemContext(self._scene_graph, ctx) + query_object = self._scene_graph.get_query_output_port().Eval(scene_graph_ctx) + + return not query_object.HasCollisions() # type: ignore[attr-defined] + + def get_min_distance(self, ctx: Context, robot_id: str) -> float: + """Get minimum signed distance (positive = clearance, negative = penetration).""" + if not self._finalized: + raise RuntimeError("World must be finalized first") + + scene_graph_ctx = self._diagram.GetSubsystemContext(self._scene_graph, ctx) + query_object = self._scene_graph.get_query_output_port().Eval(scene_graph_ctx) + + signed_distance_pairs = query_object.ComputeSignedDistancePairwiseClosestPoints() # type: ignore[attr-defined] + + if not signed_distance_pairs: + return float("inf") + + return float(min(pair.distance for pair in signed_distance_pairs)) + + # ============= Forward Kinematics (context-based) ============= + + def get_ee_pose(self, ctx: Context, robot_id: str) -> NDArray[np.float64]: + """Get end-effector pose as 4x4 transform.""" + 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] + plant_ctx = self._diagram.GetSubsystemContext(self._plant, ctx) + + ee_body = robot_data.ee_frame.body() + X_WE = self._plant.EvalBodyPoseInWorld(plant_ctx, ee_body) + + return X_WE.GetAsMatrix4() + + def get_link_pose(self, ctx: Context, robot_id: str, link_name: str) -> NDArray[np.float64]: + """Get link pose as 4x4 transform.""" + 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] + plant_ctx = self._diagram.GetSubsystemContext(self._plant, ctx) + + try: + body = self._plant.GetBodyByName(link_name, robot_data.model_instance) + except RuntimeError: + raise KeyError(f"Link '{link_name}' not found in robot '{robot_id}'") + + X_WL = self._plant.EvalBodyPoseInWorld(plant_ctx, body) + + return X_WL.GetAsMatrix4() + + def get_jacobian(self, ctx: Context, robot_id: str) -> NDArray[np.float64]: + """Get geometric Jacobian (6 x n_joints). + + 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] + plant_ctx = self._diagram.GetSubsystemContext(self._plant, ctx) + + # Compute full Jacobian + J_full = self._plant.CalcJacobianSpatialVelocity( + plant_ctx, + JacobianWrtVariable.kQDot, + robot_data.ee_frame, + np.array([0.0, 0.0, 0.0]), # 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) + J_robot = np.zeros((6, n_joints)) + + for i, joint_idx in enumerate(robot_data.joint_indices): + J_robot[:, i] = J_full[:, joint_idx] + + # Reorder rows: Drake uses [angular, linear], we want [linear, angular] + J_reordered = np.vstack([J_robot[3:6, :], J_robot[0:3, :]]) + + return J_reordered + + # ============= Visualization ============= + + def get_meshcat_url(self) -> str | None: + """Get Meshcat visualization URL.""" + if self._meshcat is not None: + return self._meshcat.web_url() + return None + + def publish_to_meshcat(self, ctx: Context | None = None) -> None: + """Force publish current state to Meshcat. + + Args: + ctx: Context to publish. Uses live context if None. + """ + if self._meshcat_visualizer is None: + return + + if ctx is None: + ctx = self._live_context + + if ctx is not None: + self._meshcat_visualizer.ForcedPublish( + self._diagram.GetSubsystemContext(self._meshcat_visualizer, ctx) + ) + + # ============= Direct Access (use with caution) ============= + + @property + def plant(self) -> MultibodyPlant: + """Get underlying MultibodyPlant.""" + return self._plant + + @property + def scene_graph(self) -> SceneGraph: + """Get underlying SceneGraph.""" + return self._scene_graph + + @property + def diagram(self) -> Any: + """Get underlying Diagram.""" + return self._diagram From b58144e9ef6d42814667f3c03cdd89eb4070684c Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Mon, 12 Jan 2026 16:03:57 -0800 Subject: [PATCH 05/31] Monitor system keeps the world model (drake planning world) in synb with real world. manages obstacle lifecycle and more --- .../manipulation/planning/monitor/__init__.py | 64 ++ .../planning/monitor/world_monitor.py | 706 ++++++++++++++++++ .../monitor/world_obstacle_monitor.py | 385 ++++++++++ .../planning/monitor/world_state_monitor.py | 326 ++++++++ 4 files changed, 1481 insertions(+) create mode 100644 dimos/manipulation/planning/monitor/__init__.py create mode 100644 dimos/manipulation/planning/monitor/world_monitor.py create mode 100644 dimos/manipulation/planning/monitor/world_obstacle_monitor.py create mode 100644 dimos/manipulation/planning/monitor/world_state_monitor.py diff --git a/dimos/manipulation/planning/monitor/__init__.py b/dimos/manipulation/planning/monitor/__init__.py new file mode 100644 index 0000000000..060579b117 --- /dev/null +++ b/dimos/manipulation/planning/monitor/__init__.py @@ -0,0 +1,64 @@ +# Copyright 2025 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. + +""" +World Monitor Module + +Provides reactive monitoring for keeping WorldSpec synchronized with the real world. + +## Components + +- WorldMonitor: Top-level monitor using WorldSpec Protocol +- WorldStateMonitor: Syncs joint state to WorldSpec +- WorldObstacleMonitor: Syncs obstacles to WorldSpec + +All monitors use the factory pattern and Protocol types. + +## Example + +```python +from dimos.manipulation.planning.monitor import WorldMonitor + +monitor = WorldMonitor(enable_viz=True) +robot_id = monitor.add_robot(config) +monitor.finalize() + +# Start monitoring +monitor.start_state_monitor(robot_id) +monitor.start_obstacle_monitor() + +# Handle joint state messages +monitor.on_joint_state(msg, robot_id) + +# Thread-safe collision checking +is_valid = monitor.is_state_valid(robot_id, q_test) +``` +""" + +from dimos.manipulation.planning.monitor.world_monitor import WorldMonitor +from dimos.manipulation.planning.monitor.world_obstacle_monitor import ( + WorldObstacleMonitor, +) +from dimos.manipulation.planning.monitor.world_state_monitor import WorldStateMonitor + +# Re-export message types from spec for convenience +from dimos.manipulation.planning.spec import CollisionObjectMessage, Detection3D + +__all__ = [ + "CollisionObjectMessage", + "Detection3D", + "WorldMonitor", + "WorldObstacleMonitor", + "WorldStateMonitor", +] diff --git a/dimos/manipulation/planning/monitor/world_monitor.py b/dimos/manipulation/planning/monitor/world_monitor.py new file mode 100644 index 0000000000..cb84c1efa5 --- /dev/null +++ b/dimos/manipulation/planning/monitor/world_monitor.py @@ -0,0 +1,706 @@ +# 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. + +""" +World Monitor + +MoveIt-like monitor that keeps a WorldSpec synchronized with the real world. +This is the WorldSpec-based replacement for PlanningSceneMonitor. + +Uses factory functions to create WorldSpec, so all code uses Protocol types. + +Example: + # Create monitor + monitor = WorldMonitor(enable_viz=True) + + # Add robot + robot_id = monitor.add_robot(robot_config) + + # Finalize + monitor.finalize() + + # Start monitoring + monitor.start_state_monitor("piper", ["joint1", "joint2", ...]) + monitor.start_obstacle_monitor() + + # Use thread-safe accessors + positions = monitor.get_current_positions(robot_id) + is_valid = monitor.is_state_valid(robot_id, target_joints) + + # Get scratch context for planning + with monitor.scratch_context() as ctx: + # Do planning operations + pass +""" + +from __future__ import annotations + +from contextlib import contextmanager +import threading +from typing import TYPE_CHECKING, Any + +import numpy as np + +from dimos.manipulation.planning.factory import create_world +from dimos.manipulation.planning.monitor.world_obstacle_monitor import WorldObstacleMonitor +from dimos.manipulation.planning.monitor.world_state_monitor import WorldStateMonitor +from dimos.manipulation.planning.spec import ( + CollisionObjectMessage, + Detection3D, + Obstacle, + ObstacleType, + RobotModelConfig, +) +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from collections.abc import Generator + + from numpy.typing import NDArray + + from dimos.manipulation.planning.spec import WorldSpec + from dimos.msgs.sensor_msgs import JointState + +logger = setup_logger() + + +class WorldMonitor: + """MoveIt-like monitor for keeping a WorldSpec synchronized. + + This class manages: + - WorldSpec instance (created via factory) + - WorldStateMonitor: Subscribes to joint states, syncs robot configurations + - WorldObstacleMonitor: Subscribes to collision objects, syncs obstacles + + ## Key Differences from PlanningSceneMonitor + + 1. Uses factory pattern - no concrete type references + 2. Uses scratch_context() for thread-safe operations + 3. State is synced via sync_from_joint_state() instead of set_robot_state() + + ## Usage Pattern + + ```python + # 1. Create monitor + monitor = WorldMonitor(enable_viz=True) + + # 2. Add robots + robot_id = monitor.add_robot(robot_config) + + # 3. Add static obstacles (optional) + monitor.add_obstacle(table_obstacle) + + # 4. Finalize (locks topology) + monitor.finalize() + + # 5. Start monitoring + monitor.start_state_monitor(robot_id, joint_names) + monitor.start_obstacle_monitor() + + # 6. Use thread-safe accessors + positions = monitor.get_current_positions(robot_id) + is_valid = monitor.is_state_valid(robot_id, target_joints) + + # 7. Get scratch context for planning + with monitor.scratch_context() as ctx: + world.set_positions(ctx, robot_id, target_joints) + if world.is_collision_free(ctx, robot_id): + # Valid configuration + pass + ``` + + ## Thread Safety + + All public methods are thread-safe via an internal RLock. + For planning operations, use scratch_context() to get an independent context. + """ + + def __init__( + self, + backend: str = "drake", + enable_viz: bool = False, + **kwargs: Any, + ): + """Create a world monitor. + + Args: + backend: Backend to use ("drake") + enable_viz: Enable visualization (Meshcat for Drake) + **kwargs: Additional arguments passed to create_world() + """ + self._backend = backend + + # Create world via factory (Protocol type) + self._world: WorldSpec = create_world( + backend=backend, + enable_viz=enable_viz, + **kwargs, + ) + + # Thread safety + self._lock = threading.RLock() + + # Robot tracking: robot_id -> joint_names + self._robot_joints: dict[str, list[str]] = {} + + # Sub-monitors + self._state_monitors: dict[str, WorldStateMonitor] = {} + self._obstacle_monitor: WorldObstacleMonitor | None = None + + # Visualization thread (publishes to Meshcat at a fixed rate) + self._viz_thread: threading.Thread | None = None + self._viz_stop_event = threading.Event() + self._viz_rate_hz: float = 10.0 # Default 10Hz + + # ============= Robot Management ============= + + def add_robot( + self, + config: RobotModelConfig, + ) -> str: + """Add a robot to the monitored world. + + Args: + config: Robot configuration + + Returns: + robot_id: Unique identifier for the robot + """ + with self._lock: + robot_id = self._world.add_robot(config) + self._robot_joints[robot_id] = config.joint_names + logger.info(f"Added robot '{config.name}' as '{robot_id}'") + return robot_id + + def get_robot_ids(self) -> list[str]: + """Get all robot IDs.""" + with self._lock: + return self._world.get_robot_ids() + + def get_robot_config(self, robot_id: str) -> RobotModelConfig: + """Get robot configuration.""" + with self._lock: + return self._world.get_robot_config(robot_id) + + def get_joint_limits(self, robot_id: str) -> tuple[NDArray[np.float64], NDArray[np.float64]]: + """Get joint limits for a robot.""" + with self._lock: + return self._world.get_joint_limits(robot_id) + + # ============= Obstacle Management ============= + + def add_obstacle(self, obstacle: Obstacle) -> str: + """Add an obstacle to the world. + + Args: + obstacle: Obstacle specification + + Returns: + obstacle_id: Unique identifier for the obstacle + """ + with self._lock: + return self._world.add_obstacle(obstacle) + + def add_box_obstacle( + self, + name: str, + pose: NDArray[np.float64], + dimensions: tuple[float, float, float], + color: tuple[float, float, float, float] = (0.8, 0.2, 0.2, 0.8), + ) -> str: + """Add a box obstacle. + + Args: + name: Obstacle name + pose: 4x4 homogeneous transform + dimensions: (width, height, depth) in meters + color: RGBA color + + Returns: + obstacle_id + """ + obstacle = Obstacle( + name=name, + obstacle_type=ObstacleType.BOX, + pose=pose, + dimensions=dimensions, + color=color, + ) + return self.add_obstacle(obstacle) + + def add_sphere_obstacle( + self, + name: str, + pose: NDArray[np.float64], + radius: float, + color: tuple[float, float, float, float] = (0.8, 0.2, 0.2, 0.8), + ) -> str: + """Add a sphere obstacle.""" + obstacle = Obstacle( + name=name, + obstacle_type=ObstacleType.SPHERE, + pose=pose, + dimensions=(radius,), + color=color, + ) + return self.add_obstacle(obstacle) + + def add_cylinder_obstacle( + self, + name: str, + pose: NDArray[np.float64], + radius: float, + height: float, + color: tuple[float, float, float, float] = (0.8, 0.2, 0.2, 0.8), + ) -> str: + """Add a cylinder obstacle.""" + obstacle = Obstacle( + name=name, + obstacle_type=ObstacleType.CYLINDER, + pose=pose, + dimensions=(radius, height), + color=color, + ) + return self.add_obstacle(obstacle) + + def remove_obstacle(self, obstacle_id: str) -> bool: + """Remove an obstacle.""" + with self._lock: + return self._world.remove_obstacle(obstacle_id) + + def clear_obstacles(self) -> None: + """Remove all obstacles.""" + with self._lock: + self._world.clear_obstacles() + + # ============= Monitor Control ============= + + def start_state_monitor( + self, + robot_id: str, + joint_names: list[str] | None = None, + joint_name_mapping: dict[str, str] | None = None, + ) -> None: + """Start monitoring joint states for a robot. + + Creates a WorldStateMonitor that will sync joint state messages + to the world's live context. + + Args: + robot_id: ID of the robot to monitor + joint_names: Joint names in URDF namespace (uses config joint_names if None) + joint_name_mapping: Maps orchestrator joint names to URDF joint names. + Example: {"left_joint1": "joint1"} means messages with "left_joint1" + will be mapped to URDF "joint1". If None, uses config mapping if available, + otherwise names must match exactly. + """ + with self._lock: + if robot_id in self._state_monitors: + logger.warning(f"State monitor for '{robot_id}' already started") + return + + # Get config for defaults + config = self._world.get_robot_config(robot_id) + + # Get joint names from config if not provided + if joint_names is None: + if robot_id in self._robot_joints: + joint_names = self._robot_joints[robot_id] + else: + joint_names = config.joint_names + + # Get joint name mapping from config if not provided + if joint_name_mapping is None and config.joint_name_mapping: + joint_name_mapping = config.joint_name_mapping + + monitor = WorldStateMonitor( + world=self._world, + lock=self._lock, + robot_id=robot_id, + joint_names=joint_names, + joint_name_mapping=joint_name_mapping, + ) + monitor.start() + self._state_monitors[robot_id] = monitor + logger.info(f"State monitor started for '{robot_id}'") + + def start_obstacle_monitor(self) -> None: + """Start monitoring obstacle updates. + + Creates a WorldObstacleMonitor that will sync obstacle messages + to the world. + """ + with self._lock: + if self._obstacle_monitor is not None: + logger.warning("Obstacle monitor already started") + return + + self._obstacle_monitor = WorldObstacleMonitor( + world=self._world, + lock=self._lock, + ) + self._obstacle_monitor.start() + logger.info("Obstacle monitor started") + + def stop_all_monitors(self) -> None: + """Stop all monitors and visualization thread.""" + # Stop visualization thread first (outside lock to avoid deadlock) + self.stop_visualization_thread() + + with self._lock: + for _robot_id, monitor in self._state_monitors.items(): + monitor.stop() + self._state_monitors.clear() + + if self._obstacle_monitor is not None: + self._obstacle_monitor.stop() + self._obstacle_monitor = None + + logger.info("All monitors stopped") + + # ============= Message Handlers ============= + + def on_joint_state(self, msg: JointState, robot_id: str | None = None) -> None: + """Handle joint state message. + + Args: + msg: JointState message + robot_id: Specific robot to update (broadcasts to all if None) + """ + try: + if robot_id is not None: + if robot_id in self._state_monitors: + self._state_monitors[robot_id].on_joint_state(msg) + else: + logger.warning(f"No state monitor for robot_id: {robot_id}") + else: + # Broadcast to all monitors + for monitor in self._state_monitors.values(): + monitor.on_joint_state(msg) + except Exception as e: + logger.error(f"[WorldMonitor] Exception in on_joint_state: {e}") + import traceback + + logger.error(traceback.format_exc()) + + def on_collision_object(self, msg: CollisionObjectMessage) -> None: + """Handle collision object message.""" + if self._obstacle_monitor is not None: + self._obstacle_monitor.on_collision_object(msg) + + def on_detections(self, detections: list[Detection3D]) -> None: + """Handle perception detections.""" + if self._obstacle_monitor is not None: + self._obstacle_monitor.on_detections(detections) + + # ============= State Access ============= + + def get_current_positions(self, robot_id: str) -> NDArray[np.float64] | None: + """Get current joint positions (thread-safe). + + Args: + robot_id: ID of the robot + + Returns: + Current positions or None if not yet received + """ + # Try state monitor first + if robot_id in self._state_monitors: + positions = self._state_monitors[robot_id].get_current_positions() + if positions is not None: + return positions + + # Fall back to world's live context + with self._lock: + ctx = self._world.get_live_context() + return self._world.get_positions(ctx, robot_id) + + def get_current_velocities(self, robot_id: str) -> NDArray[np.float64] | None: + """Get current joint velocities (thread-safe). + + Args: + robot_id: ID of the robot + + Returns: + Current velocities or None if not available + """ + if robot_id in self._state_monitors: + return self._state_monitors[robot_id].get_current_velocities() + return None + + def wait_for_state(self, robot_id: str, timeout: float = 1.0) -> bool: + """Wait until a state is received for the robot. + + Args: + robot_id: ID of the robot + timeout: Maximum time to wait + + Returns: + True if state was received, False if timeout + """ + if robot_id in self._state_monitors: + return self._state_monitors[robot_id].wait_for_state(timeout) + return False + + def is_state_stale(self, robot_id: str, max_age: float = 1.0) -> bool: + """Check if state is stale.""" + if robot_id in self._state_monitors: + return self._state_monitors[robot_id].is_state_stale(max_age) + return True + + # ============= Context Management ============= + + @contextmanager + def scratch_context(self) -> Generator[Any, None, None]: + """Get a scratch context for planning (thread-safe). + + The scratch context is a clone of the live context that can be + modified without affecting the live state. + + Example: + with monitor.scratch_context() as ctx: + world.set_positions(ctx, robot_id, q_test) + if world.is_collision_free(ctx, robot_id): + # Valid configuration + pass + + Yields: + Context that can be used with world operations + """ + with self._world.scratch_context() as ctx: + yield ctx + + def get_live_context(self) -> Any: + """Get the live context (use with caution). + + The live context is synced with real robot state. Modifications + will affect collision checking for all users. + + For planning, prefer scratch_context() instead. + + Returns: + Live context + """ + return self._world.get_live_context() + + # ============= Collision Checking ============= + + def is_state_valid( + self, + robot_id: str, + joint_positions: NDArray[np.float64], + ) -> bool: + """Check if configuration is collision-free (thread-safe). + + Args: + robot_id: ID of the robot + joint_positions: Joint configuration to check + + Returns: + True if collision-free + """ + with self._world.scratch_context() as ctx: + self._world.set_positions(ctx, robot_id, joint_positions) + return self._world.is_collision_free(ctx, robot_id) + + def is_path_valid( + self, + robot_id: str, + path: list[NDArray[np.float64]], + step_size: float = 0.05, + ) -> bool: + """Check if path is collision-free (thread-safe). + + Args: + robot_id: ID of the robot + path: List of joint configurations + step_size: Interpolation resolution + + Returns: + True if path is collision-free + """ + with self._world.scratch_context() as ctx: + for i in range(len(path) - 1): + q_start = path[i] + q_end = path[i + 1] + + # Number of interpolation steps + dist = np.linalg.norm(q_end - q_start) + num_steps = max(2, int(np.ceil(dist / step_size))) + + for j in range(num_steps): + alpha = j / (num_steps - 1) + q = q_start + alpha * (q_end - q_start) + + self._world.set_positions(ctx, robot_id, q) + if not self._world.is_collision_free(ctx, robot_id): + return False + + return True + + def get_min_distance(self, robot_id: str) -> float: + """Get minimum distance to obstacles for current state.""" + with self._world.scratch_context() as ctx: + return self._world.get_min_distance(ctx, robot_id) + + # ============= Kinematics ============= + + def get_ee_pose( + self, + robot_id: str, + joint_positions: NDArray[np.float64] | None = None, + ) -> NDArray[np.float64]: + """Get end-effector pose (thread-safe). + + Args: + robot_id: ID of the robot + joint_positions: Joint configuration (uses current if None) + + Returns: + 4x4 homogeneous transform + """ + with self._world.scratch_context() as ctx: + # If no positions provided, fetch current from state monitor + if joint_positions is None: + joint_positions = self.get_current_positions(robot_id) + + if joint_positions is not None: + self._world.set_positions(ctx, robot_id, joint_positions) + + return self._world.get_ee_pose(ctx, robot_id) + + def get_jacobian( + self, + robot_id: str, + joint_positions: NDArray[np.float64], + ) -> NDArray[np.float64]: + """Get Jacobian (thread-safe). + + Args: + robot_id: ID of the robot + joint_positions: Joint configuration + + Returns: + 6 x n_joints Jacobian matrix + """ + with self._world.scratch_context() as ctx: + self._world.set_positions(ctx, robot_id, joint_positions) + return self._world.get_jacobian(ctx, robot_id) + + # ============= Lifecycle ============= + + def finalize(self) -> None: + """Finalize world for collision checking. + + Must be called after adding all robots and before starting monitors + or performing collision checking. + """ + with self._lock: + self._world.finalize() + logger.info("World finalized") + + @property + def is_finalized(self) -> bool: + """Check if world is finalized.""" + return self._world.is_finalized + + # ============= Visualization ============= + + def get_meshcat_url(self) -> str | None: + """Get Meshcat visualization URL (Drake only). + + Returns: + URL string or None if visualization not enabled + """ + if hasattr(self._world, "get_meshcat_url"): + url = self._world.get_meshcat_url() + return str(url) if url else None + return None + + def publish_visualization(self) -> None: + """Force publish current state to visualization.""" + if hasattr(self._world, "publish_to_meshcat"): + self._world.publish_to_meshcat() + + def start_visualization_thread(self, rate_hz: float = 10.0) -> None: + """Start background thread that updates Meshcat visualization. + + This allows live visualization without blocking the LCM callback thread. + The thread publishes the current robot state to Meshcat at the specified rate. + + Args: + rate_hz: Visualization update rate in Hz (default: 10Hz) + """ + if self._viz_thread is not None and self._viz_thread.is_alive(): + logger.warning("Visualization thread already running") + return + + if not hasattr(self._world, "publish_to_meshcat"): + logger.warning("World does not support Meshcat visualization") + return + + self._viz_rate_hz = rate_hz + self._viz_stop_event.clear() + self._viz_thread = threading.Thread( + target=self._visualization_loop, + name="MeshcatVizThread", + daemon=True, + ) + self._viz_thread.start() + logger.info(f"Visualization thread started at {rate_hz}Hz") + + def stop_visualization_thread(self) -> None: + """Stop the visualization thread.""" + if self._viz_thread is None: + return + + self._viz_stop_event.set() + self._viz_thread.join(timeout=1.0) + if self._viz_thread.is_alive(): + logger.warning("Visualization thread did not stop cleanly") + self._viz_thread = None + logger.info("Visualization thread stopped") + + def _visualization_loop(self) -> None: + """Internal: Visualization update loop.""" + import time + + period = 1.0 / self._viz_rate_hz + while not self._viz_stop_event.is_set(): + try: + if hasattr(self._world, "publish_to_meshcat"): + self._world.publish_to_meshcat() + except Exception as e: + logger.debug(f"Visualization publish failed: {e}") + time.sleep(period) + + # ============= Direct World Access ============= + + @property + def world(self) -> WorldSpec: + """Get underlying WorldSpec instance. + + Direct access to the world is NOT thread-safe for state modifications. + Prefer using scratch_context() for planning operations. + """ + return self._world + + def get_state_monitor(self, robot_id: str) -> WorldStateMonitor | None: + """Get state monitor for a robot (may be None).""" + return self._state_monitors.get(robot_id) + + @property + def obstacle_monitor(self) -> WorldObstacleMonitor | None: + """Get obstacle monitor (may be None if not started).""" + return self._obstacle_monitor diff --git a/dimos/manipulation/planning/monitor/world_obstacle_monitor.py b/dimos/manipulation/planning/monitor/world_obstacle_monitor.py new file mode 100644 index 0000000000..da5a01304c --- /dev/null +++ b/dimos/manipulation/planning/monitor/world_obstacle_monitor.py @@ -0,0 +1,385 @@ +# 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. + +""" +World Obstacle Monitor + +Monitors obstacle updates and applies them to a WorldSpec instance. +This is the WorldSpec-based replacement for WorldGeometryMonitor. + +Example: + monitor = WorldObstacleMonitor(world, lock) + monitor.start() + monitor.on_collision_object(collision_msg) # Called by subscriber +""" + +from __future__ import annotations + +import time +from typing import TYPE_CHECKING + +from dimos.manipulation.planning.spec import ( + CollisionObjectMessage, + Detection3D, + Obstacle, + ObstacleType, +) +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from collections.abc import Callable + import threading + + import numpy as np + from numpy.typing import NDArray + + from dimos.manipulation.planning.spec import WorldSpec + +logger = setup_logger() + + +class WorldObstacleMonitor: + """Monitors world obstacles and updates WorldSpec. + + This class handles updates from: + - Explicit collision objects (CollisionObjectMessage) + - Perception detections (Detection3D) + + ## Thread Safety + + All obstacle operations are protected by the provided lock. + Callbacks can be called from any thread. + + ## Comparison with WorldGeometryMonitor + + - WorldGeometryMonitor: Works with PlanningScene ABC + - WorldObstacleMonitor: Works with WorldSpec Protocol + """ + + def __init__( + self, + world: WorldSpec, + lock: threading.RLock, + detection_timeout: float = 2.0, + ): + """Create a world obstacle monitor. + + Args: + world: WorldSpec instance to update + lock: Shared lock for thread-safe access + detection_timeout: Time before removing stale detections (seconds) + """ + self._world = world + self._lock = lock + self._detection_timeout = detection_timeout + + # Track obstacles from different sources + self._collision_objects: dict[str, str] = {} # msg_id -> obstacle_id + self._perception_objects: dict[str, str] = {} # detection_id -> obstacle_id + self._perception_timestamps: dict[str, float] = {} # detection_id -> timestamp + + # Running state + self._running = False + + # Callbacks + self._obstacle_callbacks: list[Callable[[str, str, Obstacle | None], None]] = [] + + def start(self) -> None: + """Start the obstacle monitor.""" + self._running = True + logger.info("World obstacle monitor started") + + def stop(self) -> None: + """Stop the obstacle monitor.""" + self._running = False + logger.info("World obstacle monitor stopped") + + def is_running(self) -> bool: + """Check if monitor is running.""" + return self._running + + def on_collision_object(self, msg: CollisionObjectMessage) -> None: + """Handle explicit collision object message. + + Args: + msg: Collision object message + """ + if not self._running: + return + + with self._lock: + if msg.operation == "add": + self._add_collision_object(msg) + elif msg.operation == "remove": + self._remove_collision_object(msg.id) + elif msg.operation == "update": + self._update_collision_object(msg) + else: + logger.warning(f"Unknown collision object operation: {msg.operation}") + + def _add_collision_object(self, msg: CollisionObjectMessage) -> None: + """Add a collision object from message.""" + if msg.id in self._collision_objects: + logger.debug(f"Collision object '{msg.id}' already exists, updating") + self._update_collision_object(msg) + return + + obstacle = self._msg_to_obstacle(msg) + if obstacle is None: + logger.warning(f"Failed to create obstacle from message: {msg.id}") + return + + obstacle_id = self._world.add_obstacle(obstacle) + self._collision_objects[msg.id] = obstacle_id + + logger.debug(f"Added collision object '{msg.id}' as '{obstacle_id}'") + + # Notify callbacks + for callback in self._obstacle_callbacks: + try: + callback("add", obstacle_id, obstacle) + except Exception as e: + logger.error(f"Obstacle callback error: {e}") + + def _remove_collision_object(self, msg_id: str) -> None: + """Remove a collision object.""" + if msg_id not in self._collision_objects: + logger.debug(f"Collision object '{msg_id}' not found") + return + + obstacle_id = self._collision_objects[msg_id] + self._world.remove_obstacle(obstacle_id) + del self._collision_objects[msg_id] + + logger.debug(f"Removed collision object '{msg_id}'") + + # Notify callbacks + for callback in self._obstacle_callbacks: + try: + callback("remove", obstacle_id, None) + except Exception as e: + logger.error(f"Obstacle callback error: {e}") + + def _update_collision_object(self, msg: CollisionObjectMessage) -> None: + """Update a collision object pose.""" + if msg.id not in self._collision_objects: + # Treat as add if doesn't exist + self._add_collision_object(msg) + return + + obstacle_id = self._collision_objects[msg.id] + + if msg.pose is not None: + self._world.update_obstacle_pose(obstacle_id, msg.pose) + logger.debug(f"Updated collision object '{msg.id}' pose") + + # Notify callbacks + for callback in self._obstacle_callbacks: + try: + callback("update", obstacle_id, None) + except Exception as e: + logger.error(f"Obstacle callback error: {e}") + + def _msg_to_obstacle(self, msg: CollisionObjectMessage) -> Obstacle | None: + """Convert collision object message to Obstacle.""" + if msg.primitive_type is None or msg.pose is None or msg.dimensions is None: + return None + + type_map = { + "box": ObstacleType.BOX, + "sphere": ObstacleType.SPHERE, + "cylinder": ObstacleType.CYLINDER, + } + + obstacle_type = type_map.get(msg.primitive_type.lower()) + if obstacle_type is None: + logger.warning(f"Unknown primitive type: {msg.primitive_type}") + return None + + return Obstacle( + name=msg.id, + obstacle_type=obstacle_type, + pose=msg.pose, + dimensions=msg.dimensions, + color=msg.color, + ) + + def on_detections(self, detections: list[Detection3D]) -> None: + """Handle perception detection results. + + Updates obstacles based on detections: + - Adds new obstacles for new detections + - Updates existing obstacles + - Removes obstacles for detections that are no longer present + + Args: + detections: List of 3D detections + """ + if not self._running: + return + + with self._lock: + current_time = time.time() + seen_ids = set() + + for detection in detections: + det_id = detection.id + seen_ids.add(det_id) + + if det_id in self._perception_objects: + # Update existing obstacle + obstacle_id = self._perception_objects[det_id] + self._world.update_obstacle_pose(obstacle_id, detection.pose) + self._perception_timestamps[det_id] = current_time + else: + # Add new obstacle + obstacle = self._detection_to_obstacle(detection) + obstacle_id = self._world.add_obstacle(obstacle) + self._perception_objects[det_id] = obstacle_id + self._perception_timestamps[det_id] = current_time + + logger.debug(f"Added perception object '{det_id}' as '{obstacle_id}'") + + # Notify callbacks + for callback in self._obstacle_callbacks: + try: + callback("add", obstacle_id, obstacle) + except Exception as e: + logger.error(f"Obstacle callback error: {e}") + + # Remove stale detections + self._cleanup_stale_detections(current_time, seen_ids) + + def _detection_to_obstacle(self, detection: Detection3D) -> Obstacle: + """Convert detection to obstacle.""" + return Obstacle( + name=f"detection_{detection.id}", + obstacle_type=ObstacleType.BOX, + pose=detection.pose, + dimensions=detection.dimensions, + color=(0.2, 0.8, 0.2, 0.6), # Green for perception objects + ) + + def _cleanup_stale_detections( + self, + current_time: float, + seen_ids: set[str], + ) -> None: + """Remove detections that haven't been seen recently.""" + stale_ids = [] + + for det_id, timestamp in self._perception_timestamps.items(): + age = current_time - timestamp + if det_id not in seen_ids and age > self._detection_timeout: + stale_ids.append(det_id) + + for det_id in stale_ids: + obstacle_id = self._perception_objects[det_id] + self._world.remove_obstacle(obstacle_id) + del self._perception_objects[det_id] + del self._perception_timestamps[det_id] + + logger.debug(f"Removed stale perception object '{det_id}'") + + # Notify callbacks + for callback in self._obstacle_callbacks: + try: + callback("remove", obstacle_id, None) + except Exception as e: + logger.error(f"Obstacle callback error: {e}") + + def add_static_obstacle( + self, + name: str, + obstacle_type: str, + pose: NDArray[np.float64], + dimensions: tuple[float, ...], + color: tuple[float, float, float, float] = (0.8, 0.2, 0.2, 0.8), + ) -> str: + """Manually add a static obstacle. + + Args: + name: Unique name for the obstacle + obstacle_type: Type ("box", "sphere", "cylinder") + pose: 4x4 homogeneous transform + dimensions: Type-specific dimensions + color: RGBA color + + Returns: + Obstacle ID + """ + msg = CollisionObjectMessage( + id=name, + operation="add", + primitive_type=obstacle_type, + pose=pose, + dimensions=dimensions, + color=color, + ) + self.on_collision_object(msg) + return self._collision_objects.get(name, "") + + def remove_static_obstacle(self, name: str) -> bool: + """Remove a static obstacle by name. + + Args: + name: Name of the obstacle + + Returns: + True if removed + """ + if name not in self._collision_objects: + return False + + msg = CollisionObjectMessage(id=name, operation="remove") + self.on_collision_object(msg) + return True + + def clear_all_obstacles(self) -> None: + """Remove all tracked obstacles.""" + with self._lock: + # Clear collision objects + for msg_id in list(self._collision_objects.keys()): + self._remove_collision_object(msg_id) + + # Clear perception objects + for det_id, obstacle_id in list(self._perception_objects.items()): + self._world.remove_obstacle(obstacle_id) + del self._perception_objects[det_id] + del self._perception_timestamps[det_id] + + def get_obstacle_count(self) -> int: + """Get total number of tracked obstacles.""" + with self._lock: + return len(self._collision_objects) + len(self._perception_objects) + + def add_obstacle_callback( + self, + callback: Callable[[str, str, Obstacle | None], None], + ) -> None: + """Add callback for obstacle changes. + + Args: + callback: Function called with (operation, obstacle_id, obstacle) + where operation is "add", "update", or "remove" + """ + self._obstacle_callbacks.append(callback) + + def remove_obstacle_callback( + self, + callback: Callable[[str, str, Obstacle | None], None], + ) -> None: + """Remove an obstacle callback.""" + if callback in self._obstacle_callbacks: + self._obstacle_callbacks.remove(callback) diff --git a/dimos/manipulation/planning/monitor/world_state_monitor.py b/dimos/manipulation/planning/monitor/world_state_monitor.py new file mode 100644 index 0000000000..38bb1575e1 --- /dev/null +++ b/dimos/manipulation/planning/monitor/world_state_monitor.py @@ -0,0 +1,326 @@ +# 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. + +""" +World State Monitor + +Monitors joint state updates and syncs them to a WorldSpec instance. +This is the WorldSpec-based replacement for StateMonitor. + +Example: + monitor = WorldStateMonitor(world, lock, robot_id, joint_names) + monitor.start() + monitor.on_joint_state(joint_state_msg) # Called by subscriber +""" + +from __future__ import annotations + +import time +from typing import TYPE_CHECKING + +import numpy as np + +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from collections.abc import Callable + import threading + + from numpy.typing import NDArray + + from dimos.manipulation.planning.spec import WorldSpec + from dimos.msgs.sensor_msgs import JointState + +logger = setup_logger() + + +class WorldStateMonitor: + """Monitors joint state updates and syncs them to WorldSpec. + + This class subscribes to joint state messages and calls + world.sync_from_joint_state() to keep the world's live context + synchronized with the real robot state. + + ## Thread Safety + + All state updates are protected by the provided lock. The on_joint_state + callback can be called from any thread. + + ## Comparison with StateMonitor + + - StateMonitor: Works with PlanningScene ABC + - WorldStateMonitor: Works with WorldSpec Protocol + """ + + def __init__( + self, + world: WorldSpec, + lock: threading.RLock, + robot_id: str, + joint_names: list[str], + joint_name_mapping: dict[str, str] | None = None, + timeout: float = 1.0, + ): + """Create a world state monitor. + + Args: + world: WorldSpec instance to sync state to + lock: Shared lock for thread-safe access + robot_id: ID of the robot to monitor + joint_names: Ordered list of joint names for this robot (URDF names) + joint_name_mapping: Maps orchestrator joint names to URDF joint names. + Example: {"left_joint1": "joint1"} means messages with "left_joint1" + will be mapped to URDF "joint1". If None, names must match exactly. + timeout: Timeout for waiting for initial state (seconds) + """ + self._world = world + self._lock = lock + self._robot_id = robot_id + self._joint_names = joint_names + self._timeout = timeout + + # Joint name mapping: orchestrator name -> URDF name + self._joint_name_mapping = joint_name_mapping or {} + # Build reverse mapping: URDF name -> orchestrator name + self._reverse_mapping = {v: k for k, v in self._joint_name_mapping.items()} + + # Latest state + self._latest_positions: NDArray[np.float64] | None = None + self._latest_velocities: NDArray[np.float64] | None = None + self._last_update_time: float | None = None + + # Running state + self._running = False + + # Callbacks + self._state_callbacks: list[Callable[[str, NDArray[np.float64]], None]] = [] + + def start(self) -> None: + """Start the state monitor.""" + self._running = True + logger.info(f"World state monitor started for robot '{self._robot_id}'") + + def stop(self) -> None: + """Stop the state monitor.""" + self._running = False + logger.info(f"World state monitor stopped for robot '{self._robot_id}'") + + def is_running(self) -> bool: + """Check if monitor is running.""" + return self._running + + @property + def robot_id(self) -> str: + """Get the robot ID being monitored.""" + return self._robot_id + + def on_joint_state(self, msg: JointState) -> None: + """Handle incoming joint state message. + + This is called by the subscriber when a new JointState message arrives. + It extracts joint positions and syncs them to the world. + + Args: + msg: JointState message with joint names and positions + """ + try: + if not self._running: + return + + # Extract positions for our robot's joints + positions = self._extract_positions(msg) + if positions is None: + logger.debug( + "[WorldStateMonitor] Failed to extract positions - joint names mismatch" + ) + logger.debug(f" Expected joints: {self._joint_names}") + logger.debug(f" Received joints: {msg.name}") + return # Not all joints present in message + + velocities = self._extract_velocities(msg) + + # Track message count for debugging + self._msg_count = getattr(self, "_msg_count", 0) + 1 + + with self._lock: + current_time = time.time() + + # Store latest state FIRST - this ensures planning always has + # current positions even if sync_from_joint_state fails + # (e.g., after dynamically adding obstacles) + self._latest_positions = positions + self._latest_velocities = velocities + self._last_update_time = current_time + + # Sync to world's live context (for visualization) + try: + self._world.sync_from_joint_state(self._robot_id, positions) + except Exception as e: + logger.error(f"Failed to sync joint state to live context: {e}") + + # Call registered callbacks + for callback in self._state_callbacks: + try: + callback(self._robot_id, positions) + except Exception as e: + logger.error(f"State callback error: {e}") + + except Exception as e: + logger.error(f"[WorldStateMonitor] Unexpected exception in on_joint_state: {e}") + import traceback + + logger.error(traceback.format_exc()) + + def _extract_positions(self, msg: JointState) -> NDArray[np.float64] | None: + """Extract positions for our joints from JointState message. + + Handles joint name translation from orchestrator namespace to URDF namespace. + If joint_name_mapping is set, message names are looked up via the reverse mapping. + + Args: + msg: JointState message (may use orchestrator joint names) + + Returns: + Array of joint positions or None if any joint is missing + """ + # Build name->index map from message (orchestrator names) + name_to_idx = {name: i for i, name in enumerate(msg.name)} + + positions = [] + for urdf_joint_name in self._joint_names: + # Try direct match first (when no mapping or names already match) + if urdf_joint_name in name_to_idx: + idx = name_to_idx[urdf_joint_name] + else: + # Try reverse mapping: URDF name -> orchestrator name -> msg index + orch_name = self._reverse_mapping.get(urdf_joint_name) + if orch_name is None or orch_name not in name_to_idx: + return None # Missing joint + idx = name_to_idx[orch_name] + + if idx >= len(msg.position): + return None # Position not available + positions.append(msg.position[idx]) + + return np.array(positions, dtype=np.float64) + + def _extract_velocities(self, msg: JointState) -> NDArray[np.float64] | None: + """Extract velocities for our joints. + + Uses same name translation as _extract_positions. + """ + if not msg.velocity or len(msg.velocity) == 0: + return None + + name_to_idx = {name: i for i, name in enumerate(msg.name)} + + velocities = [] + for urdf_joint_name in self._joint_names: + # Try direct match first + if urdf_joint_name in name_to_idx: + idx = name_to_idx[urdf_joint_name] + else: + # Try reverse mapping + orch_name = self._reverse_mapping.get(urdf_joint_name) + if orch_name is None or orch_name not in name_to_idx: + return None + idx = name_to_idx[orch_name] + + if idx >= len(msg.velocity): + return None + velocities.append(msg.velocity[idx]) + + return np.array(velocities, dtype=np.float64) + + def get_current_positions(self) -> NDArray[np.float64] | None: + """Get current joint positions (thread-safe). + + Returns: + Current positions or None if not yet received + """ + with self._lock: + return self._latest_positions.copy() if self._latest_positions is not None else None + + def get_current_velocities(self) -> NDArray[np.float64] | None: + """Get current joint velocities (thread-safe). + + Returns: + Current velocities or None if not available + """ + with self._lock: + return self._latest_velocities.copy() if self._latest_velocities is not None else None + + def wait_for_state(self, timeout: float | None = None) -> bool: + """Wait until a state is received. + + Args: + timeout: Maximum time to wait (uses default if None) + + Returns: + True if state was received, False if timeout + """ + timeout = timeout if timeout is not None else self._timeout + start_time = time.time() + + while time.time() - start_time < timeout: + with self._lock: + if self._latest_positions is not None: + return True + time.sleep(0.01) + + return False + + def get_state_age(self) -> float | None: + """Get age of the latest state in seconds. + + Returns: + Age in seconds or None if no state received + """ + with self._lock: + if self._last_update_time is None: + return None + return time.time() - self._last_update_time + + def is_state_stale(self, max_age: float = 1.0) -> bool: + """Check if state is stale (older than max_age). + + Args: + max_age: Maximum acceptable age in seconds + + Returns: + True if state is stale or not received + """ + age = self.get_state_age() + if age is None: + return True + return age > max_age + + def add_state_callback( + self, + callback: Callable[[str, NDArray[np.float64]], None], + ) -> None: + """Add callback for state updates. + + Args: + callback: Function called with (robot_id, positions) on each update + """ + self._state_callbacks.append(callback) + + def remove_state_callback( + self, + callback: Callable[[str, NDArray[np.float64]], None], + ) -> None: + """Remove a state callback.""" + if callback in self._state_callbacks: + self._state_callbacks.remove(callback) From d149497f436646a8e86f39f838865a49d068c240 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Mon, 12 Jan 2026 16:11:14 -0800 Subject: [PATCH 06/31] FK and IK solver implementation with drake --- .../planning/kinematics/__init__.py | 27 + .../planning/kinematics/drake_kinematics.py | 503 ++++++++++++++++++ 2 files changed, 530 insertions(+) create mode 100644 dimos/manipulation/planning/kinematics/__init__.py create mode 100644 dimos/manipulation/planning/kinematics/drake_kinematics.py diff --git a/dimos/manipulation/planning/kinematics/__init__.py b/dimos/manipulation/planning/kinematics/__init__.py new file mode 100644 index 0000000000..ec34b7a53a --- /dev/null +++ b/dimos/manipulation/planning/kinematics/__init__.py @@ -0,0 +1,27 @@ +# Copyright 2025 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. + +""" +Kinematics Module + +Contains IK solver implementations that use WorldSpec. + +## Implementations + +- DrakeKinematics: Uses Drake's InverseKinematics + SNOPT/IPOPT +""" + +from dimos.manipulation.planning.kinematics.drake_kinematics import DrakeKinematics + +__all__ = ["DrakeKinematics"] diff --git a/dimos/manipulation/planning/kinematics/drake_kinematics.py b/dimos/manipulation/planning/kinematics/drake_kinematics.py new file mode 100644 index 0000000000..952630ac53 --- /dev/null +++ b/dimos/manipulation/planning/kinematics/drake_kinematics.py @@ -0,0 +1,503 @@ +# 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. + +""" +Drake Kinematics Module + +Implements KinematicsSpec using Drake's optimization-based IK. + +## IK Methods + +- solve(): Full nonlinear IK via Drake's InverseKinematics + SNOPT/IPOPT +- solve_iterative(): Iterative differential IK loop +- solve_differential(): Single Jacobian step for velocity control + +## Key Design + +This module is stateless (except for configuration) and uses WorldSpec +for all FK/collision operations via scratch_context(). + +Example: + kinematics = DrakeKinematics(damping=0.01) + result = kinematics.solve(world, robot_id, target_pose, seed=q_current) + if result.is_success(): + q_goal = result.joint_positions +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +from dimos.manipulation.planning.spec import IKResult, IKStatus, WorldSpec +from dimos.manipulation.planning.utils.kinematics_utils import ( + check_singularity, + compute_error_twist, + compute_pose_error, + damped_pseudoinverse, +) +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from numpy.typing import NDArray + +try: + from pydrake.math import RigidTransform, RotationMatrix + from pydrake.multibody.inverse_kinematics import InverseKinematics + from pydrake.solvers import Solve + + DRAKE_AVAILABLE = True +except ImportError: + DRAKE_AVAILABLE = False + +logger = setup_logger() + + +class DrakeKinematics: + """Drake implementation of KinematicsSpec. + + Uses Drake's InverseKinematics class for optimization-based IK + and Jacobian-based methods for differential IK. + + This class is stateless except for configuration - it uses WorldSpec's + scratch_context() for all operations, making it thread-safe. + + ## Methods + + - solve(): Full optimization-based IK + - solve_iterative(): Iterative differential IK loop + - solve_differential(): Single Jacobian step for velocity control + """ + + def __init__( + self, + damping: float = 0.01, + max_iterations: int = 100, + singularity_threshold: float = 0.001, + ): + """Create Drake kinematics solver. + + Args: + damping: Damping factor for differential IK (higher = more stable near singularities) + max_iterations: Default max iterations for iterative IK + singularity_threshold: Manipulability threshold for singularity detection + """ + if not DRAKE_AVAILABLE: + raise ImportError("Drake is not installed. Install with: pip install drake") + + self._damping = damping + self._max_iterations = max_iterations + self._singularity_threshold = singularity_threshold + + def solve( + self, + world: WorldSpec, + robot_id: str, + target_pose: NDArray[np.float64], + seed: NDArray[np.float64] | None = None, + position_tolerance: float = 0.001, + orientation_tolerance: float = 0.01, + check_collision: bool = True, + max_attempts: int = 10, + ) -> IKResult: + """Full nonlinear IK via Drake's InverseKinematics. + + Uses Drake's optimization-based IK with multiple random restarts + for robustness. + + Args: + world: World for FK/collision checking + robot_id: Which robot + target_pose: 4x4 target transform + seed: Initial guess (uses current state if None) + position_tolerance: Position tolerance (meters) + orientation_tolerance: Orientation tolerance (radians) + check_collision: Verify solution is collision-free + max_attempts: Random restarts for robustness + + Returns: + IKResult with status, joint positions, and error metrics + """ + # Access Drake internals via world + # (DrakeWorld exposes plant/diagram for this purpose) + from dimos.manipulation.planning.world.drake_world import DrakeWorld + + if not isinstance(world, DrakeWorld): + return _create_failure_result( + IKStatus.NO_SOLUTION, + "DrakeKinematics requires DrakeWorld", + ) + + if not world.is_finalized: + return _create_failure_result( + IKStatus.NO_SOLUTION, + "World must be finalized before IK", + ) + + # Get joint limits + lower_limits, upper_limits = world.get_joint_limits(robot_id) + + # Get seed from current state if not provided + if seed is None: + with world.scratch_context() as ctx: + seed = world.get_positions(ctx, robot_id) + + # Target transform + target_transform = RigidTransform(target_pose) + + best_result: IKResult | None = None + best_error = float("inf") + + for attempt in range(max_attempts): + # Generate seed + if attempt == 0: + current_seed = seed + else: + # Random seed within joint limits + current_seed = np.random.uniform(lower_limits, upper_limits) + + # Solve IK + result = self._solve_single( + world=world, + robot_id=robot_id, + target_transform=target_transform, + seed=current_seed, + position_tolerance=position_tolerance, + orientation_tolerance=orientation_tolerance, + lower_limits=lower_limits, + upper_limits=upper_limits, + ) + + if result.is_success() and result.joint_positions is not None: + # Check collision if requested + if check_collision: + with world.scratch_context() as ctx: + world.set_positions(ctx, robot_id, result.joint_positions) + if not world.is_collision_free(ctx, robot_id): + continue # Try another seed + + # Check error + total_error = result.position_error + result.orientation_error + if total_error < best_error: + best_error = total_error + best_result = result + + # If error is within tolerance, we're done + if ( + result.position_error <= position_tolerance + and result.orientation_error <= orientation_tolerance + ): + return result + + if best_result is not None: + return best_result + + return _create_failure_result( + IKStatus.NO_SOLUTION, + f"IK failed after {max_attempts} attempts", + ) + + def _solve_single( + self, + world: WorldSpec, + robot_id: str, + target_transform: RigidTransform, + seed: NDArray[np.float64], + position_tolerance: float, + orientation_tolerance: float, + lower_limits: NDArray[np.float64], + upper_limits: NDArray[np.float64], + ) -> IKResult: + """Solve IK with a single seed.""" + # Get robot data from world internals (Drake-specific access) + robot_data = world._robots[robot_id] # type: ignore[attr-defined] + plant = world.plant # type: ignore[attr-defined] + + # Create IK problem + ik = InverseKinematics(plant) + + # Get end-effector frame + ee_frame = robot_data.ee_frame + + # Add position constraint + ik.AddPositionConstraint( + frameB=ee_frame, + p_BQ=np.array([0.0, 0.0, 0.0]), + frameA=plant.world_frame(), + p_AQ_lower=target_transform.translation() - np.array([position_tolerance] * 3), + p_AQ_upper=target_transform.translation() + np.array([position_tolerance] * 3), + ) + + # Add orientation constraint + ik.AddOrientationConstraint( + frameAbar=plant.world_frame(), + R_AbarA=target_transform.rotation(), + frameBbar=ee_frame, + R_BbarB=RotationMatrix(), + theta_bound=orientation_tolerance, + ) + + # Get program and set initial guess + prog = ik.get_mutable_prog() + q = ik.q() + + # Set initial guess (full positions vector) + full_seed = np.zeros(plant.num_positions()) + for i, joint_idx in enumerate(robot_data.joint_indices): + full_seed[joint_idx] = seed[i] + prog.SetInitialGuess(q, full_seed) + + # Solve + result = Solve(prog) + + if not result.is_success(): + return _create_failure_result( + IKStatus.NO_SOLUTION, + f"Optimization failed: {result.get_solution_result()}", + ) + + # Extract solution for this robot's joints + full_solution = result.GetSolution(q) + joint_solution = np.array([full_solution[idx] for idx in robot_data.joint_indices]) + + # Clip to limits + joint_solution = np.clip(joint_solution, lower_limits, upper_limits) + + # Compute actual error using FK + with world.scratch_context() as ctx: + world.set_positions(ctx, robot_id, joint_solution) + actual_pose = world.get_ee_pose(ctx, robot_id) + + position_error, orientation_error = compute_pose_error( + actual_pose, + target_transform.GetAsMatrix4(), + ) + + return _create_success_result( + joint_positions=joint_solution, + position_error=position_error, + orientation_error=orientation_error, + iterations=1, + ) + + def solve_iterative( + self, + world: WorldSpec, + robot_id: str, + target_pose: NDArray[np.float64], + seed: NDArray[np.float64], + max_iterations: int = 100, + position_tolerance: float = 0.001, + orientation_tolerance: float = 0.01, + ) -> IKResult: + """Iterative differential IK loop. + + Uses repeated Jacobian steps until convergence. Slower but more + predictable behavior near singularities. + + Args: + world: World for FK/Jacobian computation + robot_id: Which robot + target_pose: 4x4 target transform + seed: Initial joint configuration + max_iterations: Maximum iterations + position_tolerance: Position convergence tolerance (meters) + orientation_tolerance: Orientation convergence tolerance (radians) + + Returns: + IKResult with solution or failure info + """ + max_iterations = max_iterations or self._max_iterations + current_joints = seed.copy() + + lower_limits, upper_limits = world.get_joint_limits(robot_id) + + for iteration in range(max_iterations): + with world.scratch_context() as ctx: + # Set current position + world.set_positions(ctx, robot_id, current_joints) + + # Get current pose + current_pose = world.get_ee_pose(ctx, robot_id) + + # Compute error + pos_error, ori_error = compute_pose_error(current_pose, target_pose) + + # Check convergence + if pos_error <= position_tolerance and ori_error <= orientation_tolerance: + return _create_success_result( + joint_positions=current_joints, + position_error=pos_error, + orientation_error=ori_error, + iterations=iteration + 1, + ) + + # Compute twist to reduce error + twist = compute_error_twist(current_pose, target_pose, gain=0.5) + + # Get Jacobian + J = world.get_jacobian(ctx, robot_id) + + # Check for singularity + if check_singularity(J, threshold=self._singularity_threshold): + return _create_failure_result( + IKStatus.SINGULARITY, + f"Singularity at iteration {iteration}", + iterations=iteration + 1, + ) + + # Compute joint velocities + J_pinv = damped_pseudoinverse(J, self._damping) + q_dot = J_pinv @ twist + + # Step size (adaptive based on error) + step_size = min(0.5, pos_error + ori_error) + current_joints = current_joints + step_size * q_dot + + # Clip to limits + current_joints = np.clip(current_joints, lower_limits, upper_limits) + + # Compute final error + with world.scratch_context() as ctx: + world.set_positions(ctx, robot_id, current_joints) + final_pose = world.get_ee_pose(ctx, robot_id) + pos_error, ori_error = compute_pose_error(final_pose, target_pose) + + return _create_failure_result( + IKStatus.NO_SOLUTION, + f"Did not converge after {max_iterations} iterations (pos_err={pos_error:.4f}, ori_err={ori_error:.4f})", + iterations=max_iterations, + ) + + def solve_differential( + self, + world: WorldSpec, + robot_id: str, + current_joints: NDArray[np.float64], + twist: NDArray[np.float64], + dt: float, + ) -> NDArray[np.float64] | None: + """Single Jacobian step for velocity control. + + Computes joint velocities to achieve a desired end-effector twist. + Uses damped pseudoinverse for singularity robustness. + + Args: + world: World for Jacobian computation + robot_id: Which robot + current_joints: Current joint positions (radians) + twist: Desired end-effector twist [vx, vy, vz, wx, wy, wz] + dt: Time step (seconds) - used for velocity scaling + + Returns: + Joint velocities (rad/s) or None if near singularity + """ + with world.scratch_context() as ctx: + world.set_positions(ctx, robot_id, current_joints) + J = world.get_jacobian(ctx, robot_id) + + # Check for singularity + if check_singularity(J, threshold=self._singularity_threshold): + logger.warning("Near singularity in differential IK") + return None + + # Compute damped pseudoinverse + J_pinv = damped_pseudoinverse(J, self._damping) + + # Compute joint velocities + q_dot = J_pinv @ twist + + # Apply velocity limits if available + config = world.get_robot_config(robot_id) + if config.velocity_limits is not None: + max_ratio = np.max(np.abs(q_dot) / np.array(config.velocity_limits)) + if max_ratio > 1.0: + q_dot = q_dot / max_ratio + + return q_dot + + def solve_differential_position_only( + self, + world: WorldSpec, + robot_id: str, + current_joints: NDArray[np.float64], + linear_velocity: NDArray[np.float64], + ) -> NDArray[np.float64] | None: + """Solve differential IK for position-only control. + + Uses only the linear part of the Jacobian (first 3 rows). + + Args: + world: World for Jacobian computation + robot_id: Which robot + current_joints: Current joint positions + linear_velocity: Desired linear velocity [vx, vy, vz] + + Returns: + Joint velocities or None if near singularity + """ + with world.scratch_context() as ctx: + world.set_positions(ctx, robot_id, current_joints) + J = world.get_jacobian(ctx, robot_id) + + # Extract linear part (first 3 rows) + J_linear = J[:3, :] + + # Check for singularity + JJT = J_linear @ J_linear.T + manipulability = np.sqrt(max(0, np.linalg.det(JJT))) + if manipulability < self._singularity_threshold: + return None + + # Compute damped pseudoinverse + I = np.eye(3) + J_pinv = J_linear.T @ np.linalg.inv(JJT + self._damping**2 * I) + + # Compute joint velocities + return J_pinv @ linear_velocity + + +# ============= Result Helpers ============= + + +def _create_success_result( + joint_positions: NDArray[np.float64], + position_error: float, + orientation_error: float, + iterations: int, +) -> IKResult: + """Create a successful IK result.""" + return IKResult( + status=IKStatus.SUCCESS, + joint_positions=joint_positions, + position_error=position_error, + orientation_error=orientation_error, + iterations=iterations, + message="IK solution found", + ) + + +def _create_failure_result( + status: IKStatus, + message: str, + iterations: int = 0, +) -> IKResult: + """Create a failed IK result.""" + return IKResult( + status=status, + joint_positions=None, + iterations=iterations, + message=message, + ) From f89f91c6d1671bb13450f124fda49b8c2ce0428b Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Mon, 12 Jan 2026 16:17:02 -0800 Subject: [PATCH 07/31] RRT path planner using default drake planner --- .../planning/planners/__init__.py | 42 ++ .../planning/planners/drake_planner.py | 607 ++++++++++++++++++ 2 files changed, 649 insertions(+) create mode 100644 dimos/manipulation/planning/planners/__init__.py create mode 100644 dimos/manipulation/planning/planners/drake_planner.py diff --git a/dimos/manipulation/planning/planners/__init__.py b/dimos/manipulation/planning/planners/__init__.py new file mode 100644 index 0000000000..4323896e4e --- /dev/null +++ b/dimos/manipulation/planning/planners/__init__.py @@ -0,0 +1,42 @@ +# Copyright 2025 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. + +""" +Motion Planners Module + +Contains motion planning implementations that use WorldSpec. + +## Implementations + +- DrakePlanner: RRT-Connect planner +- DrakeRRTStarPlanner: RRT* planner (asymptotically optimal) + +## Usage + +Use factory functions to create planners: + +```python +from dimos.manipulation.planning.factory import create_planner + +planner = create_planner(name="rrt_connect") # Returns PlannerSpec +result = planner.plan_joint_path(world, robot_id, q_start, q_goal) +``` +""" + +from dimos.manipulation.planning.planners.drake_planner import ( + DrakePlanner, + DrakeRRTStarPlanner, +) + +__all__ = ["DrakePlanner", "DrakeRRTStarPlanner"] diff --git a/dimos/manipulation/planning/planners/drake_planner.py b/dimos/manipulation/planning/planners/drake_planner.py new file mode 100644 index 0000000000..89eb71ff73 --- /dev/null +++ b/dimos/manipulation/planning/planners/drake_planner.py @@ -0,0 +1,607 @@ +# 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. + +""" +Drake Planner Module + +Implements PlannerSpec using RRT-Connect algorithm with WorldSpec for collision checking. + +## Key Design + +- Uses WorldSpec.scratch_context() for thread-safe collision checking +- Joint-space planning only (use KinematicsSpec.solve() for pose goals first) +- Stateless except for configuration + +Example: + planner = DrakePlanner(step_size=0.1) + result = planner.plan_joint_path(world, robot_id, q_start, q_goal) + if result.is_success(): + waypoints = result.path +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import time +from typing import TYPE_CHECKING + +import numpy as np + +from dimos.manipulation.planning.spec import PlanningResult, PlanningStatus, WorldSpec +from dimos.manipulation.planning.utils.path_utils import ( + compute_path_length, + interpolate_path, + interpolate_segment, +) +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from numpy.typing import NDArray + +logger = setup_logger() + + +@dataclass +class TreeNode: + """Node in the RRT tree.""" + + config: NDArray[np.float64] + parent: TreeNode | None = None + children: list[TreeNode] = field(default_factory=list) + + def path_to_root(self) -> list[NDArray[np.float64]]: + """Get path from this node to root.""" + path = [] + node: TreeNode | None = self + while node is not None: + path.append(node.config) + node = node.parent + return list(reversed(path)) + + +class DrakePlanner: + """RRT-Connect planner implementing PlannerSpec. + + Uses bi-directional RRT (RRT-Connect) for joint-space path planning + with collision checking via WorldSpec.scratch_context(). + + ## Algorithm + + 1. Initialize start tree and goal tree + 2. Extend start tree toward random sample + 3. Try to connect goal tree to new node + 4. Swap trees and repeat + 5. If connected, extract path + + ## Thread Safety + + All collision checking uses world.scratch_context() for thread safety. + """ + + def __init__( + self, + step_size: float = 0.1, + connect_step_size: float = 0.05, + goal_tolerance: float = 0.1, + collision_step_size: float = 0.02, + ): + """Create RRT-Connect planner. + + Args: + step_size: Extension step size in joint space (radians) + connect_step_size: Step size for connect attempts + goal_tolerance: Distance to goal to consider success + collision_step_size: Step size for collision checking along edges + """ + self._step_size = step_size + self._connect_step_size = connect_step_size + self._goal_tolerance = goal_tolerance + self._collision_step_size = collision_step_size + + def plan_joint_path( + self, + world: WorldSpec, + robot_id: str, + q_start: NDArray[np.float64], + q_goal: NDArray[np.float64], + timeout: float = 10.0, + max_iterations: int = 5000, + ) -> PlanningResult: + """Plan collision-free joint-space path using RRT-Connect. + + Args: + world: World for collision checking + robot_id: Which robot + q_start: Start joint configuration (radians) + q_goal: Goal joint configuration (radians) + timeout: Maximum planning time (seconds) + max_iterations: Maximum iterations + + Returns: + PlanningResult with path (list of waypoints) or failure info + """ + start_time = time.time() + + # Validate inputs + error = self._validate_inputs(world, robot_id, q_start, q_goal) + if error is not None: + return error + + # Get joint limits + lower_limits, upper_limits = world.get_joint_limits(robot_id) + + # Initialize trees + start_tree = [TreeNode(config=q_start.copy())] + goal_tree = [TreeNode(config=q_goal.copy())] + + trees_swapped = False # Track if trees have been swapped an odd number of times + + for iteration in range(max_iterations): + # Check timeout + if time.time() - start_time > timeout: + return _create_failure_result( + PlanningStatus.TIMEOUT, + f"Timeout after {iteration} iterations", + planning_time=time.time() - start_time, + iterations=iteration, + ) + + # Sample random configuration + sample = np.random.uniform(lower_limits, upper_limits) + + # Extend start tree toward sample + extended_node = self._extend_tree(world, robot_id, start_tree, sample, self._step_size) + + if extended_node is not None: + # Try to connect goal tree to extended node + connected_node = self._connect_tree( + world, robot_id, goal_tree, extended_node.config, self._connect_step_size + ) + + if connected_node is not None: + # Trees connected! Extract path + path = self._extract_path(extended_node, connected_node) + + # If trees were swapped, path is from goal to start - reverse it + if trees_swapped: + path = list(reversed(path)) + + # Simplify path + path = self._simplify_path(world, robot_id, path) + + return _create_success_result( + path=path, + planning_time=time.time() - start_time, + iterations=iteration + 1, + ) + + # Swap trees + start_tree, goal_tree = goal_tree, start_tree + trees_swapped = not trees_swapped + + return _create_failure_result( + PlanningStatus.NO_SOLUTION, + f"No path found after {max_iterations} iterations", + planning_time=time.time() - start_time, + iterations=max_iterations, + ) + + def get_name(self) -> str: + """Get planner name.""" + return "RRTConnect" + + def _validate_inputs( + self, + world: WorldSpec, + robot_id: str, + q_start: NDArray[np.float64], + q_goal: NDArray[np.float64], + ) -> PlanningResult | None: + """Validate planning inputs, returns error result or None if valid.""" + # Check world is finalized + if not world.is_finalized: + return _create_failure_result( + PlanningStatus.NO_SOLUTION, + "World must be finalized before planning", + ) + + # Check robot exists + if robot_id not in world.get_robot_ids(): + return _create_failure_result( + PlanningStatus.NO_SOLUTION, + f"Robot '{robot_id}' not found", + ) + + # Check start validity + with world.scratch_context() as ctx: + world.set_positions(ctx, robot_id, q_start) + if not world.is_collision_free(ctx, robot_id): + return _create_failure_result( + PlanningStatus.COLLISION_AT_START, + "Start configuration is in collision", + ) + + # Check goal validity + with world.scratch_context() as ctx: + world.set_positions(ctx, robot_id, q_goal) + if not world.is_collision_free(ctx, robot_id): + return _create_failure_result( + PlanningStatus.COLLISION_AT_GOAL, + "Goal configuration is in collision", + ) + + # Check limits + lower, upper = world.get_joint_limits(robot_id) + if np.any(q_start < lower) or np.any(q_start > upper): + return _create_failure_result( + PlanningStatus.INVALID_START, + "Start configuration is outside joint limits", + ) + + if np.any(q_goal < lower) or np.any(q_goal > upper): + return _create_failure_result( + PlanningStatus.INVALID_GOAL, + "Goal configuration is outside joint limits", + ) + + return None + + def _extend_tree( + self, + world: WorldSpec, + robot_id: str, + tree: list[TreeNode], + target: NDArray[np.float64], + step_size: float, + ) -> TreeNode | None: + """Extend tree toward target, returns new node if successful.""" + # Find nearest node + nearest = min(tree, key=lambda n: float(np.linalg.norm(n.config - target))) + + # Compute new config + diff = target - nearest.config + dist = float(np.linalg.norm(diff)) + + if dist <= step_size: + new_config = target.copy() + else: + new_config = nearest.config + step_size * (diff / dist) + + # Check validity of edge + if self._is_edge_valid(world, robot_id, nearest.config, new_config): + new_node = TreeNode(config=new_config, parent=nearest) + nearest.children.append(new_node) + tree.append(new_node) + return new_node + + return None + + def _connect_tree( + self, + world: WorldSpec, + robot_id: str, + tree: list[TreeNode], + target: NDArray[np.float64], + step_size: float, + ) -> TreeNode | None: + """Try to connect tree to target, returns connected node if successful.""" + # Keep extending toward target + while True: + result = self._extend_tree(world, robot_id, tree, target, step_size) + + if result is None: + return None # Extension failed + + # Check if reached target + if float(np.linalg.norm(result.config - target)) < self._goal_tolerance: + return result + + def _is_edge_valid( + self, + world: WorldSpec, + robot_id: str, + q_start: NDArray[np.float64], + q_end: NDArray[np.float64], + ) -> bool: + """Check if edge between two configurations is collision-free.""" + # Interpolate and check each point + segment = interpolate_segment(q_start, q_end, self._collision_step_size) + + with world.scratch_context() as ctx: + for q in segment: + world.set_positions(ctx, robot_id, q) + if not world.is_collision_free(ctx, robot_id): + return False + + return True + + def _extract_path( + self, + start_node: TreeNode, + goal_node: TreeNode, + ) -> list[NDArray[np.float64]]: + """Extract path from two connected nodes.""" + # Path from start node to its root (reversed to be root->node) + start_path = start_node.path_to_root() + + # Path from goal node to its root + goal_path = goal_node.path_to_root() + + # Combine: start_root -> start_node -> goal_node -> goal_root + # But we need start -> goal, so reverse the goal path + full_path = start_path + list(reversed(goal_path)) + + return full_path + + def _simplify_path( + self, + world: WorldSpec, + robot_id: str, + path: list[NDArray[np.float64]], + max_iterations: int = 100, + ) -> list[NDArray[np.float64]]: + """Simplify path by random shortcutting.""" + if len(path) <= 2: + return path + + simplified = list(path) + + for _ in range(max_iterations): + if len(simplified) <= 2: + break + + # Pick two random indices (at least 2 apart) + i = np.random.randint(0, len(simplified) - 2) + j = np.random.randint(i + 2, len(simplified)) + + # Check if direct connection is valid + if self._is_edge_valid(world, robot_id, simplified[i], simplified[j]): + # Remove intermediate waypoints + simplified = simplified[: i + 1] + simplified[j:] + + return simplified + + +class DrakeRRTStarPlanner: + """RRT* (Optimal RRT) planner implementing PlannerSpec. + + Like RRT but optimizes path cost through rewiring. + Produces asymptotically optimal paths. + """ + + def __init__( + self, + step_size: float = 0.1, + goal_tolerance: float = 0.1, + rewire_radius: float = 0.5, + collision_step_size: float = 0.02, + ): + """Create RRT* planner. + + Args: + step_size: Extension step size + goal_tolerance: Distance to goal to consider success + rewire_radius: Radius for rewiring neighbors + collision_step_size: Step size for collision checking + """ + self._step_size = step_size + self._goal_tolerance = goal_tolerance + self._rewire_radius = rewire_radius + self._collision_step_size = collision_step_size + + def plan_joint_path( + self, + world: WorldSpec, + robot_id: str, + q_start: NDArray[np.float64], + q_goal: NDArray[np.float64], + timeout: float = 10.0, + max_iterations: int = 5000, + ) -> PlanningResult: + """Plan optimal collision-free joint-space path using RRT*.""" + start_time = time.time() + + # Get joint limits + lower_limits, upper_limits = world.get_joint_limits(robot_id) + + # Validate start/goal + with world.scratch_context() as ctx: + world.set_positions(ctx, robot_id, q_start) + if not world.is_collision_free(ctx, robot_id): + return _create_failure_result( + PlanningStatus.COLLISION_AT_START, + "Start configuration is in collision", + ) + + world.set_positions(ctx, robot_id, q_goal) + if not world.is_collision_free(ctx, robot_id): + return _create_failure_result( + PlanningStatus.COLLISION_AT_GOAL, + "Goal configuration is in collision", + ) + + # Initialize tree with costs + root = _RRTStarNode(config=q_start.copy(), cost=0.0) + nodes = [root] + goal_node: _RRTStarNode | None = None + best_cost = float("inf") + + for iteration in range(max_iterations): # noqa: B007 + if time.time() - start_time > timeout: + break # Return best path found so far + + # Sample + sample = np.random.uniform(lower_limits, upper_limits) + + # Find nearest + nearest = min(nodes, key=lambda n: float(np.linalg.norm(n.config - sample))) + + # Extend + diff = sample - nearest.config + dist = float(np.linalg.norm(diff)) + if dist > self._step_size: + new_config = nearest.config + self._step_size * (diff / dist) + else: + new_config = sample.copy() + + if not self._is_edge_valid(world, robot_id, nearest.config, new_config): + continue + + # Find neighbors within rewire radius + neighbors = [ + n + for n in nodes + if float(np.linalg.norm(n.config - new_config)) < self._rewire_radius + ] + + # Find best parent + best_parent = nearest + best_new_cost = nearest.cost + float(np.linalg.norm(new_config - nearest.config)) + + for neighbor in neighbors: + new_cost = neighbor.cost + float(np.linalg.norm(new_config - neighbor.config)) + if new_cost < best_new_cost: + if self._is_edge_valid(world, robot_id, neighbor.config, new_config): + best_parent = neighbor + best_new_cost = new_cost + + # Add new node + new_node = _RRTStarNode( + config=new_config, + parent=best_parent, + cost=best_new_cost, + ) + best_parent.children.append(new_node) + nodes.append(new_node) + + # Rewire neighbors through new node + for neighbor in neighbors: + if neighbor == best_parent: + continue + potential_cost = new_node.cost + float( + np.linalg.norm(neighbor.config - new_node.config) + ) + if potential_cost < neighbor.cost: + if self._is_edge_valid(world, robot_id, new_node.config, neighbor.config): + # Rewire + if neighbor.parent is not None: + neighbor.parent.children.remove(neighbor) + neighbor.parent = new_node + neighbor.cost = potential_cost + new_node.children.append(neighbor) + self._update_costs(neighbor) + + # Check goal + if float(np.linalg.norm(new_node.config - q_goal)) < self._goal_tolerance: + if new_node.cost < best_cost: + goal_node = new_node + best_cost = new_node.cost + + if goal_node is not None: + path = goal_node.path_to_root() + path.append(q_goal) + + return _create_success_result( + path=path, + planning_time=time.time() - start_time, + iterations=iteration + 1, + ) + + return _create_failure_result( + PlanningStatus.NO_SOLUTION, + f"No path found after {max_iterations} iterations", + planning_time=time.time() - start_time, + iterations=max_iterations, + ) + + def get_name(self) -> str: + return "RRTStar" + + def _is_edge_valid( + self, + world: WorldSpec, + robot_id: str, + q_start: NDArray[np.float64], + q_end: NDArray[np.float64], + ) -> bool: + """Check if edge is collision-free.""" + segment = interpolate_segment(q_start, q_end, self._collision_step_size) + + with world.scratch_context() as ctx: + for q in segment: + world.set_positions(ctx, robot_id, q) + if not world.is_collision_free(ctx, robot_id): + return False + + return True + + def _update_costs(self, node: _RRTStarNode) -> None: + """Recursively update costs after rewiring.""" + for child in node.children: + child.cost = node.cost + float(np.linalg.norm(child.config - node.config)) + self._update_costs(child) + + +@dataclass +class _RRTStarNode: + """Node for RRT* with cost tracking.""" + + config: NDArray[np.float64] + cost: float = 0.0 + parent: _RRTStarNode | None = None + children: list[_RRTStarNode] = field(default_factory=list) + + def path_to_root(self) -> list[NDArray[np.float64]]: + """Get path from this node to root.""" + path = [] + node: _RRTStarNode | None = self + while node is not None: + path.append(node.config) + node = node.parent + return list(reversed(path)) + + +# ============= Result Helpers ============= + + +def _create_success_result( + path: list[NDArray[np.float64]], + planning_time: float, + iterations: int, +) -> PlanningResult: + """Create a successful planning result.""" + return PlanningResult( + status=PlanningStatus.SUCCESS, + path=path, + planning_time=planning_time, + path_length=compute_path_length(path), + iterations=iterations, + message="Path found", + ) + + +def _create_failure_result( + status: PlanningStatus, + message: str, + planning_time: float = 0.0, + iterations: int = 0, +) -> PlanningResult: + """Create a failed planning result.""" + return PlanningResult( + status=status, + path=[], + planning_time=planning_time, + iterations=iterations, + message=message, + ) From 813db5dd240c036a0bce3b838a4035f6aff80fc7 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Mon, 12 Jan 2026 16:21:04 -0800 Subject: [PATCH 08/31] factory for building implementation of the different world, viz, planner etc specs --- dimos/manipulation/planning/__init__.py | 88 +++++++- dimos/manipulation/planning/factory.py | 265 ++++++++++++++++++++++++ 2 files changed, 351 insertions(+), 2 deletions(-) create mode 100644 dimos/manipulation/planning/factory.py diff --git a/dimos/manipulation/planning/__init__.py b/dimos/manipulation/planning/__init__.py index d197980a96..2c3be92d91 100644 --- a/dimos/manipulation/planning/__init__.py +++ b/dimos/manipulation/planning/__init__.py @@ -15,11 +15,95 @@ """ Manipulation Planning Module -Trajectory generation and motion planning for robotic manipulators. +Motion planning stack for robotic manipulators using Protocol-based architecture. + +## Core Components + +- WorldSpec: Core backend owning physics/collision (DrakeWorld) +- KinematicsSpec: Stateless IK operations (DrakeKinematics) +- PlannerSpec: Joint-space path planning (DrakePlanner) + +## Factory Functions + +Use factory functions to create components: + +```python +from dimos.manipulation.planning.factory import ( + create_world, + create_kinematics, + create_planner, +) + +world = create_world(backend="drake", enable_viz=True) +kinematics = create_kinematics(backend="drake") +planner = create_planner(name="rrt_connect") +``` + +## Monitors + +Use WorldMonitor for reactive state synchronization: + +```python +from dimos.manipulation.planning.monitor import WorldMonitor + +monitor = WorldMonitor(enable_viz=True) +robot_id = monitor.add_robot(config) +monitor.finalize() +monitor.start_state_monitor(robot_id) +``` """ +# Factory functions +from dimos.manipulation.planning.factory import ( + create_kinematics, + create_planner, + create_planning_stack, + create_world, +) + +# Data classes and Protocols +from dimos.manipulation.planning.spec import ( + CollisionObjectMessage, + Detection3D, + IKResult, + IKStatus, + KinematicsSpec, + Obstacle, + ObstacleType, + PlannerSpec, + PlanningResult, + PlanningStatus, + RobotModelConfig, + VizSpec, + WorldSpec, +) + +# Trajectory generation from dimos.manipulation.planning.trajectory_generator.joint_trajectory_generator import ( JointTrajectoryGenerator, ) -__all__ = ["JointTrajectoryGenerator"] +__all__ = [ + # Data classes + "CollisionObjectMessage", + "Detection3D", + "IKResult", + "IKStatus", + # Trajectory + "JointTrajectoryGenerator", + # Protocols + "KinematicsSpec", + "Obstacle", + "ObstacleType", + "PlannerSpec", + "PlanningResult", + "PlanningStatus", + "RobotModelConfig", + "VizSpec", + "WorldSpec", + # Factory functions + "create_kinematics", + "create_planner", + "create_planning_stack", + "create_world", +] diff --git a/dimos/manipulation/planning/factory.py b/dimos/manipulation/planning/factory.py new file mode 100644 index 0000000000..c49a9b874e --- /dev/null +++ b/dimos/manipulation/planning/factory.py @@ -0,0 +1,265 @@ +# 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. + +""" +Factory Functions for Manipulation Planning + +Provides factory functions to create instances of WorldSpec, KinematicsSpec, +PlannerSpec, and VizSpec implementations. + +Only these factory functions know about concrete implementation types. +All other code should use the Protocol types. + +Example: + from dimos.manipulation.planning.factory import ( + create_world, + create_kinematics, + create_planner, + ) + + # Create instances using factory functions + world = create_world(backend="drake", enable_viz=True) + kinematics = create_kinematics(backend="drake") + planner = create_planner(name="rrt_connect", backend="drake") + + # Use them with Protocol types + robot_id = world.add_robot(config) + world.finalize() + + result = kinematics.solve(world, robot_id, target_pose) + plan_result = planner.plan_joint_path(world, robot_id, q_start, q_goal) +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from dimos.manipulation.planning.spec import ( + KinematicsSpec, + PlannerSpec, + VizSpec, + WorldSpec, + ) + + +def create_world( + backend: str = "drake", + enable_viz: bool = False, + **kwargs: Any, +) -> WorldSpec: + """Create a world instance. + + The World owns the physics/collision backend and provides: + - Robot/obstacle management + - Collision checking + - Forward kinematics + - Context management for thread safety + + Args: + backend: Backend to use ("drake", future: "mujoco", "vamp") + enable_viz: Enable Meshcat visualization (Drake only) + **kwargs: Additional backend-specific arguments + - time_step: Simulation time step (default 0.0 for kinematics-only) + + Returns: + WorldSpec implementation + + Example: + world = create_world(backend="drake", enable_viz=True) + robot_id = world.add_robot(robot_config) + world.finalize() + + with world.scratch_context() as ctx: + world.set_positions(ctx, robot_id, q) + if world.is_collision_free(ctx, robot_id): + ee_pose = world.get_ee_pose(ctx, robot_id) + """ + if backend == "drake": + from dimos.manipulation.planning.world.drake_world import DrakeWorld + + return DrakeWorld(enable_viz=enable_viz, **kwargs) + else: + raise ValueError(f"Unknown backend: {backend}. Available: ['drake']") + + +def create_kinematics( + backend: str = "drake", + **kwargs: Any, +) -> KinematicsSpec: + """Create a kinematics solver instance. + + Kinematics solvers are stateless (except for configuration) and + use WorldSpec for all FK/collision operations. + + Args: + backend: Backend to use ("drake") + **kwargs: Additional backend-specific arguments + - damping: Damping factor for differential IK (default 0.01) + - max_iterations: Default max iterations for iterative IK (default 100) + - singularity_threshold: Manipulability threshold (default 0.001) + + Returns: + KinematicsSpec implementation + + Example: + kinematics = create_kinematics(backend="drake", damping=0.01) + result = kinematics.solve(world, robot_id, target_pose, seed=q_current) + if result.is_success(): + q_goal = result.joint_positions + """ + if backend == "drake": + from dimos.manipulation.planning.kinematics.drake_kinematics import ( + DrakeKinematics, + ) + + return DrakeKinematics(**kwargs) + else: + raise ValueError(f"Unknown backend: {backend}. Available: ['drake']") + + +def create_planner( + name: str = "rrt_connect", + backend: str = "drake", + **kwargs: Any, +) -> PlannerSpec: + """Create a motion planner instance. + + Planners find collision-free paths from start to goal configurations. + They use WorldSpec for collision checking and are stateless. + + Args: + name: Planner algorithm ("rrt_connect", "rrt_star") + backend: Backend to use ("drake") + **kwargs: Additional planner-specific arguments + RRT-Connect: + - step_size: Extension step size (default 0.1) + - connect_step_size: Connect step size (default 0.05) + - goal_tolerance: Goal tolerance (default 0.1) + - collision_step_size: Collision check step (default 0.02) + + RRT*: + - step_size: Extension step size (default 0.1) + - goal_tolerance: Goal tolerance (default 0.1) + - rewire_radius: Rewiring radius (default 0.5) + - collision_step_size: Collision check step (default 0.02) + + Returns: + PlannerSpec implementation + + Example: + planner = create_planner(name="rrt_connect", step_size=0.1) + result = planner.plan_joint_path(world, robot_id, q_start, q_goal) + if result.is_success(): + waypoints = result.path + """ + if backend == "drake": + if name == "rrt_connect": + from dimos.manipulation.planning.planners.drake_planner import DrakePlanner + + return DrakePlanner(**kwargs) + elif name == "rrt_star": + from dimos.manipulation.planning.planners.drake_planner import ( + DrakeRRTStarPlanner, + ) + + return DrakeRRTStarPlanner(**kwargs) + else: + raise ValueError( + f"Unknown planner: {name}. Available for Drake: ['rrt_connect', 'rrt_star']" + ) + else: + raise ValueError(f"Unknown backend: {backend}. Available: ['drake']") + + +def create_viz( + backend: str = "drake", + world: WorldSpec | None = None, + **kwargs: Any, +) -> VizSpec: + """Create a visualization instance. + + Provides methods to update robot/obstacle visualization. + Can be integrated with WorldSpec or standalone. + + Note: For Drake, visualization is typically integrated into DrakeWorld + via enable_viz=True. This function is for advanced use cases where + separate visualization is needed. + + Args: + backend: Backend to use ("drake") + world: Optional world to integrate with + **kwargs: Additional backend-specific arguments + + Returns: + VizSpec implementation + + Example: + # Option 1: Integrated visualization (recommended) + world = create_world(backend="drake", enable_viz=True) + print(f"Meshcat URL: {world.get_meshcat_url()}") + + # Option 2: Separate visualization (advanced) + viz = create_viz(backend="drake", world=world) + viz.set_robot_state(robot_id, q) + viz.publish() + """ + if backend == "drake": + # For Drake, visualization is integrated into DrakeWorld + # This is a placeholder for potential future standalone viz + raise NotImplementedError( + "For Drake, use create_world(enable_viz=True) instead. " + "Separate visualization is not yet implemented." + ) + else: + raise ValueError(f"Unknown backend: {backend}. Available: ['drake']") + + +# ============= Convenience Functions ============= + + +def create_planning_stack( + robot_config: Any, + enable_viz: bool = False, + planner_name: str = "rrt_connect", +) -> tuple[WorldSpec, KinematicsSpec, PlannerSpec, str]: + """Convenience function to create a complete planning stack. + + Creates world, kinematics, and planner, adds robot, and finalizes. + + Args: + robot_config: RobotModelConfig for the robot + enable_viz: Enable visualization + planner_name: Planner algorithm to use + + Returns: + Tuple of (world, kinematics, planner, robot_id) + + Example: + world, kinematics, planner, robot_id = create_planning_stack( + robot_config=piper_config, + enable_viz=True, + ) + + # Now ready to use + result = kinematics.solve(world, robot_id, target_pose) + """ + world = create_world(backend="drake", enable_viz=enable_viz) + kinematics = create_kinematics(backend="drake") + planner = create_planner(name=planner_name, backend="drake") + + robot_id = world.add_robot(robot_config) + world.finalize() + + return world, kinematics, planner, robot_id From c1f083807d1e107de5e61973db94b4fafb0dc14d Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Mon, 12 Jan 2026 18:01:42 -0800 Subject: [PATCH 09/31] utils for ik solver and path planning --- dimos/manipulation/planning/utils/__init__.py | 51 +++ .../planning/utils/kinematics_utils.py | 292 ++++++++++++++++++ .../manipulation/planning/utils/path_utils.py | 288 +++++++++++++++++ 3 files changed, 631 insertions(+) create mode 100644 dimos/manipulation/planning/utils/__init__.py create mode 100644 dimos/manipulation/planning/utils/kinematics_utils.py create mode 100644 dimos/manipulation/planning/utils/path_utils.py diff --git a/dimos/manipulation/planning/utils/__init__.py b/dimos/manipulation/planning/utils/__init__.py new file mode 100644 index 0000000000..9073cf2203 --- /dev/null +++ b/dimos/manipulation/planning/utils/__init__.py @@ -0,0 +1,51 @@ +# Copyright 2025 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. + +""" +Manipulation Planning Utilities + +Standalone utility functions for kinematics and path operations. +These are extracted from the old ABC base classes to enable composition over inheritance. + +## Modules + +- kinematics_utils: Jacobian operations, singularity detection, pose error computation +- path_utils: Path interpolation, simplification, length computation +""" + +from dimos.manipulation.planning.utils.kinematics_utils import ( + check_singularity, + compute_error_twist, + compute_pose_error, + damped_pseudoinverse, + get_manipulability, +) +from dimos.manipulation.planning.utils.path_utils import ( + compute_path_length, + interpolate_path, + interpolate_segment, +) + +__all__ = [ + "check_singularity", + "compute_error_twist", + "compute_path_length", + "compute_pose_error", + # Kinematics utilities + "damped_pseudoinverse", + "get_manipulability", + # Path utilities + "interpolate_path", + "interpolate_segment", +] diff --git a/dimos/manipulation/planning/utils/kinematics_utils.py b/dimos/manipulation/planning/utils/kinematics_utils.py new file mode 100644 index 0000000000..d86253f037 --- /dev/null +++ b/dimos/manipulation/planning/utils/kinematics_utils.py @@ -0,0 +1,292 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Kinematics Utilities + +Standalone utility functions for inverse kinematics operations. +These functions are stateless and can be used by any IK solver implementation. + +## Functions + +- damped_pseudoinverse(): Compute damped pseudoinverse of Jacobian +- check_singularity(): Check if Jacobian is near singularity +- get_manipulability(): Compute manipulability measure +- compute_pose_error(): Compute position/orientation error between poses +- compute_error_twist(): Compute error twist for differential IK +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +if TYPE_CHECKING: + from numpy.typing import NDArray + + +def damped_pseudoinverse( + J: NDArray[np.float64], + damping: float = 0.01, +) -> NDArray[np.float64]: + """Compute damped pseudoinverse of Jacobian. + + Uses the damped least-squares formula: + J_pinv = J^T @ (J @ J^T + λ²I)^(-1) + + This avoids numerical issues near singularities where J @ J^T becomes + ill-conditioned. The damping factor λ controls the trade-off between + accuracy and stability. + + Args: + J: 6 x n Jacobian matrix (rows: [vx, vy, vz, wx, wy, wz]) + damping: Damping factor λ (higher = more regularization, more stable) + + Returns: + n x 6 pseudoinverse matrix + + Example: + J = world.get_jacobian(ctx, robot_id) + J_pinv = damped_pseudoinverse(J, damping=0.01) + q_dot = J_pinv @ twist + """ + JJT = J @ J.T + I = np.eye(JJT.shape[0]) + return J.T @ np.linalg.inv(JJT + damping**2 * I) + + +def check_singularity( + J: NDArray[np.float64], + threshold: float = 0.01, +) -> bool: + """Check if Jacobian is near singularity. + + Computes the manipulability measure (sqrt(det(J @ J^T))) and checks + if it's below the threshold. Near singularities, the manipulability + approaches zero. + + Args: + J: 6 x n Jacobian matrix + threshold: Manipulability threshold (default 0.01) + + Returns: + True if near singularity (manipulability < threshold) + + Example: + J = world.get_jacobian(ctx, robot_id) + if check_singularity(J, threshold=0.001): + logger.warning("Near singularity, using damped IK") + """ + return get_manipulability(J) < threshold + + +def get_manipulability(J: NDArray[np.float64]) -> float: + """Compute manipulability measure. + + The manipulability measure w = sqrt(det(J @ J^T)) represents the + volume of the velocity ellipsoid - how well the robot can move + in all directions. + + Values: + - Higher = better manipulability + - Zero = singularity + - Lower = near singularity + + Args: + J: 6 x n Jacobian matrix + + Returns: + Manipulability measure (non-negative) + + Example: + J = world.get_jacobian(ctx, robot_id) + w = get_manipulability(J) + print(f"Manipulability: {w:.4f}") + """ + JJT = J @ J.T + det = np.linalg.det(JJT) + return float(np.sqrt(max(0, det))) + + +def compute_pose_error( + current_pose: NDArray[np.float64], + target_pose: NDArray[np.float64], +) -> tuple[float, float]: + """Compute position and orientation error between two poses. + + Position error is the Euclidean distance between origins. + Orientation error is the angle of the rotation matrix relating the two frames. + + Args: + current_pose: Current 4x4 homogeneous transform + target_pose: Target 4x4 homogeneous transform + + Returns: + Tuple of (position_error, orientation_error) in meters and radians + + Example: + current = world.get_ee_pose(ctx, robot_id) + pos_err, ori_err = compute_pose_error(current, target) + converged = pos_err < 0.001 and ori_err < 0.01 + """ + # Position error (Euclidean distance) + position_error = float(np.linalg.norm(target_pose[:3, 3] - current_pose[:3, 3])) + + # Orientation error using rotation matrices + R_current = current_pose[:3, :3] + R_target = target_pose[:3, :3] + R_error = R_target @ R_current.T + + # Convert to axis-angle for scalar error + trace = np.trace(R_error) + # Clamp to valid range for arccos (numerical stability) + cos_angle = (trace - 1) / 2 + cos_angle = np.clip(cos_angle, -1, 1) + orientation_error = float(np.arccos(cos_angle)) + + return position_error, orientation_error + + +def compute_error_twist( + current_pose: NDArray[np.float64], + target_pose: NDArray[np.float64], + gain: float = 1.0, +) -> NDArray[np.float64]: + """Compute error twist for differential IK. + + Computes the 6D twist (linear + angular velocity) that would move + from the current pose toward the target pose. Used in iterative IK. + + The twist is expressed in the world frame: + twist = [vx, vy, vz, wx, wy, wz] + + Args: + current_pose: Current 4x4 homogeneous transform + target_pose: Target 4x4 homogeneous transform + gain: Proportional gain (higher = faster convergence, less stable) + + Returns: + 6D twist vector [vx, vy, vz, wx, wy, wz] + + Example: + twist = compute_error_twist(current_pose, target_pose, gain=0.5) + q_dot = damped_pseudoinverse(J) @ twist + q_new = q + q_dot * dt + """ + # Position error (linear velocity direction) + pos_error = target_pose[:3, 3] - current_pose[:3, 3] + + # Orientation error -> angular velocity + R_current = current_pose[:3, :3] + R_target = target_pose[:3, :3] + R_error = R_target @ R_current.T + + # Extract axis-angle from rotation matrix + # Using Rodrigues' formula inverse + trace = np.trace(R_error) + cos_angle = (trace - 1) / 2 + cos_angle = np.clip(cos_angle, -1, 1) + angle = np.arccos(cos_angle) + + if angle < 1e-6: + # No rotation needed + angular_error = np.zeros(3) + elif angle > np.pi - 1e-6: + # 180 degree rotation - extract axis from diagonal + diag = np.diag(R_error) + idx = np.argmax(diag) + axis = np.zeros(3) + axis[idx] = 1.0 + # Refine axis + axis = axis * np.sqrt((diag[idx] + 1) / 2) + angular_error = axis * angle + else: + # General case: axis from skew-symmetric part + sin_angle = np.sin(angle) + axis = np.array( + [ + R_error[2, 1] - R_error[1, 2], + R_error[0, 2] - R_error[2, 0], + R_error[1, 0] - R_error[0, 1], + ] + ) / (2 * sin_angle) + angular_error = axis * angle + + # Combine into twist with gain + twist: NDArray[np.float64] = np.concatenate([pos_error * gain, angular_error * gain]) + + return twist + + +def skew_symmetric(v: NDArray[np.float64]) -> NDArray[np.float64]: + """Create skew-symmetric matrix from 3D vector. + + The skew-symmetric matrix [v]_x satisfies: [v]_x @ w = v cross w + + Args: + v: 3D vector + + Returns: + 3x3 skew-symmetric matrix + """ + return np.array( + [ + [0, -v[2], v[1]], + [v[2], 0, -v[0]], + [-v[1], v[0], 0], + ] + ) + + +def rotation_matrix_to_axis_angle(R: NDArray[np.float64]) -> tuple[NDArray[np.float64], float]: + """Convert rotation matrix to axis-angle representation. + + Args: + R: 3x3 rotation matrix + + Returns: + Tuple of (axis, angle) where axis is unit vector and angle is radians + """ + trace = np.trace(R) + cos_angle = (trace - 1) / 2 + cos_angle = np.clip(cos_angle, -1, 1) + angle = float(np.arccos(cos_angle)) + + if angle < 1e-6: + # Identity rotation + return np.array([1.0, 0.0, 0.0]), 0.0 + + if angle > np.pi - 1e-6: + # 180 degree rotation + diag = np.diag(R) + idx = int(np.argmax(diag)) + axis = np.zeros(3) + axis[idx] = np.sqrt((diag[idx] + 1) / 2) + for j in range(3): + if j != idx: + axis[j] = R[idx, j] / (2 * axis[idx]) + return axis, angle + + # General case + sin_angle = np.sin(angle) + axis = np.array( + [ + R[2, 1] - R[1, 2], + R[0, 2] - R[2, 0], + R[1, 0] - R[0, 1], + ] + ) / (2 * sin_angle) + + return axis, angle diff --git a/dimos/manipulation/planning/utils/path_utils.py b/dimos/manipulation/planning/utils/path_utils.py new file mode 100644 index 0000000000..05d63800a4 --- /dev/null +++ b/dimos/manipulation/planning/utils/path_utils.py @@ -0,0 +1,288 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Path Utilities + +Standalone utility functions for path manipulation and post-processing. +These functions are stateless and can be used by any planner implementation. + +## Functions + +- interpolate_path(): Interpolate path to uniform resolution +- interpolate_segment(): Interpolate between two configurations +- simplify_path(): Remove unnecessary waypoints (requires WorldSpec) +- compute_path_length(): Compute total path length in joint space +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +if TYPE_CHECKING: + from numpy.typing import NDArray + + from dimos.manipulation.planning.spec import WorldSpec + + +def interpolate_path( + path: list[NDArray[np.float64]], + resolution: float = 0.05, +) -> list[NDArray[np.float64]]: + """Interpolate path to have uniform resolution. + + Adds intermediate waypoints so that the maximum joint-space distance + between consecutive waypoints is at most `resolution`. + + Args: + path: Original path (list of joint configurations) + resolution: Maximum distance between waypoints (radians) + + Returns: + Interpolated path with more waypoints + + Example: + # After planning, interpolate for smoother execution + raw_path = planner.plan_joint_path(world, robot_id, q_start, q_goal).path + smooth_path = interpolate_path(raw_path, resolution=0.02) + """ + if len(path) <= 1: + return path + + interpolated = [path[0]] + + for i in range(len(path) - 1): + start = path[i] + end = path[i + 1] + + diff = end - start + max_diff = float(np.max(np.abs(diff))) + + if max_diff <= resolution: + interpolated.append(end) + else: + num_steps = int(np.ceil(max_diff / resolution)) + for step in range(1, num_steps + 1): + alpha = step / num_steps + interpolated.append(start + alpha * diff) + + return interpolated + + +def interpolate_segment( + start: NDArray[np.float64], + end: NDArray[np.float64], + step_size: float, +) -> list[NDArray[np.float64]]: + """Interpolate between two configurations. + + Returns a list of configurations from start to end (inclusive) + with at most `step_size` distance between consecutive points. + + Args: + start: Start joint configuration + end: End joint configuration + step_size: Maximum step size (radians) + + Returns: + List of interpolated configurations [start, ..., end] + + Example: + # Check collision along a segment + segment = interpolate_segment(q1, q2, step_size=0.02) + for q in segment: + if not world.is_collision_free(ctx, robot_id): + return False + """ + diff = end - start + distance = float(np.linalg.norm(diff)) + + if distance <= step_size: + return [start, end] + + num_steps = int(np.ceil(distance / step_size)) + segment = [] + + for i in range(num_steps + 1): + alpha = i / num_steps + segment.append(start + alpha * diff) + + return segment + + +def simplify_path( + world: WorldSpec, + robot_id: str, + path: list[NDArray[np.float64]], + max_iterations: int = 100, + collision_step_size: float = 0.02, +) -> list[NDArray[np.float64]]: + """Simplify path by removing unnecessary waypoints. + + Uses random shortcutting: randomly select two points and check if + the direct connection is collision-free. If so, remove intermediate + waypoints. + + Args: + world: World for collision checking + robot_id: Which robot + path: Original path + max_iterations: Maximum shortcutting attempts + collision_step_size: Step size for collision checking along shortcuts + + Returns: + Simplified path with fewer waypoints + + Example: + raw_path = planner.plan_joint_path(world, robot_id, q_start, q_goal).path + simplified = simplify_path(world, robot_id, raw_path) + """ + if len(path) <= 2: + return path + + simplified = list(path) + + for _ in range(max_iterations): + if len(simplified) <= 2: + break + + # Pick two random indices (at least 2 apart) + i = np.random.randint(0, len(simplified) - 2) + j = np.random.randint(i + 2, len(simplified)) + + # Check if direct connection is valid + with world.scratch_context() as ctx: + segment = interpolate_segment(simplified[i], simplified[j], collision_step_size) + valid = True + + for q in segment: + world.set_positions(ctx, robot_id, q) + if not world.is_collision_free(ctx, robot_id): + valid = False + break + + if valid: + # Remove intermediate waypoints + simplified = simplified[: i + 1] + simplified[j:] + + return simplified + + +def compute_path_length(path: list[NDArray[np.float64]]) -> float: + """Compute total path length in joint space. + + Sums the Euclidean distances between consecutive waypoints. + + Args: + path: Path to measure + + Returns: + Total length in radians + + Example: + length = compute_path_length(path) + print(f"Path length: {length:.2f} rad") + """ + if len(path) <= 1: + return 0.0 + + length = 0.0 + for i in range(len(path) - 1): + length += float(np.linalg.norm(path[i + 1] - path[i])) + + return length + + +def is_path_within_limits( + path: list[NDArray[np.float64]], + lower_limits: NDArray[np.float64], + upper_limits: NDArray[np.float64], +) -> bool: + """Check if all waypoints in path are within joint limits. + + Args: + path: Path to check + lower_limits: Lower joint limits (radians) + upper_limits: Upper joint limits (radians) + + Returns: + True if all waypoints are within limits + """ + for q in path: + if np.any(q < lower_limits) or np.any(q > upper_limits): + return False + return True + + +def clip_path_to_limits( + path: list[NDArray[np.float64]], + lower_limits: NDArray[np.float64], + upper_limits: NDArray[np.float64], +) -> list[NDArray[np.float64]]: + """Clip all waypoints in path to joint limits. + + Args: + path: Path to clip + lower_limits: Lower joint limits (radians) + upper_limits: Upper joint limits (radians) + + Returns: + Path with all waypoints clipped to limits + """ + return [np.clip(q, lower_limits, upper_limits) for q in path] + + +def reverse_path(path: list[NDArray[np.float64]]) -> list[NDArray[np.float64]]: + """Reverse a path (for returning to start, etc.). + + Args: + path: Path to reverse + + Returns: + Reversed path + """ + return list(reversed(path)) + + +def concatenate_paths( + *paths: list[NDArray[np.float64]], + remove_duplicates: bool = True, +) -> list[NDArray[np.float64]]: + """Concatenate multiple paths into one. + + Args: + *paths: Paths to concatenate + remove_duplicates: If True, remove duplicate waypoints at junctions + + Returns: + Single concatenated path + """ + result: list[NDArray[np.float64]] = [] + + for path in paths: + if not path: + continue + + if remove_duplicates and result: + # Check if last point matches first point + if np.allclose(result[-1], path[0]): + result.extend(path[1:]) + else: + result.extend(path) + else: + result.extend(path) + + return result From 00228b72861d29348cf734909289a076dd9917f8 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Tue, 13 Jan 2026 15:47:02 -0800 Subject: [PATCH 10/31] manipulation module manages the world monitor, planner, kinematics planner and interfaces with other modules --- dimos/manipulation/__init__.py | 29 + dimos/manipulation/manipulation_module.py | 583 ++++++++++++++++++ dimos/manipulation/planning/spec.py | 18 + .../planning/world/drake_world.py | 25 + 4 files changed, 655 insertions(+) create mode 100644 dimos/manipulation/manipulation_module.py diff --git a/dimos/manipulation/__init__.py b/dimos/manipulation/__init__.py index e69de29bb2..3ed1863092 100644 --- a/dimos/manipulation/__init__.py +++ b/dimos/manipulation/__init__.py @@ -0,0 +1,29 @@ +# Copyright 2025 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. + +"""Manipulation module for robot arm motion planning and control.""" + +from dimos.manipulation.manipulation_module import ( + ManipulationModule, + ManipulationModuleConfig, + ManipulationState, + manipulation_module, +) + +__all__ = [ + "ManipulationModule", + "ManipulationModuleConfig", + "ManipulationState", + "manipulation_module", +] diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py new file mode 100644 index 0000000000..96564ff892 --- /dev/null +++ b/dimos/manipulation/manipulation_module.py @@ -0,0 +1,583 @@ +# 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. + +"""Manipulation Module - Motion planning with ControlOrchestrator execution.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +import threading +from typing import TYPE_CHECKING + +import numpy as np + +from dimos.core import In, Module, rpc +from dimos.core.module import ModuleConfig +from dimos.manipulation.planning import ( + JointTrajectoryGenerator, + KinematicsSpec, + PlannerSpec, + RobotModelConfig, + create_kinematics, + create_planner, +) +from dimos.manipulation.planning.monitor import WorldMonitor + +# These must be imported at runtime (not TYPE_CHECKING) for In/Out port creation +from dimos.msgs.sensor_msgs import JointState # noqa: TC001 +from dimos.msgs.trajectory_msgs import JointTrajectory +from dimos.utils.logging_config import setup_logger +from dimos.utils.transform_utils import matrix_to_pose, pose_to_matrix + +if TYPE_CHECKING: + from numpy.typing import NDArray + + from dimos.control.orchestrator import ControlOrchestrator + from dimos.core.rpc_client import RPCClient + from dimos.msgs.geometry_msgs import Pose + +logger = setup_logger() + + +class ManipulationState(Enum): + """State machine for manipulation module.""" + + IDLE = 0 + PLANNING = 1 + EXECUTING = 2 + COMPLETED = 3 + FAULT = 4 + + +@dataclass +class ManipulationModuleConfig(ModuleConfig): + """Configuration for ManipulationModule.""" + + robots: list[RobotModelConfig] = field(default_factory=list) + planning_timeout: float = 10.0 + enable_viz: bool = False + + +class ManipulationModule(Module): + """Motion planning module with ControlOrchestrator execution.""" + + default_config = ManipulationModuleConfig + + # Type annotation for the config attribute (mypy uses this) + config: ManipulationModuleConfig + + # Input: Joint state from orchestrator (for world sync) + joint_state: In[JointState] + + def __init__(self, *args: object, **kwargs: object) -> None: + super().__init__(*args, **kwargs) + + # State machine + self._state = ManipulationState.IDLE + self._lock = threading.Lock() + self._error_message = "" + + # Planning components (initialized in start()) + self._world_monitor: WorldMonitor | None = None + self._planner: PlannerSpec | None = None + self._kinematics: KinematicsSpec | None = None + + # Robot registry: maps robot_name -> (world_robot_id, config, trajectory_gen) + self._robots: dict[str, tuple[str, RobotModelConfig, JointTrajectoryGenerator]] = {} + + # Stored path for plan/preview/execute workflow (per robot) + self._planned_paths: dict[str, list[NDArray[np.float64]]] = {} + self._planned_trajectories: dict[str, JointTrajectory] = {} + + # Orchestrator integration (lazy initialized) + self._orchestrator_client: RPCClient[ControlOrchestrator] | None = None + + logger.info("ManipulationModule initialized") + + @rpc + def start(self) -> None: + """Start the manipulation module.""" + super().start() + + # Initialize planning stack + self._initialize_planning() + + # Subscribe to joint state via port + if self.joint_state is not None: + self.joint_state.subscribe(self._on_joint_state) + logger.info("Subscribed to joint_state port") + + logger.info("ManipulationModule started") + + def _initialize_planning(self) -> None: + """Initialize world, planner, and trajectory generator.""" + if not self.config.robots: + logger.warning("No robots configured, planning disabled") + return + + self._world_monitor = WorldMonitor(enable_viz=self.config.enable_viz) + + for robot_config in self.config.robots: + robot_id = self._world_monitor.add_robot(robot_config) + traj_gen = JointTrajectoryGenerator( + num_joints=len(robot_config.joint_names), + max_velocity=robot_config.max_velocity, + max_acceleration=robot_config.max_acceleration, + ) + self._robots[robot_config.name] = (robot_id, robot_config, traj_gen) + + self._world_monitor.finalize() + + for _, (robot_id, _, _) in self._robots.items(): + self._world_monitor.start_state_monitor(robot_id) + + if self.config.enable_viz: + self._world_monitor.start_visualization_thread(rate_hz=10.0) + if url := self._world_monitor.get_meshcat_url(): + logger.info(f"Visualization: {url}") + + self._planner = create_planner(name="rrt_connect") + self._kinematics = create_kinematics(backend="drake") + + def _get_default_robot_name(self) -> str | None: + """Get default robot name (first robot if only one, else None).""" + if len(self._robots) == 1: + return next(iter(self._robots.keys())) + return None + + def _get_robot( + self, robot_name: str | None = None + ) -> tuple[str, str, RobotModelConfig, JointTrajectoryGenerator] | None: + """Get robot by name or default. + + Args: + robot_name: Robot name or None for default (if single robot) + + Returns: + (robot_name, robot_id, config, traj_gen) or None if not found + """ + if robot_name is None: + robot_name = self._get_default_robot_name() + if robot_name is None: + logger.error("Multiple robots configured, must specify robot_name") + return None + + if robot_name not in self._robots: + logger.error(f"Unknown robot: {robot_name}") + return None + + robot_id, config, traj_gen = self._robots[robot_name] + return (robot_name, robot_id, config, traj_gen) + + def _on_joint_state(self, msg: JointState) -> None: + """Callback when joint state received from driver.""" + try: + # Forward to world monitor for state synchronization + # For single robot, use default; for multi-robot, monitor routes by joint names + if self._world_monitor is not None: + robot = self._get_robot() + if robot is not None: + _, robot_id, _, _ = robot + self._world_monitor.on_joint_state(msg, robot_id) + + except Exception as e: + logger.error(f"Exception in _on_joint_state: {e}") + import traceback + + logger.error(traceback.format_exc()) + + # ========================================================================= + # RPC Methods + # ========================================================================= + + @rpc + def get_state(self) -> str: + """Get current manipulation state name.""" + return self._state.name + + @rpc + def get_error(self) -> str: + """Get last error message. + + Returns: + Error message or empty string + """ + return self._error_message + + @rpc + def cancel(self) -> bool: + """Cancel current motion.""" + if self._state != ManipulationState.EXECUTING: + return False + self._state = ManipulationState.IDLE + logger.info("Motion cancelled") + return True + + @rpc + def reset(self) -> bool: + """Reset to IDLE state (fails if EXECUTING).""" + if self._state == ManipulationState.EXECUTING: + return False + self._state = ManipulationState.IDLE + self._error_message = "" + return True + + @rpc + def get_current_joints(self) -> list[float] | None: + """Get current joint positions.""" + if (robot := self._get_robot()) and self._world_monitor: + if pos := self._world_monitor.get_current_positions(robot[1]): + return list(pos) + return None + + @rpc + def get_ee_pose(self) -> Pose | None: + """Get current end-effector pose.""" + if (robot := self._get_robot()) and self._world_monitor: + return matrix_to_pose(self._world_monitor.get_ee_pose(robot[1])) + return None + + @rpc + def is_collision_free(self, joints: list[float]) -> bool: + """Check if joint configuration is collision-free.""" + if (robot := self._get_robot()) and self._world_monitor: + return self._world_monitor.is_state_valid( + robot[1], np.array(joints) + ) # robot[1] is the robot_id. + return False + + # ========================================================================= + # Plan/Preview/Execute Workflow RPC Methods + # ========================================================================= + + def _begin_planning(self) -> tuple[str, str] | None: + """Check state and begin planning. Returns (robot_name, robot_id) or None.""" + if self._world_monitor is None: + logger.error("Planning not initialized") + return None + if (robot := self._get_robot()) is None: + return None + with self._lock: + if self._state not in (ManipulationState.IDLE, ManipulationState.COMPLETED): + logger.warning(f"Cannot plan: state is {self._state.name}") + return None + self._state = ManipulationState.PLANNING + return robot[0], robot[1] + + def _fail(self, msg: str) -> bool: + """Set FAULT state with error message.""" + logger.warning(msg) + self._state = ManipulationState.FAULT + self._error_message = msg + return False + + @rpc + def plan_to_pose(self, pose: Pose) -> bool: + """Plan motion to pose. Use preview_path() then execute().""" + if self._kinematics is None or (r := self._begin_planning()) is None: + return False + robot_name, robot_id = r + assert self._world_monitor # guaranteed by _begin_planning + + current = self._world_monitor.get_current_positions(robot_id) + if current is None: + return self._fail("No joint state") + + ik = self._kinematics.solve( + world=self._world_monitor.world, + robot_id=robot_id, + target_pose=pose_to_matrix(pose), + seed=current, + check_collision=True, + ) + if not ik.is_success() or ik.joint_positions is None: + return self._fail(f"IK failed: {ik.status.name}") + + logger.info(f"IK solved, error: {ik.position_error:.4f}m") + return self._plan_path_only(robot_name, robot_id, ik.joint_positions) + + @rpc + def plan_to_joints(self, joints: list[float]) -> bool: + """Plan motion to joint config. Use preview_path() then execute().""" + if (r := self._begin_planning()) is None: + return False + robot_name, robot_id = r + logger.info(f"Planning to joints: {[f'{j:.3f}' for j in joints]}") + return self._plan_path_only(robot_name, robot_id, np.array(joints)) + + def _plan_path_only(self, robot_name: str, robot_id: str, goal: NDArray[np.float64]) -> bool: + """Plan path from current position to goal, store result.""" + assert self._world_monitor and self._planner # guaranteed by _begin_planning + start = self._world_monitor.get_current_positions(robot_id) + if start is None: + return self._fail("No joint state") + + result = self._planner.plan_joint_path( + world=self._world_monitor.world, + robot_id=robot_id, + q_start=start, + q_goal=goal, + timeout=self.config.planning_timeout, + ) + if not result.is_success(): + return self._fail(f"Planning failed: {result.status.name}") + + logger.info(f"Path: {len(result.path)} waypoints") + self._planned_paths[robot_name] = result.path + + _, _, traj_gen = self._robots[robot_name] + traj = traj_gen.generate([list(q) for q in result.path]) + self._planned_trajectories[robot_name] = traj + logger.info(f"Trajectory: {traj.duration:.3f}s") + + self._state = ManipulationState.COMPLETED + return True + + @rpc + def preview_path(self, duration: float = 3.0) -> bool: + """Preview the planned path in the visualizer. + + Args: + duration: Total animation duration in seconds + """ + from dimos.manipulation.planning.utils.path_utils import interpolate_path + + if self._world_monitor is None: + return False + + robot = self._get_robot() + if robot is None: + return False + robot_name, robot_id, _, _ = robot + + planned_path = self._planned_paths.get(robot_name) + if planned_path is None or len(planned_path) == 0: + logger.warning("No planned path to preview") + return False + + # Interpolate and animate + interpolated = interpolate_path(planned_path, resolution=0.02) + self._world_monitor.world.animate_path(robot_id, interpolated, duration) + return True + + @rpc + def has_planned_path(self) -> bool: + """Check if there's a planned path ready. + + Returns: + True if a path is planned and ready + """ + robot = self._get_robot() + if robot is None: + return False + robot_name, _, _, _ = robot + + path = self._planned_paths.get(robot_name) + return path is not None and len(path) > 0 + + @rpc + def get_visualization_url(self) -> str | None: + """Get the visualization URL. + + Returns: + URL string or None if visualization not enabled + """ + if self._world_monitor is None: + return None + return self._world_monitor.get_meshcat_url() + + @rpc + def clear_planned_path(self) -> bool: + """Clear the stored planned path. + + Returns: + True if cleared + """ + robot = self._get_robot() + if robot is None: + return False + robot_name, _, _, _ = robot + + self._planned_paths.pop(robot_name, None) + self._planned_trajectories.pop(robot_name, None) + return True + + @rpc + def list_robots(self) -> list[str]: + """List all configured robot names. + + Returns: + List of robot names + """ + return list(self._robots.keys()) + + @rpc + def get_robot_info(self, robot_name: str | None = None) -> dict[str, object] | None: + """Get information about a robot. + + Args: + robot_name: Robot name (uses default if None) + + Returns: + Dict with robot info or None if not found + """ + robot = self._get_robot(robot_name) + if robot is None: + return None + + robot_name, robot_id, config, _ = robot + + return { + "name": config.name, + "world_robot_id": robot_id, + "joint_names": config.joint_names, + "end_effector_link": config.end_effector_link, + "base_link": config.base_link, + "max_velocity": config.max_velocity, + "max_acceleration": config.max_acceleration, + "has_joint_name_mapping": bool(config.joint_name_mapping), + "orchestrator_task_name": config.orchestrator_task_name, + } + + # ========================================================================= + # Orchestrator Integration RPC Methods + # ========================================================================= + + def _get_orchestrator_client(self) -> RPCClient[ControlOrchestrator] | None: + """Get or create orchestrator RPC client (lazy init).""" + if not any(c.orchestrator_task_name for _, c, _ in self._robots.values()): + return None + if self._orchestrator_client is None: + from dimos.control.orchestrator import ControlOrchestrator + from dimos.core.rpc_client import RPCClient + + self._orchestrator_client = RPCClient(None, ControlOrchestrator) + return self._orchestrator_client + + def _translate_trajectory_to_orchestrator( + self, + trajectory: JointTrajectory, + robot_config: RobotModelConfig, + ) -> JointTrajectory: + """Translate trajectory joint names from URDF to orchestrator namespace. + + Args: + trajectory: Trajectory with URDF joint names + robot_config: Robot config with joint name mapping + + Returns: + Trajectory with orchestrator joint names + """ + if not robot_config.joint_name_mapping: + return trajectory # No translation needed + + # Translate joint names + orchestrator_names = [ + robot_config.get_orchestrator_joint_name(j) for j in trajectory.joint_names + ] + + # Create new trajectory with translated names + # Note: duration is computed automatically from points in JointTrajectory.__init__ + return JointTrajectory( + joint_names=orchestrator_names, + points=trajectory.points, + timestamp=trajectory.timestamp, + ) + + @rpc + def execute(self, robot_name: str | None = None) -> bool: + """Execute planned trajectory via ControlOrchestrator.""" + if (robot := self._get_robot(robot_name)) is None: + return False + robot_name, _, config, _ = robot + + if (traj := self._planned_trajectories.get(robot_name)) is None: + logger.warning("No planned trajectory") + return False + if not config.orchestrator_task_name: + logger.error(f"No orchestrator_task_name for '{robot_name}'") + return False + if (client := self._get_orchestrator_client()) is None: + logger.error("No orchestrator client") + return False + + translated = self._translate_trajectory_to_orchestrator(traj, config) + logger.info( + f"Executing: task='{config.orchestrator_task_name}', {len(translated.points)} pts, {translated.duration:.2f}s" + ) + + self._state = ManipulationState.EXECUTING + if client.execute_trajectory(config.orchestrator_task_name, translated): + logger.info("Trajectory accepted") + self._state = ManipulationState.COMPLETED + return True + else: + return self._fail("Orchestrator rejected trajectory") + + @rpc + def get_trajectory_status(self, robot_name: str | None = None) -> dict[str, object] | None: + """Get trajectory execution status.""" + if (robot := self._get_robot(robot_name)) is None: + return None + _, _, config, _ = robot + if not config.orchestrator_task_name or (client := self._get_orchestrator_client()) is None: + return None + status = client.get_trajectory_status(config.orchestrator_task_name) + return dict(status) if status else None + + @property + def world_monitor(self) -> WorldMonitor | None: + """Access the world monitor for advanced obstacle/world operations.""" + return self._world_monitor + + @rpc + def add_obstacle(self, name: str, pose: Pose, shape: str, dimensions: list[float]) -> str: + """Add obstacle: shape='box'|'sphere'|'cylinder', dimensions=[w,h,d]|[r]|[r,len].""" + if not self._world_monitor: + return "" + p = pose_to_matrix(pose) + match shape: + case "box": + return self._world_monitor.add_box_obstacle(name, p, tuple(dimensions)) # type: ignore[arg-type] + case "sphere": + return self._world_monitor.add_sphere_obstacle(name, p, dimensions[0]) + case "cylinder": + return self._world_monitor.add_cylinder_obstacle( + name, p, dimensions[0], dimensions[1] + ) + case _: + return "" + + @rpc + def remove_obstacle(self, obstacle_id: str) -> bool: + """Remove an obstacle from the planning world.""" + if self._world_monitor is None: + return False + return self._world_monitor.remove_obstacle(obstacle_id) + + @rpc + def stop(self) -> None: + """Stop the manipulation module.""" + logger.info("Stopping ManipulationModule") + + # Stop world monitor (includes visualization thread) + if self._world_monitor is not None: + self._world_monitor.stop_all_monitors() + + super().stop() + + +# Expose blueprint for declarative composition +manipulation_module = ManipulationModule.blueprint diff --git a/dimos/manipulation/planning/spec.py b/dimos/manipulation/planning/spec.py index c8bf854b4a..66be979f30 100644 --- a/dimos/manipulation/planning/spec.py +++ b/dimos/manipulation/planning/spec.py @@ -542,6 +542,24 @@ def get_jacobian(self, ctx: Any, robot_id: str) -> NDArray[np.float64]: """ ... + # ========================================================================= + # Visualization (optional) + # ========================================================================= + + def get_meshcat_url(self) -> str | None: + """Get Meshcat visualization URL if enabled.""" + ... + + def publish_to_meshcat(self, ctx: Any | None = None) -> None: + """Publish current state to Meshcat visualization.""" + ... + + def animate_path( + self, robot_id: str, path: list[NDArray[np.float64]], duration: float = 3.0 + ) -> None: + """Animate a path in visualization.""" + ... + @runtime_checkable class KinematicsSpec(Protocol): diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index 5d5705a8b0..d7fc5b93fa 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -930,6 +930,31 @@ def publish_to_meshcat(self, ctx: Context | None = None) -> None: self._diagram.GetSubsystemContext(self._meshcat_visualizer, ctx) ) + def animate_path( + self, + robot_id: str, + path: list[NDArray[np.float64]], + duration: float = 3.0, + ) -> None: + """Animate a path in Meshcat visualization. + + Args: + robot_id: Robot to animate + path: List of joint configurations + duration: Total animation duration in seconds + """ + import time + + if self._meshcat is None or len(path) < 2: + return + + dt = duration / (len(path) - 1) + for q in path: + with self.scratch_context() as ctx: + self.set_positions(ctx, robot_id, q) + self.publish_to_meshcat(ctx) + time.sleep(dt) + # ============= Direct Access (use with caution) ============= @property From f479de07acdfa0c7175fc05aec84d97dc174b028 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Tue, 13 Jan 2026 15:59:03 -0800 Subject: [PATCH 11/31] added blueprints and manipulation client for testing --- dimos/manipulation/manipulation_blueprints.py | 203 ++++++++++++++ .../planning/examples/__init__.py | 35 +++ .../planning/examples/manipulation_client.py | 262 ++++++++++++++++++ 3 files changed, 500 insertions(+) create mode 100644 dimos/manipulation/manipulation_blueprints.py create mode 100644 dimos/manipulation/planning/examples/__init__.py create mode 100644 dimos/manipulation/planning/examples/manipulation_client.py diff --git a/dimos/manipulation/manipulation_blueprints.py b/dimos/manipulation/manipulation_blueprints.py new file mode 100644 index 0000000000..b12693a838 --- /dev/null +++ b/dimos/manipulation/manipulation_blueprints.py @@ -0,0 +1,203 @@ +# 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. + +""" +Blueprints for manipulation module integration with ControlOrchestrator. + +Usage: + # Start orchestrator first, then planner: + coordinator = xarm7_planner_orchestrator.build() + coordinator.loop() + + # Plan and execute via RPC client: + from dimos.manipulation.planning.examples.manipulation_client import ManipulationClient + client = ManipulationClient() + client.plan_to_joints([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]) + client.execute() +""" + +import numpy as np + +from dimos.core.transport import LCMTransport +from dimos.manipulation.manipulation_module import manipulation_module +from dimos.manipulation.planning.spec import RobotModelConfig +from dimos.msgs.sensor_msgs import JointState +from dimos.utils.data import get_data + +# ============================================================================= +# URDF Helpers +# ============================================================================= + + +def _get_xarm_urdf_path() -> str: + """Get path to xarm URDF.""" + return str(get_data("xarm_description") / "urdf/xarm_device.urdf.xacro") + + +def _get_xarm_package_paths() -> dict[str, str]: + """Get package paths for xarm xacro resolution.""" + return {"xarm_description": str(get_data("xarm_description"))} + + +def _get_piper_urdf_path() -> str: + """Get path to piper URDF.""" + return str(get_data("piper_description") / "urdf/piper_description.xacro") + + +def _get_piper_package_paths() -> dict[str, str]: + """Get package paths for piper xacro resolution.""" + return {"piper_description": str(get_data("piper_description"))} + + +# XArm gripper collision exclusions (parallel linkage mechanism) +XARM_GRIPPER_COLLISION_EXCLUSIONS: list[tuple[str, str]] = [ + ("right_inner_knuckle", "right_outer_knuckle"), + ("right_inner_knuckle", "right_finger"), + ("left_inner_knuckle", "left_outer_knuckle"), + ("left_inner_knuckle", "left_finger"), + ("left_finger", "right_finger"), + ("left_outer_knuckle", "right_outer_knuckle"), + ("left_inner_knuckle", "right_inner_knuckle"), +] + + +# ============================================================================= +# Robot Configs +# ============================================================================= + + +def _make_xarm6_config( + name: str = "arm", + y_offset: float = 0.0, + joint_prefix: str = "", + orchestrator_task: str | None = None, +) -> RobotModelConfig: + """Create XArm6 robot config. + + Args: + name: Robot name in Drake world + y_offset: Y-axis offset for base pose (for multi-arm setups) + joint_prefix: Prefix for joint name mapping (e.g., "left_" or "right_") + orchestrator_task: Task name for orchestrator RPC execution + """ + joint_names = ["joint1", "joint2", "joint3", "joint4", "joint5", "joint6"] + joint_mapping = {f"{joint_prefix}{j}": j for j in joint_names} if joint_prefix else {} + + return RobotModelConfig( + name=name, + urdf_path=_get_xarm_urdf_path(), + base_pose=np.array( + [ + [1, 0, 0, 0], + [0, 1, 0, y_offset], + [0, 0, 1, 0], + [0, 0, 0, 1], + ], + dtype=np.float64, + ), + joint_names=joint_names, + end_effector_link="link_tcp", + base_link="link_base", + package_paths=_get_xarm_package_paths(), + xacro_args={"dof": "6", "limited": "true", "add_gripper": "true"}, + collision_exclusion_pairs=XARM_GRIPPER_COLLISION_EXCLUSIONS, + max_velocity=1.0, + max_acceleration=2.0, + joint_name_mapping=joint_mapping, + orchestrator_task_name=orchestrator_task, + ) + + +def _make_xarm7_config( + name: str = "arm", + joint_prefix: str = "", + orchestrator_task: str | None = None, +) -> RobotModelConfig: + """Create XArm7 robot config.""" + joint_names = ["joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "joint7"] + joint_mapping = {f"{joint_prefix}{j}": j for j in joint_names} if joint_prefix else {} + + return RobotModelConfig( + name=name, + urdf_path=_get_xarm_urdf_path(), + base_pose=np.eye(4, dtype=np.float64), + joint_names=joint_names, + end_effector_link="link7", + base_link="link_base", + package_paths=_get_xarm_package_paths(), + xacro_args={"dof": "7", "limited": "true"}, + auto_convert_meshes=True, + max_velocity=1.0, + max_acceleration=2.0, + joint_name_mapping=joint_mapping, + orchestrator_task_name=orchestrator_task, + ) + + +# ============================================================================= +# Blueprints +# ============================================================================= + + +# Single XArm6 planner (standalone, no orchestrator) +xarm6_planner_only = manipulation_module( + robots=[_make_xarm6_config()], + planning_timeout=10.0, + enable_viz=True, +).transports( + { + ("joint_state", JointState): LCMTransport("/xarm/joint_states", JointState), + } +) + + +# Dual XArm6 planner with orchestrator integration +# Usage: Start with orchestrator_dual_mock, then plan/execute via RPC +dual_xarm6_planner = manipulation_module( + robots=[ + _make_xarm6_config( + "left_arm", y_offset=0.5, joint_prefix="left_", orchestrator_task="traj_left" + ), + _make_xarm6_config( + "right_arm", y_offset=-0.5, joint_prefix="right_", orchestrator_task="traj_right" + ), + ], + planning_timeout=10.0, + enable_viz=True, +).transports( + { + ("joint_state", JointState): LCMTransport("/orchestrator/joint_state", JointState), + } +) + + +# Single XArm7 planner for orchestrator-mock +# Usage: dimos run orchestrator-mock, then dimos run xarm7-planner-orchestrator +xarm7_planner_orchestrator = manipulation_module( + robots=[_make_xarm7_config("arm", joint_prefix="arm_", orchestrator_task="traj_arm")], + planning_timeout=10.0, + enable_viz=True, +).transports( + { + ("joint_state", JointState): LCMTransport("/orchestrator/joint_state", JointState), + } +) + + +__all__ = [ + "XARM_GRIPPER_COLLISION_EXCLUSIONS", + "dual_xarm6_planner", + "xarm6_planner_only", + "xarm7_planner_orchestrator", +] diff --git a/dimos/manipulation/planning/examples/__init__.py b/dimos/manipulation/planning/examples/__init__.py new file mode 100644 index 0000000000..6735749095 --- /dev/null +++ b/dimos/manipulation/planning/examples/__init__.py @@ -0,0 +1,35 @@ +# Copyright 2025 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. + +""" +Manipulation Client - RPC client for ManipulationModule + +Usage: + # Start orchestrator and planner first: + dimos run orchestrator-mock + dimos run xarm7-planner-orchestrator + + # Then run the interactive client: + python -m dimos.manipulation.planning.examples.manipulation_client + + # IPython shell with client pre-loaded: + c.joints() # Get current joint positions + c.plan([0.1, ...]) # Plan to joints + c.preview() # Preview in Meshcat + c.execute() # Execute via orchestrator +""" + +from dimos.manipulation.planning.examples.manipulation_client import ManipulationClient + +__all__ = ["ManipulationClient"] diff --git a/dimos/manipulation/planning/examples/manipulation_client.py b/dimos/manipulation/planning/examples/manipulation_client.py new file mode 100644 index 0000000000..47b73db36f --- /dev/null +++ b/dimos/manipulation/planning/examples/manipulation_client.py @@ -0,0 +1,262 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Manipulation Client - IPython interface for ManipulationModule + +Usage: + # Start orchestrator and planner first: + dimos run orchestrator-mock + dimos run xarm7-planner-orchestrator + + # Run interactive client: + python -m dimos.manipulation.planning.examples.manipulation_client + +IPython Commands (client is available as 'c'): + c.joints() # Get current joint positions + c.ee() # Get end-effector pose + c.state() # Get manipulation state + c.url() # Get Meshcat visualization URL + + c.plan([0.1, ...]) # Plan to joint config + c.plan_pose(x, y, z) # Plan to cartesian pose + c.preview() # Preview path in Meshcat + c.execute() # Execute via orchestrator + + c.box("table", ...) # Add box obstacle + c.sphere("ball", ...) # Add sphere obstacle + c.remove("obstacle_id") # Remove obstacle +""" + +from __future__ import annotations + +from typing import Any, cast + +import numpy as np + +from dimos.msgs.geometry_msgs import Pose, Quaternion, Vector3 +from dimos.protocol.rpc import LCMRPC + + +class ManipulationClient: + """RPC client for ManipulationModule with IPython-friendly API.""" + + def __init__(self) -> None: + self._rpc = LCMRPC() + self._rpc.start() + self._module = "ManipulationModule" + print("Connected to ManipulationModule via LCM RPC") + + def _call(self, method: str, *args: Any, **kwargs: Any) -> Any: + """Call RPC method.""" + try: + result, _ = self._rpc.call_sync( + f"{self._module}/{method}", (list(args), kwargs), rpc_timeout=30.0 + ) + return result + except TimeoutError: + print(f"Timeout: {method}") + return None + except Exception as e: + print(f"Error: {e}") + return None + + # ========================================================================= + # Query Methods + # ========================================================================= + + def state(self) -> str: + """Get manipulation state.""" + return cast("str", self._call("get_state")) + + def joints(self) -> list[float] | None: + """Get current joint positions.""" + return cast("list[float] | None", self._call("get_current_joints")) + + def ee(self) -> Pose | None: + """Get end-effector pose.""" + return cast("Pose | None", self._call("get_ee_pose")) + + def url(self) -> str | None: + """Get Meshcat visualization URL.""" + return cast("str | None", self._call("get_visualization_url")) + + def robots(self) -> list[str]: + """List configured robots.""" + return cast("list[str]", self._call("list_robots")) + + def info(self, robot_name: str | None = None) -> dict[str, Any] | None: + """Get robot info.""" + return cast("dict[str, Any] | None", self._call("get_robot_info", robot_name)) + + # ========================================================================= + # Planning Methods + # ========================================================================= + + def plan(self, joints: list[float]) -> bool: + """Plan to joint configuration.""" + print(f"Planning to: {[f'{j:.3f}' for j in joints]}") + return cast("bool", self._call("plan_to_joints", joints)) + + def plan_pose( + self, + x: float, + y: float, + z: float, + roll: float | None = None, + pitch: float | None = None, + yaw: float | None = None, + ) -> bool: + """Plan to cartesian pose. Uses current orientation if not specified.""" + # Get current orientation if not provided + if roll is None or pitch is None or yaw is None: + ee = self.ee() + if ee is None: + print("Cannot get current orientation") + return False + roll = roll if roll is not None else ee.roll + pitch = pitch if pitch is not None else ee.pitch + yaw = yaw if yaw is not None else ee.yaw + + print(f"Planning to: ({x:.3f}, {y:.3f}, {z:.3f}) rpy=({roll:.2f}, {pitch:.2f}, {yaw:.2f})") + pose = Pose( + position=Vector3(x, y, z), + orientation=Quaternion.from_euler(Vector3(roll, pitch, yaw)), + ) + return cast("bool", self._call("plan_to_pose", pose)) + + def preview(self, duration: float = 3.0) -> bool: + """Preview planned path in Meshcat.""" + return cast("bool", self._call("preview_path", duration)) + + def execute(self, robot_name: str | None = None) -> bool: + """Execute planned trajectory via orchestrator.""" + return cast("bool", self._call("execute", robot_name)) + + def has_plan(self) -> bool: + """Check if path is planned.""" + return cast("bool", self._call("has_planned_path")) + + def clear_plan(self) -> bool: + """Clear planned path.""" + return cast("bool", self._call("clear_planned_path")) + + # ========================================================================= + # Obstacle Methods + # ========================================================================= + + def box( + self, + name: str, + x: float, + y: float, + z: float, + w: float, + h: float, + d: float, + roll: float = 0.0, + pitch: float = 0.0, + yaw: float = 0.0, + ) -> str: + """Add box obstacle.""" + pose = Pose( + position=Vector3(x, y, z), + orientation=Quaternion.from_euler(Vector3(roll, pitch, yaw)), + ) + return cast("str", self._call("add_obstacle", name, pose, "box", [w, h, d])) + + def sphere(self, name: str, x: float, y: float, z: float, radius: float) -> str: + """Add sphere obstacle.""" + pose = Pose(position=Vector3(x, y, z)) + return cast("str", self._call("add_obstacle", name, pose, "sphere", [radius])) + + def cylinder( + self, name: str, x: float, y: float, z: float, radius: float, length: float + ) -> str: + """Add cylinder obstacle.""" + pose = Pose(position=Vector3(x, y, z)) + return cast("str", self._call("add_obstacle", name, pose, "cylinder", [radius, length])) + + def remove(self, obstacle_id: str) -> bool: + """Remove obstacle.""" + return cast("bool", self._call("remove_obstacle", obstacle_id)) + + # ========================================================================= + # Utility Methods + # ========================================================================= + + def collision(self, joints: list[float]) -> bool: + """Check if joint config is collision-free.""" + return cast("bool", self._call("is_collision_free", joints)) + + def reset(self) -> bool: + """Reset to IDLE state.""" + return cast("bool", self._call("reset")) + + def cancel(self) -> bool: + """Cancel current motion.""" + return cast("bool", self._call("cancel")) + + def status(self, robot_name: str | None = None) -> dict[str, object] | None: + """Get trajectory execution status.""" + return cast("dict[str, object] | None", self._call("get_trajectory_status", robot_name)) + + def stop(self) -> None: + """Stop RPC client.""" + self._rpc.stop() + + def __repr__(self) -> str: + return f"ManipulationClient(state={self.state()})" + + +def main() -> None: + """Start IPython shell with ManipulationClient.""" + try: + from IPython import embed + except ImportError: + print("IPython not installed. Run: pip install ipython") + return + + c = ManipulationClient() + + banner = """ +Manipulation Client - IPython Interface +======================================== + +Client available as 'c'. Examples: + c.joints() # Get joint positions + c.ee() # Get end-effector pose + c.url() # Get Meshcat URL + + c.plan([0.1, 0.2, ...]) # Plan to joints + c.plan_pose(0.4, 0, 0.3)# Plan to pose + c.preview() # Preview in Meshcat + c.execute() # Execute trajectory + + c.box("table", 0.5, 0, 0, 0.6, 0.4, 0.02) # Add obstacle + c.remove("obstacle_id") # Remove obstacle + +Type c. for all methods, or help(c.method) for details. +""" + print(banner) + + try: + embed(colors="neutral") + finally: + c.stop() + + +if __name__ == "__main__": + main() From 6d2435cc7bb5bf4fdd48549c08628b4c0018cb0a Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Tue, 13 Jan 2026 19:48:48 -0800 Subject: [PATCH 12/31] added unit test, e2e tests and readme --- dimos/e2e_tests/test_manipulation_module.py | 462 +++++++++++++++++++ dimos/manipulation/manipulation_module.py | 3 +- dimos/manipulation/planning/README.md | 130 ++++++ dimos/manipulation/test_manipulation_unit.py | 304 ++++++++++++ 4 files changed, 898 insertions(+), 1 deletion(-) create mode 100644 dimos/e2e_tests/test_manipulation_module.py create mode 100644 dimos/manipulation/planning/README.md create mode 100644 dimos/manipulation/test_manipulation_unit.py diff --git a/dimos/e2e_tests/test_manipulation_module.py b/dimos/e2e_tests/test_manipulation_module.py new file mode 100644 index 0000000000..ce3105dfa3 --- /dev/null +++ b/dimos/e2e_tests/test_manipulation_module.py @@ -0,0 +1,462 @@ +# 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. + +""" +Integration tests for ManipulationModule. + +These tests verify the full planning stack with Drake backend. +They require Drake to be installed and will be skipped otherwise. +""" + +from __future__ import annotations + +import os +import time +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest + +from dimos.manipulation.manipulation_module import ( + ManipulationModule, + ManipulationState, +) +from dimos.manipulation.planning.spec import RobotModelConfig +from dimos.msgs.geometry_msgs import Pose, Quaternion, Vector3 +from dimos.msgs.sensor_msgs import JointState +from dimos.utils.data import get_data + +# ============================================================================= +# Helper Functions +# ============================================================================= + + +def _drake_available() -> bool: + """Check if Drake is available.""" + try: + import pydrake + + return True + except ImportError: + return False + + +def _xarm_urdf_available() -> bool: + """Check if xarm URDF is available.""" + try: + desc_path = get_data("xarm_description") + urdf_path = desc_path / "urdf/xarm_device.urdf.xacro" + return urdf_path.exists() + except Exception: + return False + + +def _get_xarm7_config() -> RobotModelConfig: + """Create XArm7 robot config for testing.""" + desc_path = get_data("xarm_description") + return RobotModelConfig( + name="test_arm", + urdf_path=str(desc_path / "urdf/xarm_device.urdf.xacro"), + base_pose=np.eye(4, dtype=np.float64), + joint_names=["joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "joint7"], + end_effector_link="link7", + base_link="link_base", + package_paths={"xarm_description": str(desc_path)}, + xacro_args={"dof": "7", "limited": "true"}, + auto_convert_meshes=True, + max_velocity=1.0, + max_acceleration=2.0, + joint_name_mapping={ + "arm_joint1": "joint1", + "arm_joint2": "joint2", + "arm_joint3": "joint3", + "arm_joint4": "joint4", + "arm_joint5": "joint5", + "arm_joint6": "joint6", + "arm_joint7": "joint7", + }, + orchestrator_task_name="traj_arm", + ) + + +# ============================================================================= +# Fixtures +# ============================================================================= + + +@pytest.fixture +def xarm7_config(): + """Create XArm7 config fixture.""" + return _get_xarm7_config() + + +@pytest.fixture +def joint_state_zeros(): + """Create a JointState message with zeros for XArm7.""" + return JointState( + name=[ + "arm_joint1", + "arm_joint2", + "arm_joint3", + "arm_joint4", + "arm_joint5", + "arm_joint6", + "arm_joint7", + ], + position=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + velocity=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + effort=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ) + + +# ============================================================================= +# Integration Tests (require Drake + URDF) +# ============================================================================= + + +@pytest.mark.skipif(not _drake_available(), reason="Drake not installed") +@pytest.mark.skipif(not _xarm_urdf_available(), reason="XArm URDF not available") +@pytest.mark.skipif(bool(os.getenv("CI")), reason="Skip in CI - requires LFS data") +class TestManipulationModuleIntegration: + """Integration tests for ManipulationModule with real Drake backend.""" + + def test_module_initialization(self, xarm7_config): + """Test module initializes with real Drake world.""" + # Create module instance directly (without going through blueprints) + # Module __init__ takes config fields as kwargs, not a config object + module = ManipulationModule( + robots=[xarm7_config], + planning_timeout=10.0, + enable_viz=False, # No viz for tests + ) + + # Mock the joint_state port to avoid needing real transport + module.joint_state = None + + # Start to trigger planning initialization + module.start() + + try: + assert module._state == ManipulationState.IDLE + assert module._world_monitor is not None + assert module._planner is not None + assert module._kinematics is not None + assert "test_arm" in module._robots + finally: + module.stop() + + def test_joint_state_sync(self, xarm7_config, joint_state_zeros): + """Test joint state synchronization to Drake world.""" + module = ManipulationModule( + robots=[xarm7_config], + planning_timeout=10.0, + enable_viz=False, + ) + module.joint_state = None + module.start() + + try: + # Manually send joint state + module._on_joint_state(joint_state_zeros) + + # Verify state is synchronized + joints = module.get_current_joints() + assert joints is not None + assert len(joints) == 7 + # All zeros + assert all(abs(j) < 0.01 for j in joints) + finally: + module.stop() + + def test_collision_check(self, xarm7_config, joint_state_zeros): + """Test collision checking at a configuration.""" + module = ManipulationModule( + robots=[xarm7_config], + planning_timeout=10.0, + enable_viz=False, + ) + module.joint_state = None + module.start() + + try: + # Sync state first + module._on_joint_state(joint_state_zeros) + + # Check that zero config is collision-free + is_free = module.is_collision_free([0.0] * 7) + assert is_free is True + finally: + module.stop() + + def test_plan_to_joints(self, xarm7_config, joint_state_zeros): + """Test planning to a joint configuration.""" + module = ManipulationModule( + robots=[xarm7_config], + planning_timeout=10.0, + enable_viz=False, + ) + module.joint_state = None + module.start() + + try: + # Sync state first + module._on_joint_state(joint_state_zeros) + + # Plan to a small motion + target = [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1] + success = module.plan_to_joints(target) + + assert success is True + assert module._state == ManipulationState.COMPLETED + assert module.has_planned_path() is True + + # Verify trajectory was generated + assert "test_arm" in module._planned_trajectories + traj = module._planned_trajectories["test_arm"] + assert len(traj.points) > 1 + assert traj.duration > 0 + finally: + module.stop() + + def test_add_and_remove_obstacle(self, xarm7_config, joint_state_zeros): + """Test adding and removing obstacles.""" + module = ManipulationModule( + robots=[xarm7_config], + planning_timeout=10.0, + enable_viz=False, + ) + module.joint_state = None + module.start() + + try: + # Sync state + module._on_joint_state(joint_state_zeros) + + # Add a box obstacle + pose = Pose( + position=Vector3(0.5, 0.0, 0.3), + orientation=Quaternion(), # default is identity (w=1) + ) + obstacle_id = module.add_obstacle("test_box", pose, "box", [0.1, 0.1, 0.1]) + + assert obstacle_id != "" + assert obstacle_id is not None + + # Remove it + removed = module.remove_obstacle(obstacle_id) + assert removed is True + finally: + module.stop() + + def test_robot_info(self, xarm7_config): + """Test getting robot information.""" + module = ManipulationModule( + robots=[xarm7_config], + planning_timeout=10.0, + enable_viz=False, + ) + module.joint_state = None + module.start() + + try: + # Get robot info + info = module.get_robot_info() + + assert info is not None + assert info["name"] == "test_arm" + assert len(info["joint_names"]) == 7 + assert info["end_effector_link"] == "link7" + assert info["orchestrator_task_name"] == "traj_arm" + assert info["has_joint_name_mapping"] is True + finally: + module.stop() + + def test_ee_pose(self, xarm7_config, joint_state_zeros): + """Test getting end-effector pose.""" + module = ManipulationModule( + robots=[xarm7_config], + planning_timeout=10.0, + enable_viz=False, + ) + module.joint_state = None + module.start() + + try: + # Sync state + module._on_joint_state(joint_state_zeros) + + # Get EE pose + pose = module.get_ee_pose() + + assert pose is not None + # At zero config, EE should be at some position + assert hasattr(pose, "x") + assert hasattr(pose, "y") + assert hasattr(pose, "z") + finally: + module.stop() + + def test_trajectory_name_translation(self, xarm7_config, joint_state_zeros): + """Test that trajectory joint names are translated for orchestrator.""" + module = ManipulationModule( + robots=[xarm7_config], + planning_timeout=10.0, + enable_viz=False, + ) + module.joint_state = None + module.start() + + try: + # Sync state + module._on_joint_state(joint_state_zeros) + + # Plan a motion + success = module.plan_to_joints([0.05] * 7) + assert success is True + + # Get the trajectory + traj = module._planned_trajectories["test_arm"] + robot_config = module._robots["test_arm"][1] + + # Translate it + translated = module._translate_trajectory_to_orchestrator(traj, robot_config) + + # Verify names are translated + for name in translated.joint_names: + assert name.startswith("arm_") # Should have arm_ prefix + finally: + module.stop() + + +# ============================================================================= +# Orchestrator Integration Tests (mocked orchestrator) +# ============================================================================= + + +@pytest.mark.skipif(not _drake_available(), reason="Drake not installed") +@pytest.mark.skipif(not _xarm_urdf_available(), reason="XArm URDF not available") +@pytest.mark.skipif(bool(os.getenv("CI")), reason="Skip in CI - requires LFS data") +class TestOrchestratorIntegration: + """Test orchestrator integration with mocked RPC client.""" + + def test_execute_with_mock_orchestrator(self, xarm7_config, joint_state_zeros): + """Test execute sends trajectory to orchestrator.""" + module = ManipulationModule( + robots=[xarm7_config], + planning_timeout=10.0, + enable_viz=False, + ) + module.joint_state = None + module.start() + + try: + # Sync state + module._on_joint_state(joint_state_zeros) + + # Plan a motion + success = module.plan_to_joints([0.05] * 7) + assert success is True + + # Mock the orchestrator client + mock_client = MagicMock() + mock_client.execute_trajectory.return_value = True + module._orchestrator_client = mock_client + + # Execute + result = module.execute() + + assert result is True + assert module._state == ManipulationState.COMPLETED + + # Verify orchestrator was called + mock_client.execute_trajectory.assert_called_once() + call_args = mock_client.execute_trajectory.call_args + task_name, trajectory = call_args[0] + + assert task_name == "traj_arm" + assert len(trajectory.points) > 1 + # Joint names should be translated + assert all(n.startswith("arm_") for n in trajectory.joint_names) + finally: + module.stop() + + def test_execute_rejected_by_orchestrator(self, xarm7_config, joint_state_zeros): + """Test handling of orchestrator rejection.""" + module = ManipulationModule( + robots=[xarm7_config], + planning_timeout=10.0, + enable_viz=False, + ) + module.joint_state = None + module.start() + + try: + # Sync state + module._on_joint_state(joint_state_zeros) + + # Plan a motion + module.plan_to_joints([0.05] * 7) + + # Mock orchestrator to reject + mock_client = MagicMock() + mock_client.execute_trajectory.return_value = False + module._orchestrator_client = mock_client + + # Execute + result = module.execute() + + assert result is False + assert module._state == ManipulationState.FAULT + assert "rejected" in module._error_message.lower() + finally: + module.stop() + + def test_state_transitions_during_execution(self, xarm7_config, joint_state_zeros): + """Test state transitions during plan and execute.""" + module = ManipulationModule( + robots=[xarm7_config], + planning_timeout=10.0, + enable_viz=False, + ) + module.joint_state = None + module.start() + + try: + # Initial state + assert module._state == ManipulationState.IDLE + + # Sync state + module._on_joint_state(joint_state_zeros) + + # Plan - should go through PLANNING -> COMPLETED + module.plan_to_joints([0.05] * 7) + assert module._state == ManipulationState.COMPLETED + + # Reset works from COMPLETED + module.reset() + assert module._state == ManipulationState.IDLE + + # Plan again + module.plan_to_joints([0.05] * 7) + + # Mock orchestrator + mock_client = MagicMock() + mock_client.execute_trajectory.return_value = True + module._orchestrator_client = mock_client + + # Execute - should go to EXECUTING then COMPLETED + module.execute() + assert module._state == ManipulationState.COMPLETED + finally: + module.stop() diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index 96564ff892..c99be5d557 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -238,7 +238,8 @@ def reset(self) -> bool: def get_current_joints(self) -> list[float] | None: """Get current joint positions.""" if (robot := self._get_robot()) and self._world_monitor: - if pos := self._world_monitor.get_current_positions(robot[1]): + pos = self._world_monitor.get_current_positions(robot[1]) + if pos is not None: return list(pos) return None diff --git a/dimos/manipulation/planning/README.md b/dimos/manipulation/planning/README.md new file mode 100644 index 0000000000..d549d18a71 --- /dev/null +++ b/dimos/manipulation/planning/README.md @@ -0,0 +1,130 @@ +# Manipulation Planning Stack + +Motion planning for robotic manipulators using Drake. Integrates with Control Orchestrator for execution. + +## Quick Start (3 Terminals) + +```bash +# Terminal 1: Start mock orchestrator +dimos run orchestrator-mock + +# Terminal 2: Start manipulation planner +dimos run xarm7-planner-orchestrator + +# Terminal 3: Run IPython client +python -m dimos.manipulation.planning.examples.manipulation_client +``` + +Then in IPython: +```python +c.joints() # Get current joints +c.plan([0.1] * 7) # Plan to target +c.preview() # Preview in Meshcat (check c.url()) +c.execute() # Execute via orchestrator +``` + +## Architecture + +``` +ManipulationModule (RPC interface, state machine, multi-robot) + │ +Factory Functions (create_world, create_kinematics, create_planner) + │ +Drake Implementations (DrakeWorld, DrakeKinematics, DrakePlanner) + │ +Control Orchestrator (trajectory execution via RPC) +``` + +## Using ManipulationModule + +```python +from dimos.manipulation import ManipulationModule +from dimos.manipulation.planning.spec import RobotModelConfig + +config = RobotModelConfig( + name="xarm7", + urdf_path="/path/to/xarm7.urdf", + base_pose=np.eye(4), + joint_names=["joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "joint7"], + end_effector_link="link7", + base_link="link_base", + joint_name_mapping={"arm_joint1": "joint1", ...}, # URDF <-> orchestrator + orchestrator_task_name="traj_arm", +) + +module = ManipulationModule(robots=[config], planning_timeout=10.0, enable_viz=True) +module.start() +module.plan_to_joints([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]) +module.execute() # Sends to orchestrator +``` + +## RobotModelConfig Fields + +| Field | Description | +|-------|-------------| +| `name` | Robot identifier | +| `urdf_path` | Path to URDF/XACRO file | +| `base_pose` | 4x4 transform matrix for robot base | +| `joint_names` | List of joint names in URDF | +| `end_effector_link` | Name of EE link in URDF | +| `base_link` | Name of base link in URDF | +| `max_velocity` | Max joint velocity (rad/s) | +| `max_acceleration` | Max joint acceleration (rad/s²) | +| `joint_name_mapping` | Dict mapping orchestrator names to URDF names | +| `orchestrator_task_name` | Task name for trajectory execution RPC | +| `package_paths` | Dict of ROS package paths for mesh resolution | +| `xacro_args` | Dict of xacro arguments (e.g., `{"dof": "7"}`) | + +## Available Blueprints + +| Blueprint | Description | +|-----------|-------------| +| `xarm6-planner` | XArm 6-DOF planner (standalone) | +| `xarm7-planner-orchestrator` | XArm 7-DOF with orchestrator | +| `dual-xarm6-planner` | Dual XArm 6-DOF planner | + +## Directory Structure + +``` +planning/ +├── spec.py # Protocol definitions +├── factory.py # Factory functions +├── world/ # DrakeWorld +├── kinematics/ # DrakeKinematics +├── planners/ # RRT-Connect, RRT* +├── monitor/ # WorldMonitor, state sync +├── trajectory_generator/ +└── examples/ + ├── planning_tester.py # Standalone CLI + └── manipulation_client.py # IPython RPC client +``` + +## Protocols + +- **WorldSpec**: Physics/collision backend (DrakeWorld, MuJoCoWorld, PyBulletWorld) +- **KinematicsSpec**: IK solving (DrakeKinematics, TracIK, KDL) +- **PlannerSpec**: Path planning (DrakeRRTConnect, OMPL, cuRobo) + +## Supported Robots + +| Robot | DOF | +|-------|-----| +| `piper` | 6 | +| `xarm6` | 6 | +| `xarm7` | 7 | + +## Planners + +| Planner | Description | +|---------|-------------| +| `rrt_connect` | Bidirectional RRT (fast) | +| `rrt_star` | RRT* with rewiring (optimal) | + +## Obstacle Types + +| Type | Dimensions | +|------|------------| +| `BOX` | (w, h, d) | +| `SPHERE` | (radius,) | +| `CYLINDER` | (radius, height) | +| `MESH` | mesh_path | diff --git a/dimos/manipulation/test_manipulation_unit.py b/dimos/manipulation/test_manipulation_unit.py new file mode 100644 index 0000000000..b69f2cc286 --- /dev/null +++ b/dimos/manipulation/test_manipulation_unit.py @@ -0,0 +1,304 @@ +# 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. + +"""Unit tests for the ManipulationModule.""" + +from __future__ import annotations + +import threading +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest + +from dimos.manipulation.manipulation_module import ( + ManipulationModule, + ManipulationState, +) +from dimos.manipulation.planning.spec import RobotModelConfig +from dimos.msgs.trajectory_msgs import JointTrajectory, TrajectoryPoint + +# ============================================================================= +# Fixtures +# ============================================================================= + + +@pytest.fixture +def robot_config(): + """Create a robot config for testing.""" + return RobotModelConfig( + name="test_arm", + urdf_path="/path/to/robot.urdf", + base_pose=np.eye(4, dtype=np.float64), + joint_names=["joint1", "joint2", "joint3"], + end_effector_link="link_tcp", + base_link="link_base", + max_velocity=1.0, + max_acceleration=2.0, + orchestrator_task_name="traj_arm", + ) + + +@pytest.fixture +def robot_config_with_mapping(): + """Create a robot config with joint name mapping (dual-arm scenario).""" + return RobotModelConfig( + name="left_arm", + urdf_path="/path/to/robot.urdf", + base_pose=np.eye(4, dtype=np.float64), + joint_names=["joint1", "joint2", "joint3"], + end_effector_link="link_tcp", + base_link="link_base", + joint_name_mapping={ + "left_joint1": "joint1", + "left_joint2": "joint2", + "left_joint3": "joint3", + }, + orchestrator_task_name="traj_left", + ) + + +@pytest.fixture +def simple_trajectory(): + """Create a simple trajectory for testing.""" + return JointTrajectory( + joint_names=["joint1", "joint2", "joint3"], + points=[ + TrajectoryPoint( + positions=[0.0, 0.0, 0.0], velocities=[0.0, 0.0, 0.0], time_from_start=0.0 + ), + TrajectoryPoint( + positions=[0.5, 0.5, 0.5], velocities=[0.0, 0.0, 0.0], time_from_start=1.0 + ), + ], + ) + + +def _make_module(): + """Create a ManipulationModule instance with mocked __init__.""" + with patch.object(ManipulationModule, "__init__", lambda self: None): + module = ManipulationModule.__new__(ManipulationModule) + module._state = ManipulationState.IDLE + module._lock = threading.Lock() + module._error_message = "" + module._robots = {} + module._planned_paths = {} + module._planned_trajectories = {} + module._world_monitor = None + module._planner = None + module._kinematics = None + module._orchestrator_client = None + return module + + +# ============================================================================= +# Test State Machine +# ============================================================================= + + +class TestStateMachine: + """Test state transitions.""" + + def test_cancel_only_during_execution(self): + """Cancel only works in EXECUTING state.""" + module = _make_module() + + module._state = ManipulationState.IDLE + assert module.cancel() is False + + module._state = ManipulationState.EXECUTING + assert module.cancel() is True + assert module._state == ManipulationState.IDLE + + def test_reset_not_during_execution(self): + """Reset works in any state except EXECUTING.""" + module = _make_module() + + module._state = ManipulationState.FAULT + module._error_message = "Error" + assert module.reset() is True + assert module._state == ManipulationState.IDLE + assert module._error_message == "" + + module._state = ManipulationState.EXECUTING + assert module.reset() is False + + def test_fail_sets_fault_state(self): + """_fail helper sets FAULT state and message.""" + module = _make_module() + module._state = ManipulationState.PLANNING + + result = module._fail("Test error") + assert result is False + assert module._state == ManipulationState.FAULT + assert module._error_message == "Test error" + + def test_begin_planning_state_checks(self, robot_config): + """_begin_planning only allowed from IDLE or COMPLETED.""" + module = _make_module() + module._world_monitor = MagicMock() + module._robots = {"test_arm": ("robot_id", robot_config, MagicMock())} + + # From IDLE - OK + module._state = ManipulationState.IDLE + assert module._begin_planning() == ("test_arm", "robot_id") + assert module._state == ManipulationState.PLANNING + + # From COMPLETED - OK + module._state = ManipulationState.COMPLETED + assert module._begin_planning() == ("test_arm", "robot_id") + + # From EXECUTING - Fail + module._state = ManipulationState.EXECUTING + assert module._begin_planning() is None + + +# ============================================================================= +# Test Robot Selection +# ============================================================================= + + +class TestRobotSelection: + """Test robot selection logic.""" + + def test_single_robot_default(self, robot_config): + """Single robot is used by default.""" + module = _make_module() + module._robots = {"arm": ("id", robot_config, MagicMock())} + + result = module._get_robot() + assert result is not None + assert result[0] == "arm" + + def test_multiple_robots_require_name(self, robot_config): + """Multiple robots require explicit name.""" + module = _make_module() + module._robots = { + "left": ("id1", robot_config, MagicMock()), + "right": ("id2", robot_config, MagicMock()), + } + + # No name - fails + assert module._get_robot() is None + + # With name - works + result = module._get_robot("left") + assert result is not None + assert result[0] == "left" + + +# ============================================================================= +# Test Joint Name Translation (for orchestrator integration) +# ============================================================================= + + +class TestJointNameTranslation: + """Test trajectory joint name translation for orchestrator.""" + + def test_no_mapping_returns_original(self, robot_config, simple_trajectory): + """Without mapping, trajectory is returned unchanged.""" + module = _make_module() + + result = module._translate_trajectory_to_orchestrator(simple_trajectory, robot_config) + assert result is simple_trajectory # Same object + + def test_mapping_translates_names(self, robot_config_with_mapping, simple_trajectory): + """With mapping, joint names are translated.""" + module = _make_module() + + result = module._translate_trajectory_to_orchestrator( + simple_trajectory, robot_config_with_mapping + ) + assert result.joint_names == ["left_joint1", "left_joint2", "left_joint3"] + assert len(result.points) == 2 # Points preserved + + +# ============================================================================= +# Test Execute Method +# ============================================================================= + + +class TestExecute: + """Test orchestrator execution.""" + + def test_execute_requires_trajectory(self, robot_config): + """Execute fails without planned trajectory.""" + module = _make_module() + module._robots = {"test_arm": ("id", robot_config, MagicMock())} + module._planned_trajectories = {} + + assert module.execute() is False + + def test_execute_requires_task_name(self): + """Execute fails without orchestrator_task_name.""" + module = _make_module() + config_no_task = RobotModelConfig( + name="arm", + urdf_path="/path", + base_pose=np.eye(4), + joint_names=["j1"], + end_effector_link="ee", + ) + module._robots = {"arm": ("id", config_no_task, MagicMock())} + module._planned_trajectories = {"arm": MagicMock()} + + assert module.execute() is False + + def test_execute_success(self, robot_config, simple_trajectory): + """Successful execute calls orchestrator.""" + module = _make_module() + module._robots = {"test_arm": ("id", robot_config, MagicMock())} + module._planned_trajectories = {"test_arm": simple_trajectory} + + mock_client = MagicMock() + mock_client.execute_trajectory.return_value = True + module._orchestrator_client = mock_client + + assert module.execute() is True + assert module._state == ManipulationState.COMPLETED + mock_client.execute_trajectory.assert_called_once() + + def test_execute_rejected(self, robot_config, simple_trajectory): + """Rejected execution sets FAULT state.""" + module = _make_module() + module._robots = {"test_arm": ("id", robot_config, MagicMock())} + module._planned_trajectories = {"test_arm": simple_trajectory} + + mock_client = MagicMock() + mock_client.execute_trajectory.return_value = False + module._orchestrator_client = mock_client + + assert module.execute() is False + assert module._state == ManipulationState.FAULT + + +# ============================================================================= +# Test RobotModelConfig Mapping Helpers +# ============================================================================= + + +class TestRobotModelConfigMapping: + """Test RobotModelConfig joint name mapping helpers.""" + + def test_bidirectional_mapping(self, robot_config_with_mapping): + """Test URDF <-> orchestrator name translation.""" + config = robot_config_with_mapping + + # Orchestrator -> URDF + assert config.get_urdf_joint_name("left_joint1") == "joint1" + assert config.get_urdf_joint_name("unknown") == "unknown" + + # URDF -> Orchestrator + assert config.get_orchestrator_joint_name("joint1") == "left_joint1" + assert config.get_orchestrator_joint_name("unknown") == "unknown" From 8feaf8b28bbacadc7d658b09a585a86d9b259af1 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Wed, 14 Jan 2026 11:00:44 -0800 Subject: [PATCH 13/31] general cleanup --- dimos/manipulation/manipulation_history.py | 417 ---------------- .../planning/{ => utils}/mesh_utils.py | 0 .../planning/world/drake_world.py | 40 +- .../manipulation/test_manipulation_history.py | 458 ------------------ 4 files changed, 39 insertions(+), 876 deletions(-) delete mode 100644 dimos/manipulation/manipulation_history.py rename dimos/manipulation/planning/{ => utils}/mesh_utils.py (100%) delete mode 100644 dimos/manipulation/test_manipulation_history.py diff --git a/dimos/manipulation/manipulation_history.py b/dimos/manipulation/manipulation_history.py deleted file mode 100644 index 8d9b281d76..0000000000 --- a/dimos/manipulation/manipulation_history.py +++ /dev/null @@ -1,417 +0,0 @@ -# Copyright 2025-2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# 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](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. - -"""Module for manipulation history tracking and search.""" - -from dataclasses import dataclass, field -from datetime import datetime -import json -import os -import pickle -import time -from typing import Any - -from dimos.types.manipulation import ( - ManipulationTask, -) -from dimos.utils.logging_config import setup_logger - -logger = setup_logger() - - -@dataclass -class ManipulationHistoryEntry: - """An entry in the manipulation history. - - Attributes: - task: The manipulation task executed - timestamp: When the manipulation was performed - result: Result of the manipulation (success/failure) - manipulation_response: Response from the motion planner/manipulation executor - """ - - task: ManipulationTask - timestamp: float = field(default_factory=time.time) - result: dict[str, Any] = field(default_factory=dict) - manipulation_response: str | None = ( - None # Any elaborative response from the motion planner / manipulation executor - ) - - def __str__(self) -> str: - status = self.result.get("status", "unknown") - return f"ManipulationHistoryEntry(task='{self.task.description}', status={status}, time={datetime.fromtimestamp(self.timestamp).strftime('%H:%M:%S')})" - - -class ManipulationHistory: - """A simplified, dictionary-based storage for manipulation history. - - This class provides an efficient way to store and query manipulation tasks, - focusing on quick lookups and flexible search capabilities. - """ - - def __init__(self, output_dir: str | None = None, new_memory: bool = False) -> None: - """Initialize a new manipulation history. - - Args: - output_dir: Directory to save history to - new_memory: If True, creates a new memory instead of loading existing one - """ - self._history: list[ManipulationHistoryEntry] = [] - self._output_dir = output_dir - - if output_dir and not new_memory: - self.load_from_dir(output_dir) - elif output_dir: - os.makedirs(output_dir, exist_ok=True) - logger.info(f"Created new manipulation history at {output_dir}") - - def __len__(self) -> int: - """Return the number of entries in the history.""" - return len(self._history) - - def __str__(self) -> str: - """Return a string representation of the history.""" - if not self._history: - return "ManipulationHistory(empty)" - - return ( - f"ManipulationHistory(entries={len(self._history)}, " - f"time_range={datetime.fromtimestamp(self._history[0].timestamp).strftime('%Y-%m-%d %H:%M:%S')} to " - f"{datetime.fromtimestamp(self._history[-1].timestamp).strftime('%Y-%m-%d %H:%M:%S')})" - ) - - def clear(self) -> None: - """Clear all entries from the history.""" - self._history.clear() - logger.info("Cleared manipulation history") - - if self._output_dir: - self.save_history() - - def add_entry(self, entry: ManipulationHistoryEntry) -> None: - """Add an entry to the history. - - Args: - entry: The entry to add - """ - self._history.append(entry) - self._history.sort(key=lambda e: e.timestamp) - - if self._output_dir: - self.save_history() - - def save_history(self) -> None: - """Save the history to the output directory.""" - if not self._output_dir: - logger.warning("Cannot save history: no output directory specified") - return - - os.makedirs(self._output_dir, exist_ok=True) - history_path = os.path.join(self._output_dir, "manipulation_history.pickle") - - with open(history_path, "wb") as f: - pickle.dump(self._history, f) - - logger.info(f"Saved manipulation history to {history_path}") - - # Also save a JSON representation for easier inspection - json_path = os.path.join(self._output_dir, "manipulation_history.json") - try: - history_data = [ - { - "task": { - "description": entry.task.description, - "target_object": entry.task.target_object, - "target_point": entry.task.target_point, - "timestamp": entry.task.timestamp, - "task_id": entry.task.task_id, - "metadata": entry.task.metadata, - }, - "result": entry.result, - "timestamp": entry.timestamp, - "manipulation_response": entry.manipulation_response, - } - for entry in self._history - ] - - with open(json_path, "w") as f: - json.dump(history_data, f, indent=2) - - logger.info(f"Saved JSON representation to {json_path}") - except Exception as e: - logger.error(f"Failed to save JSON representation: {e}") - - def load_from_dir(self, directory: str) -> None: - """Load history from the specified directory. - - Args: - directory: Directory to load history from - """ - history_path = os.path.join(directory, "manipulation_history.pickle") - - if not os.path.exists(history_path): - logger.warning(f"No history found at {history_path}") - return - - try: - with open(history_path, "rb") as f: - self._history = pickle.load(f) - - logger.info( - f"Loaded manipulation history from {history_path} with {len(self._history)} entries" - ) - except Exception as e: - logger.error(f"Failed to load history: {e}") - - def get_all_entries(self) -> list[ManipulationHistoryEntry]: - """Get all entries in chronological order. - - Returns: - List of all manipulation history entries - """ - return self._history.copy() - - def get_entry_by_index(self, index: int) -> ManipulationHistoryEntry | None: - """Get an entry by its index. - - Args: - index: Index of the entry to retrieve - - Returns: - The entry at the specified index or None if index is out of bounds - """ - if 0 <= index < len(self._history): - return self._history[index] - return None - - def get_entries_by_timerange( - self, start_time: float, end_time: float - ) -> list[ManipulationHistoryEntry]: - """Get entries within a specific time range. - - Args: - start_time: Start time (UNIX timestamp) - end_time: End time (UNIX timestamp) - - Returns: - List of entries within the specified time range - """ - return [entry for entry in self._history if start_time <= entry.timestamp <= end_time] - - def get_entries_by_object(self, object_name: str) -> list[ManipulationHistoryEntry]: - """Get entries related to a specific object. - - Args: - object_name: Name of the object to search for - - Returns: - List of entries related to the specified object - """ - return [entry for entry in self._history if entry.task.target_object == object_name] - - def create_task_entry( - self, - task: ManipulationTask, - result: dict[str, Any] | None = None, - agent_response: str | None = None, - ) -> ManipulationHistoryEntry: - """Create a new manipulation history entry. - - Args: - task: The manipulation task - result: Result of the manipulation - agent_response: Response from the agent about this manipulation - - Returns: - The created history entry - """ - entry = ManipulationHistoryEntry( - task=task, result=result or {}, manipulation_response=agent_response - ) - self.add_entry(entry) - return entry - - def search(self, **kwargs) -> list[ManipulationHistoryEntry]: # type: ignore[no-untyped-def] - """Flexible search method that can search by any field in ManipulationHistoryEntry using dot notation. - - This method supports dot notation to access nested fields. String values automatically use - substring matching (contains), while all other types use exact matching. - - Examples: - # Time-based searches: - - search(**{"task.metadata.timestamp": ('>', start_time)}) - entries after start_time - - search(**{"task.metadata.timestamp": ('>=', time - 1800)}) - entries in last 30 mins - - # Constraint searches: - - search(**{"task.constraints.*.reference_point.x": 2.5}) - tasks with x=2.5 reference point - - search(**{"task.constraints.*.end_angle.x": 90}) - tasks with 90-degree x rotation - - search(**{"task.constraints.*.lock_x": True}) - tasks with x-axis translation locked - - # Object and result searches: - - search(**{"task.metadata.objects.*.label": "cup"}) - tasks involving cups - - search(**{"result.status": "success"}) - successful tasks - - search(**{"result.error": "Collision"}) - tasks that had collisions - - Args: - **kwargs: Key-value pairs for searching using dot notation for field paths. - - Returns: - List of matching entries - """ - if not kwargs: - return self._history.copy() - - results = self._history.copy() - - for key, value in kwargs.items(): - # For all searches, automatically determine if we should use contains for strings - results = [e for e in results if self._check_field_match(e, key, value)] - - return results - - def _check_field_match(self, entry, field_path, value) -> bool: # type: ignore[no-untyped-def] - """Check if a field matches the value, with special handling for strings, collections and comparisons. - - For string values, we automatically use substring matching (contains). - For collections (returned by * path), we check if any element matches. - For numeric values (like timestamps), supports >, <, >= and <= comparisons. - For all other types, we use exact matching. - - Args: - entry: The entry to check - field_path: Dot-separated path to the field - value: Value to match against. For comparisons, use tuples like: - ('>', timestamp) - greater than - ('<', timestamp) - less than - ('>=', timestamp) - greater or equal - ('<=', timestamp) - less or equal - - Returns: - True if the field matches the value, False otherwise - """ - try: - field_value = self._get_value_by_path(entry, field_path) # type: ignore[no-untyped-call] - - # Handle comparison operators for timestamps and numbers - if isinstance(value, tuple) and len(value) == 2: - op, compare_value = value - if op == ">": - return field_value > compare_value # type: ignore[no-any-return] - elif op == "<": - return field_value < compare_value # type: ignore[no-any-return] - elif op == ">=": - return field_value >= compare_value # type: ignore[no-any-return] - elif op == "<=": - return field_value <= compare_value # type: ignore[no-any-return] - - # Handle lists (from collection searches) - if isinstance(field_value, list): - for item in field_value: - # String values use contains matching - if isinstance(item, str) and isinstance(value, str): - if value in item: - return True - # All other types use exact matching - elif item == value: - return True - return False - - # String values use contains matching - elif isinstance(field_value, str) and isinstance(value, str): - return value in field_value - # All other types use exact matching - else: - return field_value == value # type: ignore[no-any-return] - - except (AttributeError, KeyError): - return False - - def _get_value_by_path(self, obj, path): # type: ignore[no-untyped-def] - """Get a value from an object using a dot-separated path. - - This method handles three special cases: - 1. Regular attribute access (obj.attr) - 2. Dictionary key access (dict[key]) - 3. Collection search (dict.*.attr) - when * is used, it searches all values in the collection - - Args: - obj: Object to get value from - path: Dot-separated path to the field (e.g., "task.metadata.robot") - - Returns: - Value at the specified path or list of values for collection searches - - Raises: - AttributeError: If an attribute in the path doesn't exist - KeyError: If a dictionary key in the path doesn't exist - """ - current = obj - parts = path.split(".") - - for i, part in enumerate(parts): - # Collection search (*.attr) - search across all items in a collection - if part == "*": - # Get remaining path parts - remaining_path = ".".join(parts[i + 1 :]) - - # Handle different collection types - if isinstance(current, dict): - items = current.values() - if not remaining_path: # If * is the last part, return all values - return list(items) - elif isinstance(current, list): - items = current # type: ignore[assignment] - if not remaining_path: # If * is the last part, return all items - return items - else: # Not a collection - raise AttributeError( - f"Cannot use wildcard on non-collection type: {type(current)}" - ) - - # Apply remaining path to each item in the collection - results = [] - for item in items: - try: - # Recursively get values from each item - value = self._get_value_by_path(item, remaining_path) # type: ignore[no-untyped-call] - if isinstance(value, list): # Flatten nested lists - results.extend(value) - else: - results.append(value) - except (AttributeError, KeyError): - # Skip items that don't have the attribute - pass - return results - - # Regular attribute/key access - elif isinstance(current, dict): - current = current[part] - else: - current = getattr(current, part) - - return current diff --git a/dimos/manipulation/planning/mesh_utils.py b/dimos/manipulation/planning/utils/mesh_utils.py similarity index 100% rename from dimos/manipulation/planning/mesh_utils.py rename to dimos/manipulation/planning/utils/mesh_utils.py diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index d7fc5b93fa..0df7bc548f 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -59,12 +59,12 @@ import numpy as np -from dimos.manipulation.planning.mesh_utils import prepare_urdf_for_drake from dimos.manipulation.planning.spec import ( Obstacle, ObstacleType, RobotModelConfig, ) +from dimos.manipulation.planning.utils.mesh_utils import prepare_urdf_for_drake from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: @@ -176,6 +176,9 @@ def __init__( self._builder, time_step=time_step ) self._parser = Parser(self._plant) + # Enable auto-renaming to avoid conflicts when adding multiple robots + # with the same URDF (e.g., 4 XArm6 arms all have model name "UF_ROBOT") + self._parser.SetAutoRenaming(True) # Visualization self._meshcat: Meshcat | None = None @@ -730,6 +733,10 @@ def scratch_context(self) -> Generator[Context, None, None]: Uses CreateDefaultContext() instead of Clone() so that dynamically added obstacles are visible to collision queries. + IMPORTANT: This copies the current live state of ALL robots into the + scratch context so that collision checking properly accounts for + inter-robot collisions (robot A planning considers robot B's position). + Usage: with world.scratch_context() as ctx: world.set_positions(ctx, robot_id, q) @@ -743,6 +750,25 @@ def scratch_context(self) -> Generator[Context, None, None]: # (Clone() doesn't see geometries added after the original context was created) ctx = self._diagram.CreateDefaultContext() + # Copy current live state of ALL robots into scratch context + # This ensures inter-robot collision checking works correctly + with self._lock: + if self._plant_context is not None: + plant_ctx = self._diagram.GetMutableSubsystemContext(self._plant, ctx) + for _robot_id, robot_data in self._robots.items(): + try: + # Get positions from live context + live_positions = self._plant.GetPositions( + self._plant_context, robot_data.model_instance + ) + # Set in scratch context + self._plant.SetPositions( + plant_ctx, robot_data.model_instance, live_positions + ) + except Exception: + # Robot not yet synced - leave at default position + pass + yield ctx # Context automatically cleaned up when exiting @@ -948,9 +974,21 @@ def animate_path( if self._meshcat is None or len(path) < 2: return + # Capture current positions of all OTHER robots so they don't snap to zero + other_robot_positions: dict[str, NDArray[np.float64]] = {} + for rid, _robot_data in self._robots.items(): + if rid != robot_id: + pos = self.get_positions(self.get_live_context(), rid) + if pos is not None: + other_robot_positions[rid] = pos + dt = duration / (len(path) - 1) for q in path: with self.scratch_context() as ctx: + # Restore other robots to their current positions + for rid, pos in other_robot_positions.items(): + self.set_positions(ctx, rid, pos) + # Set animated robot's position self.set_positions(ctx, robot_id, q) self.publish_to_meshcat(ctx) time.sleep(dt) diff --git a/dimos/manipulation/test_manipulation_history.py b/dimos/manipulation/test_manipulation_history.py deleted file mode 100644 index ec4e503bed..0000000000 --- a/dimos/manipulation/test_manipulation_history.py +++ /dev/null @@ -1,458 +0,0 @@ -# Copyright 2025-2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# 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](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. - -import os -import tempfile -import time - -import pytest - -from dimos.manipulation.manipulation_history import ManipulationHistory, ManipulationHistoryEntry -from dimos.types.manipulation import ( - ForceConstraint, - ManipulationTask, - RotationConstraint, - TranslationConstraint, -) -from dimos.types.vector import Vector - - -@pytest.fixture -def sample_task(): - """Create a sample manipulation task for testing.""" - return ManipulationTask( - description="Pick up the cup", - target_object="cup", - target_point=(100, 200), - task_id="task1", - metadata={ - "timestamp": time.time(), - "objects": { - "cup1": { - "object_id": 1, - "label": "cup", - "confidence": 0.95, - "position": {"x": 1.5, "y": 2.0, "z": 0.5}, - }, - "table1": { - "object_id": 2, - "label": "table", - "confidence": 0.98, - "position": {"x": 0.0, "y": 0.0, "z": 0.0}, - }, - }, - }, - ) - - -@pytest.fixture -def sample_task_with_constraints(): - """Create a sample manipulation task with constraints for testing.""" - task = ManipulationTask( - description="Rotate the bottle", - target_object="bottle", - target_point=(150, 250), - task_id="task2", - metadata={ - "timestamp": time.time(), - "objects": { - "bottle1": { - "object_id": 3, - "label": "bottle", - "confidence": 0.92, - "position": {"x": 2.5, "y": 1.0, "z": 0.3}, - } - }, - }, - ) - - # Add rich translation constraint - translation_constraint = TranslationConstraint( - translation_axis="y", - reference_point=Vector(2.5, 1.0, 0.3), - bounds_min=Vector(2.0, 0.5, 0.3), - bounds_max=Vector(3.0, 1.5, 0.3), - target_point=Vector(2.7, 1.2, 0.3), - description="Constrained translation along Y-axis only", - ) - task.add_constraint(translation_constraint) - - # Add rich rotation constraint - rotation_constraint = RotationConstraint( - rotation_axis="roll", - start_angle=Vector(0, 0, 0), - end_angle=Vector(90, 0, 0), - pivot_point=Vector(2.5, 1.0, 0.3), - secondary_pivot_point=Vector(2.5, 1.0, 0.5), - description="Constrained rotation around X-axis (roll only)", - ) - task.add_constraint(rotation_constraint) - - # Add force constraint - force_constraint = ForceConstraint( - min_force=2.0, - max_force=5.0, - force_direction=Vector(0, 0, -1), - description="Apply moderate downward force during manipulation", - ) - task.add_constraint(force_constraint) - - return task - - -@pytest.fixture -def temp_output_dir(): - """Create a temporary directory for testing history saving/loading.""" - with tempfile.TemporaryDirectory() as temp_dir: - yield temp_dir - - -@pytest.fixture -def populated_history(sample_task, sample_task_with_constraints): - """Create a populated history with multiple entries for testing.""" - history = ManipulationHistory() - - # Add first entry - entry1 = ManipulationHistoryEntry( - task=sample_task, - result={"status": "success", "execution_time": 2.5}, - manipulation_response="Successfully picked up the cup", - ) - history.add_entry(entry1) - - # Add second entry - entry2 = ManipulationHistoryEntry( - task=sample_task_with_constraints, - result={"status": "failure", "error": "Collision detected"}, - manipulation_response="Failed to rotate the bottle due to collision", - ) - history.add_entry(entry2) - - return history - - -def test_manipulation_history_init() -> None: - """Test initialization of ManipulationHistory.""" - # Default initialization - history = ManipulationHistory() - assert len(history) == 0 - assert str(history) == "ManipulationHistory(empty)" - - # With output directory - with tempfile.TemporaryDirectory() as temp_dir: - history = ManipulationHistory(output_dir=temp_dir, new_memory=True) - assert len(history) == 0 - assert os.path.exists(temp_dir) - - -def test_manipulation_history_add_entry(sample_task) -> None: - """Test adding entries to ManipulationHistory.""" - history = ManipulationHistory() - - # Create and add entry - entry = ManipulationHistoryEntry( - task=sample_task, result={"status": "success"}, manipulation_response="Task completed" - ) - history.add_entry(entry) - - assert len(history) == 1 - assert history.get_entry_by_index(0) == entry - - -def test_manipulation_history_create_task_entry(sample_task) -> None: - """Test creating a task entry directly.""" - history = ManipulationHistory() - - entry = history.create_task_entry( - task=sample_task, result={"status": "success"}, agent_response="Task completed" - ) - - assert len(history) == 1 - assert entry.task == sample_task - assert entry.result["status"] == "success" - assert entry.manipulation_response == "Task completed" - - -def test_manipulation_history_save_load(temp_output_dir, sample_task) -> None: - """Test saving and loading history from disk.""" - # Create history and add entry - history = ManipulationHistory(output_dir=temp_output_dir) - history.create_task_entry( - task=sample_task, result={"status": "success"}, agent_response="Task completed" - ) - - # Check that files were created - pickle_path = os.path.join(temp_output_dir, "manipulation_history.pickle") - json_path = os.path.join(temp_output_dir, "manipulation_history.json") - assert os.path.exists(pickle_path) - assert os.path.exists(json_path) - - # Create new history that loads from the saved files - loaded_history = ManipulationHistory(output_dir=temp_output_dir) - assert len(loaded_history) == 1 - assert loaded_history.get_entry_by_index(0).task.description == sample_task.description - - -def test_manipulation_history_clear(populated_history) -> None: - """Test clearing the history.""" - assert len(populated_history) > 0 - - populated_history.clear() - assert len(populated_history) == 0 - assert str(populated_history) == "ManipulationHistory(empty)" - - -def test_manipulation_history_get_methods(populated_history) -> None: - """Test various getter methods of ManipulationHistory.""" - # get_all_entries - entries = populated_history.get_all_entries() - assert len(entries) == 2 - - # get_entry_by_index - entry = populated_history.get_entry_by_index(0) - assert entry.task.task_id == "task1" - - # Out of bounds index - assert populated_history.get_entry_by_index(100) is None - - # get_entries_by_timerange - start_time = time.time() - 3600 # 1 hour ago - end_time = time.time() + 3600 # 1 hour from now - entries = populated_history.get_entries_by_timerange(start_time, end_time) - assert len(entries) == 2 - - # get_entries_by_object - cup_entries = populated_history.get_entries_by_object("cup") - assert len(cup_entries) == 1 - assert cup_entries[0].task.task_id == "task1" - - bottle_entries = populated_history.get_entries_by_object("bottle") - assert len(bottle_entries) == 1 - assert bottle_entries[0].task.task_id == "task2" - - -def test_manipulation_history_search_basic(populated_history) -> None: - """Test basic search functionality.""" - # Search by exact match on top-level fields - results = populated_history.search(timestamp=populated_history.get_entry_by_index(0).timestamp) - assert len(results) == 1 - - # Search by task fields - results = populated_history.search(**{"task.task_id": "task1"}) - assert len(results) == 1 - assert results[0].task.target_object == "cup" - - # Search by result fields - results = populated_history.search(**{"result.status": "success"}) - assert len(results) == 1 - assert results[0].task.task_id == "task1" - - # Search by manipulation_response (substring match for strings) - results = populated_history.search(manipulation_response="picked up") - assert len(results) == 1 - assert results[0].task.task_id == "task1" - - -def test_manipulation_history_search_nested(populated_history) -> None: - """Test search with nested field paths.""" - # Search by nested metadata fields - results = populated_history.search( - **{ - "task.metadata.timestamp": populated_history.get_entry_by_index(0).task.metadata[ - "timestamp" - ] - } - ) - assert len(results) == 1 - - # Search by nested object fields - results = populated_history.search(**{"task.metadata.objects.cup1.label": "cup"}) - assert len(results) == 1 - assert results[0].task.task_id == "task1" - - # Search by position values - results = populated_history.search(**{"task.metadata.objects.cup1.position.x": 1.5}) - assert len(results) == 1 - assert results[0].task.task_id == "task1" - - -def test_manipulation_history_search_wildcards(populated_history) -> None: - """Test search with wildcard patterns.""" - # Search for any object with label "cup" - results = populated_history.search(**{"task.metadata.objects.*.label": "cup"}) - assert len(results) == 1 - assert results[0].task.task_id == "task1" - - # Search for any object with confidence > 0.95 - results = populated_history.search(**{"task.metadata.objects.*.confidence": 0.98}) - assert len(results) == 1 - assert results[0].task.task_id == "task1" - - # Search for any object position with x=2.5 - results = populated_history.search(**{"task.metadata.objects.*.position.x": 2.5}) - assert len(results) == 1 - assert results[0].task.task_id == "task2" - - -def test_manipulation_history_search_constraints(populated_history) -> None: - """Test search by constraint properties.""" - # Find entries with any TranslationConstraint with y-axis - results = populated_history.search(**{"task.constraints.*.translation_axis": "y"}) - assert len(results) == 1 - assert results[0].task.task_id == "task2" - - # Find entries with any RotationConstraint with roll axis - results = populated_history.search(**{"task.constraints.*.rotation_axis": "roll"}) - assert len(results) == 1 - assert results[0].task.task_id == "task2" - - -def test_manipulation_history_search_string_contains(populated_history) -> None: - """Test string contains searching.""" - # Basic string contains - results = populated_history.search(**{"task.description": "Pick"}) - assert len(results) == 1 - assert results[0].task.task_id == "task1" - - # Nested string contains - results = populated_history.search(manipulation_response="collision") - assert len(results) == 1 - assert results[0].task.task_id == "task2" - - -def test_manipulation_history_search_multiple_criteria(populated_history) -> None: - """Test search with multiple criteria.""" - # Multiple criteria - all must match - results = populated_history.search(**{"task.target_object": "cup", "result.status": "success"}) - assert len(results) == 1 - assert results[0].task.task_id == "task1" - - # Multiple criteria with no matches - results = populated_history.search(**{"task.target_object": "cup", "result.status": "failure"}) - assert len(results) == 0 - - # Combination of direct and wildcard paths - results = populated_history.search( - **{"task.target_object": "bottle", "task.metadata.objects.*.position.z": 0.3} - ) - assert len(results) == 1 - assert results[0].task.task_id == "task2" - - -def test_manipulation_history_search_nonexistent_fields(populated_history) -> None: - """Test search with fields that don't exist.""" - # Search by nonexistent field - results = populated_history.search(nonexistent_field="value") - assert len(results) == 0 - - # Search by nonexistent nested field - results = populated_history.search(**{"task.nonexistent_field": "value"}) - assert len(results) == 0 - - # Search by nonexistent object - results = populated_history.search(**{"task.metadata.objects.nonexistent_object": "value"}) - assert len(results) == 0 - - -def test_manipulation_history_search_timestamp_ranges(populated_history) -> None: - """Test searching by timestamp ranges.""" - # Get reference timestamps - entry1_time = populated_history.get_entry_by_index(0).task.metadata["timestamp"] - entry2_time = populated_history.get_entry_by_index(1).task.metadata["timestamp"] - mid_time = (entry1_time + entry2_time) / 2 - - # Search for timestamps before second entry - results = populated_history.search(**{"task.metadata.timestamp": ("<", entry2_time)}) - assert len(results) == 1 - assert results[0].task.task_id == "task1" - - # Search for timestamps after first entry - results = populated_history.search(**{"task.metadata.timestamp": (">", entry1_time)}) - assert len(results) == 1 - assert results[0].task.task_id == "task2" - - # Search within a time window using >= and <= - results = populated_history.search(**{"task.metadata.timestamp": (">=", mid_time - 1800)}) - assert len(results) == 2 - assert results[0].task.task_id == "task1" - assert results[1].task.task_id == "task2" - - -def test_manipulation_history_search_vector_fields(populated_history) -> None: - """Test searching by vector components in constraints.""" - # Search by reference point components - results = populated_history.search(**{"task.constraints.*.reference_point.x": 2.5}) - assert len(results) == 1 - assert results[0].task.task_id == "task2" - - # Search by target point components - results = populated_history.search(**{"task.constraints.*.target_point.z": 0.3}) - assert len(results) == 1 - assert results[0].task.task_id == "task2" - - # Search by rotation angles - results = populated_history.search(**{"task.constraints.*.end_angle.x": 90}) - assert len(results) == 1 - assert results[0].task.task_id == "task2" - - -def test_manipulation_history_search_execution_details(populated_history) -> None: - """Test searching by execution time and error patterns.""" - # Search by execution time - results = populated_history.search(**{"result.execution_time": 2.5}) - assert len(results) == 1 - assert results[0].task.task_id == "task1" - - # Search by error message pattern - results = populated_history.search(**{"result.error": "Collision"}) - assert len(results) == 1 - assert results[0].task.task_id == "task2" - - # Search by status - results = populated_history.search(**{"result.status": "success"}) - assert len(results) == 1 - assert results[0].task.task_id == "task1" - - -def test_manipulation_history_search_multiple_criteria(populated_history) -> None: - """Test search with multiple criteria.""" - # Multiple criteria - all must match - results = populated_history.search(**{"task.target_object": "cup", "result.status": "success"}) - assert len(results) == 1 - assert results[0].task.task_id == "task1" - - # Multiple criteria with no matches - results = populated_history.search(**{"task.target_object": "cup", "result.status": "failure"}) - assert len(results) == 0 - - # Combination of direct and wildcard paths - results = populated_history.search( - **{"task.target_object": "bottle", "task.metadata.objects.*.position.z": 0.3} - ) - assert len(results) == 1 - assert results[0].task.task_id == "task2" From 8545bd1abf4ec2b3a456027506ad8699ec3fdb7f Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Wed, 14 Jan 2026 14:43:32 -0800 Subject: [PATCH 14/31] added multi robot management and control to the manipulation module --- dimos/manipulation/manipulation_module.py | 82 ++++++++++++------- .../planning/examples/manipulation_client.py | 78 +++++++++++++----- 2 files changed, 111 insertions(+), 49 deletions(-) diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index c99be5d557..94c90cd468 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -184,13 +184,11 @@ def _get_robot( def _on_joint_state(self, msg: JointState) -> None: """Callback when joint state received from driver.""" try: - # Forward to world monitor for state synchronization - # For single robot, use default; for multi-robot, monitor routes by joint names + # Forward to world monitor for state synchronization. + # Pass robot_id=None to broadcast to all monitors - each monitor + # extracts only its robot's joints based on joint_name_mapping. if self._world_monitor is not None: - robot = self._get_robot() - if robot is not None: - _, robot_id, _, _ = robot - self._world_monitor.on_joint_state(msg, robot_id) + self._world_monitor.on_joint_state(msg, robot_id=None) except Exception as e: logger.error(f"Exception in _on_joint_state: {e}") @@ -235,25 +233,38 @@ def reset(self) -> bool: return True @rpc - def get_current_joints(self) -> list[float] | None: - """Get current joint positions.""" - if (robot := self._get_robot()) and self._world_monitor: + def get_current_joints(self, robot_name: str | None = None) -> list[float] | None: + """Get current joint positions. + + Args: + robot_name: Robot to query (required if multiple robots configured) + """ + if (robot := self._get_robot(robot_name)) and self._world_monitor: pos = self._world_monitor.get_current_positions(robot[1]) if pos is not None: return list(pos) return None @rpc - def get_ee_pose(self) -> Pose | None: - """Get current end-effector pose.""" - if (robot := self._get_robot()) and self._world_monitor: + def get_ee_pose(self, robot_name: str | None = None) -> Pose | None: + """Get current end-effector pose. + + Args: + robot_name: Robot to query (required if multiple robots configured) + """ + if (robot := self._get_robot(robot_name)) and self._world_monitor: return matrix_to_pose(self._world_monitor.get_ee_pose(robot[1])) return None @rpc - def is_collision_free(self, joints: list[float]) -> bool: - """Check if joint configuration is collision-free.""" - if (robot := self._get_robot()) and self._world_monitor: + def is_collision_free(self, joints: list[float], robot_name: str | None = None) -> bool: + """Check if joint configuration is collision-free. + + Args: + joints: Joint configuration to check + robot_name: Robot to check (required if multiple robots configured) + """ + if (robot := self._get_robot(robot_name)) and self._world_monitor: return self._world_monitor.is_state_valid( robot[1], np.array(joints) ) # robot[1] is the robot_id. @@ -263,12 +274,16 @@ def is_collision_free(self, joints: list[float]) -> bool: # Plan/Preview/Execute Workflow RPC Methods # ========================================================================= - def _begin_planning(self) -> tuple[str, str] | None: - """Check state and begin planning. Returns (robot_name, robot_id) or None.""" + def _begin_planning(self, robot_name: str | None = None) -> tuple[str, str] | None: + """Check state and begin planning. Returns (robot_name, robot_id) or None. + + Args: + robot_name: Robot to plan for (required if multiple robots configured) + """ if self._world_monitor is None: logger.error("Planning not initialized") return None - if (robot := self._get_robot()) is None: + if (robot := self._get_robot(robot_name)) is None: return None with self._lock: if self._state not in (ManipulationState.IDLE, ManipulationState.COMPLETED): @@ -285,9 +300,14 @@ def _fail(self, msg: str) -> bool: return False @rpc - def plan_to_pose(self, pose: Pose) -> bool: - """Plan motion to pose. Use preview_path() then execute().""" - if self._kinematics is None or (r := self._begin_planning()) is None: + def plan_to_pose(self, pose: Pose, robot_name: str | None = None) -> bool: + """Plan motion to pose. Use preview_path() then execute(). + + Args: + pose: Target end-effector pose + robot_name: Robot to plan for (required if multiple robots configured) + """ + if self._kinematics is None or (r := self._begin_planning(robot_name)) is None: return False robot_name, robot_id = r assert self._world_monitor # guaranteed by _begin_planning @@ -310,12 +330,17 @@ def plan_to_pose(self, pose: Pose) -> bool: return self._plan_path_only(robot_name, robot_id, ik.joint_positions) @rpc - def plan_to_joints(self, joints: list[float]) -> bool: - """Plan motion to joint config. Use preview_path() then execute().""" - if (r := self._begin_planning()) is None: + def plan_to_joints(self, joints: list[float], robot_name: str | None = None) -> bool: + """Plan motion to joint config. Use preview_path() then execute(). + + Args: + joints: Target joint configuration + robot_name: Robot to plan for (required if multiple robots configured) + """ + if (r := self._begin_planning(robot_name)) is None: return False robot_name, robot_id = r - logger.info(f"Planning to joints: {[f'{j:.3f}' for j in joints]}") + logger.info(f"Planning to joints for {robot_name}: {[f'{j:.3f}' for j in joints]}") return self._plan_path_only(robot_name, robot_id, np.array(joints)) def _plan_path_only(self, robot_name: str, robot_id: str, goal: NDArray[np.float64]) -> bool: @@ -347,25 +372,26 @@ def _plan_path_only(self, robot_name: str, robot_id: str, goal: NDArray[np.float return True @rpc - def preview_path(self, duration: float = 3.0) -> bool: + def preview_path(self, duration: float = 3.0, robot_name: str | None = None) -> bool: """Preview the planned path in the visualizer. Args: duration: Total animation duration in seconds + robot_name: Robot to preview (required if multiple robots configured) """ from dimos.manipulation.planning.utils.path_utils import interpolate_path if self._world_monitor is None: return False - robot = self._get_robot() + robot = self._get_robot(robot_name) if robot is None: return False robot_name, robot_id, _, _ = robot planned_path = self._planned_paths.get(robot_name) if planned_path is None or len(planned_path) == 0: - logger.warning("No planned path to preview") + logger.warning(f"No planned path to preview for {robot_name}") return False # Interpolate and animate diff --git a/dimos/manipulation/planning/examples/manipulation_client.py b/dimos/manipulation/planning/examples/manipulation_client.py index 47b73db36f..08f79f6e36 100644 --- a/dimos/manipulation/planning/examples/manipulation_client.py +++ b/dimos/manipulation/planning/examples/manipulation_client.py @@ -81,13 +81,21 @@ def state(self) -> str: """Get manipulation state.""" return cast("str", self._call("get_state")) - def joints(self) -> list[float] | None: - """Get current joint positions.""" - return cast("list[float] | None", self._call("get_current_joints")) + def joints(self, robot_name: str | None = None) -> list[float] | None: + """Get current joint positions. - def ee(self) -> Pose | None: - """Get end-effector pose.""" - return cast("Pose | None", self._call("get_ee_pose")) + Args: + robot_name: Robot to query (required if multiple robots configured) + """ + return cast("list[float] | None", self._call("get_current_joints", robot_name)) + + def ee(self, robot_name: str | None = None) -> Pose | None: + """Get end-effector pose. + + Args: + robot_name: Robot to query (required if multiple robots configured) + """ + return cast("Pose | None", self._call("get_ee_pose", robot_name)) def url(self) -> str | None: """Get Meshcat visualization URL.""" @@ -105,10 +113,18 @@ def info(self, robot_name: str | None = None) -> dict[str, Any] | None: # Planning Methods # ========================================================================= - def plan(self, joints: list[float]) -> bool: - """Plan to joint configuration.""" - print(f"Planning to: {[f'{j:.3f}' for j in joints]}") - return cast("bool", self._call("plan_to_joints", joints)) + def plan(self, joints: list[float], robot_name: str | None = None) -> bool: + """Plan to joint configuration. + + Args: + joints: Target joint configuration + robot_name: Robot to plan for (required if multiple robots configured) + """ + print( + f"Planning to: {[f'{j:.3f}' for j in joints]}" + + (f" for {robot_name}" if robot_name else "") + ) + return cast("bool", self._call("plan_to_joints", joints, robot_name=robot_name)) def plan_pose( self, @@ -118,28 +134,43 @@ def plan_pose( roll: float | None = None, pitch: float | None = None, yaw: float | None = None, + robot_name: str | None = None, ) -> bool: - """Plan to cartesian pose. Uses current orientation if not specified.""" + """Plan to cartesian pose. Uses current orientation if not specified. + + Args: + x, y, z: Target position + roll, pitch, yaw: Target orientation (uses current if not specified) + robot_name: Robot to plan for (required if multiple robots configured) + """ # Get current orientation if not provided if roll is None or pitch is None or yaw is None: - ee = self.ee() + ee = self.ee(robot_name) if ee is None: - print("Cannot get current orientation") + print("Cannot get current orientation - specify roll, pitch, yaw explicitly") return False roll = roll if roll is not None else ee.roll pitch = pitch if pitch is not None else ee.pitch yaw = yaw if yaw is not None else ee.yaw - print(f"Planning to: ({x:.3f}, {y:.3f}, {z:.3f}) rpy=({roll:.2f}, {pitch:.2f}, {yaw:.2f})") + print( + f"Planning to: ({x:.3f}, {y:.3f}, {z:.3f}) rpy=({roll:.2f}, {pitch:.2f}, {yaw:.2f})" + + (f" for {robot_name}" if robot_name else "") + ) pose = Pose( position=Vector3(x, y, z), orientation=Quaternion.from_euler(Vector3(roll, pitch, yaw)), ) - return cast("bool", self._call("plan_to_pose", pose)) + return cast("bool", self._call("plan_to_pose", pose, robot_name=robot_name)) - def preview(self, duration: float = 3.0) -> bool: - """Preview planned path in Meshcat.""" - return cast("bool", self._call("preview_path", duration)) + def preview(self, duration: float = 3.0, robot_name: str | None = None) -> bool: + """Preview planned path in Meshcat. + + Args: + duration: Animation duration in seconds + robot_name: Robot to preview (required if multiple robots configured) + """ + return cast("bool", self._call("preview_path", duration, robot_name=robot_name)) def execute(self, robot_name: str | None = None) -> bool: """Execute planned trajectory via orchestrator.""" @@ -197,9 +228,14 @@ def remove(self, obstacle_id: str) -> bool: # Utility Methods # ========================================================================= - def collision(self, joints: list[float]) -> bool: - """Check if joint config is collision-free.""" - return cast("bool", self._call("is_collision_free", joints)) + def collision(self, joints: list[float], robot_name: str | None = None) -> bool: + """Check if joint config is collision-free. + + Args: + joints: Joint configuration to check + robot_name: Robot to check (required if multiple robots configured) + """ + return cast("bool", self._call("is_collision_free", joints, robot_name)) def reset(self) -> bool: """Reset to IDLE state.""" From b5605fc3157afe7d1b551f7640c6085988409f88 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Wed, 14 Jan 2026 14:46:57 -0800 Subject: [PATCH 15/31] refactored planning stack --- dimos/manipulation/manipulation_blueprints.py | 216 +++++- .../planning/kinematics/drake_kinematics.py | 143 +--- .../planning/monitor/world_monitor.py | 376 +-------- .../planning/planners/drake_planner.py | 275 ++----- dimos/manipulation/planning/spec.py | 733 ------------------ dimos/manipulation/planning/spec/__init__.py | 83 ++ dimos/manipulation/planning/spec/config.py | 91 +++ dimos/manipulation/planning/spec/enums.py | 49 ++ dimos/manipulation/planning/spec/protocols.py | 261 +++++++ dimos/manipulation/planning/spec/types.py | 154 ++++ .../planning/world/drake_world.py | 375 +++------ 11 files changed, 1092 insertions(+), 1664 deletions(-) delete mode 100644 dimos/manipulation/planning/spec.py create mode 100644 dimos/manipulation/planning/spec/__init__.py create mode 100644 dimos/manipulation/planning/spec/config.py create mode 100644 dimos/manipulation/planning/spec/enums.py create mode 100644 dimos/manipulation/planning/spec/protocols.py create mode 100644 dimos/manipulation/planning/spec/types.py diff --git a/dimos/manipulation/manipulation_blueprints.py b/dimos/manipulation/manipulation_blueprints.py index b12693a838..f3334bcbc0 100644 --- a/dimos/manipulation/manipulation_blueprints.py +++ b/dimos/manipulation/manipulation_blueprints.py @@ -29,6 +29,7 @@ import numpy as np +from dimos.core.blueprints import autoconnect from dimos.core.transport import LCMTransport from dimos.manipulation.manipulation_module import manipulation_module from dimos.manipulation.planning.spec import RobotModelConfig @@ -60,15 +61,41 @@ def _get_piper_package_paths() -> dict[str, str]: return {"piper_description": str(get_data("piper_description"))} +# Piper gripper collision exclusions (parallel jaw gripper) +# The gripper fingers (link7, link8) can touch each other and gripper_base +PIPER_GRIPPER_COLLISION_EXCLUSIONS: list[tuple[str, str]] = [ + ("gripper_base", "link7"), + ("gripper_base", "link8"), + ("link7", "link8"), + ("link6", "gripper_base"), +] + + # XArm gripper collision exclusions (parallel linkage mechanism) +# The gripper uses mimic joints where non-adjacent links can overlap legitimately XARM_GRIPPER_COLLISION_EXCLUSIONS: list[tuple[str, str]] = [ + # Inner knuckle <-> outer knuckle (parallel linkage) ("right_inner_knuckle", "right_outer_knuckle"), - ("right_inner_knuckle", "right_finger"), ("left_inner_knuckle", "left_outer_knuckle"), + # Inner knuckle <-> finger (parallel linkage) + ("right_inner_knuckle", "right_finger"), ("left_inner_knuckle", "left_finger"), + # Cross-finger pairs (mimic joint symmetry) ("left_finger", "right_finger"), ("left_outer_knuckle", "right_outer_knuckle"), ("left_inner_knuckle", "right_inner_knuckle"), + # Outer knuckle <-> opposite finger + ("left_outer_knuckle", "right_finger"), + ("right_outer_knuckle", "left_finger"), + # Gripper base <-> all moving parts (can touch at limits) + ("xarm_gripper_base_link", "left_inner_knuckle"), + ("xarm_gripper_base_link", "right_inner_knuckle"), + ("xarm_gripper_base_link", "left_finger"), + ("xarm_gripper_base_link", "right_finger"), + # Arm link6 <-> gripper (attached via fixed joint, can touch) + ("link6", "xarm_gripper_base_link"), + ("link6", "left_outer_knuckle"), + ("link6", "right_outer_knuckle"), ] @@ -110,8 +137,15 @@ def _make_xarm6_config( end_effector_link="link_tcp", base_link="link_base", package_paths=_get_xarm_package_paths(), - xacro_args={"dof": "6", "limited": "true", "add_gripper": "true"}, + # Pass attach_xyz to position each robot; xacro creates world_joint at this offset + xacro_args={ + "dof": "6", + "limited": "true", + "add_gripper": "true", + "attach_xyz": f"0 {y_offset} 0", + }, collision_exclusion_pairs=XARM_GRIPPER_COLLISION_EXCLUSIONS, + auto_convert_meshes=True, max_velocity=1.0, max_acceleration=2.0, joint_name_mapping=joint_mapping, @@ -121,22 +155,97 @@ def _make_xarm6_config( def _make_xarm7_config( name: str = "arm", + y_offset: float = 0.0, joint_prefix: str = "", orchestrator_task: str | None = None, + add_gripper: bool = False, ) -> RobotModelConfig: - """Create XArm7 robot config.""" + """Create XArm7 robot config. + + Args: + name: Robot name in Drake world + y_offset: Y-axis offset for base pose (for multi-arm setups) + joint_prefix: Prefix for joint name mapping (e.g., "left_" or "right_") + orchestrator_task: Task name for orchestrator RPC execution + add_gripper: Whether to add the xarm gripper + """ joint_names = ["joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "joint7"] joint_mapping = {f"{joint_prefix}{j}": j for j in joint_names} if joint_prefix else {} + xacro_args: dict[str, str] = { + "dof": "7", + "limited": "true", + "attach_xyz": f"0 {y_offset} 0", + } + if add_gripper: + xacro_args["add_gripper"] = "true" + return RobotModelConfig( name=name, urdf_path=_get_xarm_urdf_path(), - base_pose=np.eye(4, dtype=np.float64), + base_pose=np.array( + [ + [1, 0, 0, 0], + [0, 1, 0, y_offset], + [0, 0, 1, 0], + [0, 0, 0, 1], + ], + dtype=np.float64, + ), joint_names=joint_names, - end_effector_link="link7", + end_effector_link="link_tcp" if add_gripper else "link7", base_link="link_base", package_paths=_get_xarm_package_paths(), - xacro_args={"dof": "7", "limited": "true"}, + xacro_args=xacro_args, + collision_exclusion_pairs=XARM_GRIPPER_COLLISION_EXCLUSIONS if add_gripper else [], + auto_convert_meshes=True, + max_velocity=1.0, + max_acceleration=2.0, + joint_name_mapping=joint_mapping, + orchestrator_task_name=orchestrator_task, + ) + + +def _make_piper_config( + name: str = "piper", + y_offset: float = 0.0, + joint_prefix: str = "", + orchestrator_task: str | None = None, +) -> RobotModelConfig: + """Create Piper robot config. + + Args: + name: Robot name in Drake world + y_offset: Y-axis offset for base pose (for multi-arm setups) + joint_prefix: Prefix for joint name mapping (e.g., "piper_") + orchestrator_task: Task name for orchestrator RPC execution + + Note: + Piper has 6 revolute joints (joint1-joint6) for the arm and 2 prismatic + joints (joint7, joint8) for the parallel jaw gripper. + """ + # Piper arm joints (6-DOF) + joint_names = ["joint1", "joint2", "joint3", "joint4", "joint5", "joint6"] + joint_mapping = {f"{joint_prefix}{j}": j for j in joint_names} if joint_prefix else {} + + return RobotModelConfig( + name=name, + urdf_path=_get_piper_urdf_path(), + base_pose=np.array( + [ + [1, 0, 0, 0], + [0, 1, 0, y_offset], + [0, 0, 1, 0], + [0, 0, 0, 1], + ], + dtype=np.float64, + ), + joint_names=joint_names, + end_effector_link="gripper_base", # End of arm, before gripper fingers + base_link="arm_base", + package_paths=_get_piper_package_paths(), + xacro_args={}, # Piper xacro doesn't need special args + collision_exclusion_pairs=PIPER_GRIPPER_COLLISION_EXCLUSIONS, auto_convert_meshes=True, max_velocity=1.0, max_acceleration=2.0, @@ -195,9 +304,104 @@ def _make_xarm7_config( ) +# ============================================================================= +# Combined Blueprints (Orchestrator + Planner in one process) +# ============================================================================= + + +# Combined: Quad mock orchestrator + Quad manipulation planner +# Usage: dimos run quad-mock-with-planner +def _make_quad_mock_with_planner(): + """Create combined quad-arm blueprint (lazy import to avoid circular deps).""" + from dimos.control.blueprints import orchestrator_quad_mock + + return autoconnect( + orchestrator_quad_mock, + manipulation_module( + robots=[ + _make_xarm6_config( + "arm1", y_offset=0.75, joint_prefix="arm1_", orchestrator_task="traj_arm1" + ), + _make_xarm6_config( + "arm2", y_offset=0.25, joint_prefix="arm2_", orchestrator_task="traj_arm2" + ), + _make_xarm6_config( + "arm3", y_offset=-0.25, joint_prefix="arm3_", orchestrator_task="traj_arm3" + ), + _make_xarm6_config( + "arm4", y_offset=-0.75, joint_prefix="arm4_", orchestrator_task="traj_arm4" + ), + ], + planning_timeout=10.0, + enable_viz=True, + ), + ).transports( + { + ("joint_state", JointState): LCMTransport("/orchestrator/joint_state", JointState), + } + ) + + +# Lazy-loaded to avoid circular import +quad_mock_with_planner = _make_quad_mock_with_planner() + + +# Combined: Mixed-arm orchestrator (XArm7 + XArm6 + Piper) + manipulation planner +# Usage: dimos run mixed-mock-with-planner +def _make_mixed_mock_with_planner(): + """Create combined mixed-arm blueprint (XArm7 + XArm6 + Piper). + + Layout (top view, looking down at table): + y=0.5: XArm7 (7-DOF with gripper) + y=0.0: XArm6 (6-DOF with gripper) + y=-0.5: Piper (6-DOF with parallel jaw gripper) + """ + from dimos.control.blueprints import orchestrator_mixed_mock + + return autoconnect( + orchestrator_mixed_mock, + manipulation_module( + robots=[ + _make_xarm7_config( + "xarm7", + y_offset=0.5, + joint_prefix="xarm7_", + orchestrator_task="traj_xarm7", + add_gripper=True, + ), + _make_xarm6_config( + "xarm6", + y_offset=0.0, + joint_prefix="xarm6_", + orchestrator_task="traj_xarm6", + ), + _make_piper_config( + "piper", + y_offset=-0.5, + joint_prefix="piper_", + orchestrator_task="traj_piper", + ), + ], + planning_timeout=10.0, + enable_viz=True, + ), + ).transports( + { + ("joint_state", JointState): LCMTransport("/orchestrator/joint_state", JointState), + } + ) + + +# Lazy-loaded to avoid circular import +mixed_mock_with_planner = _make_mixed_mock_with_planner() + + __all__ = [ + "PIPER_GRIPPER_COLLISION_EXCLUSIONS", "XARM_GRIPPER_COLLISION_EXCLUSIONS", "dual_xarm6_planner", + "mixed_mock_with_planner", + "quad_mock_with_planner", "xarm6_planner_only", "xarm7_planner_orchestrator", ] diff --git a/dimos/manipulation/planning/kinematics/drake_kinematics.py b/dimos/manipulation/planning/kinematics/drake_kinematics.py index 952630ac53..7c6bfad218 100644 --- a/dimos/manipulation/planning/kinematics/drake_kinematics.py +++ b/dimos/manipulation/planning/kinematics/drake_kinematics.py @@ -12,28 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -""" -Drake Kinematics Module - -Implements KinematicsSpec using Drake's optimization-based IK. - -## IK Methods - -- solve(): Full nonlinear IK via Drake's InverseKinematics + SNOPT/IPOPT -- solve_iterative(): Iterative differential IK loop -- solve_differential(): Single Jacobian step for velocity control - -## Key Design - -This module is stateless (except for configuration) and uses WorldSpec -for all FK/collision operations via scratch_context(). - -Example: - kinematics = DrakeKinematics(damping=0.01) - result = kinematics.solve(world, robot_id, target_pose, seed=q_current) - if result.is_success(): - q_goal = result.joint_positions -""" +"""Drake Kinematics - optimization-based and differential IK using Drake.""" from __future__ import annotations @@ -66,20 +45,7 @@ class DrakeKinematics: - """Drake implementation of KinematicsSpec. - - Uses Drake's InverseKinematics class for optimization-based IK - and Jacobian-based methods for differential IK. - - This class is stateless except for configuration - it uses WorldSpec's - scratch_context() for all operations, making it thread-safe. - - ## Methods - - - solve(): Full optimization-based IK - - solve_iterative(): Iterative differential IK loop - - solve_differential(): Single Jacobian step for velocity control - """ + """Drake IK solver - optimization-based and Jacobian-based methods.""" def __init__( self, @@ -87,20 +53,24 @@ def __init__( max_iterations: int = 100, singularity_threshold: float = 0.001, ): - """Create Drake kinematics solver. - - Args: - damping: Damping factor for differential IK (higher = more stable near singularities) - max_iterations: Default max iterations for iterative IK - singularity_threshold: Manipulability threshold for singularity detection - """ if not DRAKE_AVAILABLE: raise ImportError("Drake is not installed. Install with: pip install drake") - self._damping = damping self._max_iterations = max_iterations self._singularity_threshold = singularity_threshold + def _validate_world(self, world: WorldSpec) -> IKResult | None: + """Validate world is DrakeWorld and finalized. Returns error or None.""" + from dimos.manipulation.planning.world.drake_world import DrakeWorld + + if not isinstance(world, DrakeWorld): + return _create_failure_result( + IKStatus.NO_SOLUTION, "DrakeKinematics requires DrakeWorld" + ) + if not world.is_finalized: + return _create_failure_result(IKStatus.NO_SOLUTION, "World must be finalized before IK") + return None + def solve( self, world: WorldSpec, @@ -112,39 +82,10 @@ def solve( check_collision: bool = True, max_attempts: int = 10, ) -> IKResult: - """Full nonlinear IK via Drake's InverseKinematics. - - Uses Drake's optimization-based IK with multiple random restarts - for robustness. - - Args: - world: World for FK/collision checking - robot_id: Which robot - target_pose: 4x4 target transform - seed: Initial guess (uses current state if None) - position_tolerance: Position tolerance (meters) - orientation_tolerance: Orientation tolerance (radians) - check_collision: Verify solution is collision-free - max_attempts: Random restarts for robustness - - Returns: - IKResult with status, joint positions, and error metrics - """ - # Access Drake internals via world - # (DrakeWorld exposes plant/diagram for this purpose) - from dimos.manipulation.planning.world.drake_world import DrakeWorld - - if not isinstance(world, DrakeWorld): - return _create_failure_result( - IKStatus.NO_SOLUTION, - "DrakeKinematics requires DrakeWorld", - ) - - if not world.is_finalized: - return _create_failure_result( - IKStatus.NO_SOLUTION, - "World must be finalized before IK", - ) + """Full nonlinear IK with multiple random restarts.""" + error = self._validate_world(world) + if error is not None: + return error # Get joint limits lower_limits, upper_limits = world.get_joint_limits(robot_id) @@ -302,23 +243,7 @@ def solve_iterative( position_tolerance: float = 0.001, orientation_tolerance: float = 0.01, ) -> IKResult: - """Iterative differential IK loop. - - Uses repeated Jacobian steps until convergence. Slower but more - predictable behavior near singularities. - - Args: - world: World for FK/Jacobian computation - robot_id: Which robot - target_pose: 4x4 target transform - seed: Initial joint configuration - max_iterations: Maximum iterations - position_tolerance: Position convergence tolerance (meters) - orientation_tolerance: Orientation convergence tolerance (radians) - - Returns: - IKResult with solution or failure info - """ + """Iterative Jacobian-based IK until convergence.""" max_iterations = max_iterations or self._max_iterations current_joints = seed.copy() @@ -389,21 +314,7 @@ def solve_differential( twist: NDArray[np.float64], dt: float, ) -> NDArray[np.float64] | None: - """Single Jacobian step for velocity control. - - Computes joint velocities to achieve a desired end-effector twist. - Uses damped pseudoinverse for singularity robustness. - - Args: - world: World for Jacobian computation - robot_id: Which robot - current_joints: Current joint positions (radians) - twist: Desired end-effector twist [vx, vy, vz, wx, wy, wz] - dt: Time step (seconds) - used for velocity scaling - - Returns: - Joint velocities (rad/s) or None if near singularity - """ + """Single Jacobian step for velocity control. Returns None if near singularity.""" with world.scratch_context() as ctx: world.set_positions(ctx, robot_id, current_joints) J = world.get_jacobian(ctx, robot_id) @@ -435,19 +346,7 @@ def solve_differential_position_only( current_joints: NDArray[np.float64], linear_velocity: NDArray[np.float64], ) -> NDArray[np.float64] | None: - """Solve differential IK for position-only control. - - Uses only the linear part of the Jacobian (first 3 rows). - - Args: - world: World for Jacobian computation - robot_id: Which robot - current_joints: Current joint positions - linear_velocity: Desired linear velocity [vx, vy, vz] - - Returns: - Joint velocities or None if near singularity - """ + """Position-only differential IK using linear Jacobian. Returns None if singular.""" with world.scratch_context() as ctx: world.set_positions(ctx, robot_id, current_joints) J = world.get_jacobian(ctx, robot_id) diff --git a/dimos/manipulation/planning/monitor/world_monitor.py b/dimos/manipulation/planning/monitor/world_monitor.py index cb84c1efa5..f1be2f338d 100644 --- a/dimos/manipulation/planning/monitor/world_monitor.py +++ b/dimos/manipulation/planning/monitor/world_monitor.py @@ -12,37 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -""" -World Monitor - -MoveIt-like monitor that keeps a WorldSpec synchronized with the real world. -This is the WorldSpec-based replacement for PlanningSceneMonitor. - -Uses factory functions to create WorldSpec, so all code uses Protocol types. - -Example: - # Create monitor - monitor = WorldMonitor(enable_viz=True) - - # Add robot - robot_id = monitor.add_robot(robot_config) - - # Finalize - monitor.finalize() - - # Start monitoring - monitor.start_state_monitor("piper", ["joint1", "joint2", ...]) - monitor.start_obstacle_monitor() - - # Use thread-safe accessors - positions = monitor.get_current_positions(robot_id) - is_valid = monitor.is_state_valid(robot_id, target_joints) - - # Get scratch context for planning - with monitor.scratch_context() as ctx: - # Do planning operations - pass -""" +"""World Monitor - keeps WorldSpec synchronized with real robot state and obstacles.""" from __future__ import annotations @@ -55,13 +25,6 @@ from dimos.manipulation.planning.factory import create_world from dimos.manipulation.planning.monitor.world_obstacle_monitor import WorldObstacleMonitor from dimos.manipulation.planning.monitor.world_state_monitor import WorldStateMonitor -from dimos.manipulation.planning.spec import ( - CollisionObjectMessage, - Detection3D, - Obstacle, - ObstacleType, - RobotModelConfig, -) from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: @@ -69,62 +32,20 @@ from numpy.typing import NDArray - from dimos.manipulation.planning.spec import WorldSpec + from dimos.manipulation.planning.spec import ( + CollisionObjectMessage, + Detection3D, + Obstacle, + RobotModelConfig, + WorldSpec, + ) from dimos.msgs.sensor_msgs import JointState logger = setup_logger() class WorldMonitor: - """MoveIt-like monitor for keeping a WorldSpec synchronized. - - This class manages: - - WorldSpec instance (created via factory) - - WorldStateMonitor: Subscribes to joint states, syncs robot configurations - - WorldObstacleMonitor: Subscribes to collision objects, syncs obstacles - - ## Key Differences from PlanningSceneMonitor - - 1. Uses factory pattern - no concrete type references - 2. Uses scratch_context() for thread-safe operations - 3. State is synced via sync_from_joint_state() instead of set_robot_state() - - ## Usage Pattern - - ```python - # 1. Create monitor - monitor = WorldMonitor(enable_viz=True) - - # 2. Add robots - robot_id = monitor.add_robot(robot_config) - - # 3. Add static obstacles (optional) - monitor.add_obstacle(table_obstacle) - - # 4. Finalize (locks topology) - monitor.finalize() - - # 5. Start monitoring - monitor.start_state_monitor(robot_id, joint_names) - monitor.start_obstacle_monitor() - - # 6. Use thread-safe accessors - positions = monitor.get_current_positions(robot_id) - is_valid = monitor.is_state_valid(robot_id, target_joints) - - # 7. Get scratch context for planning - with monitor.scratch_context() as ctx: - world.set_positions(ctx, robot_id, target_joints) - if world.is_collision_free(ctx, robot_id): - # Valid configuration - pass - ``` - - ## Thread Safety - - All public methods are thread-safe via an internal RLock. - For planning operations, use scratch_context() to get an independent context. - """ + """Manages WorldSpec with state/obstacle monitors. Thread-safe via RLock.""" def __init__( self, @@ -132,51 +53,20 @@ def __init__( enable_viz: bool = False, **kwargs: Any, ): - """Create a world monitor. - - Args: - backend: Backend to use ("drake") - enable_viz: Enable visualization (Meshcat for Drake) - **kwargs: Additional arguments passed to create_world() - """ self._backend = backend - - # Create world via factory (Protocol type) - self._world: WorldSpec = create_world( - backend=backend, - enable_viz=enable_viz, - **kwargs, - ) - - # Thread safety + self._world: WorldSpec = create_world(backend=backend, enable_viz=enable_viz, **kwargs) self._lock = threading.RLock() - - # Robot tracking: robot_id -> joint_names self._robot_joints: dict[str, list[str]] = {} - - # Sub-monitors self._state_monitors: dict[str, WorldStateMonitor] = {} self._obstacle_monitor: WorldObstacleMonitor | None = None - - # Visualization thread (publishes to Meshcat at a fixed rate) self._viz_thread: threading.Thread | None = None self._viz_stop_event = threading.Event() - self._viz_rate_hz: float = 10.0 # Default 10Hz + self._viz_rate_hz: float = 10.0 # ============= Robot Management ============= - def add_robot( - self, - config: RobotModelConfig, - ) -> str: - """Add a robot to the monitored world. - - Args: - config: Robot configuration - - Returns: - robot_id: Unique identifier for the robot - """ + def add_robot(self, config: RobotModelConfig) -> str: + """Add a robot. Returns robot_id.""" with self._lock: robot_id = self._world.add_robot(config) self._robot_joints[robot_id] = config.joint_names @@ -201,79 +91,10 @@ def get_joint_limits(self, robot_id: str) -> tuple[NDArray[np.float64], NDArray[ # ============= Obstacle Management ============= def add_obstacle(self, obstacle: Obstacle) -> str: - """Add an obstacle to the world. - - Args: - obstacle: Obstacle specification - - Returns: - obstacle_id: Unique identifier for the obstacle - """ + """Add an obstacle. Returns obstacle_id.""" with self._lock: return self._world.add_obstacle(obstacle) - def add_box_obstacle( - self, - name: str, - pose: NDArray[np.float64], - dimensions: tuple[float, float, float], - color: tuple[float, float, float, float] = (0.8, 0.2, 0.2, 0.8), - ) -> str: - """Add a box obstacle. - - Args: - name: Obstacle name - pose: 4x4 homogeneous transform - dimensions: (width, height, depth) in meters - color: RGBA color - - Returns: - obstacle_id - """ - obstacle = Obstacle( - name=name, - obstacle_type=ObstacleType.BOX, - pose=pose, - dimensions=dimensions, - color=color, - ) - return self.add_obstacle(obstacle) - - def add_sphere_obstacle( - self, - name: str, - pose: NDArray[np.float64], - radius: float, - color: tuple[float, float, float, float] = (0.8, 0.2, 0.2, 0.8), - ) -> str: - """Add a sphere obstacle.""" - obstacle = Obstacle( - name=name, - obstacle_type=ObstacleType.SPHERE, - pose=pose, - dimensions=(radius,), - color=color, - ) - return self.add_obstacle(obstacle) - - def add_cylinder_obstacle( - self, - name: str, - pose: NDArray[np.float64], - radius: float, - height: float, - color: tuple[float, float, float, float] = (0.8, 0.2, 0.2, 0.8), - ) -> str: - """Add a cylinder obstacle.""" - obstacle = Obstacle( - name=name, - obstacle_type=ObstacleType.CYLINDER, - pose=pose, - dimensions=(radius, height), - color=color, - ) - return self.add_obstacle(obstacle) - def remove_obstacle(self, obstacle_id: str) -> bool: """Remove an obstacle.""" with self._lock: @@ -292,19 +113,7 @@ def start_state_monitor( joint_names: list[str] | None = None, joint_name_mapping: dict[str, str] | None = None, ) -> None: - """Start monitoring joint states for a robot. - - Creates a WorldStateMonitor that will sync joint state messages - to the world's live context. - - Args: - robot_id: ID of the robot to monitor - joint_names: Joint names in URDF namespace (uses config joint_names if None) - joint_name_mapping: Maps orchestrator joint names to URDF joint names. - Example: {"left_joint1": "joint1"} means messages with "left_joint1" - will be mapped to URDF "joint1". If None, uses config mapping if available, - otherwise names must match exactly. - """ + """Start monitoring joint states. Uses config defaults if args are None.""" with self._lock: if robot_id in self._state_monitors: logger.warning(f"State monitor for '{robot_id}' already started") @@ -336,11 +145,7 @@ def start_state_monitor( logger.info(f"State monitor started for '{robot_id}'") def start_obstacle_monitor(self) -> None: - """Start monitoring obstacle updates. - - Creates a WorldObstacleMonitor that will sync obstacle messages - to the world. - """ + """Start monitoring obstacle updates.""" with self._lock: if self._obstacle_monitor is not None: logger.warning("Obstacle monitor already started") @@ -372,12 +177,7 @@ def stop_all_monitors(self) -> None: # ============= Message Handlers ============= def on_joint_state(self, msg: JointState, robot_id: str | None = None) -> None: - """Handle joint state message. - - Args: - msg: JointState message - robot_id: Specific robot to update (broadcasts to all if None) - """ + """Handle joint state message. Broadcasts to all monitors if robot_id is None.""" try: if robot_id is not None: if robot_id in self._state_monitors: @@ -407,14 +207,7 @@ def on_detections(self, detections: list[Detection3D]) -> None: # ============= State Access ============= def get_current_positions(self, robot_id: str) -> NDArray[np.float64] | None: - """Get current joint positions (thread-safe). - - Args: - robot_id: ID of the robot - - Returns: - Current positions or None if not yet received - """ + """Get current joint positions. Returns None if not yet received.""" # Try state monitor first if robot_id in self._state_monitors: positions = self._state_monitors[robot_id].get_current_positions() @@ -427,28 +220,13 @@ def get_current_positions(self, robot_id: str) -> NDArray[np.float64] | None: return self._world.get_positions(ctx, robot_id) def get_current_velocities(self, robot_id: str) -> NDArray[np.float64] | None: - """Get current joint velocities (thread-safe). - - Args: - robot_id: ID of the robot - - Returns: - Current velocities or None if not available - """ + """Get current joint velocities. Returns None if not available.""" if robot_id in self._state_monitors: return self._state_monitors[robot_id].get_current_velocities() return None def wait_for_state(self, robot_id: str, timeout: float = 1.0) -> bool: - """Wait until a state is received for the robot. - - Args: - robot_id: ID of the robot - timeout: Maximum time to wait - - Returns: - True if state was received, False if timeout - """ + """Wait until state is received. Returns False on timeout.""" if robot_id in self._state_monitors: return self._state_monitors[robot_id].wait_for_state(timeout) return False @@ -463,73 +241,26 @@ def is_state_stale(self, robot_id: str, max_age: float = 1.0) -> bool: @contextmanager def scratch_context(self) -> Generator[Any, None, None]: - """Get a scratch context for planning (thread-safe). - - The scratch context is a clone of the live context that can be - modified without affecting the live state. - - Example: - with monitor.scratch_context() as ctx: - world.set_positions(ctx, robot_id, q_test) - if world.is_collision_free(ctx, robot_id): - # Valid configuration - pass - - Yields: - Context that can be used with world operations - """ + """Thread-safe scratch context for planning.""" with self._world.scratch_context() as ctx: yield ctx def get_live_context(self) -> Any: - """Get the live context (use with caution). - - The live context is synced with real robot state. Modifications - will affect collision checking for all users. - - For planning, prefer scratch_context() instead. - - Returns: - Live context - """ + """Get live context. Prefer scratch_context() for planning.""" return self._world.get_live_context() # ============= Collision Checking ============= - def is_state_valid( - self, - robot_id: str, - joint_positions: NDArray[np.float64], - ) -> bool: - """Check if configuration is collision-free (thread-safe). - - Args: - robot_id: ID of the robot - joint_positions: Joint configuration to check - - Returns: - True if collision-free - """ + def is_state_valid(self, robot_id: str, joint_positions: NDArray[np.float64]) -> bool: + """Check if configuration is collision-free.""" with self._world.scratch_context() as ctx: self._world.set_positions(ctx, robot_id, joint_positions) return self._world.is_collision_free(ctx, robot_id) def is_path_valid( - self, - robot_id: str, - path: list[NDArray[np.float64]], - step_size: float = 0.05, + self, robot_id: str, path: list[NDArray[np.float64]], step_size: float = 0.05 ) -> bool: - """Check if path is collision-free (thread-safe). - - Args: - robot_id: ID of the robot - path: List of joint configurations - step_size: Interpolation resolution - - Returns: - True if path is collision-free - """ + """Check if path is collision-free with interpolation.""" with self._world.scratch_context() as ctx: for i in range(len(path) - 1): q_start = path[i] @@ -557,19 +288,9 @@ def get_min_distance(self, robot_id: str) -> float: # ============= Kinematics ============= def get_ee_pose( - self, - robot_id: str, - joint_positions: NDArray[np.float64] | None = None, + self, robot_id: str, joint_positions: NDArray[np.float64] | None = None ) -> NDArray[np.float64]: - """Get end-effector pose (thread-safe). - - Args: - robot_id: ID of the robot - joint_positions: Joint configuration (uses current if None) - - Returns: - 4x4 homogeneous transform - """ + """Get end-effector pose as 4x4 transform. Uses current state if positions is None.""" with self._world.scratch_context() as ctx: # If no positions provided, fetch current from state monitor if joint_positions is None: @@ -581,19 +302,9 @@ def get_ee_pose( return self._world.get_ee_pose(ctx, robot_id) def get_jacobian( - self, - robot_id: str, - joint_positions: NDArray[np.float64], + self, robot_id: str, joint_positions: NDArray[np.float64] ) -> NDArray[np.float64]: - """Get Jacobian (thread-safe). - - Args: - robot_id: ID of the robot - joint_positions: Joint configuration - - Returns: - 6 x n_joints Jacobian matrix - """ + """Get 6xN Jacobian matrix.""" with self._world.scratch_context() as ctx: self._world.set_positions(ctx, robot_id, joint_positions) return self._world.get_jacobian(ctx, robot_id) @@ -601,11 +312,7 @@ def get_jacobian( # ============= Lifecycle ============= def finalize(self) -> None: - """Finalize world for collision checking. - - Must be called after adding all robots and before starting monitors - or performing collision checking. - """ + """Finalize world. Must be called before collision checking.""" with self._lock: self._world.finalize() logger.info("World finalized") @@ -618,11 +325,7 @@ def is_finalized(self) -> bool: # ============= Visualization ============= def get_meshcat_url(self) -> str | None: - """Get Meshcat visualization URL (Drake only). - - Returns: - URL string or None if visualization not enabled - """ + """Get Meshcat URL or None if not enabled.""" if hasattr(self._world, "get_meshcat_url"): url = self._world.get_meshcat_url() return str(url) if url else None @@ -634,14 +337,7 @@ def publish_visualization(self) -> None: self._world.publish_to_meshcat() def start_visualization_thread(self, rate_hz: float = 10.0) -> None: - """Start background thread that updates Meshcat visualization. - - This allows live visualization without blocking the LCM callback thread. - The thread publishes the current robot state to Meshcat at the specified rate. - - Args: - rate_hz: Visualization update rate in Hz (default: 10Hz) - """ + """Start background thread for Meshcat updates at given rate.""" if self._viz_thread is not None and self._viz_thread.is_alive(): logger.warning("Visualization thread already running") return @@ -689,11 +385,7 @@ def _visualization_loop(self) -> None: @property def world(self) -> WorldSpec: - """Get underlying WorldSpec instance. - - Direct access to the world is NOT thread-safe for state modifications. - Prefer using scratch_context() for planning operations. - """ + """Get underlying WorldSpec. Not thread-safe for modifications.""" return self._world def get_state_monitor(self, robot_id: str) -> WorldStateMonitor | None: diff --git a/dimos/manipulation/planning/planners/drake_planner.py b/dimos/manipulation/planning/planners/drake_planner.py index 89eb71ff73..95289f379a 100644 --- a/dimos/manipulation/planning/planners/drake_planner.py +++ b/dimos/manipulation/planning/planners/drake_planner.py @@ -12,23 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -""" -Drake Planner Module - -Implements PlannerSpec using RRT-Connect algorithm with WorldSpec for collision checking. - -## Key Design - -- Uses WorldSpec.scratch_context() for thread-safe collision checking -- Joint-space planning only (use KinematicsSpec.solve() for pose goals first) -- Stateless except for configuration - -Example: - planner = DrakePlanner(step_size=0.1) - result = planner.plan_joint_path(world, robot_id, q_start, q_goal) - if result.is_success(): - waypoints = result.path -""" +"""RRT-Connect and RRT* motion planners implementing PlannerSpec.""" from __future__ import annotations @@ -41,7 +25,6 @@ from dimos.manipulation.planning.spec import PlanningResult, PlanningStatus, WorldSpec from dimos.manipulation.planning.utils.path_utils import ( compute_path_length, - interpolate_path, interpolate_segment, ) from dimos.utils.logging_config import setup_logger @@ -54,11 +37,12 @@ @dataclass class TreeNode: - """Node in the RRT tree.""" + """Node in RRT tree with optional cost tracking (for RRT*).""" config: NDArray[np.float64] parent: TreeNode | None = None children: list[TreeNode] = field(default_factory=list) + cost: float = 0.0 def path_to_root(self) -> list[NDArray[np.float64]]: """Get path from this node to root.""" @@ -71,23 +55,7 @@ def path_to_root(self) -> list[NDArray[np.float64]]: class DrakePlanner: - """RRT-Connect planner implementing PlannerSpec. - - Uses bi-directional RRT (RRT-Connect) for joint-space path planning - with collision checking via WorldSpec.scratch_context(). - - ## Algorithm - - 1. Initialize start tree and goal tree - 2. Extend start tree toward random sample - 3. Try to connect goal tree to new node - 4. Swap trees and repeat - 5. If connected, extract path - - ## Thread Safety - - All collision checking uses world.scratch_context() for thread safety. - """ + """Bi-directional RRT-Connect planner.""" def __init__( self, @@ -96,14 +64,6 @@ def __init__( goal_tolerance: float = 0.1, collision_step_size: float = 0.02, ): - """Create RRT-Connect planner. - - Args: - step_size: Extension step size in joint space (radians) - connect_step_size: Step size for connect attempts - goal_tolerance: Distance to goal to consider success - collision_step_size: Step size for collision checking along edges - """ self._step_size = step_size self._connect_step_size = connect_step_size self._goal_tolerance = goal_tolerance @@ -118,83 +78,49 @@ def plan_joint_path( timeout: float = 10.0, max_iterations: int = 5000, ) -> PlanningResult: - """Plan collision-free joint-space path using RRT-Connect. - - Args: - world: World for collision checking - robot_id: Which robot - q_start: Start joint configuration (radians) - q_goal: Goal joint configuration (radians) - timeout: Maximum planning time (seconds) - max_iterations: Maximum iterations - - Returns: - PlanningResult with path (list of waypoints) or failure info - """ + """Plan collision-free path using bi-directional RRT.""" start_time = time.time() - # Validate inputs error = self._validate_inputs(world, robot_id, q_start, q_goal) if error is not None: return error - # Get joint limits - lower_limits, upper_limits = world.get_joint_limits(robot_id) - - # Initialize trees + lower, upper = world.get_joint_limits(robot_id) start_tree = [TreeNode(config=q_start.copy())] goal_tree = [TreeNode(config=q_goal.copy())] - - trees_swapped = False # Track if trees have been swapped an odd number of times + trees_swapped = False for iteration in range(max_iterations): - # Check timeout if time.time() - start_time > timeout: return _create_failure_result( PlanningStatus.TIMEOUT, f"Timeout after {iteration} iterations", - planning_time=time.time() - start_time, - iterations=iteration, + time.time() - start_time, + iteration, ) - # Sample random configuration - sample = np.random.uniform(lower_limits, upper_limits) - - # Extend start tree toward sample - extended_node = self._extend_tree(world, robot_id, start_tree, sample, self._step_size) + sample = np.random.uniform(lower, upper) + extended = self._extend_tree(world, robot_id, start_tree, sample, self._step_size) - if extended_node is not None: - # Try to connect goal tree to extended node - connected_node = self._connect_tree( - world, robot_id, goal_tree, extended_node.config, self._connect_step_size + if extended is not None: + connected = self._connect_tree( + world, robot_id, goal_tree, extended.config, self._connect_step_size ) - - if connected_node is not None: - # Trees connected! Extract path - path = self._extract_path(extended_node, connected_node) - - # If trees were swapped, path is from goal to start - reverse it + if connected is not None: + path = self._extract_path(extended, connected) if trees_swapped: path = list(reversed(path)) - - # Simplify path path = self._simplify_path(world, robot_id, path) + return _create_success_result(path, time.time() - start_time, iteration + 1) - return _create_success_result( - path=path, - planning_time=time.time() - start_time, - iterations=iteration + 1, - ) - - # Swap trees start_tree, goal_tree = goal_tree, start_tree trees_swapped = not trees_swapped return _create_failure_result( PlanningStatus.NO_SOLUTION, f"No path found after {max_iterations} iterations", - planning_time=time.time() - start_time, - iterations=max_iterations, + time.time() - start_time, + max_iterations, ) def get_name(self) -> str: @@ -409,120 +335,86 @@ def plan_joint_path( timeout: float = 10.0, max_iterations: int = 5000, ) -> PlanningResult: - """Plan optimal collision-free joint-space path using RRT*.""" + """Plan optimal path using RRT* with rewiring.""" start_time = time.time() - - # Get joint limits - lower_limits, upper_limits = world.get_joint_limits(robot_id) + lower, upper = world.get_joint_limits(robot_id) # Validate start/goal with world.scratch_context() as ctx: world.set_positions(ctx, robot_id, q_start) if not world.is_collision_free(ctx, robot_id): return _create_failure_result( - PlanningStatus.COLLISION_AT_START, - "Start configuration is in collision", + PlanningStatus.COLLISION_AT_START, "Start in collision" ) - world.set_positions(ctx, robot_id, q_goal) if not world.is_collision_free(ctx, robot_id): - return _create_failure_result( - PlanningStatus.COLLISION_AT_GOAL, - "Goal configuration is in collision", - ) + return _create_failure_result(PlanningStatus.COLLISION_AT_GOAL, "Goal in collision") - # Initialize tree with costs - root = _RRTStarNode(config=q_start.copy(), cost=0.0) - nodes = [root] - goal_node: _RRTStarNode | None = None + nodes = [TreeNode(config=q_start.copy(), cost=0.0)] + goal_node: TreeNode | None = None best_cost = float("inf") + iterations_run = 0 - for iteration in range(max_iterations): # noqa: B007 + for _ in range(max_iterations): + iterations_run += 1 if time.time() - start_time > timeout: - break # Return best path found so far - - # Sample - sample = np.random.uniform(lower_limits, upper_limits) + break - # Find nearest + # Sample and find nearest + sample = np.random.uniform(lower, upper) nearest = min(nodes, key=lambda n: float(np.linalg.norm(n.config - sample))) - # Extend + # Extend toward sample diff = sample - nearest.config dist = float(np.linalg.norm(diff)) - if dist > self._step_size: - new_config = nearest.config + self._step_size * (diff / dist) - else: - new_config = sample.copy() + new_config = ( + nearest.config + min(dist, self._step_size) * (diff / dist) + if dist > 0 + else sample.copy() + ) if not self._is_edge_valid(world, robot_id, nearest.config, new_config): continue - # Find neighbors within rewire radius + # Find best parent among neighbors neighbors = [ n for n in nodes if float(np.linalg.norm(n.config - new_config)) < self._rewire_radius ] - - # Find best parent - best_parent = nearest - best_new_cost = nearest.cost + float(np.linalg.norm(new_config - nearest.config)) - - for neighbor in neighbors: - new_cost = neighbor.cost + float(np.linalg.norm(new_config - neighbor.config)) - if new_cost < best_new_cost: - if self._is_edge_valid(world, robot_id, neighbor.config, new_config): - best_parent = neighbor - best_new_cost = new_cost - - # Add new node - new_node = _RRTStarNode( - config=new_config, - parent=best_parent, - cost=best_new_cost, + best_parent, best_cost_to_new = ( + nearest, + nearest.cost + float(np.linalg.norm(new_config - nearest.config)), ) + for n in neighbors: + cost = n.cost + float(np.linalg.norm(new_config - n.config)) + if cost < best_cost_to_new and self._is_edge_valid( + world, robot_id, n.config, new_config + ): + best_parent, best_cost_to_new = n, cost + + # Add node and rewire neighbors + new_node = TreeNode(config=new_config, parent=best_parent, cost=best_cost_to_new) best_parent.children.append(new_node) nodes.append(new_node) - - # Rewire neighbors through new node - for neighbor in neighbors: - if neighbor == best_parent: - continue - potential_cost = new_node.cost + float( - np.linalg.norm(neighbor.config - new_node.config) - ) - if potential_cost < neighbor.cost: - if self._is_edge_valid(world, robot_id, new_node.config, neighbor.config): - # Rewire - if neighbor.parent is not None: - neighbor.parent.children.remove(neighbor) - neighbor.parent = new_node - neighbor.cost = potential_cost - new_node.children.append(neighbor) - self._update_costs(neighbor) + self._rewire_neighbors(world, robot_id, new_node, neighbors) # Check goal - if float(np.linalg.norm(new_node.config - q_goal)) < self._goal_tolerance: - if new_node.cost < best_cost: - goal_node = new_node - best_cost = new_node.cost + if ( + float(np.linalg.norm(new_node.config - q_goal)) < self._goal_tolerance + and new_node.cost < best_cost + ): + goal_node, best_cost = new_node, new_node.cost if goal_node is not None: - path = goal_node.path_to_root() - path.append(q_goal) - - return _create_success_result( - path=path, - planning_time=time.time() - start_time, - iterations=iteration + 1, - ) + path = [*goal_node.path_to_root(), q_goal] + return _create_success_result(path, time.time() - start_time, iterations_run) return _create_failure_result( PlanningStatus.NO_SOLUTION, - f"No path found after {max_iterations} iterations", - planning_time=time.time() - start_time, - iterations=max_iterations, + f"No path after {max_iterations} iterations", + time.time() - start_time, + max_iterations, ) def get_name(self) -> str: @@ -546,32 +438,37 @@ def _is_edge_valid( return True - def _update_costs(self, node: _RRTStarNode) -> None: + def _rewire_neighbors( + self, + world: WorldSpec, + robot_id: str, + new_node: TreeNode, + neighbors: list[TreeNode], + ) -> None: + """Rewire neighbors through new node if it provides a shorter path.""" + for neighbor in neighbors: + if neighbor == new_node.parent: + continue + potential_cost = new_node.cost + float( + np.linalg.norm(neighbor.config - new_node.config) + ) + if potential_cost < neighbor.cost and self._is_edge_valid( + world, robot_id, new_node.config, neighbor.config + ): + if neighbor.parent is not None: + neighbor.parent.children.remove(neighbor) + neighbor.parent = new_node + neighbor.cost = potential_cost + new_node.children.append(neighbor) + self._update_costs(neighbor) + + def _update_costs(self, node: TreeNode) -> None: """Recursively update costs after rewiring.""" for child in node.children: child.cost = node.cost + float(np.linalg.norm(child.config - node.config)) self._update_costs(child) -@dataclass -class _RRTStarNode: - """Node for RRT* with cost tracking.""" - - config: NDArray[np.float64] - cost: float = 0.0 - parent: _RRTStarNode | None = None - children: list[_RRTStarNode] = field(default_factory=list) - - def path_to_root(self) -> list[NDArray[np.float64]]: - """Get path from this node to root.""" - path = [] - node: _RRTStarNode | None = self - while node is not None: - path.append(node.config) - node = node.parent - return list(reversed(path)) - - # ============= Result Helpers ============= diff --git a/dimos/manipulation/planning/spec.py b/dimos/manipulation/planning/spec.py deleted file mode 100644 index 66be979f30..0000000000 --- a/dimos/manipulation/planning/spec.py +++ /dev/null @@ -1,733 +0,0 @@ -# Copyright 2025-2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -Manipulation Planning Specifications - -Contains Protocol definitions and data classes for the manipulation planning stack. - -## Protocols - -- WorldSpec: Core backend owning physics/collision (implemented by DrakeWorld) -- KinematicsSpec: Stateless IK operations (implemented by DrakeKinematics) -- PlannerSpec: Joint-space path planning (implemented by DrakePlanner) -- VizSpec: Visualization backend (future) - -## Data Classes - -- RobotModelConfig: Robot configuration for adding to world -- Obstacle: Obstacle specification -- IKResult: Result of IK solve -- PlanningResult: Result of path planning - -## Usage - -All code should use Protocol types (not concrete classes): - -```python -from dimos.manipulation.planning.spec import WorldSpec, KinematicsSpec - -def plan_motion(world: WorldSpec, kinematics: KinematicsSpec, ...): - # Works with any conforming implementation - pass -``` - -Use factory functions to create concrete instances: - -```python -from dimos.manipulation.planning.factory import create_world, create_kinematics - -world = create_world(backend="drake") # Returns WorldSpec -kinematics = create_kinematics() # Returns KinematicsSpec -``` -""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from enum import Enum, auto -from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable - -if TYPE_CHECKING: - from contextlib import AbstractContextManager - - import numpy as np - from numpy.typing import NDArray - - -# ============================================================================= -# Enums -# ============================================================================= - - -class ObstacleType(Enum): - """Type of obstacle geometry.""" - - BOX = auto() - SPHERE = auto() - CYLINDER = auto() - MESH = auto() - - -class IKStatus(Enum): - """Status of IK solution.""" - - SUCCESS = auto() - NO_SOLUTION = auto() - SINGULARITY = auto() - JOINT_LIMITS = auto() - COLLISION = auto() - TIMEOUT = auto() - - -class PlanningStatus(Enum): - """Status of motion planning.""" - - SUCCESS = auto() - NO_SOLUTION = auto() - TIMEOUT = auto() - INVALID_START = auto() - INVALID_GOAL = auto() - COLLISION_AT_START = auto() - COLLISION_AT_GOAL = auto() - - -# ============================================================================= -# Data Classes -# ============================================================================= - - -@dataclass -class RobotModelConfig: - """Configuration for adding a robot to the world. - - Attributes: - name: Human-readable robot name - urdf_path: Path to URDF file (can be .urdf or .xacro) - base_pose: 4x4 homogeneous transform for robot base - joint_names: Ordered list of controlled joint names (in URDF namespace) - end_effector_link: Name of the end-effector link for FK/IK - base_link: Name of the base link (default: "base_link") - package_paths: Dict mapping package names to filesystem paths - joint_limits_lower: Lower joint limits (radians) - joint_limits_upper: Upper joint limits (radians) - velocity_limits: Joint velocity limits (rad/s) - auto_convert_meshes: Auto-convert DAE/STL meshes to OBJ for Drake - xacro_args: Arguments to pass to xacro processor (for .xacro files) - collision_exclusion_pairs: List of (link1, link2) pairs to exclude from collision. - Useful for parallel linkage mechanisms like grippers where non-adjacent - links may legitimately overlap (e.g., mimic joints). - max_velocity: Maximum joint velocity for trajectory generation (rad/s) - max_acceleration: Maximum joint acceleration for trajectory generation (rad/s^2) - joint_name_mapping: Maps orchestrator joint names to URDF joint names. - Example: {"left_joint1": "joint1"} means orchestrator's "left_joint1" - corresponds to URDF's "joint1". If empty, names are assumed to match. - orchestrator_task_name: Task name for executing trajectories via orchestrator RPC. - If set, trajectories can be executed via execute_trajectory() RPC. - """ - - name: str - urdf_path: str - base_pose: NDArray[np.float64] - joint_names: list[str] - end_effector_link: str - base_link: str = "base_link" - package_paths: dict[str, str] = field(default_factory=dict) - joint_limits_lower: NDArray[np.float64] | None = None - joint_limits_upper: NDArray[np.float64] | None = None - velocity_limits: NDArray[np.float64] | None = None - auto_convert_meshes: bool = False - xacro_args: dict[str, str] = field(default_factory=dict) - collision_exclusion_pairs: list[tuple[str, str]] = field(default_factory=list) - # Motion constraints for trajectory generation - max_velocity: float = 1.0 - max_acceleration: float = 2.0 - # Orchestrator integration - joint_name_mapping: dict[str, str] = field(default_factory=dict) - orchestrator_task_name: str | None = None - - def get_urdf_joint_name(self, orchestrator_name: str) -> str: - """Translate orchestrator joint name to URDF joint name. - - Args: - orchestrator_name: Joint name from orchestrator (e.g., "left_joint1") - - Returns: - Corresponding URDF joint name (e.g., "joint1") - """ - return self.joint_name_mapping.get(orchestrator_name, orchestrator_name) - - def get_orchestrator_joint_name(self, urdf_name: str) -> str: - """Translate URDF joint name to orchestrator joint name. - - Args: - urdf_name: Joint name from URDF (e.g., "joint1") - - Returns: - Corresponding orchestrator joint name (e.g., "left_joint1") - """ - # Reverse lookup - for orch_name, u_name in self.joint_name_mapping.items(): - if u_name == urdf_name: - return orch_name - return urdf_name - - def get_orchestrator_joint_names(self) -> list[str]: - """Get joint names in orchestrator namespace. - - Returns: - List of joint names as they appear in orchestrator messages. - If no mapping, returns the URDF joint names. - """ - if not self.joint_name_mapping: - return self.joint_names - return [self.get_orchestrator_joint_name(j) for j in self.joint_names] - - -@dataclass -class Obstacle: - """Obstacle specification for collision avoidance. - - Attributes: - name: Unique name for the obstacle - obstacle_type: Type of geometry (BOX, SPHERE, CYLINDER, MESH) - pose: 4x4 homogeneous transform - dimensions: Type-specific dimensions: - - BOX: (width, height, depth) - - SPHERE: (radius,) - - CYLINDER: (radius, height) - - MESH: Not used - color: RGBA color tuple (0-1 range) - mesh_path: Path to mesh file (for MESH type) - """ - - name: str - obstacle_type: ObstacleType - pose: NDArray[np.float64] - dimensions: tuple[float, ...] = () - color: tuple[float, float, float, float] = (0.8, 0.2, 0.2, 0.8) - mesh_path: str | None = None - - -@dataclass -class IKResult: - """Result of an IK solve. - - Attributes: - status: Solution status - joint_positions: Solution joint positions (None if failed) - position_error: Cartesian position error (meters) - orientation_error: Orientation error (radians) - iterations: Number of iterations taken - message: Human-readable status message - """ - - status: IKStatus - joint_positions: NDArray[np.float64] | None = None - position_error: float = 0.0 - orientation_error: float = 0.0 - iterations: int = 0 - message: str = "" - - def is_success(self) -> bool: - """Check if IK was successful.""" - return self.status == IKStatus.SUCCESS - - -@dataclass -class PlanningResult: - """Result of motion planning. - - Attributes: - status: Planning status - path: List of joint configurations (empty if failed) - planning_time: Time taken to plan (seconds) - path_length: Total path length in joint space (radians) - iterations: Number of iterations/nodes expanded - message: Human-readable status message - """ - - status: PlanningStatus - path: list[NDArray[np.float64]] = field(default_factory=list) - planning_time: float = 0.0 - path_length: float = 0.0 - iterations: int = 0 - message: str = "" - - def is_success(self) -> bool: - """Check if planning was successful.""" - return self.status == PlanningStatus.SUCCESS - - -@dataclass -class CollisionObjectMessage: - """Message for adding/updating/removing obstacles. - - Used by monitors to handle obstacle updates from external sources. - - Attributes: - id: Unique identifier for the object - operation: "add", "update", or "remove" - primitive_type: "box", "sphere", or "cylinder" (for add/update) - pose: 4x4 transform (for add/update) - dimensions: Type-specific dimensions (for add/update) - color: RGBA color tuple - """ - - id: str - operation: str # "add", "update", "remove" - primitive_type: str | None = None - pose: NDArray[np.float64] | None = None - dimensions: tuple[float, ...] | None = None - color: tuple[float, float, float, float] = (0.8, 0.2, 0.2, 0.8) - - -@dataclass -class Detection3D: - """3D detection result from perception pipeline. - - Used by monitors to handle perception updates. - - Attributes: - id: Unique detection ID - label: Object class label - pose: 4x4 transform - dimensions: (width, height, depth) - confidence: Detection confidence (0-1) - """ - - id: str - label: str - pose: NDArray[np.float64] - dimensions: tuple[float, float, float] - confidence: float = 1.0 - - -# ============================================================================= -# Protocols -# ============================================================================= - - -@runtime_checkable -class WorldSpec(Protocol): - """Protocol for the world/scene backend. - - The world owns the physics/collision backend and provides: - - Robot/obstacle management - - Collision checking - - Forward kinematics - - Context management for thread safety - - ## Context Management - - The world maintains: - - Live context: Mirrors current robot state (synced from driver) - - Scratch contexts: Thread-safe clones for planning/IK operations - - All state-dependent operations (collision checking, FK, Jacobian) take - a context parameter to ensure thread safety. - - ## Implementations - - - DrakeWorld: Uses Drake's MultibodyPlant and SceneGraph - """ - - # ========================================================================= - # Robot Management - # ========================================================================= - - def add_robot(self, config: RobotModelConfig) -> str: - """Add a robot to the world. - - Args: - config: Robot configuration - - Returns: - Unique robot ID - """ - ... - - def get_robot_ids(self) -> list[str]: - """Get all robot IDs.""" - ... - - def get_robot_config(self, robot_id: str) -> RobotModelConfig: - """Get robot configuration.""" - ... - - def get_joint_limits(self, robot_id: str) -> tuple[NDArray[np.float64], NDArray[np.float64]]: - """Get joint limits (lower, upper) for a robot.""" - ... - - # ========================================================================= - # Obstacle Management - # ========================================================================= - - def add_obstacle(self, obstacle: Obstacle) -> str: - """Add an obstacle to the world. - - Args: - obstacle: Obstacle specification - - Returns: - Unique obstacle ID - """ - ... - - def remove_obstacle(self, obstacle_id: str) -> bool: - """Remove an obstacle. - - Returns: - True if obstacle was removed - """ - ... - - def update_obstacle_pose(self, obstacle_id: str, pose: NDArray[np.float64]) -> bool: - """Update obstacle pose. - - Returns: - True if obstacle was updated - """ - ... - - def clear_obstacles(self) -> None: - """Remove all obstacles.""" - ... - - # ========================================================================= - # Lifecycle - # ========================================================================= - - def finalize(self) -> None: - """Finalize the world for simulation/collision checking. - - Must be called after adding all robots and before any operations. - """ - ... - - @property - def is_finalized(self) -> bool: - """Check if world is finalized.""" - ... - - # ========================================================================= - # Context Management - # ========================================================================= - - def get_live_context(self) -> Any: - """Get the live context (mirrors real robot state). - - The live context is synced with real robot state via - sync_from_joint_state(). Use with caution - modifications - affect all users. - - For planning, prefer scratch_context() instead. - """ - ... - - def scratch_context(self) -> AbstractContextManager[Any]: - """Get a scratch context for planning (thread-safe clone). - - Example: - with world.scratch_context() as ctx: - world.set_positions(ctx, robot_id, q_test) - if world.is_collision_free(ctx, robot_id): - # Valid configuration - pass - - Yields: - Context that can be used for state operations - """ - ... - - def sync_from_joint_state(self, robot_id: str, positions: NDArray[np.float64]) -> None: - """Sync live context from joint state (called by driver/monitor). - - Args: - robot_id: Robot to update - positions: Joint positions from driver - """ - ... - - # ========================================================================= - # State Operations (require context) - # ========================================================================= - - def set_positions(self, ctx: Any, robot_id: str, positions: NDArray[np.float64]) -> None: - """Set robot joint positions in a context. - - Args: - ctx: Context from scratch_context() or get_live_context() - robot_id: Robot to update - positions: Joint positions (radians) - """ - ... - - def get_positions(self, ctx: Any, robot_id: str) -> NDArray[np.float64]: - """Get robot joint positions from a context. - - Args: - ctx: Context - robot_id: Robot to query - - Returns: - Joint positions (radians) - """ - ... - - # ========================================================================= - # Collision Checking (require context) - # ========================================================================= - - def is_collision_free(self, ctx: Any, robot_id: str) -> bool: - """Check if robot configuration is collision-free. - - Args: - ctx: Context with robot state set - robot_id: Robot to check - - Returns: - True if no collisions - """ - ... - - def get_min_distance(self, ctx: Any, robot_id: str) -> float: - """Get minimum distance to obstacles. - - Args: - ctx: Context with robot state set - robot_id: Robot to check - - Returns: - Minimum distance (meters), negative if in collision - """ - ... - - # ========================================================================= - # Forward Kinematics (require context) - # ========================================================================= - - def get_ee_pose(self, ctx: Any, robot_id: str) -> NDArray[np.float64]: - """Get end-effector pose. - - Args: - ctx: Context with robot state set - robot_id: Robot to query - - Returns: - 4x4 homogeneous transform - """ - ... - - def get_jacobian(self, ctx: Any, robot_id: str) -> NDArray[np.float64]: - """Get end-effector Jacobian. - - Args: - ctx: Context with robot state set - robot_id: Robot to query - - Returns: - 6 x n_joints Jacobian matrix [linear; angular] - """ - ... - - # ========================================================================= - # Visualization (optional) - # ========================================================================= - - def get_meshcat_url(self) -> str | None: - """Get Meshcat visualization URL if enabled.""" - ... - - def publish_to_meshcat(self, ctx: Any | None = None) -> None: - """Publish current state to Meshcat visualization.""" - ... - - def animate_path( - self, robot_id: str, path: list[NDArray[np.float64]], duration: float = 3.0 - ) -> None: - """Animate a path in visualization.""" - ... - - -@runtime_checkable -class KinematicsSpec(Protocol): - """Protocol for inverse kinematics solver. - - Kinematics solvers are stateless (except for configuration) and - use WorldSpec for all FK/collision operations. - - ## Methods - - - solve(): Full optimization-based IK with collision checking - - solve_iterative(): Iterative Jacobian-based IK - - solve_differential(): Single Jacobian step for velocity control - - ## Implementations - - - DrakeKinematics: Uses Drake's InverseKinematics + SNOPT/IPOPT - """ - - def solve( - self, - world: WorldSpec, - robot_id: str, - target_pose: NDArray[np.float64], - seed: NDArray[np.float64] | None = None, - position_tolerance: float = 0.001, - orientation_tolerance: float = 0.01, - check_collision: bool = True, - max_attempts: int = 10, - ) -> IKResult: - """Solve full IK with optional collision checking. - - Args: - world: World for FK/collision - robot_id: Robot to solve for - target_pose: 4x4 target end-effector transform - seed: Initial guess (uses current state if None) - position_tolerance: Position tolerance (meters) - orientation_tolerance: Orientation tolerance (radians) - check_collision: Check solution is collision-free - max_attempts: Random restarts for robustness - - Returns: - IKResult with status and solution - """ - ... - - def solve_iterative( - self, - world: WorldSpec, - robot_id: str, - target_pose: NDArray[np.float64], - seed: NDArray[np.float64], - max_iterations: int = 100, - position_tolerance: float = 0.001, - orientation_tolerance: float = 0.01, - ) -> IKResult: - """Solve IK iteratively using Jacobian method. - - Slower but more predictable behavior near singularities. - - Args: - world: World for FK/Jacobian - robot_id: Robot to solve for - target_pose: 4x4 target transform - seed: Initial joint configuration - max_iterations: Maximum iterations - position_tolerance: Convergence tolerance (meters) - orientation_tolerance: Convergence tolerance (radians) - - Returns: - IKResult with status and solution - """ - ... - - def solve_differential( - self, - world: WorldSpec, - robot_id: str, - current_joints: NDArray[np.float64], - twist: NDArray[np.float64], - dt: float, - ) -> NDArray[np.float64] | None: - """Single Jacobian step for velocity control. - - Args: - world: World for Jacobian computation - robot_id: Robot to control - current_joints: Current joint positions - twist: Desired end-effector twist [vx, vy, vz, wx, wy, wz] - dt: Time step (seconds) - - Returns: - Joint velocities (rad/s) or None if near singularity - """ - ... - - -@runtime_checkable -class PlannerSpec(Protocol): - """Protocol for motion planner. - - Planners find collision-free paths from start to goal configurations. - They use WorldSpec for collision checking and are stateless. - - ## Implementations - - - DrakePlanner: RRT-Connect planner - - DrakeRRTStarPlanner: RRT* planner (asymptotically optimal) - """ - - def plan_joint_path( - self, - world: WorldSpec, - robot_id: str, - q_start: NDArray[np.float64], - q_goal: NDArray[np.float64], - timeout: float = 10.0, - ) -> PlanningResult: - """Plan a collision-free joint-space path. - - Args: - world: World for collision checking - robot_id: Robot to plan for - q_start: Start configuration - q_goal: Goal configuration - timeout: Planning timeout (seconds) - - Returns: - PlanningResult with status and path - """ - ... - - def get_name(self) -> str: - """Get planner name.""" - ... - - -@runtime_checkable -class VizSpec(Protocol): - """Protocol for visualization backend. - - Provides methods to update robot/obstacle visualization. - - Note: For Drake, visualization is typically integrated into DrakeWorld - via enable_viz=True. This protocol is for advanced use cases. - - ## Implementations - - - DrakeWorld (integrated): Use create_world(enable_viz=True) - """ - - def set_robot_state(self, robot_id: str, positions: NDArray[np.float64]) -> None: - """Update robot visualization state.""" - ... - - def add_obstacle(self, obstacle: Obstacle) -> str: - """Add obstacle to visualization.""" - ... - - def remove_obstacle(self, obstacle_id: str) -> None: - """Remove obstacle from visualization.""" - ... - - def get_url(self) -> str | None: - """Get visualization URL (e.g., Meshcat URL).""" - ... - - def publish(self) -> None: - """Force publish current state to visualization.""" - ... diff --git a/dimos/manipulation/planning/spec/__init__.py b/dimos/manipulation/planning/spec/__init__.py new file mode 100644 index 0000000000..5d079c6d87 --- /dev/null +++ b/dimos/manipulation/planning/spec/__init__.py @@ -0,0 +1,83 @@ +# 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. + +"""Manipulation Planning Specifications. + +Protocols: + - WorldSpec: Core backend owning physics/collision (implemented by DrakeWorld) + - KinematicsSpec: Stateless IK operations (implemented by DrakeKinematics) + - PlannerSpec: Joint-space path planning (implemented by DrakePlanner) + - VizSpec: Visualization backend + +Data Classes: + - RobotModelConfig: Robot configuration for adding to world + - Obstacle: Obstacle specification + - IKResult: Result of IK solve + - PlanningResult: Result of path planning + +Usage: + All code should use Protocol types (not concrete classes): + + ```python + from dimos.manipulation.planning.spec import WorldSpec, KinematicsSpec + + def plan_motion(world: WorldSpec, kinematics: KinematicsSpec, ...): + pass + ``` + + Use factory functions to create concrete instances: + + ```python + from dimos.manipulation.planning.factory import create_world + world = create_world(backend="drake") + ``` +""" + +from dimos.manipulation.planning.spec.config import RobotModelConfig +from dimos.manipulation.planning.spec.enums import IKStatus, ObstacleType, PlanningStatus +from dimos.manipulation.planning.spec.protocols import ( + KinematicsSpec, + PlannerSpec, + VizSpec, + WorldSpec, +) +from dimos.manipulation.planning.spec.types import ( + CollisionObjectMessage, + Detection3D, + IKResult, + Obstacle, + PerceptionDetection, + PlanningResult, +) + +__all__ = [ + "CollisionObjectMessage", + "Detection3D", # Backwards compat alias for PerceptionDetection + "IKResult", + "IKStatus", + "KinematicsSpec", + # Types + "Obstacle", + # Enums + "ObstacleType", + "PerceptionDetection", + "PlannerSpec", + "PlanningResult", + "PlanningStatus", + # Config + "RobotModelConfig", + "VizSpec", + # Protocols + "WorldSpec", +] diff --git a/dimos/manipulation/planning/spec/config.py b/dimos/manipulation/planning/spec/config.py new file mode 100644 index 0000000000..bb8706e6f2 --- /dev/null +++ b/dimos/manipulation/planning/spec/config.py @@ -0,0 +1,91 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Robot configuration for manipulation planning.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import numpy as np + from numpy.typing import NDArray + + +@dataclass +class RobotModelConfig: + """Configuration for adding a robot to the world. + + Attributes: + name: Human-readable robot name + urdf_path: Path to URDF file (can be .urdf or .xacro) + base_pose: 4x4 homogeneous transform for robot base + joint_names: Ordered list of controlled joint names (in URDF namespace) + end_effector_link: Name of the end-effector link for FK/IK + base_link: Name of the base link (default: "base_link") + package_paths: Dict mapping package names to filesystem paths + joint_limits_lower: Lower joint limits (radians) + joint_limits_upper: Upper joint limits (radians) + velocity_limits: Joint velocity limits (rad/s) + auto_convert_meshes: Auto-convert DAE/STL meshes to OBJ for Drake + xacro_args: Arguments to pass to xacro processor (for .xacro files) + collision_exclusion_pairs: List of (link1, link2) pairs to exclude from collision. + Useful for parallel linkage mechanisms like grippers where non-adjacent + links may legitimately overlap (e.g., mimic joints). + max_velocity: Maximum joint velocity for trajectory generation (rad/s) + max_acceleration: Maximum joint acceleration for trajectory generation (rad/s^2) + joint_name_mapping: Maps orchestrator joint names to URDF joint names. + Example: {"left_joint1": "joint1"} means orchestrator's "left_joint1" + corresponds to URDF's "joint1". If empty, names are assumed to match. + orchestrator_task_name: Task name for executing trajectories via orchestrator RPC. + If set, trajectories can be executed via execute_trajectory() RPC. + """ + + name: str + urdf_path: str + base_pose: NDArray[np.float64] + joint_names: list[str] + end_effector_link: str + base_link: str = "base_link" + package_paths: dict[str, str] = field(default_factory=dict) + joint_limits_lower: NDArray[np.float64] | None = None + joint_limits_upper: NDArray[np.float64] | None = None + velocity_limits: NDArray[np.float64] | None = None + auto_convert_meshes: bool = False + xacro_args: dict[str, str] = field(default_factory=dict) + collision_exclusion_pairs: list[tuple[str, str]] = field(default_factory=list) + # Motion constraints for trajectory generation + max_velocity: float = 1.0 + max_acceleration: float = 2.0 + # Orchestrator integration + joint_name_mapping: dict[str, str] = field(default_factory=dict) + orchestrator_task_name: str | None = None + + def get_urdf_joint_name(self, orchestrator_name: str) -> str: + """Translate orchestrator joint name to URDF joint name.""" + return self.joint_name_mapping.get(orchestrator_name, orchestrator_name) + + def get_orchestrator_joint_name(self, urdf_name: str) -> str: + """Translate URDF joint name to orchestrator joint name.""" + for orch_name, u_name in self.joint_name_mapping.items(): + if u_name == urdf_name: + return orch_name + return urdf_name + + def get_orchestrator_joint_names(self) -> list[str]: + """Get joint names in orchestrator namespace.""" + if not self.joint_name_mapping: + return self.joint_names + return [self.get_orchestrator_joint_name(j) for j in self.joint_names] diff --git a/dimos/manipulation/planning/spec/enums.py b/dimos/manipulation/planning/spec/enums.py new file mode 100644 index 0000000000..66a17ee199 --- /dev/null +++ b/dimos/manipulation/planning/spec/enums.py @@ -0,0 +1,49 @@ +# 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. + +"""Enumerations for manipulation planning.""" + +from enum import Enum, auto + + +class ObstacleType(Enum): + """Type of obstacle geometry.""" + + BOX = auto() + SPHERE = auto() + CYLINDER = auto() + MESH = auto() + + +class IKStatus(Enum): + """Status of IK solution.""" + + SUCCESS = auto() + NO_SOLUTION = auto() + SINGULARITY = auto() + JOINT_LIMITS = auto() + COLLISION = auto() + TIMEOUT = auto() + + +class PlanningStatus(Enum): + """Status of motion planning.""" + + SUCCESS = auto() + NO_SOLUTION = auto() + TIMEOUT = auto() + INVALID_START = auto() + INVALID_GOAL = auto() + COLLISION_AT_START = auto() + COLLISION_AT_GOAL = auto() diff --git a/dimos/manipulation/planning/spec/protocols.py b/dimos/manipulation/planning/spec/protocols.py new file mode 100644 index 0000000000..6108ae60c5 --- /dev/null +++ b/dimos/manipulation/planning/spec/protocols.py @@ -0,0 +1,261 @@ +# 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. + +"""Protocol definitions for manipulation planning. + +All code should use these Protocol types (not concrete classes). +Use factory functions from dimos.manipulation.planning.factory to create instances. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable + +if TYPE_CHECKING: + from contextlib import AbstractContextManager + + import numpy as np + from numpy.typing import NDArray + + from dimos.manipulation.planning.spec.config import RobotModelConfig + from dimos.manipulation.planning.spec.types import IKResult, Obstacle, PlanningResult + + +@runtime_checkable +class WorldSpec(Protocol): + """Protocol for the world/scene backend. + + The world owns the physics/collision backend and provides: + - Robot/obstacle management + - Collision checking + - Forward kinematics + - Context management for thread safety + + Context Management: + - Live context: Mirrors current robot state (synced from driver) + - Scratch contexts: Thread-safe clones for planning/IK operations + + Implementations: + - DrakeWorld: Uses Drake's MultibodyPlant and SceneGraph + """ + + # Robot Management + def add_robot(self, config: RobotModelConfig) -> str: + """Add a robot to the world. Returns unique robot ID.""" + ... + + def get_robot_ids(self) -> list[str]: + """Get all robot IDs.""" + ... + + def get_robot_config(self, robot_id: str) -> RobotModelConfig: + """Get robot configuration.""" + ... + + def get_joint_limits(self, robot_id: str) -> tuple[NDArray[np.float64], NDArray[np.float64]]: + """Get joint limits (lower, upper) for a robot.""" + ... + + # Obstacle Management + def add_obstacle(self, obstacle: Obstacle) -> str: + """Add an obstacle to the world. Returns unique obstacle ID.""" + ... + + def remove_obstacle(self, obstacle_id: str) -> bool: + """Remove an obstacle. Returns True if removed.""" + ... + + def update_obstacle_pose(self, obstacle_id: str, pose: NDArray[np.float64]) -> bool: + """Update obstacle pose. Returns True if updated.""" + ... + + def clear_obstacles(self) -> None: + """Remove all obstacles.""" + ... + + # Lifecycle + def finalize(self) -> None: + """Finalize the world. Must be called after adding robots.""" + ... + + @property + def is_finalized(self) -> bool: + """Check if world is finalized.""" + ... + + # Context Management + def get_live_context(self) -> Any: + """Get the live context (mirrors real robot state).""" + ... + + def scratch_context(self) -> AbstractContextManager[Any]: + """Get a scratch context for planning (thread-safe clone).""" + ... + + def sync_from_joint_state(self, robot_id: str, positions: NDArray[np.float64]) -> None: + """Sync live context from joint state.""" + ... + + # State Operations (require context) + def set_positions(self, ctx: Any, robot_id: str, positions: NDArray[np.float64]) -> None: + """Set robot joint positions in a context.""" + ... + + def get_positions(self, ctx: Any, robot_id: str) -> NDArray[np.float64]: + """Get robot joint positions from a context.""" + ... + + # Collision Checking (require context) + def is_collision_free(self, ctx: Any, robot_id: str) -> bool: + """Check if robot configuration is collision-free.""" + ... + + def get_min_distance(self, ctx: Any, robot_id: str) -> float: + """Get minimum distance to obstacles (negative if collision).""" + ... + + # Forward Kinematics (require context) + def get_ee_pose(self, ctx: Any, robot_id: str) -> NDArray[np.float64]: + """Get end-effector pose (4x4 transform).""" + ... + + def get_jacobian(self, ctx: Any, robot_id: str) -> NDArray[np.float64]: + """Get end-effector Jacobian (6 x n_joints).""" + ... + + # Visualization (optional) + def get_meshcat_url(self) -> str | None: + """Get Meshcat visualization URL if enabled.""" + ... + + def publish_to_meshcat(self, ctx: Any | None = None) -> None: + """Publish current state to Meshcat visualization.""" + ... + + def animate_path( + self, robot_id: str, path: list[NDArray[np.float64]], duration: float = 3.0 + ) -> None: + """Animate a path in visualization.""" + ... + + +@runtime_checkable +class KinematicsSpec(Protocol): + """Protocol for inverse kinematics solver. + + Kinematics solvers are stateless and use WorldSpec for FK/collision. + + Methods: + - solve(): Full optimization-based IK with collision checking + - solve_iterative(): Iterative Jacobian-based IK + - solve_differential(): Single Jacobian step for velocity control + + Implementations: + - DrakeKinematics: Uses Drake's InverseKinematics + SNOPT/IPOPT + """ + + def solve( + self, + world: WorldSpec, + robot_id: str, + target_pose: NDArray[np.float64], + seed: NDArray[np.float64] | None = None, + position_tolerance: float = 0.001, + orientation_tolerance: float = 0.01, + check_collision: bool = True, + max_attempts: int = 10, + ) -> IKResult: + """Solve full IK with optional collision checking.""" + ... + + def solve_iterative( + self, + world: WorldSpec, + robot_id: str, + target_pose: NDArray[np.float64], + seed: NDArray[np.float64], + max_iterations: int = 100, + position_tolerance: float = 0.001, + orientation_tolerance: float = 0.01, + ) -> IKResult: + """Solve IK iteratively using Jacobian method.""" + ... + + def solve_differential( + self, + world: WorldSpec, + robot_id: str, + current_joints: NDArray[np.float64], + twist: NDArray[np.float64], + dt: float, + ) -> NDArray[np.float64] | None: + """Single Jacobian step for velocity control.""" + ... + + +@runtime_checkable +class PlannerSpec(Protocol): + """Protocol for motion planner. + + Planners find collision-free paths from start to goal configurations. + They use WorldSpec for collision checking and are stateless. + + Implementations: + - DrakePlanner: RRT-Connect planner + - DrakeRRTStarPlanner: RRT* planner (asymptotically optimal) + """ + + def plan_joint_path( + self, + world: WorldSpec, + robot_id: str, + q_start: NDArray[np.float64], + q_goal: NDArray[np.float64], + timeout: float = 10.0, + ) -> PlanningResult: + """Plan a collision-free joint-space path.""" + ... + + def get_name(self) -> str: + """Get planner name.""" + ... + + +@runtime_checkable +class VizSpec(Protocol): + """Protocol for visualization backend. + + Note: For Drake, visualization is typically integrated into DrakeWorld + via enable_viz=True. This protocol is for advanced use cases. + """ + + def set_robot_state(self, robot_id: str, positions: NDArray[np.float64]) -> None: + """Update robot visualization state.""" + ... + + def add_obstacle(self, obstacle: Obstacle) -> str: + """Add obstacle to visualization.""" + ... + + def remove_obstacle(self, obstacle_id: str) -> None: + """Remove obstacle from visualization.""" + ... + + def get_url(self) -> str | None: + """Get visualization URL (e.g., Meshcat URL).""" + ... + + def publish(self) -> None: + """Force publish current state to visualization.""" + ... diff --git a/dimos/manipulation/planning/spec/types.py b/dimos/manipulation/planning/spec/types.py new file mode 100644 index 0000000000..c698f37b87 --- /dev/null +++ b/dimos/manipulation/planning/spec/types.py @@ -0,0 +1,154 @@ +# 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. + +"""Data types for manipulation planning.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +from dimos.manipulation.planning.spec.enums import ( + IKStatus, + ObstacleType, + PlanningStatus, +) + +if TYPE_CHECKING: + import numpy as np + from numpy.typing import NDArray + + +@dataclass +class Obstacle: + """Obstacle specification for collision avoidance. + + Attributes: + name: Unique name for the obstacle + obstacle_type: Type of geometry (BOX, SPHERE, CYLINDER, MESH) + pose: 4x4 homogeneous transform + dimensions: Type-specific dimensions: + - BOX: (width, height, depth) + - SPHERE: (radius,) + - CYLINDER: (radius, height) + - MESH: Not used + color: RGBA color tuple (0-1 range) + mesh_path: Path to mesh file (for MESH type) + """ + + name: str + obstacle_type: ObstacleType + pose: NDArray[np.float64] + dimensions: tuple[float, ...] = () + color: tuple[float, float, float, float] = (0.8, 0.2, 0.2, 0.8) + mesh_path: str | None = None + + +@dataclass +class IKResult: + """Result of an IK solve. + + Attributes: + status: Solution status + joint_positions: Solution joint positions (None if failed) + position_error: Cartesian position error (meters) + orientation_error: Orientation error (radians) + iterations: Number of iterations taken + message: Human-readable status message + """ + + status: IKStatus + joint_positions: NDArray[np.float64] | None = None + position_error: float = 0.0 + orientation_error: float = 0.0 + iterations: int = 0 + message: str = "" + + def is_success(self) -> bool: + """Check if IK was successful.""" + return self.status == IKStatus.SUCCESS + + +@dataclass +class PlanningResult: + """Result of motion planning. + + Attributes: + status: Planning status + path: List of joint configurations (empty if failed) + planning_time: Time taken to plan (seconds) + path_length: Total path length in joint space (radians) + iterations: Number of iterations/nodes expanded + message: Human-readable status message + """ + + status: PlanningStatus + path: list[NDArray[np.float64]] = field(default_factory=list) + planning_time: float = 0.0 + path_length: float = 0.0 + iterations: int = 0 + message: str = "" + + def is_success(self) -> bool: + """Check if planning was successful.""" + return self.status == PlanningStatus.SUCCESS + + +@dataclass +class CollisionObjectMessage: + """Message for adding/updating/removing obstacles. + + Used by monitors to handle obstacle updates from external sources. + + Attributes: + id: Unique identifier for the object + operation: "add", "update", or "remove" + primitive_type: "box", "sphere", or "cylinder" (for add/update) + pose: 4x4 transform (for add/update) + dimensions: Type-specific dimensions (for add/update) + color: RGBA color tuple + """ + + id: str + operation: str # "add", "update", "remove" + primitive_type: str | None = None + pose: NDArray[np.float64] | None = None + dimensions: tuple[float, ...] | None = None + color: tuple[float, float, float, float] = (0.8, 0.2, 0.2, 0.8) + + +@dataclass +class PerceptionDetection: + """3D detection for planning obstacle updates. + + Internal type for the planning system using numpy transforms. + For LCM messages, use dimos_lcm.vision_msgs.Detection3D instead. + + Attributes: + id: Unique detection ID + label: Object class label + pose: 4x4 homogeneous transform + dimensions: (width, height, depth) + confidence: Detection confidence (0-1) + """ + + id: str + label: str + pose: NDArray[np.float64] + dimensions: tuple[float, float, float] + confidence: float = 1.0 + + +# Backwards compatibility alias - will be removed in future version +Detection3D = PerceptionDetection diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index 0df7bc548f..1bac856671 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -12,41 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -""" -Drake World Implementation - -Implements WorldSpec using Drake's MultibodyPlant and SceneGraph. - -## Architecture - -Uses Drake's DiagramBuilder pattern: -``` -DiagramBuilder -├── MultibodyPlant (kinematics, dynamics) -├── SceneGraph (collision geometry) -└── MeshcatVisualizer (optional) -``` - -## Context Management - -- Live context: Mirrors current robot state (synced from driver) -- Scratch contexts: Thread-safe clones for planning/IK operations - -Example: - world = DrakeWorld(enable_viz=True) - robot_id = world.add_robot(config) - world.add_obstacle(table) - world.finalize() - - # Sync live state from driver - world.sync_from_joint_state(robot_id, joint_positions) - - # Planning uses scratch contexts - with world.scratch_context() as ctx: - world.set_positions(ctx, robot_id, q_test) - if world.is_collision_free(ctx, robot_id): - ee_pose = world.get_ee_pose(ctx, robot_id) -""" +"""Drake World Implementation - WorldSpec using Drake's MultibodyPlant and SceneGraph.""" from __future__ import annotations @@ -131,36 +97,9 @@ class _ObstacleData: class DrakeWorld: - """Drake implementation of WorldSpec. - - Owns a single Drake Diagram containing: - - MultibodyPlant (kinematics, dynamics) - - SceneGraph (collision geometry) - - Optional MeshcatVisualizer + """Drake implementation of WorldSpec with MultibodyPlant, SceneGraph, optional Meshcat.""" - ## Context Management - - - _live_context: Mirrors current robot state (synced from driver) - - scratch_context(): Returns cloned context for planning/IK - - ## Thread Safety - - - Live context writes are protected by RLock - - scratch_context() returns independent clones - - All public methods that modify state are thread-safe - """ - - def __init__( - self, - time_step: float = 0.0, - enable_viz: bool = False, - ): - """Create a Drake world. - - Args: - time_step: Simulation time step (0 for kinematics-only) - enable_viz: If True, enable Meshcat visualization - """ + def __init__(self, time_step: float = 0.0, enable_viz: bool = False): if not DRAKE_AVAILABLE: raise ImportError("Drake is not installed. Install with: pip install drake") @@ -205,118 +144,94 @@ def __init__( # Obstacle source for dynamic obstacles self._obstacle_source_id: Any = None - # ============= Robot Management ============= - def add_robot(self, config: RobotModelConfig) -> str: - """Add a robot to the world. - - Args: - config: Robot configuration including URDF path and joint names - - Returns: - robot_id: Unique identifier for the robot - - Raises: - RuntimeError: If world is already finalized - """ + """Add a robot to the world. Returns robot_id.""" if self._finalized: raise RuntimeError("Cannot add robot after world is finalized") with self._lock: - # Generate unique robot ID self._robot_counter += 1 robot_id = f"robot_{self._robot_counter}" - # Prepare URDF: process xacro and convert STL meshes if needed - original_path = Path(config.urdf_path).resolve() - if not original_path.exists(): - raise FileNotFoundError(f"URDF/xacro not found: {original_path}") - - urdf_path = prepare_urdf_for_drake( - urdf_path=original_path, - package_paths=config.package_paths, - xacro_args=config.xacro_args, - # Always convert meshes when requested - Drake needs OBJ for collision - convert_meshes=config.auto_convert_meshes, - ) - - urdf_path_obj = Path(urdf_path) - logger.info(f"Using prepared URDF: {urdf_path_obj}") - - # Register package paths for mesh resolution - if config.package_paths: - for pkg_name, pkg_path in config.package_paths.items(): - self._parser.package_map().Add(pkg_name, Path(pkg_path)) - else: - # Fallback: register URDF directory as package - package_dir = urdf_path_obj.parent - self._parser.package_map().Add( - package_name=f"{config.name}_description", - package_path=package_dir, - ) - - # Parse the URDF - model_instances = self._parser.AddModels(urdf_path_obj) - if not model_instances: - raise ValueError(f"Failed to parse URDF: {urdf_path}") - - model_instance = model_instances[0] - - # Get base body - base_body = self._plant.GetBodyByName(config.base_link, model_instance) - - # Check if URDF already welded base to world - base_already_welded = False - try: - world_joint = self._plant.GetJointByName("world_joint", model_instance) - if ( - world_joint.parent_body().name() == "world" - and world_joint.child_body().name() == config.base_link - ): - base_already_welded = True - logger.info("URDF has 'world_joint', skipping weld") - except RuntimeError: - pass - - # Weld base to world if needed - if not base_already_welded: - world_frame = self._plant.world_frame() - base_transform = RigidTransform(config.base_pose) - self._plant.WeldFrames( - world_frame, - base_body.body_frame(), - base_transform, - ) - - # Verify joints exist - for joint_name in config.joint_names: - try: - self._plant.GetJointByName(joint_name, model_instance) - except RuntimeError: - raise ValueError(f"Joint '{joint_name}' not found in URDF") + model_instance = self._load_urdf(config) + self._weld_base_if_needed(config, model_instance) + self._validate_joints(config, model_instance) - # Get end-effector frame - try: - ee_body = self._plant.GetBodyByName(config.end_effector_link, model_instance) - ee_frame = ee_body.body_frame() - except RuntimeError: - raise ValueError( - f"End-effector link '{config.end_effector_link}' not found in URDF" - ) + ee_frame = self._plant.GetBodyByName( + config.end_effector_link, model_instance + ).body_frame() + base_frame = self._plant.GetBodyByName(config.base_link, model_instance).body_frame() - # Store robot data (joint_indices computed after finalize) self._robots[robot_id] = _RobotData( robot_id=robot_id, config=config, model_instance=model_instance, joint_indices=[], ee_frame=ee_frame, - base_frame=base_body.body_frame(), + base_frame=base_frame, ) logger.info(f"Added robot '{robot_id}' ({config.name})") return robot_id + def _load_urdf(self, config: RobotModelConfig) -> Any: + """Load URDF/xacro and return model instance.""" + original_path = Path(config.urdf_path).resolve() + if not original_path.exists(): + raise FileNotFoundError(f"URDF/xacro not found: {original_path}") + + urdf_path = prepare_urdf_for_drake( + urdf_path=original_path, + package_paths=config.package_paths, + xacro_args=config.xacro_args, + convert_meshes=config.auto_convert_meshes, + ) + urdf_path_obj = Path(urdf_path) + logger.info(f"Using prepared URDF: {urdf_path_obj}") + + # Register package paths + if config.package_paths: + for pkg_name, pkg_path in config.package_paths.items(): + self._parser.package_map().Add(pkg_name, Path(pkg_path)) + else: + self._parser.package_map().Add(f"{config.name}_description", urdf_path_obj.parent) + + model_instances = self._parser.AddModels(urdf_path_obj) + if not model_instances: + raise ValueError(f"Failed to parse URDF: {urdf_path}") + return model_instances[0] + + def _weld_base_if_needed(self, config: RobotModelConfig, model_instance: Any) -> None: + """Weld robot base to world if not already welded in URDF.""" + base_body = self._plant.GetBodyByName(config.base_link, model_instance) + + # Check if URDF already has world_joint + try: + world_joint = self._plant.GetJointByName("world_joint", model_instance) + if ( + world_joint.parent_body().name() == "world" + and world_joint.child_body().name() == config.base_link + ): + logger.info("URDF has 'world_joint', skipping weld") + return + except RuntimeError: + pass + + # Weld base to world + self._plant.WeldFrames( + self._plant.world_frame(), + base_body.body_frame(), + RigidTransform(config.base_pose), + ) + + def _validate_joints(self, config: RobotModelConfig, model_instance: Any) -> None: + """Validate that all configured joints exist in URDF.""" + for joint_name in config.joint_names: + try: + self._plant.GetJointByName(joint_name, model_instance) + except RuntimeError: + raise ValueError(f"Joint '{joint_name}' not found in URDF") + def get_robot_ids(self) -> list[str]: """Get all robot IDs in the world.""" return list(self._robots.keys()) @@ -622,98 +537,37 @@ def is_finalized(self) -> bool: return self._finalized def _setup_collision_filters(self) -> None: - """Filter out collisions between adjacent links in kinematic chain.""" - scene_graph = self._scene_graph - - for _robot_id, robot_data in self._robots.items(): - model_instance = robot_data.model_instance - body_indices = self._plant.GetBodyIndices(model_instance) - bodies = [self._plant.get_body(bi) for bi in body_indices] - - # Build parent-child relationships - parent_map: dict[str, str] = {} - for joint_idx in self._plant.GetJointIndices(model_instance): + """Filter collisions between adjacent links and user-specified pairs.""" + for robot_data in self._robots.values(): + # Filter parent-child pairs (adjacent links always "collide") + for joint_idx in self._plant.GetJointIndices(robot_data.model_instance): joint = self._plant.get_joint(joint_idx) - parent_body = joint.parent_body() - child_body = joint.child_body() - if parent_body.index() != self._plant.world_body().index(): - parent_map[child_body.name()] = parent_body.name() - - # Compute kinematic distance - def get_chain_distance(name1: str, name2: str) -> int: - def path_to_root(name: str) -> list[str]: - path = [name] - while name in parent_map: - name = parent_map[name] - path.append(name) - return path - - path1 = path_to_root(name1) - path2 = path_to_root(name2) - - set1 = set(path1) - for i, node in enumerate(path2): - if node in set1: - idx1 = path1.index(node) - return idx1 + i - - return len(path1) + len(path2) - - # Filter collisions within 2 joints - filter_distance = 2 - filtered_pairs: set[tuple[str, str]] = set() - - for i, body1 in enumerate(bodies): - for body2 in bodies[i + 1 :]: - name1, name2 = body1.name(), body2.name() - dist = get_chain_distance(name1, name2) - - if dist <= filter_distance: - pair = (min(name1, name2), max(name1, name2)) - if pair not in filtered_pairs: - filtered_pairs.add(pair) - - try: - geoms1 = self._plant.GetCollisionGeometriesForBody(body1) - geoms2 = self._plant.GetCollisionGeometriesForBody(body2) - - if geoms1 and geoms2: - scene_graph.collision_filter_manager().Apply( - CollisionFilterDeclaration().ExcludeBetween( - GeometrySet(geoms1), GeometrySet(geoms2) - ) - ) - except Exception as e: - logger.warning(f"Could not filter {name1}<->{name2}: {e}") - - # Add user-specified collision exclusions from config - # Useful for parallel linkage mechanisms (grippers) where non-adjacent - # links may legitimately overlap due to mimic joints - exclusion_pairs = robot_data.config.collision_exclusion_pairs - if exclusion_pairs: - body_name_map = {body.name(): body for body in bodies} - for name1, name2 in exclusion_pairs: - if name1 in body_name_map and name2 in body_name_map: - body1 = body_name_map[name1] - body2 = body_name_map[name2] - try: - geoms1 = self._plant.GetCollisionGeometriesForBody(body1) - geoms2 = self._plant.GetCollisionGeometriesForBody(body2) - if geoms1 and geoms2: - scene_graph.collision_filter_manager().Apply( - CollisionFilterDeclaration().ExcludeBetween( - GeometrySet(geoms1), GeometrySet(geoms2) - ) - ) - logger.debug(f"Excluded collision pair: {name1}<->{name2}") - except Exception as e: - logger.warning(f"Could not filter pair {name1}<->{name2}: {e}") - else: - missing = [n for n in (name1, name2) if n not in body_name_map] - logger.warning(f"Collision exclusion: links not found: {missing}") + parent, child = joint.parent_body(), joint.child_body() + if parent.index() != self._plant.world_body().index(): + self._exclude_body_pair(parent, child) + + # Filter user-specified pairs (e.g., parallel linkage grippers) + for name1, name2 in robot_data.config.collision_exclusion_pairs: + try: + body1 = self._plant.GetBodyByName(name1, robot_data.model_instance) + body2 = self._plant.GetBodyByName(name2, robot_data.model_instance) + self._exclude_body_pair(body1, body2) + except RuntimeError: + logger.warning(f"Collision exclusion: link not found: {name1} or {name2}") logger.info("Collision filters applied") + def _exclude_body_pair(self, body1: Any, body2: Any) -> None: + """Exclude collision between two bodies.""" + geoms1 = self._plant.GetCollisionGeometriesForBody(body1) + geoms2 = self._plant.GetCollisionGeometriesForBody(body2) + if geoms1 and geoms2: + self._scene_graph.collision_filter_manager().Apply( + CollisionFilterDeclaration().ExcludeBetween( + GeometrySet(geoms1), GeometrySet(geoms2) + ) + ) + # ============= Context Management ============= def get_live_context(self) -> Context: @@ -728,49 +582,26 @@ def get_live_context(self) -> Context: @contextmanager def scratch_context(self) -> Generator[Context, None, None]: - """Get a scratch context for planning (thread-safe). - - Uses CreateDefaultContext() instead of Clone() so that dynamically - added obstacles are visible to collision queries. - - IMPORTANT: This copies the current live state of ALL robots into the - scratch context so that collision checking properly accounts for - inter-robot collisions (robot A planning considers robot B's position). - - Usage: - with world.scratch_context() as ctx: - world.set_positions(ctx, robot_id, q) - if world.is_collision_free(ctx, robot_id): - ... - """ + """Thread-safe context for planning. Copies current robot states for inter-robot collision checking.""" if not self._finalized: raise RuntimeError("World must be finalized first") - # Create fresh context - this sees all geometries including dynamic obstacles - # (Clone() doesn't see geometries added after the original context was created) ctx = self._diagram.CreateDefaultContext() - # Copy current live state of ALL robots into scratch context - # This ensures inter-robot collision checking works correctly + # Copy live robot states so inter-robot collision checking works with self._lock: if self._plant_context is not None: plant_ctx = self._diagram.GetMutableSubsystemContext(self._plant, ctx) - for _robot_id, robot_data in self._robots.items(): + for robot_data in self._robots.values(): try: - # Get positions from live context - live_positions = self._plant.GetPositions( + positions = self._plant.GetPositions( self._plant_context, robot_data.model_instance ) - # Set in scratch context - self._plant.SetPositions( - plant_ctx, robot_data.model_instance, live_positions - ) + self._plant.SetPositions(plant_ctx, robot_data.model_instance, positions) except Exception: - # Robot not yet synced - leave at default position - pass + pass # Robot not yet synced yield ctx - # Context automatically cleaned up when exiting def sync_from_joint_state(self, robot_id: str, positions: NDArray[np.float64]) -> None: """Sync live context from driver's joint state. From e38b12cd4a61be51600ff91a88660640e2bb9510 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Tue, 20 Jan 2026 22:34:12 -0800 Subject: [PATCH 16/31] Refactored manipulation planning stack: seperated planners from world implementation, split kinematics into JacobianIK and DrakeOptimizationIK --- dimos/manipulation/manipulation_blueprints.py | 95 ----- dimos/manipulation/manipulation_module.py | 61 ++- dimos/manipulation/planning/__init__.py | 24 +- .../planning/examples/manipulation_client.py | 74 ++-- dimos/manipulation/planning/factory.py | 222 ++-------- .../planning/kinematics/__init__.py | 28 +- ...kinematics.py => drake_optimization_ik.py} | 210 ++++------ .../planning/kinematics/jacobian_ik.py | 387 ++++++++++++++++++ .../manipulation/planning/monitor/__init__.py | 3 +- .../planning/monitor/world_monitor.py | 26 +- .../monitor/world_obstacle_monitor.py | 23 +- .../planning/planners/__init__.py | 15 +- .../{drake_planner.py => rrt_planner.py} | 250 ++--------- dimos/manipulation/planning/spec/__init__.py | 42 +- dimos/manipulation/planning/spec/protocols.py | 60 ++- dimos/manipulation/planning/spec/types.py | 26 -- .../planning/world/drake_world.py | 51 ++- 17 files changed, 751 insertions(+), 846 deletions(-) rename dimos/manipulation/planning/kinematics/{drake_kinematics.py => drake_optimization_ik.py} (62%) create mode 100644 dimos/manipulation/planning/kinematics/jacobian_ik.py rename dimos/manipulation/planning/planners/{drake_planner.py => rrt_planner.py} (53%) diff --git a/dimos/manipulation/manipulation_blueprints.py b/dimos/manipulation/manipulation_blueprints.py index f3334bcbc0..1d85490e7b 100644 --- a/dimos/manipulation/manipulation_blueprints.py +++ b/dimos/manipulation/manipulation_blueprints.py @@ -29,7 +29,6 @@ import numpy as np -from dimos.core.blueprints import autoconnect from dimos.core.transport import LCMTransport from dimos.manipulation.manipulation_module import manipulation_module from dimos.manipulation.planning.spec import RobotModelConfig @@ -304,104 +303,10 @@ def _make_piper_config( ) -# ============================================================================= -# Combined Blueprints (Orchestrator + Planner in one process) -# ============================================================================= - - -# Combined: Quad mock orchestrator + Quad manipulation planner -# Usage: dimos run quad-mock-with-planner -def _make_quad_mock_with_planner(): - """Create combined quad-arm blueprint (lazy import to avoid circular deps).""" - from dimos.control.blueprints import orchestrator_quad_mock - - return autoconnect( - orchestrator_quad_mock, - manipulation_module( - robots=[ - _make_xarm6_config( - "arm1", y_offset=0.75, joint_prefix="arm1_", orchestrator_task="traj_arm1" - ), - _make_xarm6_config( - "arm2", y_offset=0.25, joint_prefix="arm2_", orchestrator_task="traj_arm2" - ), - _make_xarm6_config( - "arm3", y_offset=-0.25, joint_prefix="arm3_", orchestrator_task="traj_arm3" - ), - _make_xarm6_config( - "arm4", y_offset=-0.75, joint_prefix="arm4_", orchestrator_task="traj_arm4" - ), - ], - planning_timeout=10.0, - enable_viz=True, - ), - ).transports( - { - ("joint_state", JointState): LCMTransport("/orchestrator/joint_state", JointState), - } - ) - - -# Lazy-loaded to avoid circular import -quad_mock_with_planner = _make_quad_mock_with_planner() - - -# Combined: Mixed-arm orchestrator (XArm7 + XArm6 + Piper) + manipulation planner -# Usage: dimos run mixed-mock-with-planner -def _make_mixed_mock_with_planner(): - """Create combined mixed-arm blueprint (XArm7 + XArm6 + Piper). - - Layout (top view, looking down at table): - y=0.5: XArm7 (7-DOF with gripper) - y=0.0: XArm6 (6-DOF with gripper) - y=-0.5: Piper (6-DOF with parallel jaw gripper) - """ - from dimos.control.blueprints import orchestrator_mixed_mock - - return autoconnect( - orchestrator_mixed_mock, - manipulation_module( - robots=[ - _make_xarm7_config( - "xarm7", - y_offset=0.5, - joint_prefix="xarm7_", - orchestrator_task="traj_xarm7", - add_gripper=True, - ), - _make_xarm6_config( - "xarm6", - y_offset=0.0, - joint_prefix="xarm6_", - orchestrator_task="traj_xarm6", - ), - _make_piper_config( - "piper", - y_offset=-0.5, - joint_prefix="piper_", - orchestrator_task="traj_piper", - ), - ], - planning_timeout=10.0, - enable_viz=True, - ), - ).transports( - { - ("joint_state", JointState): LCMTransport("/orchestrator/joint_state", JointState), - } - ) - - -# Lazy-loaded to avoid circular import -mixed_mock_with_planner = _make_mixed_mock_with_planner() - - __all__ = [ "PIPER_GRIPPER_COLLISION_EXCLUSIONS", "XARM_GRIPPER_COLLISION_EXCLUSIONS", "dual_xarm6_planner", - "mixed_mock_with_planner", - "quad_mock_with_planner", "xarm6_planner_only", "xarm7_planner_orchestrator", ] diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index 94c90cd468..95c4b94778 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -28,6 +28,8 @@ from dimos.manipulation.planning import ( JointTrajectoryGenerator, KinematicsSpec, + Obstacle, + ObstacleType, PlannerSpec, RobotModelConfig, create_kinematics, @@ -68,6 +70,8 @@ class ManipulationModuleConfig(ModuleConfig): robots: list[RobotModelConfig] = field(default_factory=list) planning_timeout: float = 10.0 enable_viz: bool = False + planner_name: str = "rrt_connect" # "rrt_connect" + kinematics_name: str = "jacobian" # "jacobian" or "drake_optimization" class ManipulationModule(Module): @@ -145,11 +149,11 @@ def _initialize_planning(self) -> None: if self.config.enable_viz: self._world_monitor.start_visualization_thread(rate_hz=10.0) - if url := self._world_monitor.get_meshcat_url(): + if url := self._world_monitor.get_visualization_url(): logger.info(f"Visualization: {url}") - self._planner = create_planner(name="rrt_connect") - self._kinematics = create_kinematics(backend="drake") + self._planner = create_planner(name=self.config.planner_name) + self._kinematics = create_kinematics(name=self.config.kinematics_name) def _get_default_robot_name(self) -> str | None: """Get default robot name (first robot if only one, else None).""" @@ -423,7 +427,7 @@ def get_visualization_url(self) -> str | None: """ if self._world_monitor is None: return None - return self._world_monitor.get_meshcat_url() + return self._world_monitor.get_visualization_url() @rpc def clear_planned_path(self) -> bool: @@ -570,22 +574,43 @@ def world_monitor(self) -> WorldMonitor | None: return self._world_monitor @rpc - def add_obstacle(self, name: str, pose: Pose, shape: str, dimensions: list[float]) -> str: - """Add obstacle: shape='box'|'sphere'|'cylinder', dimensions=[w,h,d]|[r]|[r,len].""" + def add_obstacle( + self, + name: str, + pose: Pose, + shape: str, + dimensions: list[float] | None = None, + mesh_path: str | None = None, + ) -> str: + """Add obstacle: shape='box'|'sphere'|'cylinder'|'mesh'. Returns obstacle_id.""" if not self._world_monitor: return "" - p = pose_to_matrix(pose) - match shape: - case "box": - return self._world_monitor.add_box_obstacle(name, p, tuple(dimensions)) # type: ignore[arg-type] - case "sphere": - return self._world_monitor.add_sphere_obstacle(name, p, dimensions[0]) - case "cylinder": - return self._world_monitor.add_cylinder_obstacle( - name, p, dimensions[0], dimensions[1] - ) - case _: - return "" + + # Map shape string to ObstacleType + shape_map = { + "box": ObstacleType.BOX, + "sphere": ObstacleType.SPHERE, + "cylinder": ObstacleType.CYLINDER, + "mesh": ObstacleType.MESH, + } + obstacle_type = shape_map.get(shape) + if obstacle_type is None: + logger.warning(f"Unknown obstacle shape: {shape}") + return "" + + # Validate mesh_path for mesh type + if obstacle_type == ObstacleType.MESH and not mesh_path: + logger.warning("mesh_path required for mesh obstacles") + return "" + + obstacle = Obstacle( + name=name, + obstacle_type=obstacle_type, + pose=pose_to_matrix(pose), + dimensions=tuple(dimensions) if dimensions else (), + mesh_path=mesh_path, + ) + return self._world_monitor.add_obstacle(obstacle) @rpc def remove_obstacle(self, obstacle_id: str) -> bool: diff --git a/dimos/manipulation/planning/__init__.py b/dimos/manipulation/planning/__init__.py index 2c3be92d91..71638f376b 100644 --- a/dimos/manipulation/planning/__init__.py +++ b/dimos/manipulation/planning/__init__.py @@ -17,11 +17,15 @@ Motion planning stack for robotic manipulators using Protocol-based architecture. -## Core Components +## Architecture -- WorldSpec: Core backend owning physics/collision (DrakeWorld) -- KinematicsSpec: Stateless IK operations (DrakeKinematics) -- PlannerSpec: Joint-space path planning (DrakePlanner) +- WorldSpec: Core backend owning physics/collision (DrakeWorld, future: MuJoCoWorld) +- KinematicsSpec: IK solvers + - JacobianIK: Backend-agnostic iterative/differential IK + - DrakeOptimizationIK: Drake-specific nonlinear optimization IK +- PlannerSpec: Backend-agnostic joint-space path planning + - RRTConnectPlanner: Bi-directional RRT-Connect + - RRTStarPlanner: RRT* (asymptotically optimal) ## Factory Functions @@ -35,8 +39,8 @@ ) world = create_world(backend="drake", enable_viz=True) -kinematics = create_kinematics(backend="drake") -planner = create_planner(name="rrt_connect") +kinematics = create_kinematics(name="jacobian") # or "drake_optimization" +planner = create_planner(name="rrt_connect") # backend-agnostic ``` ## Monitors @@ -64,7 +68,6 @@ # Data classes and Protocols from dimos.manipulation.planning.spec import ( CollisionObjectMessage, - Detection3D, IKResult, IKStatus, KinematicsSpec, @@ -74,7 +77,6 @@ PlanningResult, PlanningStatus, RobotModelConfig, - VizSpec, WorldSpec, ) @@ -84,14 +86,10 @@ ) __all__ = [ - # Data classes "CollisionObjectMessage", - "Detection3D", "IKResult", "IKStatus", - # Trajectory "JointTrajectoryGenerator", - # Protocols "KinematicsSpec", "Obstacle", "ObstacleType", @@ -99,9 +97,7 @@ "PlanningResult", "PlanningStatus", "RobotModelConfig", - "VizSpec", "WorldSpec", - # Factory functions "create_kinematics", "create_planner", "create_planning_stack", diff --git a/dimos/manipulation/planning/examples/manipulation_client.py b/dimos/manipulation/planning/examples/manipulation_client.py index 08f79f6e36..a749690629 100644 --- a/dimos/manipulation/planning/examples/manipulation_client.py +++ b/dimos/manipulation/planning/examples/manipulation_client.py @@ -24,20 +24,21 @@ # Run interactive client: python -m dimos.manipulation.planning.examples.manipulation_client -IPython Commands (client is available as 'c'): - c.joints() # Get current joint positions - c.ee() # Get end-effector pose - c.state() # Get manipulation state - c.url() # Get Meshcat visualization URL - - c.plan([0.1, ...]) # Plan to joint config - c.plan_pose(x, y, z) # Plan to cartesian pose - c.preview() # Preview path in Meshcat - c.execute() # Execute via orchestrator - - c.box("table", ...) # Add box obstacle - c.sphere("ball", ...) # Add sphere obstacle - c.remove("obstacle_id") # Remove obstacle +Commands (call directly, no prefix needed): + joints() # Get current joint positions + ee() # Get end-effector pose + state() # Get manipulation state + url() # Get Meshcat visualization URL + + plan([0.1, ...]) # Plan to joint config + plan_pose(x, y, z) # Plan to cartesian pose + preview() # Preview path in Meshcat + execute() # Execute via orchestrator + + box("name", x, y, z, w, h, d) # Add box obstacle + sphere("name", x, y, z, r) # Add sphere obstacle + cylinder("name", x, y, z, r, l) # Add cylinder obstacle + remove("obstacle_id") # Remove obstacle """ from __future__ import annotations @@ -267,29 +268,44 @@ def main() -> None: c = ManipulationClient() + # Expose methods directly (no 'c.' prefix needed) + banner = """ Manipulation Client - IPython Interface ======================================== -Client available as 'c'. Examples: - c.joints() # Get joint positions - c.ee() # Get end-effector pose - c.url() # Get Meshcat URL - - c.plan([0.1, 0.2, ...]) # Plan to joints - c.plan_pose(0.4, 0, 0.3)# Plan to pose - c.preview() # Preview in Meshcat - c.execute() # Execute trajectory - - c.box("table", 0.5, 0, 0, 0.6, 0.4, 0.02) # Add obstacle - c.remove("obstacle_id") # Remove obstacle - -Type c. for all methods, or help(c.method) for details. +Commands (no prefix needed): + joints() # Get joint positions + ee() # Get end-effector pose + url() # Get Meshcat URL + state() # Get manipulation state + robots() # List configured robots + +Planning: + plan([0.1, 0.2, ...]) # Plan to joint config + plan_pose(0.4, 0, 0.3) # Plan to cartesian pose (keeps orientation) + plan_pose(0.4, 0, 0.3, roll=0, pitch=3.14, yaw=0) # With orientation + preview() # Preview path in Meshcat + execute() # Execute via orchestrator + +Obstacles: + box("name", x, y, z, width, height, depth) # Add box + box("table", 0.5, 0, -0.02, 1.0, 0.6, 0.04) # Example: table + sphere("name", x, y, z, radius) # Add sphere + cylinder("name", x, y, z, radius, length) # Add cylinder + remove("obstacle_id") # Remove obstacle + +Utility: + collision([0.1, ...]) # Check if config is collision-free + reset() # Reset to IDLE state + cancel() # Cancel current motion + +Type help(command) for details, e.g. help(box) """ print(banner) try: - embed(colors="neutral") + embed(colors="neutral") # type: ignore[no-untyped-call] finally: c.stop() diff --git a/dimos/manipulation/planning/factory.py b/dimos/manipulation/planning/factory.py index c49a9b874e..d392bac563 100644 --- a/dimos/manipulation/planning/factory.py +++ b/dimos/manipulation/planning/factory.py @@ -12,34 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -""" -Factory Functions for Manipulation Planning - -Provides factory functions to create instances of WorldSpec, KinematicsSpec, -PlannerSpec, and VizSpec implementations. - -Only these factory functions know about concrete implementation types. -All other code should use the Protocol types. - -Example: - from dimos.manipulation.planning.factory import ( - create_world, - create_kinematics, - create_planner, - ) - - # Create instances using factory functions - world = create_world(backend="drake", enable_viz=True) - kinematics = create_kinematics(backend="drake") - planner = create_planner(name="rrt_connect", backend="drake") - - # Use them with Protocol types - robot_id = world.add_robot(config) - world.finalize() - - result = kinematics.solve(world, robot_id, target_pose) - plan_result = planner.plan_joint_path(world, robot_id, q_start, q_goal) -""" +"""Factory functions for manipulation planning components.""" from __future__ import annotations @@ -49,7 +22,6 @@ from dimos.manipulation.planning.spec import ( KinematicsSpec, PlannerSpec, - VizSpec, WorldSpec, ) @@ -59,33 +31,7 @@ def create_world( enable_viz: bool = False, **kwargs: Any, ) -> WorldSpec: - """Create a world instance. - - The World owns the physics/collision backend and provides: - - Robot/obstacle management - - Collision checking - - Forward kinematics - - Context management for thread safety - - Args: - backend: Backend to use ("drake", future: "mujoco", "vamp") - enable_viz: Enable Meshcat visualization (Drake only) - **kwargs: Additional backend-specific arguments - - time_step: Simulation time step (default 0.0 for kinematics-only) - - Returns: - WorldSpec implementation - - Example: - world = create_world(backend="drake", enable_viz=True) - robot_id = world.add_robot(robot_config) - world.finalize() - - with world.scratch_context() as ctx: - world.set_positions(ctx, robot_id, q) - if world.is_collision_free(ctx, robot_id): - ee_pose = world.get_ee_pose(ctx, robot_id) - """ + """Create a world instance. backend='drake', enable_viz for Meshcat.""" if backend == "drake": from dimos.manipulation.planning.world.drake_world import DrakeWorld @@ -95,169 +41,49 @@ def create_world( def create_kinematics( - backend: str = "drake", + name: str = "jacobian", **kwargs: Any, ) -> KinematicsSpec: - """Create a kinematics solver instance. - - Kinematics solvers are stateless (except for configuration) and - use WorldSpec for all FK/collision operations. - - Args: - backend: Backend to use ("drake") - **kwargs: Additional backend-specific arguments - - damping: Damping factor for differential IK (default 0.01) - - max_iterations: Default max iterations for iterative IK (default 100) - - singularity_threshold: Manipulability threshold (default 0.001) - - Returns: - KinematicsSpec implementation - - Example: - kinematics = create_kinematics(backend="drake", damping=0.01) - result = kinematics.solve(world, robot_id, target_pose, seed=q_current) - if result.is_success(): - q_goal = result.joint_positions - """ - if backend == "drake": - from dimos.manipulation.planning.kinematics.drake_kinematics import ( - DrakeKinematics, + """Create IK solver. name='jacobian'|'drake_optimization'.""" + if name == "jacobian": + from dimos.manipulation.planning.kinematics.jacobian_ik import JacobianIK + + return JacobianIK(**kwargs) + elif name == "drake_optimization": + from dimos.manipulation.planning.kinematics.drake_optimization_ik import ( + DrakeOptimizationIK, ) - return DrakeKinematics(**kwargs) + return DrakeOptimizationIK(**kwargs) else: - raise ValueError(f"Unknown backend: {backend}. Available: ['drake']") + raise ValueError( + f"Unknown kinematics solver: {name}. Available: ['jacobian', 'drake_optimization']" + ) def create_planner( name: str = "rrt_connect", - backend: str = "drake", **kwargs: Any, ) -> PlannerSpec: - """Create a motion planner instance. - - Planners find collision-free paths from start to goal configurations. - They use WorldSpec for collision checking and are stateless. - - Args: - name: Planner algorithm ("rrt_connect", "rrt_star") - backend: Backend to use ("drake") - **kwargs: Additional planner-specific arguments - RRT-Connect: - - step_size: Extension step size (default 0.1) - - connect_step_size: Connect step size (default 0.05) - - goal_tolerance: Goal tolerance (default 0.1) - - collision_step_size: Collision check step (default 0.02) - - RRT*: - - step_size: Extension step size (default 0.1) - - goal_tolerance: Goal tolerance (default 0.1) - - rewire_radius: Rewiring radius (default 0.5) - - collision_step_size: Collision check step (default 0.02) - - Returns: - PlannerSpec implementation + """Create motion planner. name='rrt_connect'.""" + if name == "rrt_connect": + from dimos.manipulation.planning.planners.rrt_planner import RRTConnectPlanner - Example: - planner = create_planner(name="rrt_connect", step_size=0.1) - result = planner.plan_joint_path(world, robot_id, q_start, q_goal) - if result.is_success(): - waypoints = result.path - """ - if backend == "drake": - if name == "rrt_connect": - from dimos.manipulation.planning.planners.drake_planner import DrakePlanner - - return DrakePlanner(**kwargs) - elif name == "rrt_star": - from dimos.manipulation.planning.planners.drake_planner import ( - DrakeRRTStarPlanner, - ) - - return DrakeRRTStarPlanner(**kwargs) - else: - raise ValueError( - f"Unknown planner: {name}. Available for Drake: ['rrt_connect', 'rrt_star']" - ) - else: - raise ValueError(f"Unknown backend: {backend}. Available: ['drake']") - - -def create_viz( - backend: str = "drake", - world: WorldSpec | None = None, - **kwargs: Any, -) -> VizSpec: - """Create a visualization instance. - - Provides methods to update robot/obstacle visualization. - Can be integrated with WorldSpec or standalone. - - Note: For Drake, visualization is typically integrated into DrakeWorld - via enable_viz=True. This function is for advanced use cases where - separate visualization is needed. - - Args: - backend: Backend to use ("drake") - world: Optional world to integrate with - **kwargs: Additional backend-specific arguments - - Returns: - VizSpec implementation - - Example: - # Option 1: Integrated visualization (recommended) - world = create_world(backend="drake", enable_viz=True) - print(f"Meshcat URL: {world.get_meshcat_url()}") - - # Option 2: Separate visualization (advanced) - viz = create_viz(backend="drake", world=world) - viz.set_robot_state(robot_id, q) - viz.publish() - """ - if backend == "drake": - # For Drake, visualization is integrated into DrakeWorld - # This is a placeholder for potential future standalone viz - raise NotImplementedError( - "For Drake, use create_world(enable_viz=True) instead. " - "Separate visualization is not yet implemented." - ) + return RRTConnectPlanner(**kwargs) else: - raise ValueError(f"Unknown backend: {backend}. Available: ['drake']") - - -# ============= Convenience Functions ============= + raise ValueError(f"Unknown planner: {name}. Available: ['rrt_connect']") def create_planning_stack( robot_config: Any, enable_viz: bool = False, planner_name: str = "rrt_connect", + kinematics_name: str = "jacobian", ) -> tuple[WorldSpec, KinematicsSpec, PlannerSpec, str]: - """Convenience function to create a complete planning stack. - - Creates world, kinematics, and planner, adds robot, and finalizes. - - Args: - robot_config: RobotModelConfig for the robot - enable_viz: Enable visualization - planner_name: Planner algorithm to use - - Returns: - Tuple of (world, kinematics, planner, robot_id) - - Example: - world, kinematics, planner, robot_id = create_planning_stack( - robot_config=piper_config, - enable_viz=True, - ) - - # Now ready to use - result = kinematics.solve(world, robot_id, target_pose) - """ + """Create complete planning stack. Returns (world, kinematics, planner, robot_id).""" world = create_world(backend="drake", enable_viz=enable_viz) - kinematics = create_kinematics(backend="drake") - planner = create_planner(name=planner_name, backend="drake") + kinematics = create_kinematics(name=kinematics_name) + planner = create_planner(name=planner_name) robot_id = world.add_robot(robot_config) world.finalize() diff --git a/dimos/manipulation/planning/kinematics/__init__.py b/dimos/manipulation/planning/kinematics/__init__.py index ec34b7a53a..588fccbd72 100644 --- a/dimos/manipulation/planning/kinematics/__init__.py +++ b/dimos/manipulation/planning/kinematics/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2025 Dimensional Inc. +# 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. @@ -19,9 +19,29 @@ ## Implementations -- DrakeKinematics: Uses Drake's InverseKinematics + SNOPT/IPOPT +- JacobianIK: Backend-agnostic iterative/differential IK (works with any WorldSpec) +- DrakeOptimizationIK: Drake-specific nonlinear optimization IK (requires DrakeWorld) + +## Usage + +Use factory functions to create IK solvers: + +```python +from dimos.manipulation.planning.factory import create_kinematics + +# Backend-agnostic (works with any WorldSpec) +kinematics = create_kinematics(name="jacobian") + +# Drake-specific (requires DrakeWorld, more accurate) +kinematics = create_kinematics(name="drake_optimization") + +result = kinematics.solve(world, robot_id, target_pose) +``` """ -from dimos.manipulation.planning.kinematics.drake_kinematics import DrakeKinematics +from dimos.manipulation.planning.kinematics.drake_optimization_ik import ( + DrakeOptimizationIK, +) +from dimos.manipulation.planning.kinematics.jacobian_ik import JacobianIK -__all__ = ["DrakeKinematics"] +__all__ = ["DrakeOptimizationIK", "JacobianIK"] diff --git a/dimos/manipulation/planning/kinematics/drake_kinematics.py b/dimos/manipulation/planning/kinematics/drake_optimization_ik.py similarity index 62% rename from dimos/manipulation/planning/kinematics/drake_kinematics.py rename to dimos/manipulation/planning/kinematics/drake_optimization_ik.py index 7c6bfad218..7637c03c30 100644 --- a/dimos/manipulation/planning/kinematics/drake_kinematics.py +++ b/dimos/manipulation/planning/kinematics/drake_optimization_ik.py @@ -12,7 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Drake Kinematics - optimization-based and differential IK using Drake.""" +"""Drake-specific optimization-based inverse kinematics. + +DrakeOptimizationIK uses Drake's InverseKinematics class with nonlinear optimization +(SNOPT/IPOPT) to find accurate IK solutions. It requires DrakeWorld and cannot work +with other physics backends. + +For backend-agnostic Jacobian-based IK, use JacobianIK. +""" from __future__ import annotations @@ -21,12 +28,7 @@ import numpy as np from dimos.manipulation.planning.spec import IKResult, IKStatus, WorldSpec -from dimos.manipulation.planning.utils.kinematics_utils import ( - check_singularity, - compute_error_twist, - compute_pose_error, - damped_pseudoinverse, -) +from dimos.manipulation.planning.utils.kinematics_utils import compute_pose_error from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: @@ -44,20 +46,40 @@ logger = setup_logger() -class DrakeKinematics: - """Drake IK solver - optimization-based and Jacobian-based methods.""" +class DrakeOptimizationIK: + """Drake-specific optimization-based IK solver. - def __init__( - self, - damping: float = 0.01, - max_iterations: int = 100, - singularity_threshold: float = 0.001, - ): + Uses Drake's InverseKinematics class with SNOPT/IPOPT to solve constrained + nonlinear optimization for IK. This produces highly accurate solutions but + requires DrakeWorld. + + For backend-agnostic IK, use JacobianIK instead. + + Methods: + - solve(): Full nonlinear IK with multiple random restarts + - solve_single(): Single optimization attempt with given seed + + Example: + from dimos.manipulation.planning.world import DrakeWorld + + world = DrakeWorld(enable_viz=True) + world.add_robot(config) + world.finalize() + + ik = DrakeOptimizationIK() + result = ik.solve(world, robot_id, target_pose) + if result.is_success(): + print(f"Solution: {result.joint_positions}") + """ + + def __init__(self) -> None: + """Create Drake optimization IK solver.""" if not DRAKE_AVAILABLE: raise ImportError("Drake is not installed. Install with: pip install drake") - self._damping = damping - self._max_iterations = max_iterations - self._singularity_threshold = singularity_threshold + # Internal JacobianIK for iterative/differential methods + from dimos.manipulation.planning.kinematics.jacobian_ik import JacobianIK + + self._jacobian_ik = JacobianIK() def _validate_world(self, world: WorldSpec) -> IKResult | None: """Validate world is DrakeWorld and finalized. Returns error or None.""" @@ -65,7 +87,7 @@ def _validate_world(self, world: WorldSpec) -> IKResult | None: if not isinstance(world, DrakeWorld): return _create_failure_result( - IKStatus.NO_SOLUTION, "DrakeKinematics requires DrakeWorld" + IKStatus.NO_SOLUTION, "DrakeOptimizationIK requires DrakeWorld" ) if not world.is_finalized: return _create_failure_result(IKStatus.NO_SOLUTION, "World must be finalized before IK") @@ -82,7 +104,25 @@ def solve( check_collision: bool = True, max_attempts: int = 10, ) -> IKResult: - """Full nonlinear IK with multiple random restarts.""" + """Full nonlinear IK with multiple random restarts. + + Solves constrained nonlinear optimization using Drake's InverseKinematics + and SNOPT/IPOPT solvers. Tries multiple starting configurations to find + a collision-free solution. + + Args: + world: DrakeWorld instance (NOT other WorldSpec implementations) + robot_id: Robot to solve IK for + target_pose: Target end-effector pose (4x4 transform) + seed: Initial guess (uses current position if None) + position_tolerance: Required position accuracy (meters) + orientation_tolerance: Required orientation accuracy (radians) + check_collision: Whether to check collision of solution + max_attempts: Maximum random restart attempts + + Returns: + IKResult with solution or failure status + """ error = self._validate_world(world) if error is not None: return error @@ -243,67 +283,15 @@ def solve_iterative( position_tolerance: float = 0.001, orientation_tolerance: float = 0.01, ) -> IKResult: - """Iterative Jacobian-based IK until convergence.""" - max_iterations = max_iterations or self._max_iterations - current_joints = seed.copy() - - lower_limits, upper_limits = world.get_joint_limits(robot_id) - - for iteration in range(max_iterations): - with world.scratch_context() as ctx: - # Set current position - world.set_positions(ctx, robot_id, current_joints) - - # Get current pose - current_pose = world.get_ee_pose(ctx, robot_id) - - # Compute error - pos_error, ori_error = compute_pose_error(current_pose, target_pose) - - # Check convergence - if pos_error <= position_tolerance and ori_error <= orientation_tolerance: - return _create_success_result( - joint_positions=current_joints, - position_error=pos_error, - orientation_error=ori_error, - iterations=iteration + 1, - ) - - # Compute twist to reduce error - twist = compute_error_twist(current_pose, target_pose, gain=0.5) - - # Get Jacobian - J = world.get_jacobian(ctx, robot_id) - - # Check for singularity - if check_singularity(J, threshold=self._singularity_threshold): - return _create_failure_result( - IKStatus.SINGULARITY, - f"Singularity at iteration {iteration}", - iterations=iteration + 1, - ) - - # Compute joint velocities - J_pinv = damped_pseudoinverse(J, self._damping) - q_dot = J_pinv @ twist - - # Step size (adaptive based on error) - step_size = min(0.5, pos_error + ori_error) - current_joints = current_joints + step_size * q_dot - - # Clip to limits - current_joints = np.clip(current_joints, lower_limits, upper_limits) - - # Compute final error - with world.scratch_context() as ctx: - world.set_positions(ctx, robot_id, current_joints) - final_pose = world.get_ee_pose(ctx, robot_id) - pos_error, ori_error = compute_pose_error(final_pose, target_pose) - - return _create_failure_result( - IKStatus.NO_SOLUTION, - f"Did not converge after {max_iterations} iterations (pos_err={pos_error:.4f}, ori_err={ori_error:.4f})", - iterations=max_iterations, + """Iterative Jacobian-based IK. Delegates to JacobianIK.""" + return self._jacobian_ik.solve_iterative( + world=world, + robot_id=robot_id, + target_pose=target_pose, + seed=seed, + max_iterations=max_iterations, + position_tolerance=position_tolerance, + orientation_tolerance=orientation_tolerance, ) def solve_differential( @@ -314,58 +302,14 @@ def solve_differential( twist: NDArray[np.float64], dt: float, ) -> NDArray[np.float64] | None: - """Single Jacobian step for velocity control. Returns None if near singularity.""" - with world.scratch_context() as ctx: - world.set_positions(ctx, robot_id, current_joints) - J = world.get_jacobian(ctx, robot_id) - - # Check for singularity - if check_singularity(J, threshold=self._singularity_threshold): - logger.warning("Near singularity in differential IK") - return None - - # Compute damped pseudoinverse - J_pinv = damped_pseudoinverse(J, self._damping) - - # Compute joint velocities - q_dot = J_pinv @ twist - - # Apply velocity limits if available - config = world.get_robot_config(robot_id) - if config.velocity_limits is not None: - max_ratio = np.max(np.abs(q_dot) / np.array(config.velocity_limits)) - if max_ratio > 1.0: - q_dot = q_dot / max_ratio - - return q_dot - - def solve_differential_position_only( - self, - world: WorldSpec, - robot_id: str, - current_joints: NDArray[np.float64], - linear_velocity: NDArray[np.float64], - ) -> NDArray[np.float64] | None: - """Position-only differential IK using linear Jacobian. Returns None if singular.""" - with world.scratch_context() as ctx: - world.set_positions(ctx, robot_id, current_joints) - J = world.get_jacobian(ctx, robot_id) - - # Extract linear part (first 3 rows) - J_linear = J[:3, :] - - # Check for singularity - JJT = J_linear @ J_linear.T - manipulability = np.sqrt(max(0, np.linalg.det(JJT))) - if manipulability < self._singularity_threshold: - return None - - # Compute damped pseudoinverse - I = np.eye(3) - J_pinv = J_linear.T @ np.linalg.inv(JJT + self._damping**2 * I) - - # Compute joint velocities - return J_pinv @ linear_velocity + """Single Jacobian step for velocity control. Delegates to JacobianIK.""" + return self._jacobian_ik.solve_differential( + world=world, + robot_id=robot_id, + current_joints=current_joints, + twist=twist, + dt=dt, + ) # ============= Result Helpers ============= diff --git a/dimos/manipulation/planning/kinematics/jacobian_ik.py b/dimos/manipulation/planning/kinematics/jacobian_ik.py new file mode 100644 index 0000000000..1ab26c9fd5 --- /dev/null +++ b/dimos/manipulation/planning/kinematics/jacobian_ik.py @@ -0,0 +1,387 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Backend-agnostic Jacobian-based inverse kinematics. + +JacobianIK provides iterative and differential IK methods that work with any +WorldSpec implementation. It only uses the standard WorldSpec interface methods +(get_jacobian, get_ee_pose, get_joint_limits) and doesn't depend on any specific +physics backend. + +For full nonlinear optimization IK with Drake, use DrakeOptimizationIK. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +from dimos.manipulation.planning.spec import IKResult, IKStatus, WorldSpec +from dimos.manipulation.planning.utils.kinematics_utils import ( + check_singularity, + compute_error_twist, + compute_pose_error, + damped_pseudoinverse, +) +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from numpy.typing import NDArray + +logger = setup_logger() + + +class JacobianIK: + """Backend-agnostic Jacobian-based IK solver. + + This class provides iterative and differential IK methods using only + the standard WorldSpec interface. It works with any physics backend + (Drake, MuJoCo, PyBullet, etc.). + + Methods: + - solve_iterative(): Iterative Jacobian-based IK until convergence + - solve_differential(): Single Jacobian step for velocity control + - solve_differential_position_only(): Position-only differential IK + - solve(): Wrapper for solve_iterative with multiple random restarts + + Example: + ik = JacobianIK(damping=0.01) + result = ik.solve_iterative( + world, robot_id, + target_pose=target, + seed=current_joints, + ) + if result.is_success(): + print(f"Solution: {result.joint_positions}") + """ + + def __init__( + self, + damping: float = 0.05, + max_iterations: int = 200, + singularity_threshold: float = 1e-6, + ): + """Create Jacobian IK solver. + + Args: + damping: Damping factor for pseudoinverse (higher = more stable near singularities) + max_iterations: Default maximum iterations for iterative IK + singularity_threshold: Manipulability threshold for singularity detection + """ + self._damping = damping + self._max_iterations = max_iterations + self._singularity_threshold = singularity_threshold + + def solve( + self, + world: WorldSpec, + robot_id: str, + target_pose: NDArray[np.float64], + seed: NDArray[np.float64] | None = None, + position_tolerance: float = 0.001, + orientation_tolerance: float = 0.01, + check_collision: bool = True, + max_attempts: int = 10, + ) -> IKResult: + """Solve IK with multiple random restarts. + + Tries iterative IK from multiple starting configurations to find + a collision-free solution. + + Args: + world: World for FK/collision checking + robot_id: Robot to solve IK for + target_pose: Target end-effector pose (4x4 transform) + seed: Initial guess (uses current position if None) + position_tolerance: Required position accuracy (meters) + orientation_tolerance: Required orientation accuracy (radians) + check_collision: Whether to check collision of solution + max_attempts: Maximum random restart attempts + + Returns: + IKResult with solution or failure status + """ + if not world.is_finalized: + return _create_failure_result(IKStatus.NO_SOLUTION, "World must be finalized before IK") + + lower_limits, upper_limits = world.get_joint_limits(robot_id) + + # Get seed from current state if not provided + if seed is None: + with world.scratch_context() as ctx: + seed = world.get_positions(ctx, robot_id) + + best_result: IKResult | None = None + best_error = float("inf") + + for attempt in range(max_attempts): + # Generate seed + if attempt == 0: + current_seed = seed + else: + # Random seed within joint limits + current_seed = np.random.uniform(lower_limits, upper_limits) + + # Solve iterative IK + result = self.solve_iterative( + world=world, + robot_id=robot_id, + target_pose=target_pose, + seed=current_seed, + max_iterations=self._max_iterations, + position_tolerance=position_tolerance, + orientation_tolerance=orientation_tolerance, + ) + + if result.is_success() and result.joint_positions is not None: + # Check collision if requested + if check_collision: + if not world.check_config_collision_free(robot_id, result.joint_positions): + continue # Try another seed + + # Check error + total_error = result.position_error + result.orientation_error + if total_error < best_error: + best_error = total_error + best_result = result + + # If error is within tolerance, we're done + if ( + result.position_error <= position_tolerance + and result.orientation_error <= orientation_tolerance + ): + return result + + if best_result is not None: + return best_result + + return _create_failure_result( + IKStatus.NO_SOLUTION, + f"IK failed after {max_attempts} attempts", + ) + + def solve_iterative( + self, + world: WorldSpec, + robot_id: str, + target_pose: NDArray[np.float64], + seed: NDArray[np.float64], + max_iterations: int = 100, + position_tolerance: float = 0.001, + orientation_tolerance: float = 0.01, + ) -> IKResult: + """Iterative Jacobian-based IK until convergence. + + Uses the damped pseudoinverse method with adaptive step size. + Converges when both position and orientation errors are within tolerance. + + Args: + world: World for FK/Jacobian computation + robot_id: Robot to solve IK for + target_pose: Target end-effector pose (4x4 transform) + seed: Initial joint configuration + max_iterations: Maximum iterations before giving up + position_tolerance: Required position accuracy (meters) + orientation_tolerance: Required orientation accuracy (radians) + + Returns: + IKResult with solution or failure status + """ + max_iterations = max_iterations or self._max_iterations + current_joints = seed.copy() + + lower_limits, upper_limits = world.get_joint_limits(robot_id) + + for iteration in range(max_iterations): + with world.scratch_context() as ctx: + # Set current position + world.set_positions(ctx, robot_id, current_joints) + + # Get current pose + current_pose = world.get_ee_pose(ctx, robot_id) + + # Compute error + pos_error, ori_error = compute_pose_error(current_pose, target_pose) + + # Check convergence + if pos_error <= position_tolerance and ori_error <= orientation_tolerance: + return _create_success_result( + joint_positions=current_joints, + position_error=pos_error, + orientation_error=ori_error, + iterations=iteration + 1, + ) + + # Compute twist to reduce error + twist = compute_error_twist(current_pose, target_pose, gain=0.5) + + # Get Jacobian + J = world.get_jacobian(ctx, robot_id) + + # Adaptive damping near singularities + if check_singularity(J, threshold=self._singularity_threshold): + # Increase damping near singularity instead of failing + effective_damping = self._damping * 10.0 + else: + effective_damping = self._damping + + # Compute joint velocities + J_pinv = damped_pseudoinverse(J, effective_damping) + q_dot = J_pinv @ twist + + # Clamp maximum joint change per iteration (like reference implementations) + max_delta = 0.1 # radians per iteration + max_change = np.max(np.abs(q_dot)) + if max_change > max_delta: + q_dot = q_dot * (max_delta / max_change) + + current_joints = current_joints + q_dot + + # Clip to limits + current_joints = np.clip(current_joints, lower_limits, upper_limits) + + # Compute final error + with world.scratch_context() as ctx: + world.set_positions(ctx, robot_id, current_joints) + final_pose = world.get_ee_pose(ctx, robot_id) + pos_error, ori_error = compute_pose_error(final_pose, target_pose) + + return _create_failure_result( + IKStatus.NO_SOLUTION, + f"Did not converge after {max_iterations} iterations (pos_err={pos_error:.4f}, ori_err={ori_error:.4f})", + iterations=max_iterations, + ) + + def solve_differential( + self, + world: WorldSpec, + robot_id: str, + current_joints: NDArray[np.float64], + twist: NDArray[np.float64], + dt: float, + ) -> NDArray[np.float64] | None: + """Single Jacobian step for velocity control. + + Computes joint velocities from desired end-effector twist using + the damped pseudoinverse method. Returns None if near singularity. + + Args: + world: World for Jacobian computation + robot_id: Robot to compute for + current_joints: Current joint configuration + twist: Desired 6D end-effector twist [vx, vy, vz, wx, wy, wz] + dt: Time step (not used, but kept for interface compatibility) + + Returns: + Joint velocities, or None if near singularity + """ + with world.scratch_context() as ctx: + world.set_positions(ctx, robot_id, current_joints) + J = world.get_jacobian(ctx, robot_id) + + # Check for singularity + if check_singularity(J, threshold=self._singularity_threshold): + logger.warning("Near singularity in differential IK") + return None + + # Compute damped pseudoinverse + J_pinv = damped_pseudoinverse(J, self._damping) + + # Compute joint velocities + q_dot = J_pinv @ twist + + # Apply velocity limits if available + config = world.get_robot_config(robot_id) + if config.velocity_limits is not None: + max_ratio = np.max(np.abs(q_dot) / np.array(config.velocity_limits)) + if max_ratio > 1.0: + q_dot = q_dot / max_ratio + + return q_dot + + def solve_differential_position_only( + self, + world: WorldSpec, + robot_id: str, + current_joints: NDArray[np.float64], + linear_velocity: NDArray[np.float64], + ) -> NDArray[np.float64] | None: + """Position-only differential IK using linear Jacobian. + + Computes joint velocities from desired linear velocity, ignoring + orientation. Returns None if near singularity. + + Args: + world: World for Jacobian computation + robot_id: Robot to compute for + current_joints: Current joint configuration + linear_velocity: Desired linear velocity [vx, vy, vz] + + Returns: + Joint velocities, or None if singular + """ + with world.scratch_context() as ctx: + world.set_positions(ctx, robot_id, current_joints) + J = world.get_jacobian(ctx, robot_id) + + # Extract linear part (first 3 rows) + J_linear = J[:3, :] + + # Check for singularity + JJT = J_linear @ J_linear.T + manipulability = np.sqrt(max(0, np.linalg.det(JJT))) + if manipulability < self._singularity_threshold: + return None + + # Compute damped pseudoinverse + I = np.eye(3) + J_pinv = J_linear.T @ np.linalg.inv(JJT + self._damping**2 * I) + + # Compute joint velocities + return J_pinv @ linear_velocity + + +# ============= Result Helpers ============= + + +def _create_success_result( + joint_positions: NDArray[np.float64], + position_error: float, + orientation_error: float, + iterations: int, +) -> IKResult: + """Create a successful IK result.""" + return IKResult( + status=IKStatus.SUCCESS, + joint_positions=joint_positions, + position_error=position_error, + orientation_error=orientation_error, + iterations=iterations, + message="IK solution found", + ) + + +def _create_failure_result( + status: IKStatus, + message: str, + iterations: int = 0, +) -> IKResult: + """Create a failed IK result.""" + return IKResult( + status=status, + joint_positions=None, + iterations=iterations, + message=message, + ) diff --git a/dimos/manipulation/planning/monitor/__init__.py b/dimos/manipulation/planning/monitor/__init__.py index 060579b117..c280bd4d56 100644 --- a/dimos/manipulation/planning/monitor/__init__.py +++ b/dimos/manipulation/planning/monitor/__init__.py @@ -53,11 +53,10 @@ from dimos.manipulation.planning.monitor.world_state_monitor import WorldStateMonitor # Re-export message types from spec for convenience -from dimos.manipulation.planning.spec import CollisionObjectMessage, Detection3D +from dimos.manipulation.planning.spec import CollisionObjectMessage __all__ = [ "CollisionObjectMessage", - "Detection3D", "WorldMonitor", "WorldObstacleMonitor", "WorldStateMonitor", diff --git a/dimos/manipulation/planning/monitor/world_monitor.py b/dimos/manipulation/planning/monitor/world_monitor.py index f1be2f338d..18de1f352c 100644 --- a/dimos/manipulation/planning/monitor/world_monitor.py +++ b/dimos/manipulation/planning/monitor/world_monitor.py @@ -34,12 +34,12 @@ from dimos.manipulation.planning.spec import ( CollisionObjectMessage, - Detection3D, Obstacle, RobotModelConfig, WorldSpec, ) from dimos.msgs.sensor_msgs import JointState + from dimos.msgs.vision_msgs import Detection3D logger = setup_logger() @@ -200,7 +200,7 @@ def on_collision_object(self, msg: CollisionObjectMessage) -> None: self._obstacle_monitor.on_collision_object(msg) def on_detections(self, detections: list[Detection3D]) -> None: - """Handle perception detections.""" + """Handle perception detections (Detection3D from dimos.msgs.vision_msgs).""" if self._obstacle_monitor is not None: self._obstacle_monitor.on_detections(detections) @@ -324,26 +324,26 @@ def is_finalized(self) -> bool: # ============= Visualization ============= - def get_meshcat_url(self) -> str | None: - """Get Meshcat URL or None if not enabled.""" - if hasattr(self._world, "get_meshcat_url"): - url = self._world.get_meshcat_url() + def get_visualization_url(self) -> str | None: + """Get visualization URL or None if not enabled.""" + if hasattr(self._world, "get_visualization_url"): + url = self._world.get_visualization_url() return str(url) if url else None return None def publish_visualization(self) -> None: """Force publish current state to visualization.""" - if hasattr(self._world, "publish_to_meshcat"): - self._world.publish_to_meshcat() + if hasattr(self._world, "publish_visualization"): + self._world.publish_visualization() def start_visualization_thread(self, rate_hz: float = 10.0) -> None: - """Start background thread for Meshcat updates at given rate.""" + """Start background thread for visualization updates at given rate.""" if self._viz_thread is not None and self._viz_thread.is_alive(): logger.warning("Visualization thread already running") return - if not hasattr(self._world, "publish_to_meshcat"): - logger.warning("World does not support Meshcat visualization") + if not hasattr(self._world, "publish_visualization"): + logger.warning("World does not support visualization") return self._viz_rate_hz = rate_hz @@ -375,8 +375,8 @@ def _visualization_loop(self) -> None: period = 1.0 / self._viz_rate_hz while not self._viz_stop_event.is_set(): try: - if hasattr(self._world, "publish_to_meshcat"): - self._world.publish_to_meshcat() + if hasattr(self._world, "publish_visualization"): + self._world.publish_visualization() except Exception as e: logger.debug(f"Visualization publish failed: {e}") time.sleep(period) diff --git a/dimos/manipulation/planning/monitor/world_obstacle_monitor.py b/dimos/manipulation/planning/monitor/world_obstacle_monitor.py index da5a01304c..725095090e 100644 --- a/dimos/manipulation/planning/monitor/world_obstacle_monitor.py +++ b/dimos/manipulation/planning/monitor/world_obstacle_monitor.py @@ -31,11 +31,11 @@ from dimos.manipulation.planning.spec import ( CollisionObjectMessage, - Detection3D, Obstacle, ObstacleType, ) from dimos.utils.logging_config import setup_logger +from dimos.utils.transform_utils import pose_to_matrix if TYPE_CHECKING: from collections.abc import Callable @@ -45,6 +45,7 @@ from numpy.typing import NDArray from dimos.manipulation.planning.spec import WorldSpec + from dimos.msgs.vision_msgs import Detection3D logger = setup_logger() @@ -54,7 +55,7 @@ class WorldObstacleMonitor: This class handles updates from: - Explicit collision objects (CollisionObjectMessage) - - Perception detections (Detection3D) + - Perception detections (Detection3D from dimos.msgs.vision_msgs) ## Thread Safety @@ -224,7 +225,7 @@ def on_detections(self, detections: list[Detection3D]) -> None: - Removes obstacles for detections that are no longer present Args: - detections: List of 3D detections + detections: List of Detection3D messages from dimos.msgs.vision_msgs """ if not self._running: return @@ -237,10 +238,12 @@ def on_detections(self, detections: list[Detection3D]) -> None: det_id = detection.id seen_ids.add(det_id) + pose = self._detection3d_to_pose(detection) + if det_id in self._perception_objects: # Update existing obstacle obstacle_id = self._perception_objects[det_id] - self._world.update_obstacle_pose(obstacle_id, detection.pose) + self._world.update_obstacle_pose(obstacle_id, pose) self._perception_timestamps[det_id] = current_time else: # Add new obstacle @@ -261,13 +264,19 @@ def on_detections(self, detections: list[Detection3D]) -> None: # Remove stale detections self._cleanup_stale_detections(current_time, seen_ids) + def _detection3d_to_pose(self, detection: Detection3D) -> NDArray[np.float64]: + """Convert Detection3D bbox.center to 4x4 transform.""" + return pose_to_matrix(detection.bbox.center) + def _detection_to_obstacle(self, detection: Detection3D) -> Obstacle: - """Convert detection to obstacle.""" + """Convert Detection3D to Obstacle.""" + pose = self._detection3d_to_pose(detection) + size = detection.bbox.size return Obstacle( name=f"detection_{detection.id}", obstacle_type=ObstacleType.BOX, - pose=detection.pose, - dimensions=detection.dimensions, + pose=pose, + dimensions=(size.x, size.y, size.z), color=(0.2, 0.8, 0.2, 0.6), # Green for perception objects ) diff --git a/dimos/manipulation/planning/planners/__init__.py b/dimos/manipulation/planning/planners/__init__.py index 4323896e4e..8fb8ae042b 100644 --- a/dimos/manipulation/planning/planners/__init__.py +++ b/dimos/manipulation/planning/planners/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2025 Dimensional Inc. +# 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. @@ -17,10 +17,12 @@ Contains motion planning implementations that use WorldSpec. +All planners are backend-agnostic - they only use WorldSpec methods and +work with any physics backend (Drake, MuJoCo, PyBullet, etc.). + ## Implementations -- DrakePlanner: RRT-Connect planner -- DrakeRRTStarPlanner: RRT* planner (asymptotically optimal) +- RRTConnectPlanner: Bi-directional RRT-Connect planner (fast, reliable) ## Usage @@ -34,9 +36,6 @@ ``` """ -from dimos.manipulation.planning.planners.drake_planner import ( - DrakePlanner, - DrakeRRTStarPlanner, -) +from dimos.manipulation.planning.planners.rrt_planner import RRTConnectPlanner -__all__ = ["DrakePlanner", "DrakeRRTStarPlanner"] +__all__ = ["RRTConnectPlanner"] diff --git a/dimos/manipulation/planning/planners/drake_planner.py b/dimos/manipulation/planning/planners/rrt_planner.py similarity index 53% rename from dimos/manipulation/planning/planners/drake_planner.py rename to dimos/manipulation/planning/planners/rrt_planner.py index 95289f379a..9ad2cad57f 100644 --- a/dimos/manipulation/planning/planners/drake_planner.py +++ b/dimos/manipulation/planning/planners/rrt_planner.py @@ -12,7 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""RRT-Connect and RRT* motion planners implementing PlannerSpec.""" +"""RRT-Connect and RRT* motion planners implementing PlannerSpec. + +These planners are backend-agnostic - they only use WorldSpec methods and can work +with any physics backend (Drake, MuJoCo, PyBullet, etc.). +""" from __future__ import annotations @@ -23,10 +27,7 @@ import numpy as np from dimos.manipulation.planning.spec import PlanningResult, PlanningStatus, WorldSpec -from dimos.manipulation.planning.utils.path_utils import ( - compute_path_length, - interpolate_segment, -) +from dimos.manipulation.planning.utils.path_utils import compute_path_length from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: @@ -35,7 +36,7 @@ logger = setup_logger() -@dataclass +@dataclass(eq=False) class TreeNode: """Node in RRT tree with optional cost tracking (for RRT*).""" @@ -54,8 +55,12 @@ def path_to_root(self) -> list[NDArray[np.float64]]: return list(reversed(path)) -class DrakePlanner: - """Bi-directional RRT-Connect planner.""" +class RRTConnectPlanner: + """Bi-directional RRT-Connect planner. + + This planner is backend-agnostic - it only uses WorldSpec methods for + collision checking and can work with any physics backend. + """ def __init__( self, @@ -149,23 +154,19 @@ def _validate_inputs( f"Robot '{robot_id}' not found", ) - # Check start validity - with world.scratch_context() as ctx: - world.set_positions(ctx, robot_id, q_start) - if not world.is_collision_free(ctx, robot_id): - return _create_failure_result( - PlanningStatus.COLLISION_AT_START, - "Start configuration is in collision", - ) + # Check start validity using context-free method + if not world.check_config_collision_free(robot_id, q_start): + return _create_failure_result( + PlanningStatus.COLLISION_AT_START, + "Start configuration is in collision", + ) - # Check goal validity - with world.scratch_context() as ctx: - world.set_positions(ctx, robot_id, q_goal) - if not world.is_collision_free(ctx, robot_id): - return _create_failure_result( - PlanningStatus.COLLISION_AT_GOAL, - "Goal configuration is in collision", - ) + # Check goal validity using context-free method + if not world.check_config_collision_free(robot_id, q_goal): + return _create_failure_result( + PlanningStatus.COLLISION_AT_GOAL, + "Goal configuration is in collision", + ) # Check limits lower, upper = world.get_joint_limits(robot_id) @@ -204,8 +205,10 @@ def _extend_tree( else: new_config = nearest.config + step_size * (diff / dist) - # Check validity of edge - if self._is_edge_valid(world, robot_id, nearest.config, new_config): + # Check validity of edge using context-free method + if world.check_edge_collision_free( + robot_id, nearest.config, new_config, self._collision_step_size + ): new_node = TreeNode(config=new_config, parent=nearest) nearest.children.append(new_node) tree.append(new_node) @@ -233,25 +236,6 @@ def _connect_tree( if float(np.linalg.norm(result.config - target)) < self._goal_tolerance: return result - def _is_edge_valid( - self, - world: WorldSpec, - robot_id: str, - q_start: NDArray[np.float64], - q_end: NDArray[np.float64], - ) -> bool: - """Check if edge between two configurations is collision-free.""" - # Interpolate and check each point - segment = interpolate_segment(q_start, q_end, self._collision_step_size) - - with world.scratch_context() as ctx: - for q in segment: - world.set_positions(ctx, robot_id, q) - if not world.is_collision_free(ctx, robot_id): - return False - - return True - def _extract_path( self, start_node: TreeNode, @@ -291,184 +275,16 @@ def _simplify_path( i = np.random.randint(0, len(simplified) - 2) j = np.random.randint(i + 2, len(simplified)) - # Check if direct connection is valid - if self._is_edge_valid(world, robot_id, simplified[i], simplified[j]): + # Check if direct connection is valid using context-free method + if world.check_edge_collision_free( + robot_id, simplified[i], simplified[j], self._collision_step_size + ): # Remove intermediate waypoints simplified = simplified[: i + 1] + simplified[j:] return simplified -class DrakeRRTStarPlanner: - """RRT* (Optimal RRT) planner implementing PlannerSpec. - - Like RRT but optimizes path cost through rewiring. - Produces asymptotically optimal paths. - """ - - def __init__( - self, - step_size: float = 0.1, - goal_tolerance: float = 0.1, - rewire_radius: float = 0.5, - collision_step_size: float = 0.02, - ): - """Create RRT* planner. - - Args: - step_size: Extension step size - goal_tolerance: Distance to goal to consider success - rewire_radius: Radius for rewiring neighbors - collision_step_size: Step size for collision checking - """ - self._step_size = step_size - self._goal_tolerance = goal_tolerance - self._rewire_radius = rewire_radius - self._collision_step_size = collision_step_size - - def plan_joint_path( - self, - world: WorldSpec, - robot_id: str, - q_start: NDArray[np.float64], - q_goal: NDArray[np.float64], - timeout: float = 10.0, - max_iterations: int = 5000, - ) -> PlanningResult: - """Plan optimal path using RRT* with rewiring.""" - start_time = time.time() - lower, upper = world.get_joint_limits(robot_id) - - # Validate start/goal - with world.scratch_context() as ctx: - world.set_positions(ctx, robot_id, q_start) - if not world.is_collision_free(ctx, robot_id): - return _create_failure_result( - PlanningStatus.COLLISION_AT_START, "Start in collision" - ) - world.set_positions(ctx, robot_id, q_goal) - if not world.is_collision_free(ctx, robot_id): - return _create_failure_result(PlanningStatus.COLLISION_AT_GOAL, "Goal in collision") - - nodes = [TreeNode(config=q_start.copy(), cost=0.0)] - goal_node: TreeNode | None = None - best_cost = float("inf") - iterations_run = 0 - - for _ in range(max_iterations): - iterations_run += 1 - if time.time() - start_time > timeout: - break - - # Sample and find nearest - sample = np.random.uniform(lower, upper) - nearest = min(nodes, key=lambda n: float(np.linalg.norm(n.config - sample))) - - # Extend toward sample - diff = sample - nearest.config - dist = float(np.linalg.norm(diff)) - new_config = ( - nearest.config + min(dist, self._step_size) * (diff / dist) - if dist > 0 - else sample.copy() - ) - - if not self._is_edge_valid(world, robot_id, nearest.config, new_config): - continue - - # Find best parent among neighbors - neighbors = [ - n - for n in nodes - if float(np.linalg.norm(n.config - new_config)) < self._rewire_radius - ] - best_parent, best_cost_to_new = ( - nearest, - nearest.cost + float(np.linalg.norm(new_config - nearest.config)), - ) - for n in neighbors: - cost = n.cost + float(np.linalg.norm(new_config - n.config)) - if cost < best_cost_to_new and self._is_edge_valid( - world, robot_id, n.config, new_config - ): - best_parent, best_cost_to_new = n, cost - - # Add node and rewire neighbors - new_node = TreeNode(config=new_config, parent=best_parent, cost=best_cost_to_new) - best_parent.children.append(new_node) - nodes.append(new_node) - self._rewire_neighbors(world, robot_id, new_node, neighbors) - - # Check goal - if ( - float(np.linalg.norm(new_node.config - q_goal)) < self._goal_tolerance - and new_node.cost < best_cost - ): - goal_node, best_cost = new_node, new_node.cost - - if goal_node is not None: - path = [*goal_node.path_to_root(), q_goal] - return _create_success_result(path, time.time() - start_time, iterations_run) - - return _create_failure_result( - PlanningStatus.NO_SOLUTION, - f"No path after {max_iterations} iterations", - time.time() - start_time, - max_iterations, - ) - - def get_name(self) -> str: - return "RRTStar" - - def _is_edge_valid( - self, - world: WorldSpec, - robot_id: str, - q_start: NDArray[np.float64], - q_end: NDArray[np.float64], - ) -> bool: - """Check if edge is collision-free.""" - segment = interpolate_segment(q_start, q_end, self._collision_step_size) - - with world.scratch_context() as ctx: - for q in segment: - world.set_positions(ctx, robot_id, q) - if not world.is_collision_free(ctx, robot_id): - return False - - return True - - def _rewire_neighbors( - self, - world: WorldSpec, - robot_id: str, - new_node: TreeNode, - neighbors: list[TreeNode], - ) -> None: - """Rewire neighbors through new node if it provides a shorter path.""" - for neighbor in neighbors: - if neighbor == new_node.parent: - continue - potential_cost = new_node.cost + float( - np.linalg.norm(neighbor.config - new_node.config) - ) - if potential_cost < neighbor.cost and self._is_edge_valid( - world, robot_id, new_node.config, neighbor.config - ): - if neighbor.parent is not None: - neighbor.parent.children.remove(neighbor) - neighbor.parent = new_node - neighbor.cost = potential_cost - new_node.children.append(neighbor) - self._update_costs(neighbor) - - def _update_costs(self, node: TreeNode) -> None: - """Recursively update costs after rewiring.""" - for child in node.children: - child.cost = node.cost + float(np.linalg.norm(child.config - node.config)) - self._update_costs(child) - - # ============= Result Helpers ============= diff --git a/dimos/manipulation/planning/spec/__init__.py b/dimos/manipulation/planning/spec/__init__.py index 5d079c6d87..345f408e03 100644 --- a/dimos/manipulation/planning/spec/__init__.py +++ b/dimos/manipulation/planning/spec/__init__.py @@ -12,72 +12,32 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Manipulation Planning Specifications. - -Protocols: - - WorldSpec: Core backend owning physics/collision (implemented by DrakeWorld) - - KinematicsSpec: Stateless IK operations (implemented by DrakeKinematics) - - PlannerSpec: Joint-space path planning (implemented by DrakePlanner) - - VizSpec: Visualization backend - -Data Classes: - - RobotModelConfig: Robot configuration for adding to world - - Obstacle: Obstacle specification - - IKResult: Result of IK solve - - PlanningResult: Result of path planning - -Usage: - All code should use Protocol types (not concrete classes): - - ```python - from dimos.manipulation.planning.spec import WorldSpec, KinematicsSpec - - def plan_motion(world: WorldSpec, kinematics: KinematicsSpec, ...): - pass - ``` - - Use factory functions to create concrete instances: - - ```python - from dimos.manipulation.planning.factory import create_world - world = create_world(backend="drake") - ``` -""" +"""Manipulation Planning Specifications.""" from dimos.manipulation.planning.spec.config import RobotModelConfig from dimos.manipulation.planning.spec.enums import IKStatus, ObstacleType, PlanningStatus from dimos.manipulation.planning.spec.protocols import ( KinematicsSpec, PlannerSpec, - VizSpec, WorldSpec, ) from dimos.manipulation.planning.spec.types import ( CollisionObjectMessage, - Detection3D, IKResult, Obstacle, - PerceptionDetection, PlanningResult, ) __all__ = [ "CollisionObjectMessage", - "Detection3D", # Backwards compat alias for PerceptionDetection "IKResult", "IKStatus", "KinematicsSpec", - # Types "Obstacle", - # Enums "ObstacleType", - "PerceptionDetection", "PlannerSpec", "PlanningResult", "PlanningStatus", - # Config "RobotModelConfig", - "VizSpec", - # Protocols "WorldSpec", ] diff --git a/dimos/manipulation/planning/spec/protocols.py b/dimos/manipulation/planning/spec/protocols.py index 6108ae60c5..fac97f2c49 100644 --- a/dimos/manipulation/planning/spec/protocols.py +++ b/dimos/manipulation/planning/spec/protocols.py @@ -125,6 +125,21 @@ def get_min_distance(self, ctx: Any, robot_id: str) -> float: """Get minimum distance to obstacles (negative if collision).""" ... + # Collision Checking (context-free, for planning) + def check_config_collision_free(self, robot_id: str, q: NDArray[np.float64]) -> bool: + """Check if a configuration is collision-free (manages context internally).""" + ... + + def check_edge_collision_free( + self, + robot_id: str, + q_start: NDArray[np.float64], + q_end: NDArray[np.float64], + step_size: float = 0.05, + ) -> bool: + """Check if the entire edge between two configurations is collision-free.""" + ... + # Forward Kinematics (require context) def get_ee_pose(self, ctx: Any, robot_id: str) -> NDArray[np.float64]: """Get end-effector pose (4x4 transform).""" @@ -135,12 +150,12 @@ def get_jacobian(self, ctx: Any, robot_id: str) -> NDArray[np.float64]: ... # Visualization (optional) - def get_meshcat_url(self) -> str | None: - """Get Meshcat visualization URL if enabled.""" + def get_visualization_url(self) -> str | None: + """Get visualization URL if enabled.""" ... - def publish_to_meshcat(self, ctx: Any | None = None) -> None: - """Publish current state to Meshcat visualization.""" + def publish_visualization(self, ctx: Any | None = None) -> None: + """Publish current state to visualization.""" ... def animate_path( @@ -162,7 +177,8 @@ class KinematicsSpec(Protocol): - solve_differential(): Single Jacobian step for velocity control Implementations: - - DrakeKinematics: Uses Drake's InverseKinematics + SNOPT/IPOPT + - JacobianIK: Backend-agnostic iterative/differential IK + - DrakeOptimizationIK: Uses Drake's InverseKinematics + SNOPT/IPOPT """ def solve( @@ -210,10 +226,11 @@ class PlannerSpec(Protocol): Planners find collision-free paths from start to goal configurations. They use WorldSpec for collision checking and are stateless. + All planners are backend-agnostic - they only use WorldSpec methods. Implementations: - - DrakePlanner: RRT-Connect planner - - DrakeRRTStarPlanner: RRT* planner (asymptotically optimal) + - RRTConnectPlanner: Bi-directional RRT-Connect planner + - RRTStarPlanner: RRT* planner (asymptotically optimal) """ def plan_joint_path( @@ -230,32 +247,3 @@ def plan_joint_path( def get_name(self) -> str: """Get planner name.""" ... - - -@runtime_checkable -class VizSpec(Protocol): - """Protocol for visualization backend. - - Note: For Drake, visualization is typically integrated into DrakeWorld - via enable_viz=True. This protocol is for advanced use cases. - """ - - def set_robot_state(self, robot_id: str, positions: NDArray[np.float64]) -> None: - """Update robot visualization state.""" - ... - - def add_obstacle(self, obstacle: Obstacle) -> str: - """Add obstacle to visualization.""" - ... - - def remove_obstacle(self, obstacle_id: str) -> None: - """Remove obstacle from visualization.""" - ... - - def get_url(self) -> str | None: - """Get visualization URL (e.g., Meshcat URL).""" - ... - - def publish(self) -> None: - """Force publish current state to visualization.""" - ... diff --git a/dimos/manipulation/planning/spec/types.py b/dimos/manipulation/planning/spec/types.py index c698f37b87..c776b57299 100644 --- a/dimos/manipulation/planning/spec/types.py +++ b/dimos/manipulation/planning/spec/types.py @@ -126,29 +126,3 @@ class CollisionObjectMessage: pose: NDArray[np.float64] | None = None dimensions: tuple[float, ...] | None = None color: tuple[float, float, float, float] = (0.8, 0.2, 0.2, 0.8) - - -@dataclass -class PerceptionDetection: - """3D detection for planning obstacle updates. - - Internal type for the planning system using numpy transforms. - For LCM messages, use dimos_lcm.vision_msgs.Detection3D instead. - - Attributes: - id: Unique detection ID - label: Object class label - pose: 4x4 homogeneous transform - dimensions: (width, height, depth) - confidence: Detection confidence (0-1) - """ - - id: str - label: str - pose: NDArray[np.float64] - dimensions: tuple[float, float, float] - confidence: float = 1.0 - - -# Backwards compatibility alias - will be removed in future version -Detection3D = PerceptionDetection diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index 1bac856671..cf318c4cf1 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -688,6 +688,47 @@ def get_min_distance(self, ctx: Context, robot_id: str) -> float: return float(min(pair.distance for pair in signed_distance_pairs)) + # ============= Collision Checking (context-free, for planning) ============= + + def check_config_collision_free(self, robot_id: str, q: NDArray[np.float64]) -> bool: + """Check if a configuration is collision-free (manages context internally). + + This is a convenience method for planners that don't need to manage contexts. + """ + with self.scratch_context() as ctx: + self.set_positions(ctx, robot_id, q) + return self.is_collision_free(ctx, robot_id) + + def check_edge_collision_free( + self, + robot_id: str, + q_start: NDArray[np.float64], + q_end: NDArray[np.float64], + step_size: float = 0.05, + ) -> bool: + """Check if the entire edge between two configurations is collision-free. + + Interpolates between q_start and q_end at the given step_size and checks + each configuration for collisions. This is more efficient than checking + each configuration separately as it uses a single scratch context. + """ + # Compute number of steps needed + dist = float(np.linalg.norm(q_end - q_start)) + if dist < 1e-8: + return self.check_config_collision_free(robot_id, q_start) + + n_steps = max(2, int(np.ceil(dist / step_size)) + 1) + + with self.scratch_context() as ctx: + for i in range(n_steps): + t = i / (n_steps - 1) + q = q_start + t * (q_end - q_start) + self.set_positions(ctx, robot_id, q) + if not self.is_collision_free(ctx, robot_id): + return False + + return True + # ============= Forward Kinematics (context-based) ============= def get_ee_pose(self, ctx: Context, robot_id: str) -> NDArray[np.float64]: @@ -764,14 +805,14 @@ def get_jacobian(self, ctx: Context, robot_id: str) -> NDArray[np.float64]: # ============= Visualization ============= - def get_meshcat_url(self) -> str | None: - """Get Meshcat visualization URL.""" + def get_visualization_url(self) -> str | None: + """Get visualization URL if enabled.""" if self._meshcat is not None: return self._meshcat.web_url() return None - def publish_to_meshcat(self, ctx: Context | None = None) -> None: - """Force publish current state to Meshcat. + def publish_visualization(self, ctx: Context | None = None) -> None: + """Publish current state to visualization. Args: ctx: Context to publish. Uses live context if None. @@ -821,7 +862,7 @@ def animate_path( self.set_positions(ctx, rid, pos) # Set animated robot's position self.set_positions(ctx, robot_id, q) - self.publish_to_meshcat(ctx) + self.publish_visualization(ctx) time.sleep(dt) # ============= Direct Access (use with caution) ============= From d4f5ccd294d01fc70a0b42bb24777c4e477a8d1d Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Tue, 20 Jan 2026 22:36:38 -0800 Subject: [PATCH 17/31] updated README --- dimos/manipulation/planning/README.md | 161 +++++++++++++++++--------- 1 file changed, 104 insertions(+), 57 deletions(-) diff --git a/dimos/manipulation/planning/README.md b/dimos/manipulation/planning/README.md index d549d18a71..7ca4e9cd2d 100644 --- a/dimos/manipulation/planning/README.md +++ b/dimos/manipulation/planning/README.md @@ -1,38 +1,57 @@ # Manipulation Planning Stack -Motion planning for robotic manipulators using Drake. Integrates with Control Orchestrator for execution. +Motion planning for robotic manipulators. Backend-agnostic design with Drake implementation. -## Quick Start (3 Terminals) +## Quick Start ```bash -# Terminal 1: Start mock orchestrator +# Terminal 1: Mock orchestrator dimos run orchestrator-mock -# Terminal 2: Start manipulation planner +# Terminal 2: Manipulation planner dimos run xarm7-planner-orchestrator -# Terminal 3: Run IPython client +# Terminal 3: IPython client python -m dimos.manipulation.planning.examples.manipulation_client ``` -Then in IPython: +In IPython: ```python -c.joints() # Get current joints -c.plan([0.1] * 7) # Plan to target -c.preview() # Preview in Meshcat (check c.url()) -c.execute() # Execute via orchestrator +joints() # Get current joints +plan([0.1] * 7) # Plan to target +preview() # Preview in Meshcat (url() for link) +execute() # Execute via orchestrator ``` ## Architecture ``` -ManipulationModule (RPC interface, state machine, multi-robot) - │ -Factory Functions (create_world, create_kinematics, create_planner) - │ -Drake Implementations (DrakeWorld, DrakeKinematics, DrakePlanner) - │ -Control Orchestrator (trajectory execution via RPC) +┌─────────────────────────────────────────────────────────────┐ +│ ManipulationModule │ +│ (RPC interface, state machine, multi-robot) │ +└─────────────────────────────────────────────────────────────┘ + │ +┌─────────────────────────────────────────────────────────────┐ +│ Backend-Agnostic Components │ +│ ┌──────────────────┐ ┌─────────────────────────────┐ │ +│ │ RRTConnectPlanner│ │ JacobianIK │ │ +│ │ (rrt_planner.py) │ │ (iterative & differential) │ │ +│ └──────────────────┘ └─────────────────────────────┘ │ +│ Uses only WorldSpec interface │ +└─────────────────────────────────────────────────────────────┘ + │ +┌─────────────────────────────────────────────────────────────┐ +│ WorldSpec Protocol │ +│ Context management, collision checking, FK, Jacobian │ +└─────────────────────────────────────────────────────────────┘ + │ +┌─────────────────────────────────────────────────────────────┐ +│ Backend-Specific Implementations │ +│ ┌──────────────────┐ ┌─────────────────────────────┐ │ +│ │ DrakeWorld │ │ DrakeOptimizationIK │ │ +│ │ (physics/viz) │ │ (nonlinear IK) │ │ +│ └──────────────────┘ └─────────────────────────────┘ │ +└─────────────────────────────────────────────────────────────┘ ``` ## Using ManipulationModule @@ -48,11 +67,17 @@ config = RobotModelConfig( joint_names=["joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "joint7"], end_effector_link="link7", base_link="link_base", - joint_name_mapping={"arm_joint1": "joint1", ...}, # URDF <-> orchestrator + joint_name_mapping={"arm_joint1": "joint1", ...}, # orchestrator <-> URDF orchestrator_task_name="traj_arm", ) -module = ManipulationModule(robots=[config], planning_timeout=10.0, enable_viz=True) +module = ManipulationModule( + robots=[config], + planning_timeout=10.0, + enable_viz=True, + planner_name="rrt_connect", # Only option + kinematics_name="drake_optimization", # Or "jacobian" +) module.start() module.plan_to_joints([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]) module.execute() # Sends to orchestrator @@ -64,46 +89,74 @@ module.execute() # Sends to orchestrator |-------|-------------| | `name` | Robot identifier | | `urdf_path` | Path to URDF/XACRO file | -| `base_pose` | 4x4 transform matrix for robot base | -| `joint_names` | List of joint names in URDF | -| `end_effector_link` | Name of EE link in URDF | -| `base_link` | Name of base link in URDF | +| `base_pose` | 4x4 transform for robot base | +| `joint_names` | Joint names in URDF | +| `end_effector_link` | EE link name | +| `base_link` | Base link name | | `max_velocity` | Max joint velocity (rad/s) | -| `max_acceleration` | Max joint acceleration (rad/s²) | -| `joint_name_mapping` | Dict mapping orchestrator names to URDF names | -| `orchestrator_task_name` | Task name for trajectory execution RPC | -| `package_paths` | Dict of ROS package paths for mesh resolution | -| `xacro_args` | Dict of xacro arguments (e.g., `{"dof": "7"}`) | +| `max_acceleration` | Max acceleration (rad/s²) | +| `joint_name_mapping` | Orchestrator → URDF name mapping | +| `orchestrator_task_name` | Task name for execution RPC | +| `package_paths` | ROS package paths for meshes | +| `xacro_args` | Xacro arguments (e.g., `{"dof": "7"}`) | + +## Components + +### Planners (Backend-Agnostic) + +| Planner | Description | +|---------|-------------| +| `RRTConnectPlanner` | Bi-directional RRT-Connect (fast, reliable) | + +### IK Solvers -## Available Blueprints +| Solver | Type | Description | +|--------|------|-------------| +| `JacobianIK` | Backend-agnostic | Iterative damped least-squares | +| `DrakeOptimizationIK` | Drake-specific | Full nonlinear optimization | + +### World Backends + +| Backend | Description | +|---------|-------------| +| `DrakeWorld` | Drake physics with Meshcat visualization | + +## Blueprints | Blueprint | Description | |-----------|-------------| -| `xarm6-planner` | XArm 6-DOF planner (standalone) | +| `xarm6_planner_only` | XArm 6-DOF standalone (no orchestrator) | | `xarm7-planner-orchestrator` | XArm 7-DOF with orchestrator | -| `dual-xarm6-planner` | Dual XArm 6-DOF planner | +| `dual-xarm6-planner` | Dual XArm 6-DOF | ## Directory Structure ``` planning/ -├── spec.py # Protocol definitions -├── factory.py # Factory functions -├── world/ # DrakeWorld -├── kinematics/ # DrakeKinematics -├── planners/ # RRT-Connect, RRT* -├── monitor/ # WorldMonitor, state sync -├── trajectory_generator/ +├── spec.py # Protocols (WorldSpec, KinematicsSpec, PlannerSpec) +├── factory.py # create_world, create_kinematics, create_planner +├── world/ +│ └── drake_world.py # DrakeWorld implementation +├── kinematics/ +│ ├── jacobian_ik.py # Backend-agnostic Jacobian IK +│ └── drake_optimization_ik.py # Drake nonlinear IK +├── planners/ +│ └── rrt_planner.py # RRTConnectPlanner +├── monitor/ # WorldMonitor (live state sync) +├── trajectory_generator/ # Time-parameterized trajectories └── examples/ - ├── planning_tester.py # Standalone CLI - └── manipulation_client.py # IPython RPC client + ├── planning_tester.py # Standalone CLI tester + └── manipulation_client.py # IPython RPC client ``` -## Protocols +## Obstacle Types -- **WorldSpec**: Physics/collision backend (DrakeWorld, MuJoCoWorld, PyBulletWorld) -- **KinematicsSpec**: IK solving (DrakeKinematics, TracIK, KDL) -- **PlannerSpec**: Path planning (DrakeRRTConnect, OMPL, cuRobo) +| Type | Dimensions | +|------|------------| +| `BOX` | (width, height, depth) | +| `SPHERE` | (radius,) | +| `CYLINDER` | (radius, height) | +| `MESH` | mesh_path | ## Supported Robots @@ -113,18 +166,12 @@ planning/ | `xarm6` | 6 | | `xarm7` | 7 | -## Planners +## Testing -| Planner | Description | -|---------|-------------| -| `rrt_connect` | Bidirectional RRT (fast) | -| `rrt_star` | RRT* with rewiring (optimal) | - -## Obstacle Types +```bash +# Unit tests (fast, no Drake) +pytest dimos/manipulation/test_manipulation_unit.py -v -| Type | Dimensions | -|------|------------| -| `BOX` | (w, h, d) | -| `SPHERE` | (radius,) | -| `CYLINDER` | (radius, height) | -| `MESH` | mesh_path | +# Integration tests (requires Drake) +pytest dimos/e2e_tests/test_manipulation_module.py -v +``` From f8043e57aad0503de3219c50d2d7bd8efa8ba803 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Wed, 21 Jan 2026 10:53:11 -0800 Subject: [PATCH 18/31] Address greptile comments: manipulation_blueprint - Added optional add_gripper: bool = True to make xarm6 ans xarm7 config consistent. world_obstacle_monitor - Added warning log when obstacle not found during cleanup manipulation_client - Removed unused numpy impor path_utils - Added explicit tolerances atol=1e-6, rtol=0 to np.allclose() for stricter joint-space duplicate detection jacobian_ik - Added division-by-zero protection for velocity limits using nonzero_mask to skip zero-valued limits drake_world.py - added more specific exception handling (avoids hiding bugs) mesh_utils - regex fixes --- dimos/manipulation/manipulation_blueprints.py | 22 +++++++----- dimos/manipulation/manipulation_interface.py | 2 +- dimos/manipulation/manipulation_module.py | 5 ++- .../planning/examples/manipulation_client.py | 34 ++++++++++++++++--- .../kinematics/drake_optimization_ik.py | 8 +++-- .../planning/kinematics/jacobian_ik.py | 10 ++++-- .../monitor/world_obstacle_monitor.py | 4 ++- .../planning/utils/kinematics_utils.py | 7 ++-- .../manipulation/planning/utils/mesh_utils.py | 11 +++--- .../manipulation/planning/utils/path_utils.py | 4 +-- .../planning/world/drake_world.py | 27 ++++++++------- 11 files changed, 86 insertions(+), 48 deletions(-) diff --git a/dimos/manipulation/manipulation_blueprints.py b/dimos/manipulation/manipulation_blueprints.py index 1d85490e7b..593ac94555 100644 --- a/dimos/manipulation/manipulation_blueprints.py +++ b/dimos/manipulation/manipulation_blueprints.py @@ -108,6 +108,7 @@ def _make_xarm6_config( y_offset: float = 0.0, joint_prefix: str = "", orchestrator_task: str | None = None, + add_gripper: bool = True, ) -> RobotModelConfig: """Create XArm6 robot config. @@ -116,10 +117,19 @@ def _make_xarm6_config( y_offset: Y-axis offset for base pose (for multi-arm setups) joint_prefix: Prefix for joint name mapping (e.g., "left_" or "right_") orchestrator_task: Task name for orchestrator RPC execution + add_gripper: Whether to add the xarm gripper """ joint_names = ["joint1", "joint2", "joint3", "joint4", "joint5", "joint6"] joint_mapping = {f"{joint_prefix}{j}": j for j in joint_names} if joint_prefix else {} + xacro_args: dict[str, str] = { + "dof": "6", + "limited": "true", + "attach_xyz": f"0 {y_offset} 0", + } + if add_gripper: + xacro_args["add_gripper"] = "true" + return RobotModelConfig( name=name, urdf_path=_get_xarm_urdf_path(), @@ -133,17 +143,11 @@ def _make_xarm6_config( dtype=np.float64, ), joint_names=joint_names, - end_effector_link="link_tcp", + end_effector_link="link_tcp" if add_gripper else "link6", base_link="link_base", package_paths=_get_xarm_package_paths(), - # Pass attach_xyz to position each robot; xacro creates world_joint at this offset - xacro_args={ - "dof": "6", - "limited": "true", - "add_gripper": "true", - "attach_xyz": f"0 {y_offset} 0", - }, - collision_exclusion_pairs=XARM_GRIPPER_COLLISION_EXCLUSIONS, + xacro_args=xacro_args, + collision_exclusion_pairs=XARM_GRIPPER_COLLISION_EXCLUSIONS if add_gripper else [], auto_convert_meshes=True, max_velocity=1.0, max_acceleration=2.0, diff --git a/dimos/manipulation/manipulation_interface.py b/dimos/manipulation/manipulation_interface.py index 10e71fbc66..24b0f823e7 100644 --- a/dimos/manipulation/manipulation_interface.py +++ b/dimos/manipulation/manipulation_interface.py @@ -23,7 +23,7 @@ import os from typing import TYPE_CHECKING, Any -from dimos.manipulation.manipulation_history import ( +from dimos.manipulation.manipulation_history import ( # type: ignore[import-not-found] ManipulationHistory, ) from dimos.types.manipulation import ( diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index 95c4b94778..b2a1bfeb2a 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -46,7 +46,6 @@ if TYPE_CHECKING: from numpy.typing import NDArray - from dimos.control.orchestrator import ControlOrchestrator from dimos.core.rpc_client import RPCClient from dimos.msgs.geometry_msgs import Pose @@ -106,7 +105,7 @@ def __init__(self, *args: object, **kwargs: object) -> None: self._planned_trajectories: dict[str, JointTrajectory] = {} # Orchestrator integration (lazy initialized) - self._orchestrator_client: RPCClient[ControlOrchestrator] | None = None + self._orchestrator_client: RPCClient | None = None logger.info("ManipulationModule initialized") @@ -486,7 +485,7 @@ def get_robot_info(self, robot_name: str | None = None) -> dict[str, object] | N # Orchestrator Integration RPC Methods # ========================================================================= - def _get_orchestrator_client(self) -> RPCClient[ControlOrchestrator] | None: + def _get_orchestrator_client(self) -> RPCClient | None: """Get or create orchestrator RPC client (lazy init).""" if not any(c.orchestrator_task_name for _, c, _ in self._robots.values()): return None diff --git a/dimos/manipulation/planning/examples/manipulation_client.py b/dimos/manipulation/planning/examples/manipulation_client.py index a749690629..aa207af948 100644 --- a/dimos/manipulation/planning/examples/manipulation_client.py +++ b/dimos/manipulation/planning/examples/manipulation_client.py @@ -45,8 +45,6 @@ from typing import Any, cast -import numpy as np - from dimos.msgs.geometry_msgs import Pose, Quaternion, Vector3 from dimos.protocol.rpc import LCMRPC @@ -268,7 +266,35 @@ def main() -> None: c = ManipulationClient() - # Expose methods directly (no 'c.' prefix needed) + # Expose methods directly in user namespace (no 'c.' prefix needed) + user_ns = { + "c": c, + "client": c, + # Query methods + "state": c.state, + "joints": c.joints, + "ee": c.ee, + "url": c.url, + "robots": c.robots, + "info": c.info, + # Planning methods + "plan": c.plan, + "plan_pose": c.plan_pose, + "preview": c.preview, + "execute": c.execute, + "has_plan": c.has_plan, + "clear_plan": c.clear_plan, + # Obstacle methods + "box": c.box, + "sphere": c.sphere, + "cylinder": c.cylinder, + "remove": c.remove, + # Utility methods + "collision": c.collision, + "reset": c.reset, + "cancel": c.cancel, + "status": c.status, + } banner = """ Manipulation Client - IPython Interface @@ -305,7 +331,7 @@ def main() -> None: print(banner) try: - embed(colors="neutral") # type: ignore[no-untyped-call] + embed(user_ns=user_ns, colors="neutral") # type: ignore[no-untyped-call] finally: c.stop() diff --git a/dimos/manipulation/planning/kinematics/drake_optimization_ik.py b/dimos/manipulation/planning/kinematics/drake_optimization_ik.py index 7637c03c30..853372f3b2 100644 --- a/dimos/manipulation/planning/kinematics/drake_optimization_ik.py +++ b/dimos/manipulation/planning/kinematics/drake_optimization_ik.py @@ -35,9 +35,11 @@ from numpy.typing import NDArray try: - from pydrake.math import RigidTransform, RotationMatrix - from pydrake.multibody.inverse_kinematics import InverseKinematics - from pydrake.solvers import Solve + from pydrake.math import RigidTransform, RotationMatrix # type: ignore[import-not-found] + from pydrake.multibody.inverse_kinematics import ( # type: ignore[import-not-found] + InverseKinematics, + ) + from pydrake.solvers import Solve # type: ignore[import-not-found] DRAKE_AVAILABLE = True except ImportError: diff --git a/dimos/manipulation/planning/kinematics/jacobian_ik.py b/dimos/manipulation/planning/kinematics/jacobian_ik.py index 1ab26c9fd5..aa5e1eae1e 100644 --- a/dimos/manipulation/planning/kinematics/jacobian_ik.py +++ b/dimos/manipulation/planning/kinematics/jacobian_ik.py @@ -305,9 +305,13 @@ def solve_differential( # Apply velocity limits if available config = world.get_robot_config(robot_id) if config.velocity_limits is not None: - max_ratio = np.max(np.abs(q_dot) / np.array(config.velocity_limits)) - if max_ratio > 1.0: - q_dot = q_dot / max_ratio + velocity_limits = np.array(config.velocity_limits) + # Only consider joints with non-zero velocity limits + nonzero_mask = velocity_limits > 0 + if np.any(nonzero_mask): + max_ratio = np.max(np.abs(q_dot[nonzero_mask]) / velocity_limits[nonzero_mask]) + if max_ratio > 1.0: + q_dot = q_dot / max_ratio return q_dot diff --git a/dimos/manipulation/planning/monitor/world_obstacle_monitor.py b/dimos/manipulation/planning/monitor/world_obstacle_monitor.py index 725095090e..5a879e34fb 100644 --- a/dimos/manipulation/planning/monitor/world_obstacle_monitor.py +++ b/dimos/manipulation/planning/monitor/world_obstacle_monitor.py @@ -295,7 +295,9 @@ def _cleanup_stale_detections( for det_id in stale_ids: obstacle_id = self._perception_objects[det_id] - self._world.remove_obstacle(obstacle_id) + removed = self._world.remove_obstacle(obstacle_id) + if not removed: + logger.warning(f"Obstacle '{obstacle_id}' not found in world during cleanup") del self._perception_objects[det_id] del self._perception_timestamps[det_id] diff --git a/dimos/manipulation/planning/utils/kinematics_utils.py b/dimos/manipulation/planning/utils/kinematics_utils.py index d86253f037..fa6b74e61c 100644 --- a/dimos/manipulation/planning/utils/kinematics_utils.py +++ b/dimos/manipulation/planning/utils/kinematics_utils.py @@ -274,9 +274,10 @@ def rotation_matrix_to_axis_angle(R: NDArray[np.float64]) -> tuple[NDArray[np.fl idx = int(np.argmax(diag)) axis = np.zeros(3) axis[idx] = np.sqrt((diag[idx] + 1) / 2) - for j in range(3): - if j != idx: - axis[j] = R[idx, j] / (2 * axis[idx]) + if axis[idx] > 1e-12: + for j in range(3): + if j != idx: + axis[j] = R[idx, j] / (2 * axis[idx]) return axis, angle # General case diff --git a/dimos/manipulation/planning/utils/mesh_utils.py b/dimos/manipulation/planning/utils/mesh_utils.py index d7645f359e..48f6452e15 100644 --- a/dimos/manipulation/planning/utils/mesh_utils.py +++ b/dimos/manipulation/planning/utils/mesh_utils.py @@ -32,7 +32,6 @@ from __future__ import annotations import hashlib -import os from pathlib import Path import re import shutil @@ -135,7 +134,7 @@ def _process_xacro( ) -> str: """Process xacro file to URDF.""" try: - import xacro # type: ignore[import-untyped] + import xacro # type: ignore[import-not-found] except ImportError: raise ImportError( "xacro is required for processing .xacro files. Install with: pip install xacro" @@ -184,9 +183,9 @@ def _strip_transmission_blocks(urdf_content: str) -> str: Returns: URDF content with transmission blocks removed """ - # Pattern to match ... blocks (including nested content) + # Pattern to match ... blocks and self-closing # Uses non-greedy matching and handles nested tags - pattern = r"]*>.*?" + pattern = r"]*(?:/>|>.*?)" # Remove transmission blocks (with flags for multiline and dotall) result = re.sub(pattern, "", urdf_content, flags=re.DOTALL | re.MULTILINE) @@ -205,8 +204,8 @@ def _resolve_package_uris( output_dir: Path, ) -> str: """Resolve package:// URIs to filesystem paths.""" - # Pattern for package:// URIs - pattern = r'package://([^/]+)/(.+?)(["\s<>])' + # Pattern for package:// URIs (handles both single and double quotes) + pattern = r'package://([^/]+)/(.+?)(["\'\\s<>])' def replace_uri(match: re.Match[str]) -> str: pkg_name = match.group(1) diff --git a/dimos/manipulation/planning/utils/path_utils.py b/dimos/manipulation/planning/utils/path_utils.py index 05d63800a4..cd1360598a 100644 --- a/dimos/manipulation/planning/utils/path_utils.py +++ b/dimos/manipulation/planning/utils/path_utils.py @@ -277,8 +277,8 @@ def concatenate_paths( continue if remove_duplicates and result: - # Check if last point matches first point - if np.allclose(result[-1], path[0]): + # Check if last point matches first point (tight tolerance for joint space) + if np.allclose(result[-1], path[0], atol=1e-6, rtol=0): result.extend(path[1:]) else: result.extend(path) diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index cf318c4cf1..05495251f5 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -39,7 +39,7 @@ from numpy.typing import NDArray try: - from pydrake.geometry import ( + from pydrake.geometry import ( # type: ignore[import-not-found] AddContactMaterial, Box, CollisionFilterDeclaration, @@ -57,15 +57,15 @@ SceneGraph, Sphere, ) - from pydrake.math import RigidTransform - from pydrake.multibody.parsing import Parser - from pydrake.multibody.plant import ( + from pydrake.math import RigidTransform # type: ignore[import-not-found] + from pydrake.multibody.parsing import Parser # type: ignore[import-not-found] + from pydrake.multibody.plant import ( # type: ignore[import-not-found] AddMultibodyPlantSceneGraph, CoulombFriction, MultibodyPlant, ) - from pydrake.multibody.tree import JacobianWrtVariable - from pydrake.systems.framework import Context, DiagramBuilder + from pydrake.multibody.tree import JacobianWrtVariable # type: ignore[import-not-found] + from pydrake.systems.framework import Context, DiagramBuilder # type: ignore[import-not-found] DRAKE_AVAILABLE = True except ImportError: @@ -598,7 +598,7 @@ def scratch_context(self) -> Generator[Context, None, None]: self._plant_context, robot_data.model_instance ) self._plant.SetPositions(plant_ctx, robot_data.model_instance, positions) - except Exception: + except RuntimeError: pass # Robot not yet synced yield ctx @@ -745,7 +745,8 @@ def get_ee_pose(self, ctx: Context, robot_id: str) -> NDArray[np.float64]: ee_body = robot_data.ee_frame.body() X_WE = self._plant.EvalBodyPoseInWorld(plant_ctx, ee_body) - return X_WE.GetAsMatrix4() + result: NDArray[np.float64] = X_WE.GetAsMatrix4() + return result def get_link_pose(self, ctx: Context, robot_id: str, link_name: str) -> NDArray[np.float64]: """Get link pose as 4x4 transform.""" @@ -765,7 +766,8 @@ def get_link_pose(self, ctx: Context, robot_id: str, link_name: str) -> NDArray[ X_WL = self._plant.EvalBodyPoseInWorld(plant_ctx, body) - return X_WL.GetAsMatrix4() + result: NDArray[np.float64] = X_WL.GetAsMatrix4() + return result def get_jacobian(self, ctx: Context, robot_id: str) -> NDArray[np.float64]: """Get geometric Jacobian (6 x n_joints). @@ -808,7 +810,8 @@ def get_jacobian(self, ctx: Context, robot_id: str) -> NDArray[np.float64]: def get_visualization_url(self) -> str | None: """Get visualization URL if enabled.""" if self._meshcat is not None: - return self._meshcat.web_url() + url: str = self._meshcat.web_url() + return url return None def publish_visualization(self, ctx: Context | None = None) -> None: @@ -850,9 +853,7 @@ def animate_path( other_robot_positions: dict[str, NDArray[np.float64]] = {} for rid, _robot_data in self._robots.items(): if rid != robot_id: - pos = self.get_positions(self.get_live_context(), rid) - if pos is not None: - other_robot_positions[rid] = pos + other_robot_positions[rid] = self.get_positions(self.get_live_context(), rid) dt = duration / (len(path) - 1) for q in path: From e3cac082d1a1edce7db987678b4a830c5dc8a4a8 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Wed, 21 Jan 2026 11:18:27 -0800 Subject: [PATCH 19/31] fix mypy import error with manipulation interface (file is deprecated) --- dimos/manipulation/manipulation_interface.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dimos/manipulation/manipulation_interface.py b/dimos/manipulation/manipulation_interface.py index 24b0f823e7..eb528a3a43 100644 --- a/dimos/manipulation/manipulation_interface.py +++ b/dimos/manipulation/manipulation_interface.py @@ -23,7 +23,7 @@ import os from typing import TYPE_CHECKING, Any -from dimos.manipulation.manipulation_history import ( # type: ignore[import-not-found] +from dimos.manipulation.manipulation_history import ( # type: ignore[import-not-found,import-untyped] ManipulationHistory, ) from dimos.types.manipulation import ( From c7c4d59aaf4838f5f61ab3f1912a2cdced0df6aa Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Mon, 26 Jan 2026 15:57:05 -0800 Subject: [PATCH 20/31] protocol now only requires solve() method. --- .../kinematics/drake_optimization_ik.py | 103 +----------------- dimos/manipulation/planning/spec/protocols.py | 40 +------ 2 files changed, 6 insertions(+), 137 deletions(-) diff --git a/dimos/manipulation/planning/kinematics/drake_optimization_ik.py b/dimos/manipulation/planning/kinematics/drake_optimization_ik.py index 853372f3b2..7d079f9058 100644 --- a/dimos/manipulation/planning/kinematics/drake_optimization_ik.py +++ b/dimos/manipulation/planning/kinematics/drake_optimization_ik.py @@ -12,14 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Drake-specific optimization-based inverse kinematics. - -DrakeOptimizationIK uses Drake's InverseKinematics class with nonlinear optimization -(SNOPT/IPOPT) to find accurate IK solutions. It requires DrakeWorld and cannot work -with other physics backends. - -For backend-agnostic Jacobian-based IK, use JacobianIK. -""" +"""Drake optimization-based IK using SNOPT/IPOPT. Requires DrakeWorld.""" from __future__ import annotations @@ -49,42 +42,16 @@ class DrakeOptimizationIK: - """Drake-specific optimization-based IK solver. - - Uses Drake's InverseKinematics class with SNOPT/IPOPT to solve constrained - nonlinear optimization for IK. This produces highly accurate solutions but - requires DrakeWorld. - - For backend-agnostic IK, use JacobianIK instead. - - Methods: - - solve(): Full nonlinear IK with multiple random restarts - - solve_single(): Single optimization attempt with given seed - - Example: - from dimos.manipulation.planning.world import DrakeWorld - - world = DrakeWorld(enable_viz=True) - world.add_robot(config) - world.finalize() + """Drake optimization-based IK solver using constrained nonlinear optimization. - ik = DrakeOptimizationIK() - result = ik.solve(world, robot_id, target_pose) - if result.is_success(): - print(f"Solution: {result.joint_positions}") + Requires DrakeWorld. For backend-agnostic IK, use JacobianIK. """ def __init__(self) -> None: - """Create Drake optimization IK solver.""" if not DRAKE_AVAILABLE: raise ImportError("Drake is not installed. Install with: pip install drake") - # Internal JacobianIK for iterative/differential methods - from dimos.manipulation.planning.kinematics.jacobian_ik import JacobianIK - - self._jacobian_ik = JacobianIK() def _validate_world(self, world: WorldSpec) -> IKResult | None: - """Validate world is DrakeWorld and finalized. Returns error or None.""" from dimos.manipulation.planning.world.drake_world import DrakeWorld if not isinstance(world, DrakeWorld): @@ -106,25 +73,7 @@ def solve( check_collision: bool = True, max_attempts: int = 10, ) -> IKResult: - """Full nonlinear IK with multiple random restarts. - - Solves constrained nonlinear optimization using Drake's InverseKinematics - and SNOPT/IPOPT solvers. Tries multiple starting configurations to find - a collision-free solution. - - Args: - world: DrakeWorld instance (NOT other WorldSpec implementations) - robot_id: Robot to solve IK for - target_pose: Target end-effector pose (4x4 transform) - seed: Initial guess (uses current position if None) - position_tolerance: Required position accuracy (meters) - orientation_tolerance: Required orientation accuracy (radians) - check_collision: Whether to check collision of solution - max_attempts: Maximum random restart attempts - - Returns: - IKResult with solution or failure status - """ + """Solve IK with multiple random restarts, returning the best collision-free solution.""" error = self._validate_world(world) if error is not None: return error @@ -203,7 +152,6 @@ def _solve_single( lower_limits: NDArray[np.float64], upper_limits: NDArray[np.float64], ) -> IKResult: - """Solve IK with a single seed.""" # Get robot data from world internals (Drake-specific access) robot_data = world._robots[robot_id] # type: ignore[attr-defined] plant = world.plant # type: ignore[attr-defined] @@ -275,47 +223,6 @@ def _solve_single( iterations=1, ) - def solve_iterative( - self, - world: WorldSpec, - robot_id: str, - target_pose: NDArray[np.float64], - seed: NDArray[np.float64], - max_iterations: int = 100, - position_tolerance: float = 0.001, - orientation_tolerance: float = 0.01, - ) -> IKResult: - """Iterative Jacobian-based IK. Delegates to JacobianIK.""" - return self._jacobian_ik.solve_iterative( - world=world, - robot_id=robot_id, - target_pose=target_pose, - seed=seed, - max_iterations=max_iterations, - position_tolerance=position_tolerance, - orientation_tolerance=orientation_tolerance, - ) - - def solve_differential( - self, - world: WorldSpec, - robot_id: str, - current_joints: NDArray[np.float64], - twist: NDArray[np.float64], - dt: float, - ) -> NDArray[np.float64] | None: - """Single Jacobian step for velocity control. Delegates to JacobianIK.""" - return self._jacobian_ik.solve_differential( - world=world, - robot_id=robot_id, - current_joints=current_joints, - twist=twist, - dt=dt, - ) - - -# ============= Result Helpers ============= - def _create_success_result( joint_positions: NDArray[np.float64], @@ -323,7 +230,6 @@ def _create_success_result( orientation_error: float, iterations: int, ) -> IKResult: - """Create a successful IK result.""" return IKResult( status=IKStatus.SUCCESS, joint_positions=joint_positions, @@ -339,7 +245,6 @@ def _create_failure_result( message: str, iterations: int = 0, ) -> IKResult: - """Create a failed IK result.""" return IKResult( status=status, joint_positions=None, diff --git a/dimos/manipulation/planning/spec/protocols.py b/dimos/manipulation/planning/spec/protocols.py index fac97f2c49..b8c2a22be5 100644 --- a/dimos/manipulation/planning/spec/protocols.py +++ b/dimos/manipulation/planning/spec/protocols.py @@ -167,19 +167,7 @@ def animate_path( @runtime_checkable class KinematicsSpec(Protocol): - """Protocol for inverse kinematics solver. - - Kinematics solvers are stateless and use WorldSpec for FK/collision. - - Methods: - - solve(): Full optimization-based IK with collision checking - - solve_iterative(): Iterative Jacobian-based IK - - solve_differential(): Single Jacobian step for velocity control - - Implementations: - - JacobianIK: Backend-agnostic iterative/differential IK - - DrakeOptimizationIK: Uses Drake's InverseKinematics + SNOPT/IPOPT - """ + """Protocol for inverse kinematics solvers. Stateless, uses WorldSpec for FK/collision.""" def solve( self, @@ -192,31 +180,7 @@ def solve( check_collision: bool = True, max_attempts: int = 10, ) -> IKResult: - """Solve full IK with optional collision checking.""" - ... - - def solve_iterative( - self, - world: WorldSpec, - robot_id: str, - target_pose: NDArray[np.float64], - seed: NDArray[np.float64], - max_iterations: int = 100, - position_tolerance: float = 0.001, - orientation_tolerance: float = 0.01, - ) -> IKResult: - """Solve IK iteratively using Jacobian method.""" - ... - - def solve_differential( - self, - world: WorldSpec, - robot_id: str, - current_joints: NDArray[np.float64], - twist: NDArray[np.float64], - dt: float, - ) -> NDArray[np.float64] | None: - """Single Jacobian step for velocity control.""" + """Solve IK with optional collision checking.""" ... From f8b9e5f6736ac3412dcb88eca54f7501089e4e58 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Mon, 26 Jan 2026 17:28:55 -0800 Subject: [PATCH 21/31] moved to using standard LCM data types and added type alias for readability --- dimos/manipulation/manipulation_module.py | 78 ++++++++++++------ dimos/manipulation/planning/__init__.py | 6 ++ .../kinematics/drake_optimization_ik.py | 18 +++-- .../planning/kinematics/jacobian_ik.py | 24 ++++-- .../planning/monitor/world_monitor.py | 38 +++++---- .../monitor/world_obstacle_monitor.py | 19 ++--- .../planning/planners/rrt_planner.py | 24 +++--- dimos/manipulation/planning/spec/__init__.py | 6 ++ dimos/manipulation/planning/spec/protocols.py | 51 +++++++----- dimos/manipulation/planning/spec/types.py | 30 +++++-- .../manipulation/planning/utils/mesh_utils.py | 3 +- .../manipulation/planning/utils/path_utils.py | 26 +++--- .../planning/world/drake_world.py | 79 +++++++++++-------- dimos/robot/all_blueprints.py | 4 + 14 files changed, 265 insertions(+), 141 deletions(-) diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index b2a1bfeb2a..bc5b4bd87a 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -19,19 +19,22 @@ from dataclasses import dataclass, field from enum import Enum import threading -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, TypeAlias import numpy as np from dimos.core import In, Module, rpc from dimos.core.module import ModuleConfig from dimos.manipulation.planning import ( + JointPath, JointTrajectoryGenerator, KinematicsSpec, Obstacle, ObstacleType, PlannerSpec, RobotModelConfig, + RobotName, + WorldRobotID, create_kinematics, create_planner, ) @@ -41,16 +44,29 @@ from dimos.msgs.sensor_msgs import JointState # noqa: TC001 from dimos.msgs.trajectory_msgs import JointTrajectory from dimos.utils.logging_config import setup_logger -from dimos.utils.transform_utils import matrix_to_pose, pose_to_matrix +from dimos.utils.transform_utils import matrix_to_pose if TYPE_CHECKING: from numpy.typing import NDArray from dimos.core.rpc_client import RPCClient - from dimos.msgs.geometry_msgs import Pose + from dimos.msgs.geometry_msgs import Pose, PoseStamped logger = setup_logger() +# Composite type aliases for readability (using semantic IDs from planning.spec) +RobotEntry: TypeAlias = tuple[WorldRobotID, RobotModelConfig, JointTrajectoryGenerator] +"""(world_robot_id, config, trajectory_generator)""" + +RobotRegistry: TypeAlias = dict[RobotName, RobotEntry] +"""Maps robot_name -> RobotEntry""" + +PlannedPaths: TypeAlias = dict[RobotName, JointPath] +"""Maps robot_name -> planned joint path""" + +PlannedTrajectories: TypeAlias = dict[RobotName, JointTrajectory] +"""Maps robot_name -> planned trajectory""" + class ManipulationState(Enum): """State machine for manipulation module.""" @@ -98,11 +114,11 @@ def __init__(self, *args: object, **kwargs: object) -> None: self._kinematics: KinematicsSpec | None = None # Robot registry: maps robot_name -> (world_robot_id, config, trajectory_gen) - self._robots: dict[str, tuple[str, RobotModelConfig, JointTrajectoryGenerator]] = {} + self._robots: RobotRegistry = {} # Stored path for plan/preview/execute workflow (per robot) - self._planned_paths: dict[str, list[NDArray[np.float64]]] = {} - self._planned_trajectories: dict[str, JointTrajectory] = {} + self._planned_paths: PlannedPaths = {} + self._planned_trajectories: PlannedTrajectories = {} # Orchestrator integration (lazy initialized) self._orchestrator_client: RPCClient | None = None @@ -154,15 +170,15 @@ def _initialize_planning(self) -> None: self._planner = create_planner(name=self.config.planner_name) self._kinematics = create_kinematics(name=self.config.kinematics_name) - def _get_default_robot_name(self) -> str | None: + def _get_default_robot_name(self) -> RobotName | None: """Get default robot name (first robot if only one, else None).""" if len(self._robots) == 1: return next(iter(self._robots.keys())) return None def _get_robot( - self, robot_name: str | None = None - ) -> tuple[str, str, RobotModelConfig, JointTrajectoryGenerator] | None: + self, robot_name: RobotName | None = None + ) -> tuple[RobotName, WorldRobotID, RobotModelConfig, JointTrajectoryGenerator] | None: """Get robot by name or default. Args: @@ -236,7 +252,7 @@ def reset(self) -> bool: return True @rpc - def get_current_joints(self, robot_name: str | None = None) -> list[float] | None: + def get_current_joints(self, robot_name: RobotName | None = None) -> list[float] | None: """Get current joint positions. Args: @@ -249,7 +265,7 @@ def get_current_joints(self, robot_name: str | None = None) -> list[float] | Non return None @rpc - def get_ee_pose(self, robot_name: str | None = None) -> Pose | None: + def get_ee_pose(self, robot_name: RobotName | None = None) -> Pose | None: """Get current end-effector pose. Args: @@ -260,7 +276,7 @@ def get_ee_pose(self, robot_name: str | None = None) -> Pose | None: return None @rpc - def is_collision_free(self, joints: list[float], robot_name: str | None = None) -> bool: + def is_collision_free(self, joints: list[float], robot_name: RobotName | None = None) -> bool: """Check if joint configuration is collision-free. Args: @@ -277,7 +293,9 @@ def is_collision_free(self, joints: list[float], robot_name: str | None = None) # Plan/Preview/Execute Workflow RPC Methods # ========================================================================= - def _begin_planning(self, robot_name: str | None = None) -> tuple[str, str] | None: + def _begin_planning( + self, robot_name: RobotName | None = None + ) -> tuple[RobotName, WorldRobotID] | None: """Check state and begin planning. Returns (robot_name, robot_id) or None. Args: @@ -303,7 +321,7 @@ def _fail(self, msg: str) -> bool: return False @rpc - def plan_to_pose(self, pose: Pose, robot_name: str | None = None) -> bool: + def plan_to_pose(self, pose: Pose, robot_name: RobotName | None = None) -> bool: """Plan motion to pose. Use preview_path() then execute(). Args: @@ -319,10 +337,19 @@ def plan_to_pose(self, pose: Pose, robot_name: str | None = None) -> bool: if current is None: return self._fail("No joint state") + # Convert Pose to PoseStamped for the IK solver + from dimos.msgs.geometry_msgs import PoseStamped + + target_pose = PoseStamped( + frame_id="world", + position=pose.position, + orientation=pose.orientation, + ) + ik = self._kinematics.solve( world=self._world_monitor.world, robot_id=robot_id, - target_pose=pose_to_matrix(pose), + target_pose=target_pose, seed=current, check_collision=True, ) @@ -333,7 +360,7 @@ def plan_to_pose(self, pose: Pose, robot_name: str | None = None) -> bool: return self._plan_path_only(robot_name, robot_id, ik.joint_positions) @rpc - def plan_to_joints(self, joints: list[float], robot_name: str | None = None) -> bool: + def plan_to_joints(self, joints: list[float], robot_name: RobotName | None = None) -> bool: """Plan motion to joint config. Use preview_path() then execute(). Args: @@ -346,7 +373,9 @@ def plan_to_joints(self, joints: list[float], robot_name: str | None = None) -> logger.info(f"Planning to joints for {robot_name}: {[f'{j:.3f}' for j in joints]}") return self._plan_path_only(robot_name, robot_id, np.array(joints)) - def _plan_path_only(self, robot_name: str, robot_id: str, goal: NDArray[np.float64]) -> bool: + def _plan_path_only( + self, robot_name: RobotName, robot_id: WorldRobotID, goal: NDArray[np.float64] + ) -> bool: """Plan path from current position to goal, store result.""" assert self._world_monitor and self._planner # guaranteed by _begin_planning start = self._world_monitor.get_current_positions(robot_id) @@ -375,7 +404,7 @@ def _plan_path_only(self, robot_name: str, robot_id: str, goal: NDArray[np.float return True @rpc - def preview_path(self, duration: float = 3.0, robot_name: str | None = None) -> bool: + def preview_path(self, duration: float = 3.0, robot_name: RobotName | None = None) -> bool: """Preview the planned path in the visualizer. Args: @@ -454,7 +483,7 @@ def list_robots(self) -> list[str]: return list(self._robots.keys()) @rpc - def get_robot_info(self, robot_name: str | None = None) -> dict[str, object] | None: + def get_robot_info(self, robot_name: RobotName | None = None) -> dict[str, object] | None: """Get information about a robot. Args: @@ -527,7 +556,7 @@ def _translate_trajectory_to_orchestrator( ) @rpc - def execute(self, robot_name: str | None = None) -> bool: + def execute(self, robot_name: RobotName | None = None) -> bool: """Execute planned trajectory via ControlOrchestrator.""" if (robot := self._get_robot(robot_name)) is None: return False @@ -557,7 +586,9 @@ def execute(self, robot_name: str | None = None) -> bool: return self._fail("Orchestrator rejected trajectory") @rpc - def get_trajectory_status(self, robot_name: str | None = None) -> dict[str, object] | None: + def get_trajectory_status( + self, robot_name: RobotName | None = None + ) -> dict[str, object] | None: """Get trajectory execution status.""" if (robot := self._get_robot(robot_name)) is None: return None @@ -602,10 +633,13 @@ def add_obstacle( logger.warning("mesh_path required for mesh obstacles") return "" + # Import PoseStamped here to avoid circular imports + from dimos.msgs.geometry_msgs import PoseStamped + obstacle = Obstacle( name=name, obstacle_type=obstacle_type, - pose=pose_to_matrix(pose), + pose=PoseStamped(position=pose.position, orientation=pose.orientation), dimensions=tuple(dimensions) if dimensions else (), mesh_path=mesh_path, ) diff --git a/dimos/manipulation/planning/__init__.py b/dimos/manipulation/planning/__init__.py index 71638f376b..a9a89a96cd 100644 --- a/dimos/manipulation/planning/__init__.py +++ b/dimos/manipulation/planning/__init__.py @@ -70,6 +70,7 @@ CollisionObjectMessage, IKResult, IKStatus, + JointPath, KinematicsSpec, Obstacle, ObstacleType, @@ -77,6 +78,8 @@ PlanningResult, PlanningStatus, RobotModelConfig, + RobotName, + WorldRobotID, WorldSpec, ) @@ -89,6 +92,7 @@ "CollisionObjectMessage", "IKResult", "IKStatus", + "JointPath", "JointTrajectoryGenerator", "KinematicsSpec", "Obstacle", @@ -97,6 +101,8 @@ "PlanningResult", "PlanningStatus", "RobotModelConfig", + "RobotName", + "WorldRobotID", "WorldSpec", "create_kinematics", "create_planner", diff --git a/dimos/manipulation/planning/kinematics/drake_optimization_ik.py b/dimos/manipulation/planning/kinematics/drake_optimization_ik.py index 7d079f9058..73766f00a8 100644 --- a/dimos/manipulation/planning/kinematics/drake_optimization_ik.py +++ b/dimos/manipulation/planning/kinematics/drake_optimization_ik.py @@ -20,13 +20,15 @@ import numpy as np -from dimos.manipulation.planning.spec import IKResult, IKStatus, WorldSpec +from dimos.manipulation.planning.spec import IKResult, IKStatus, WorldRobotID, WorldSpec from dimos.manipulation.planning.utils.kinematics_utils import compute_pose_error from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: from numpy.typing import NDArray +from dimos.msgs.geometry_msgs import PoseStamped, Transform + try: from pydrake.math import RigidTransform, RotationMatrix # type: ignore[import-not-found] from pydrake.multibody.inverse_kinematics import ( # type: ignore[import-not-found] @@ -65,8 +67,8 @@ def _validate_world(self, world: WorldSpec) -> IKResult | None: def solve( self, world: WorldSpec, - robot_id: str, - target_pose: NDArray[np.float64], + robot_id: WorldRobotID, + target_pose: PoseStamped, seed: NDArray[np.float64] | None = None, position_tolerance: float = 0.001, orientation_tolerance: float = 0.01, @@ -78,6 +80,12 @@ def solve( if error is not None: return error + # Convert PoseStamped to 4x4 matrix via Transform + target_matrix = Transform( + translation=target_pose.position, + rotation=target_pose.orientation, + ).to_matrix() + # Get joint limits lower_limits, upper_limits = world.get_joint_limits(robot_id) @@ -87,7 +95,7 @@ def solve( seed = world.get_positions(ctx, robot_id) # Target transform - target_transform = RigidTransform(target_pose) + target_transform = RigidTransform(target_matrix) best_result: IKResult | None = None best_error = float("inf") @@ -144,7 +152,7 @@ def solve( def _solve_single( self, world: WorldSpec, - robot_id: str, + robot_id: WorldRobotID, target_transform: RigidTransform, seed: NDArray[np.float64], position_tolerance: float, diff --git a/dimos/manipulation/planning/kinematics/jacobian_ik.py b/dimos/manipulation/planning/kinematics/jacobian_ik.py index aa5e1eae1e..e31b16372c 100644 --- a/dimos/manipulation/planning/kinematics/jacobian_ik.py +++ b/dimos/manipulation/planning/kinematics/jacobian_ik.py @@ -28,7 +28,7 @@ import numpy as np -from dimos.manipulation.planning.spec import IKResult, IKStatus, WorldSpec +from dimos.manipulation.planning.spec import IKResult, IKStatus, WorldRobotID, WorldSpec from dimos.manipulation.planning.utils.kinematics_utils import ( check_singularity, compute_error_twist, @@ -40,6 +40,8 @@ if TYPE_CHECKING: from numpy.typing import NDArray +from dimos.msgs.geometry_msgs import PoseStamped, Transform + logger = setup_logger() @@ -87,8 +89,8 @@ def __init__( def solve( self, world: WorldSpec, - robot_id: str, - target_pose: NDArray[np.float64], + robot_id: WorldRobotID, + target_pose: PoseStamped, seed: NDArray[np.float64] | None = None, position_tolerance: float = 0.001, orientation_tolerance: float = 0.01, @@ -103,7 +105,7 @@ def solve( Args: world: World for FK/collision checking robot_id: Robot to solve IK for - target_pose: Target end-effector pose (4x4 transform) + target_pose: Target end-effector pose seed: Initial guess (uses current position if None) position_tolerance: Required position accuracy (meters) orientation_tolerance: Required orientation accuracy (radians) @@ -113,6 +115,12 @@ def solve( Returns: IKResult with solution or failure status """ + # Convert PoseStamped to 4x4 matrix via Transform + target_matrix = Transform( + translation=target_pose.position, + rotation=target_pose.orientation, + ).to_matrix() + if not world.is_finalized: return _create_failure_result(IKStatus.NO_SOLUTION, "World must be finalized before IK") @@ -138,7 +146,7 @@ def solve( result = self.solve_iterative( world=world, robot_id=robot_id, - target_pose=target_pose, + target_pose=target_matrix, seed=current_seed, max_iterations=self._max_iterations, position_tolerance=position_tolerance, @@ -175,7 +183,7 @@ def solve( def solve_iterative( self, world: WorldSpec, - robot_id: str, + robot_id: WorldRobotID, target_pose: NDArray[np.float64], seed: NDArray[np.float64], max_iterations: int = 100, @@ -267,7 +275,7 @@ def solve_iterative( def solve_differential( self, world: WorldSpec, - robot_id: str, + robot_id: WorldRobotID, current_joints: NDArray[np.float64], twist: NDArray[np.float64], dt: float, @@ -318,7 +326,7 @@ def solve_differential( def solve_differential_position_only( self, world: WorldSpec, - robot_id: str, + robot_id: WorldRobotID, current_joints: NDArray[np.float64], linear_velocity: NDArray[np.float64], ) -> NDArray[np.float64] | None: diff --git a/dimos/manipulation/planning/monitor/world_monitor.py b/dimos/manipulation/planning/monitor/world_monitor.py index 18de1f352c..f7f9e002a2 100644 --- a/dimos/manipulation/planning/monitor/world_monitor.py +++ b/dimos/manipulation/planning/monitor/world_monitor.py @@ -34,8 +34,10 @@ from dimos.manipulation.planning.spec import ( CollisionObjectMessage, + JointPath, Obstacle, RobotModelConfig, + WorldRobotID, WorldSpec, ) from dimos.msgs.sensor_msgs import JointState @@ -56,8 +58,8 @@ def __init__( self._backend = backend self._world: WorldSpec = create_world(backend=backend, enable_viz=enable_viz, **kwargs) self._lock = threading.RLock() - self._robot_joints: dict[str, list[str]] = {} - self._state_monitors: dict[str, WorldStateMonitor] = {} + self._robot_joints: dict[WorldRobotID, list[str]] = {} + self._state_monitors: dict[WorldRobotID, WorldStateMonitor] = {} self._obstacle_monitor: WorldObstacleMonitor | None = None self._viz_thread: threading.Thread | None = None self._viz_stop_event = threading.Event() @@ -65,7 +67,7 @@ def __init__( # ============= Robot Management ============= - def add_robot(self, config: RobotModelConfig) -> str: + def add_robot(self, config: RobotModelConfig) -> WorldRobotID: """Add a robot. Returns robot_id.""" with self._lock: robot_id = self._world.add_robot(config) @@ -73,17 +75,19 @@ def add_robot(self, config: RobotModelConfig) -> str: logger.info(f"Added robot '{config.name}' as '{robot_id}'") return robot_id - def get_robot_ids(self) -> list[str]: + def get_robot_ids(self) -> list[WorldRobotID]: """Get all robot IDs.""" with self._lock: return self._world.get_robot_ids() - def get_robot_config(self, robot_id: str) -> RobotModelConfig: + def get_robot_config(self, robot_id: WorldRobotID) -> RobotModelConfig: """Get robot configuration.""" with self._lock: return self._world.get_robot_config(robot_id) - def get_joint_limits(self, robot_id: str) -> tuple[NDArray[np.float64], NDArray[np.float64]]: + def get_joint_limits( + self, robot_id: WorldRobotID + ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: """Get joint limits for a robot.""" with self._lock: return self._world.get_joint_limits(robot_id) @@ -109,7 +113,7 @@ def clear_obstacles(self) -> None: def start_state_monitor( self, - robot_id: str, + robot_id: WorldRobotID, joint_names: list[str] | None = None, joint_name_mapping: dict[str, str] | None = None, ) -> None: @@ -176,7 +180,7 @@ def stop_all_monitors(self) -> None: # ============= Message Handlers ============= - def on_joint_state(self, msg: JointState, robot_id: str | None = None) -> None: + def on_joint_state(self, msg: JointState, robot_id: WorldRobotID | None = None) -> None: """Handle joint state message. Broadcasts to all monitors if robot_id is None.""" try: if robot_id is not None: @@ -206,7 +210,7 @@ def on_detections(self, detections: list[Detection3D]) -> None: # ============= State Access ============= - def get_current_positions(self, robot_id: str) -> NDArray[np.float64] | None: + def get_current_positions(self, robot_id: WorldRobotID) -> NDArray[np.float64] | None: """Get current joint positions. Returns None if not yet received.""" # Try state monitor first if robot_id in self._state_monitors: @@ -219,19 +223,19 @@ def get_current_positions(self, robot_id: str) -> NDArray[np.float64] | None: ctx = self._world.get_live_context() return self._world.get_positions(ctx, robot_id) - def get_current_velocities(self, robot_id: str) -> NDArray[np.float64] | None: + def get_current_velocities(self, robot_id: WorldRobotID) -> NDArray[np.float64] | None: """Get current joint velocities. Returns None if not available.""" if robot_id in self._state_monitors: return self._state_monitors[robot_id].get_current_velocities() return None - def wait_for_state(self, robot_id: str, timeout: float = 1.0) -> bool: + def wait_for_state(self, robot_id: WorldRobotID, timeout: float = 1.0) -> bool: """Wait until state is received. Returns False on timeout.""" if robot_id in self._state_monitors: return self._state_monitors[robot_id].wait_for_state(timeout) return False - def is_state_stale(self, robot_id: str, max_age: float = 1.0) -> bool: + def is_state_stale(self, robot_id: WorldRobotID, max_age: float = 1.0) -> bool: """Check if state is stale.""" if robot_id in self._state_monitors: return self._state_monitors[robot_id].is_state_stale(max_age) @@ -251,14 +255,14 @@ def get_live_context(self) -> Any: # ============= Collision Checking ============= - def is_state_valid(self, robot_id: str, joint_positions: NDArray[np.float64]) -> bool: + def is_state_valid(self, robot_id: WorldRobotID, joint_positions: NDArray[np.float64]) -> bool: """Check if configuration is collision-free.""" with self._world.scratch_context() as ctx: self._world.set_positions(ctx, robot_id, joint_positions) return self._world.is_collision_free(ctx, robot_id) def is_path_valid( - self, robot_id: str, path: list[NDArray[np.float64]], step_size: float = 0.05 + self, robot_id: WorldRobotID, path: JointPath, step_size: float = 0.05 ) -> bool: """Check if path is collision-free with interpolation.""" with self._world.scratch_context() as ctx: @@ -280,7 +284,7 @@ def is_path_valid( return True - def get_min_distance(self, robot_id: str) -> float: + def get_min_distance(self, robot_id: WorldRobotID) -> float: """Get minimum distance to obstacles for current state.""" with self._world.scratch_context() as ctx: return self._world.get_min_distance(ctx, robot_id) @@ -288,7 +292,7 @@ def get_min_distance(self, robot_id: str) -> float: # ============= Kinematics ============= def get_ee_pose( - self, robot_id: str, joint_positions: NDArray[np.float64] | None = None + self, robot_id: WorldRobotID, joint_positions: NDArray[np.float64] | None = None ) -> NDArray[np.float64]: """Get end-effector pose as 4x4 transform. Uses current state if positions is None.""" with self._world.scratch_context() as ctx: @@ -302,7 +306,7 @@ def get_ee_pose( return self._world.get_ee_pose(ctx, robot_id) def get_jacobian( - self, robot_id: str, joint_positions: NDArray[np.float64] + self, robot_id: WorldRobotID, joint_positions: NDArray[np.float64] ) -> NDArray[np.float64]: """Get 6xN Jacobian matrix.""" with self._world.scratch_context() as ctx: diff --git a/dimos/manipulation/planning/monitor/world_obstacle_monitor.py b/dimos/manipulation/planning/monitor/world_obstacle_monitor.py index 5a879e34fb..4401149282 100644 --- a/dimos/manipulation/planning/monitor/world_obstacle_monitor.py +++ b/dimos/manipulation/planning/monitor/world_obstacle_monitor.py @@ -34,16 +34,13 @@ Obstacle, ObstacleType, ) +from dimos.msgs.geometry_msgs import PoseStamped from dimos.utils.logging_config import setup_logger -from dimos.utils.transform_utils import pose_to_matrix if TYPE_CHECKING: from collections.abc import Callable import threading - import numpy as np - from numpy.typing import NDArray - from dimos.manipulation.planning.spec import WorldSpec from dimos.msgs.vision_msgs import Detection3D @@ -264,9 +261,13 @@ def on_detections(self, detections: list[Detection3D]) -> None: # Remove stale detections self._cleanup_stale_detections(current_time, seen_ids) - def _detection3d_to_pose(self, detection: Detection3D) -> NDArray[np.float64]: - """Convert Detection3D bbox.center to 4x4 transform.""" - return pose_to_matrix(detection.bbox.center) + def _detection3d_to_pose(self, detection: Detection3D) -> PoseStamped: + """Convert Detection3D bbox.center to PoseStamped.""" + center = detection.bbox.center + return PoseStamped( + position=center.position, + orientation=center.orientation, + ) def _detection_to_obstacle(self, detection: Detection3D) -> Obstacle: """Convert Detection3D to Obstacle.""" @@ -314,7 +315,7 @@ def add_static_obstacle( self, name: str, obstacle_type: str, - pose: NDArray[np.float64], + pose: PoseStamped, dimensions: tuple[float, ...], color: tuple[float, float, float, float] = (0.8, 0.2, 0.2, 0.8), ) -> str: @@ -323,7 +324,7 @@ def add_static_obstacle( Args: name: Unique name for the obstacle obstacle_type: Type ("box", "sphere", "cylinder") - pose: 4x4 homogeneous transform + pose: Pose of the obstacle in world frame dimensions: Type-specific dimensions color: RGBA color diff --git a/dimos/manipulation/planning/planners/rrt_planner.py b/dimos/manipulation/planning/planners/rrt_planner.py index 9ad2cad57f..ee81b64808 100644 --- a/dimos/manipulation/planning/planners/rrt_planner.py +++ b/dimos/manipulation/planning/planners/rrt_planner.py @@ -26,7 +26,13 @@ import numpy as np -from dimos.manipulation.planning.spec import PlanningResult, PlanningStatus, WorldSpec +from dimos.manipulation.planning.spec import ( + JointPath, + PlanningResult, + PlanningStatus, + WorldRobotID, + WorldSpec, +) from dimos.manipulation.planning.utils.path_utils import compute_path_length from dimos.utils.logging_config import setup_logger @@ -77,7 +83,7 @@ def __init__( def plan_joint_path( self, world: WorldSpec, - robot_id: str, + robot_id: WorldRobotID, q_start: NDArray[np.float64], q_goal: NDArray[np.float64], timeout: float = 10.0, @@ -135,7 +141,7 @@ def get_name(self) -> str: def _validate_inputs( self, world: WorldSpec, - robot_id: str, + robot_id: WorldRobotID, q_start: NDArray[np.float64], q_goal: NDArray[np.float64], ) -> PlanningResult | None: @@ -187,7 +193,7 @@ def _validate_inputs( def _extend_tree( self, world: WorldSpec, - robot_id: str, + robot_id: WorldRobotID, tree: list[TreeNode], target: NDArray[np.float64], step_size: float, @@ -219,7 +225,7 @@ def _extend_tree( def _connect_tree( self, world: WorldSpec, - robot_id: str, + robot_id: WorldRobotID, tree: list[TreeNode], target: NDArray[np.float64], step_size: float, @@ -240,7 +246,7 @@ def _extract_path( self, start_node: TreeNode, goal_node: TreeNode, - ) -> list[NDArray[np.float64]]: + ) -> JointPath: """Extract path from two connected nodes.""" # Path from start node to its root (reversed to be root->node) start_path = start_node.path_to_root() @@ -257,10 +263,10 @@ def _extract_path( def _simplify_path( self, world: WorldSpec, - robot_id: str, - path: list[NDArray[np.float64]], + robot_id: WorldRobotID, + path: JointPath, max_iterations: int = 100, - ) -> list[NDArray[np.float64]]: + ) -> JointPath: """Simplify path by random shortcutting.""" if len(path) <= 2: return path diff --git a/dimos/manipulation/planning/spec/__init__.py b/dimos/manipulation/planning/spec/__init__.py index 345f408e03..790b4ecf3f 100644 --- a/dimos/manipulation/planning/spec/__init__.py +++ b/dimos/manipulation/planning/spec/__init__.py @@ -24,14 +24,18 @@ from dimos.manipulation.planning.spec.types import ( CollisionObjectMessage, IKResult, + JointPath, Obstacle, PlanningResult, + RobotName, + WorldRobotID, ) __all__ = [ "CollisionObjectMessage", "IKResult", "IKStatus", + "JointPath", "KinematicsSpec", "Obstacle", "ObstacleType", @@ -39,5 +43,7 @@ "PlanningResult", "PlanningStatus", "RobotModelConfig", + "RobotName", + "WorldRobotID", "WorldSpec", ] diff --git a/dimos/manipulation/planning/spec/protocols.py b/dimos/manipulation/planning/spec/protocols.py index b8c2a22be5..04322711e6 100644 --- a/dimos/manipulation/planning/spec/protocols.py +++ b/dimos/manipulation/planning/spec/protocols.py @@ -29,7 +29,14 @@ from numpy.typing import NDArray from dimos.manipulation.planning.spec.config import RobotModelConfig - from dimos.manipulation.planning.spec.types import IKResult, Obstacle, PlanningResult + from dimos.manipulation.planning.spec.types import ( + IKResult, + JointPath, + Obstacle, + PlanningResult, + WorldRobotID, + ) + from dimos.msgs.geometry_msgs import Pose, PoseStamped @runtime_checkable @@ -51,19 +58,21 @@ class WorldSpec(Protocol): """ # Robot Management - def add_robot(self, config: RobotModelConfig) -> str: + def add_robot(self, config: RobotModelConfig) -> WorldRobotID: """Add a robot to the world. Returns unique robot ID.""" ... - def get_robot_ids(self) -> list[str]: + def get_robot_ids(self) -> list[WorldRobotID]: """Get all robot IDs.""" ... - def get_robot_config(self, robot_id: str) -> RobotModelConfig: + def get_robot_config(self, robot_id: WorldRobotID) -> RobotModelConfig: """Get robot configuration.""" ... - def get_joint_limits(self, robot_id: str) -> tuple[NDArray[np.float64], NDArray[np.float64]]: + def get_joint_limits( + self, robot_id: WorldRobotID + ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: """Get joint limits (lower, upper) for a robot.""" ... @@ -76,7 +85,7 @@ def remove_obstacle(self, obstacle_id: str) -> bool: """Remove an obstacle. Returns True if removed.""" ... - def update_obstacle_pose(self, obstacle_id: str, pose: NDArray[np.float64]) -> bool: + def update_obstacle_pose(self, obstacle_id: str, pose: PoseStamped) -> bool: """Update obstacle pose. Returns True if updated.""" ... @@ -103,36 +112,38 @@ def scratch_context(self) -> AbstractContextManager[Any]: """Get a scratch context for planning (thread-safe clone).""" ... - def sync_from_joint_state(self, robot_id: str, positions: NDArray[np.float64]) -> None: + def sync_from_joint_state(self, robot_id: WorldRobotID, positions: NDArray[np.float64]) -> None: """Sync live context from joint state.""" ... # State Operations (require context) - def set_positions(self, ctx: Any, robot_id: str, positions: NDArray[np.float64]) -> None: + def set_positions( + self, ctx: Any, robot_id: WorldRobotID, positions: NDArray[np.float64] + ) -> None: """Set robot joint positions in a context.""" ... - def get_positions(self, ctx: Any, robot_id: str) -> NDArray[np.float64]: + def get_positions(self, ctx: Any, robot_id: WorldRobotID) -> NDArray[np.float64]: """Get robot joint positions from a context.""" ... # Collision Checking (require context) - def is_collision_free(self, ctx: Any, robot_id: str) -> bool: + def is_collision_free(self, ctx: Any, robot_id: WorldRobotID) -> bool: """Check if robot configuration is collision-free.""" ... - def get_min_distance(self, ctx: Any, robot_id: str) -> float: + def get_min_distance(self, ctx: Any, robot_id: WorldRobotID) -> float: """Get minimum distance to obstacles (negative if collision).""" ... # Collision Checking (context-free, for planning) - def check_config_collision_free(self, robot_id: str, q: NDArray[np.float64]) -> bool: + def check_config_collision_free(self, robot_id: WorldRobotID, q: NDArray[np.float64]) -> bool: """Check if a configuration is collision-free (manages context internally).""" ... def check_edge_collision_free( self, - robot_id: str, + robot_id: WorldRobotID, q_start: NDArray[np.float64], q_end: NDArray[np.float64], step_size: float = 0.05, @@ -141,11 +152,11 @@ def check_edge_collision_free( ... # Forward Kinematics (require context) - def get_ee_pose(self, ctx: Any, robot_id: str) -> NDArray[np.float64]: + def get_ee_pose(self, ctx: Any, robot_id: WorldRobotID) -> NDArray[np.float64]: """Get end-effector pose (4x4 transform).""" ... - def get_jacobian(self, ctx: Any, robot_id: str) -> NDArray[np.float64]: + def get_jacobian(self, ctx: Any, robot_id: WorldRobotID) -> NDArray[np.float64]: """Get end-effector Jacobian (6 x n_joints).""" ... @@ -158,9 +169,7 @@ def publish_visualization(self, ctx: Any | None = None) -> None: """Publish current state to visualization.""" ... - def animate_path( - self, robot_id: str, path: list[NDArray[np.float64]], duration: float = 3.0 - ) -> None: + def animate_path(self, robot_id: WorldRobotID, path: JointPath, duration: float = 3.0) -> None: """Animate a path in visualization.""" ... @@ -172,8 +181,8 @@ class KinematicsSpec(Protocol): def solve( self, world: WorldSpec, - robot_id: str, - target_pose: NDArray[np.float64], + robot_id: WorldRobotID, + target_pose: PoseStamped, seed: NDArray[np.float64] | None = None, position_tolerance: float = 0.001, orientation_tolerance: float = 0.01, @@ -200,7 +209,7 @@ class PlannerSpec(Protocol): def plan_joint_path( self, world: WorldSpec, - robot_id: str, + robot_id: WorldRobotID, q_start: NDArray[np.float64], q_goal: NDArray[np.float64], timeout: float = 10.0, diff --git a/dimos/manipulation/planning/spec/types.py b/dimos/manipulation/planning/spec/types.py index c776b57299..8ce1515ae9 100644 --- a/dimos/manipulation/planning/spec/types.py +++ b/dimos/manipulation/planning/spec/types.py @@ -17,7 +17,7 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, TypeAlias from dimos.manipulation.planning.spec.enums import ( IKStatus, @@ -29,6 +29,26 @@ import numpy as np from numpy.typing import NDArray + from dimos.msgs.geometry_msgs import PoseStamped + +# ============================================================================= +# Semantic ID Types (documentation only, not enforced at runtime) +# ============================================================================= + +RobotName: TypeAlias = str +"""User-facing robot name (e.g., 'left_arm', 'right_arm')""" + +WorldRobotID: TypeAlias = str +"""Internal Drake world robot ID""" + +JointPath: TypeAlias = "list[NDArray[np.float64]]" +"""List of joint configurations forming a path""" + + +# ============================================================================= +# Data Classes +# ============================================================================= + @dataclass class Obstacle: @@ -37,7 +57,7 @@ class Obstacle: Attributes: name: Unique name for the obstacle obstacle_type: Type of geometry (BOX, SPHERE, CYLINDER, MESH) - pose: 4x4 homogeneous transform + pose: Pose of the obstacle in world frame dimensions: Type-specific dimensions: - BOX: (width, height, depth) - SPHERE: (radius,) @@ -49,7 +69,7 @@ class Obstacle: name: str obstacle_type: ObstacleType - pose: NDArray[np.float64] + pose: PoseStamped dimensions: tuple[float, ...] = () color: tuple[float, float, float, float] = (0.8, 0.2, 0.2, 0.8) mesh_path: str | None = None @@ -115,7 +135,7 @@ class CollisionObjectMessage: id: Unique identifier for the object operation: "add", "update", or "remove" primitive_type: "box", "sphere", or "cylinder" (for add/update) - pose: 4x4 transform (for add/update) + pose: Pose of the obstacle (for add/update) dimensions: Type-specific dimensions (for add/update) color: RGBA color tuple """ @@ -123,6 +143,6 @@ class CollisionObjectMessage: id: str operation: str # "add", "update", "remove" primitive_type: str | None = None - pose: NDArray[np.float64] | None = None + pose: PoseStamped | None = None dimensions: tuple[float, ...] | None = None color: tuple[float, float, float, float] = (0.8, 0.2, 0.2, 0.8) diff --git a/dimos/manipulation/planning/utils/mesh_utils.py b/dimos/manipulation/planning/utils/mesh_utils.py index 48f6452e15..05bc000739 100644 --- a/dimos/manipulation/planning/utils/mesh_utils.py +++ b/dimos/manipulation/planning/utils/mesh_utils.py @@ -205,7 +205,8 @@ def _resolve_package_uris( ) -> str: """Resolve package:// URIs to filesystem paths.""" # Pattern for package:// URIs (handles both single and double quotes) - pattern = r'package://([^/]+)/(.+?)(["\'\\s<>])' + # Note: Use triple quotes so \s is correctly interpreted as whitespace, not literal 's' + pattern = r"""package://([^/]+)/(.+?)(["'<>\s])""" def replace_uri(match: re.Match[str]) -> str: pkg_name = match.group(1) diff --git a/dimos/manipulation/planning/utils/path_utils.py b/dimos/manipulation/planning/utils/path_utils.py index cd1360598a..03eada3f34 100644 --- a/dimos/manipulation/planning/utils/path_utils.py +++ b/dimos/manipulation/planning/utils/path_utils.py @@ -35,13 +35,13 @@ if TYPE_CHECKING: from numpy.typing import NDArray - from dimos.manipulation.planning.spec import WorldSpec + from dimos.manipulation.planning.spec import JointPath, WorldRobotID, WorldSpec def interpolate_path( - path: list[NDArray[np.float64]], + path: JointPath, resolution: float = 0.05, -) -> list[NDArray[np.float64]]: +) -> JointPath: """Interpolate path to have uniform resolution. Adds intermediate waypoints so that the maximum joint-space distance @@ -125,11 +125,11 @@ def interpolate_segment( def simplify_path( world: WorldSpec, - robot_id: str, - path: list[NDArray[np.float64]], + robot_id: WorldRobotID, + path: JointPath, max_iterations: int = 100, collision_step_size: float = 0.02, -) -> list[NDArray[np.float64]]: +) -> JointPath: """Simplify path by removing unnecessary waypoints. Uses random shortcutting: randomly select two points and check if @@ -181,7 +181,7 @@ def simplify_path( return simplified -def compute_path_length(path: list[NDArray[np.float64]]) -> float: +def compute_path_length(path: JointPath) -> float: """Compute total path length in joint space. Sums the Euclidean distances between consecutive waypoints. @@ -207,7 +207,7 @@ def compute_path_length(path: list[NDArray[np.float64]]) -> float: def is_path_within_limits( - path: list[NDArray[np.float64]], + path: JointPath, lower_limits: NDArray[np.float64], upper_limits: NDArray[np.float64], ) -> bool: @@ -228,10 +228,10 @@ def is_path_within_limits( def clip_path_to_limits( - path: list[NDArray[np.float64]], + path: JointPath, lower_limits: NDArray[np.float64], upper_limits: NDArray[np.float64], -) -> list[NDArray[np.float64]]: +) -> JointPath: """Clip all waypoints in path to joint limits. Args: @@ -245,7 +245,7 @@ def clip_path_to_limits( return [np.clip(q, lower_limits, upper_limits) for q in path] -def reverse_path(path: list[NDArray[np.float64]]) -> list[NDArray[np.float64]]: +def reverse_path(path: JointPath) -> JointPath: """Reverse a path (for returning to start, etc.). Args: @@ -258,9 +258,9 @@ def reverse_path(path: list[NDArray[np.float64]]) -> list[NDArray[np.float64]]: def concatenate_paths( - *paths: list[NDArray[np.float64]], + *paths: JointPath, remove_duplicates: bool = True, -) -> list[NDArray[np.float64]]: +) -> JointPath: """Concatenate multiple paths into one. Args: diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index 05495251f5..8959c66e25 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -26,9 +26,12 @@ import numpy as np from dimos.manipulation.planning.spec import ( + JointPath, Obstacle, ObstacleType, RobotModelConfig, + WorldRobotID, + WorldSpec, ) from dimos.manipulation.planning.utils.mesh_utils import prepare_urdf_for_drake from dimos.utils.logging_config import setup_logger @@ -38,6 +41,8 @@ from numpy.typing import NDArray +from dimos.msgs.geometry_msgs import PoseStamped, Transform + try: from pydrake.geometry import ( # type: ignore[import-not-found] AddContactMaterial, @@ -78,7 +83,7 @@ class _RobotData: """Internal data for tracking a robot in the world.""" - robot_id: str + robot_id: WorldRobotID config: RobotModelConfig model_instance: Any # ModelInstanceIndex joint_indices: list[int] # Indices into plant's position vector @@ -96,7 +101,7 @@ class _ObstacleData: source_id: Any # SourceId -class DrakeWorld: +class DrakeWorld(WorldSpec): """Drake implementation of WorldSpec with MultibodyPlant, SceneGraph, optional Meshcat.""" def __init__(self, time_step: float = 0.0, enable_viz: bool = False): @@ -129,7 +134,7 @@ def __init__(self, time_step: float = 0.0, enable_viz: bool = False): self._obstacles_model_instance = self._plant.AddModelInstance("obstacles") # Tracking data - self._robots: dict[str, _RobotData] = {} + self._robots: dict[WorldRobotID, _RobotData] = {} self._obstacles: dict[str, _ObstacleData] = {} self._robot_counter = 0 self._obstacle_counter = 0 @@ -144,7 +149,7 @@ def __init__(self, time_step: float = 0.0, enable_viz: bool = False): # Obstacle source for dynamic obstacles self._obstacle_source_id: Any = None - def add_robot(self, config: RobotModelConfig) -> str: + def add_robot(self, config: RobotModelConfig) -> WorldRobotID: """Add a robot to the world. Returns robot_id.""" if self._finalized: raise RuntimeError("Cannot add robot after world is finalized") @@ -232,17 +237,19 @@ def _validate_joints(self, config: RobotModelConfig, model_instance: Any) -> Non except RuntimeError: raise ValueError(f"Joint '{joint_name}' not found in URDF") - def get_robot_ids(self) -> list[str]: + def get_robot_ids(self) -> list[WorldRobotID]: """Get all robot IDs in the world.""" return list(self._robots.keys()) - def get_robot_config(self, robot_id: str) -> RobotModelConfig: + def get_robot_config(self, robot_id: WorldRobotID) -> RobotModelConfig: """Get robot configuration by ID.""" if robot_id not in self._robots: raise KeyError(f"Robot '{robot_id}' not found") return self._robots[robot_id].config - def get_joint_limits(self, robot_id: str) -> tuple[NDArray[np.float64], NDArray[np.float64]]: + def get_joint_limits( + self, robot_id: WorldRobotID + ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: """Get joint limits (lower, upper) in radians.""" if robot_id not in self._robots: raise KeyError(f"Robot '{robot_id}' not found") @@ -313,7 +320,7 @@ def _add_obstacle_to_plant(self, obstacle: Obstacle, obstacle_id: str) -> Any: self._obstacles_model_instance, ) - transform = RigidTransform(obstacle.pose) + transform = self._pose_to_rigid_transform(obstacle.pose) geometry_id = self._plant.RegisterCollisionGeometry( body, RigidTransform(), @@ -345,7 +352,7 @@ def _add_obstacle_to_scene_graph(self, obstacle: Obstacle, obstacle_id: str) -> raise RuntimeError("Obstacle source not initialized") shape = self._create_shape(obstacle) - transform = RigidTransform(obstacle.pose) + transform = self._pose_to_rigid_transform(obstacle.pose) # MakePhongIllustrationProperties expects numpy array, not Rgba rgba_array = np.array(obstacle.color, dtype=np.float64) @@ -387,7 +394,7 @@ def _add_obstacle_to_meshcat(self, obstacle: Obstacle, obstacle_id: str) -> None # Use Drake's geometry types for Meshcat path = f"obstacles/{obstacle_id}" - transform = RigidTransform(obstacle.pose) + transform = self._pose_to_rigid_transform(obstacle.pose) rgba = Rgba(*obstacle.color) # Create Drake shape and add to Meshcat @@ -406,6 +413,14 @@ def _add_obstacle_to_meshcat(self, obstacle: Obstacle, obstacle_id: str) -> None self._meshcat.SetObject(path, shape, rgba) self._meshcat.SetTransform(path, transform) + def _pose_to_rigid_transform(self, pose: PoseStamped) -> Any: + """Convert PoseStamped to Drake RigidTransform.""" + pose_matrix = Transform( + translation=pose.position, + rotation=pose.orientation, + ).to_matrix() + return RigidTransform(pose_matrix) + def _create_shape(self, obstacle: Obstacle) -> Any: """Create Drake shape from obstacle specification.""" if obstacle.obstacle_type == ObstacleType.BOX: @@ -440,21 +455,19 @@ def remove_obstacle(self, obstacle_id: str) -> bool: logger.debug(f"Removed obstacle '{obstacle_id}'") return True - def update_obstacle_pose(self, obstacle_id: str, pose: NDArray[np.float64]) -> bool: - """Update obstacle pose (4x4 transform).""" + def update_obstacle_pose(self, obstacle_id: str, pose: PoseStamped) -> bool: + """Update obstacle pose.""" with self._lock: if obstacle_id not in self._obstacles: return False - if pose.shape != (4, 4): - raise ValueError(f"Pose must be 4x4, got {pose.shape}") - - self._obstacles[obstacle_id].obstacle.pose = pose.copy() + # Store PoseStamped directly + self._obstacles[obstacle_id].obstacle.pose = pose # Update Meshcat visualization if self._meshcat is not None: path = f"obstacles/{obstacle_id}" - transform = RigidTransform(pose) + transform = self._pose_to_rigid_transform(pose) self._meshcat.SetTransform(path, transform) # Note: SceneGraph geometry pose is fixed after registration @@ -603,7 +616,7 @@ def scratch_context(self) -> Generator[Context, None, None]: yield ctx - def sync_from_joint_state(self, robot_id: str, positions: NDArray[np.float64]) -> None: + def sync_from_joint_state(self, robot_id: WorldRobotID, positions: NDArray[np.float64]) -> None: """Sync live context from driver's joint state. Called by StateMonitor when new JointState arrives. @@ -620,7 +633,9 @@ def sync_from_joint_state(self, robot_id: str, positions: NDArray[np.float64]) - # ============= State Operations (context-based) ============= - def set_positions(self, ctx: Context, robot_id: str, positions: NDArray[np.float64]) -> None: + def set_positions( + self, ctx: Context, robot_id: WorldRobotID, positions: NDArray[np.float64] + ) -> None: """Set robot positions in given context.""" if not self._finalized: raise RuntimeError("World must be finalized first") @@ -630,7 +645,7 @@ def set_positions(self, ctx: Context, robot_id: str, positions: NDArray[np.float self._set_positions_internal(plant_ctx, robot_id, positions) def _set_positions_internal( - self, plant_ctx: Context, robot_id: str, positions: NDArray[np.float64] + self, plant_ctx: Context, robot_id: WorldRobotID, positions: NDArray[np.float64] ) -> None: """Internal: Set positions in a plant context.""" if robot_id not in self._robots: @@ -644,7 +659,7 @@ def _set_positions_internal( self._plant.SetPositions(plant_ctx, full_positions) - def get_positions(self, ctx: Context, robot_id: str) -> NDArray[np.float64]: + def get_positions(self, ctx: Context, robot_id: WorldRobotID) -> NDArray[np.float64]: """Get robot positions from given context.""" if not self._finalized: raise RuntimeError("World must be finalized first") @@ -660,7 +675,7 @@ def get_positions(self, ctx: Context, robot_id: str) -> NDArray[np.float64]: # ============= Collision Checking (context-based) ============= - def is_collision_free(self, ctx: Context, robot_id: str) -> bool: + def is_collision_free(self, ctx: Context, robot_id: WorldRobotID) -> bool: """Check if current configuration in context is collision-free.""" if not self._finalized: raise RuntimeError("World must be finalized first") @@ -673,7 +688,7 @@ def is_collision_free(self, ctx: Context, robot_id: str) -> bool: return not query_object.HasCollisions() # type: ignore[attr-defined] - def get_min_distance(self, ctx: Context, robot_id: str) -> float: + def get_min_distance(self, ctx: Context, robot_id: WorldRobotID) -> float: """Get minimum signed distance (positive = clearance, negative = penetration).""" if not self._finalized: raise RuntimeError("World must be finalized first") @@ -690,7 +705,7 @@ def get_min_distance(self, ctx: Context, robot_id: str) -> float: # ============= Collision Checking (context-free, for planning) ============= - def check_config_collision_free(self, robot_id: str, q: NDArray[np.float64]) -> bool: + def check_config_collision_free(self, robot_id: WorldRobotID, q: NDArray[np.float64]) -> bool: """Check if a configuration is collision-free (manages context internally). This is a convenience method for planners that don't need to manage contexts. @@ -701,7 +716,7 @@ def check_config_collision_free(self, robot_id: str, q: NDArray[np.float64]) -> def check_edge_collision_free( self, - robot_id: str, + robot_id: WorldRobotID, q_start: NDArray[np.float64], q_end: NDArray[np.float64], step_size: float = 0.05, @@ -731,7 +746,7 @@ def check_edge_collision_free( # ============= Forward Kinematics (context-based) ============= - def get_ee_pose(self, ctx: Context, robot_id: str) -> NDArray[np.float64]: + def get_ee_pose(self, ctx: Context, robot_id: WorldRobotID) -> NDArray[np.float64]: """Get end-effector pose as 4x4 transform.""" if not self._finalized: raise RuntimeError("World must be finalized first") @@ -748,7 +763,9 @@ def get_ee_pose(self, ctx: Context, robot_id: str) -> NDArray[np.float64]: result: NDArray[np.float64] = X_WE.GetAsMatrix4() return result - def get_link_pose(self, ctx: Context, robot_id: str, link_name: str) -> NDArray[np.float64]: + def get_link_pose( + self, ctx: Context, robot_id: WorldRobotID, link_name: str + ) -> NDArray[np.float64]: """Get link pose as 4x4 transform.""" if not self._finalized: raise RuntimeError("World must be finalized first") @@ -769,7 +786,7 @@ def get_link_pose(self, ctx: Context, robot_id: str, link_name: str) -> NDArray[ result: NDArray[np.float64] = X_WL.GetAsMatrix4() return result - def get_jacobian(self, ctx: Context, robot_id: str) -> NDArray[np.float64]: + def get_jacobian(self, ctx: Context, robot_id: WorldRobotID) -> NDArray[np.float64]: """Get geometric Jacobian (6 x n_joints). Rows: [vx, vy, vz, wx, wy, wz] (linear, then angular) @@ -833,8 +850,8 @@ def publish_visualization(self, ctx: Context | None = None) -> None: def animate_path( self, - robot_id: str, - path: list[NDArray[np.float64]], + robot_id: WorldRobotID, + path: JointPath, duration: float = 3.0, ) -> None: """Animate a path in Meshcat visualization. @@ -850,7 +867,7 @@ def animate_path( return # Capture current positions of all OTHER robots so they don't snap to zero - other_robot_positions: dict[str, NDArray[np.float64]] = {} + other_robot_positions: dict[WorldRobotID, NDArray[np.float64]] = {} for rid, _robot_data in self._robots.items(): if rid != robot_id: other_robot_positions[rid] = self.get_positions(self.get_live_context(), rid) diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index ccd7ef1850..e593a72472 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -25,6 +25,7 @@ "coordinator-piper-xarm": "dimos.control.blueprints:coordinator_piper_xarm", "coordinator-xarm6": "dimos.control.blueprints:coordinator_xarm6", "coordinator-xarm7": "dimos.control.blueprints:coordinator_xarm7", + "dual-xarm6-planner": "dimos.manipulation.manipulation_blueprints:dual_xarm6_planner", "demo-camera": "dimos.hardware.sensors.camera.module:demo_camera", "demo-error-on-name-conflicts": "dimos.robot.unitree_webrtc.demo_error_on_name_conflicts:demo_error_on_name_conflicts", "demo-google-maps-skill": "dimos.agents.skills.demo_google_maps_skill:demo_google_maps_skill", @@ -53,6 +54,8 @@ "unitree-go2-spatial": "dimos.robot.unitree_webrtc.unitree_go2_blueprints:unitree_go2_spatial", "unitree-go2-temporal-memory": "dimos.robot.unitree_webrtc.unitree_go2_blueprints:unitree_go2_temporal_memory", "unitree-go2-vlm-stream-test": "dimos.robot.unitree_webrtc.unitree_go2_blueprints:unitree_go2_vlm_stream_test", + "xarm6-planner-only": "dimos.manipulation.manipulation_blueprints:xarm6_planner_only", + "xarm7-planner-orchestrator": "dimos.manipulation.manipulation_blueprints:xarm7_planner_orchestrator", } @@ -77,6 +80,7 @@ "joint_trajectory_controller": "dimos.manipulation.control.trajectory_controller.joint_trajectory_controller", "keyboard_teleop": "dimos.robot.unitree_webrtc.keyboard_teleop", "llm_agent": "dimos.agents.agent", + "manipulation_module": "dimos.manipulation.manipulation_module", "mapper": "dimos.robot.unitree_webrtc.type.map", "navigation_skill": "dimos.agents.skills.navigation", "object_scene_registration_module": "dimos.perception.object_scene_registration", From 12ffac2b2d6d2ba16ae6bb4a23362a772ed8ece3 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Mon, 26 Jan 2026 17:43:32 -0800 Subject: [PATCH 22/31] updated all coordinator references --- dimos/e2e_tests/test_manipulation_module.py | 36 +++++------ dimos/manipulation/manipulation_blueprints.py | 48 +++++++------- dimos/manipulation/manipulation_module.py | 62 +++++++++---------- dimos/manipulation/planning/README.md | 22 +++---- .../planning/examples/__init__.py | 8 +-- .../planning/examples/manipulation_client.py | 12 ++-- .../planning/monitor/world_state_monitor.py | 14 ++--- dimos/manipulation/planning/spec/config.py | 30 ++++----- dimos/manipulation/test_manipulation_unit.py | 34 +++++----- dimos/robot/all_blueprints.py | 2 +- 10 files changed, 134 insertions(+), 134 deletions(-) diff --git a/dimos/e2e_tests/test_manipulation_module.py b/dimos/e2e_tests/test_manipulation_module.py index ce3105dfa3..7ae83d47be 100644 --- a/dimos/e2e_tests/test_manipulation_module.py +++ b/dimos/e2e_tests/test_manipulation_module.py @@ -86,7 +86,7 @@ def _get_xarm7_config() -> RobotModelConfig: "arm_joint6": "joint6", "arm_joint7": "joint7", }, - orchestrator_task_name="traj_arm", + coordinator_task_name="traj_arm", ) @@ -277,7 +277,7 @@ def test_robot_info(self, xarm7_config): assert info["name"] == "test_arm" assert len(info["joint_names"]) == 7 assert info["end_effector_link"] == "link7" - assert info["orchestrator_task_name"] == "traj_arm" + assert info["coordinator_task_name"] == "traj_arm" assert info["has_joint_name_mapping"] is True finally: module.stop() @@ -308,7 +308,7 @@ def test_ee_pose(self, xarm7_config, joint_state_zeros): module.stop() def test_trajectory_name_translation(self, xarm7_config, joint_state_zeros): - """Test that trajectory joint names are translated for orchestrator.""" + """Test that trajectory joint names are translated for coordinator.""" module = ManipulationModule( robots=[xarm7_config], planning_timeout=10.0, @@ -330,7 +330,7 @@ def test_trajectory_name_translation(self, xarm7_config, joint_state_zeros): robot_config = module._robots["test_arm"][1] # Translate it - translated = module._translate_trajectory_to_orchestrator(traj, robot_config) + translated = module._translate_trajectory_to_coordinator(traj, robot_config) # Verify names are translated for name in translated.joint_names: @@ -340,18 +340,18 @@ def test_trajectory_name_translation(self, xarm7_config, joint_state_zeros): # ============================================================================= -# Orchestrator Integration Tests (mocked orchestrator) +# Coordinator Integration Tests (mocked coordinator) # ============================================================================= @pytest.mark.skipif(not _drake_available(), reason="Drake not installed") @pytest.mark.skipif(not _xarm_urdf_available(), reason="XArm URDF not available") @pytest.mark.skipif(bool(os.getenv("CI")), reason="Skip in CI - requires LFS data") -class TestOrchestratorIntegration: - """Test orchestrator integration with mocked RPC client.""" +class TestCoordinatorIntegration: + """Test coordinator integration with mocked RPC client.""" - def test_execute_with_mock_orchestrator(self, xarm7_config, joint_state_zeros): - """Test execute sends trajectory to orchestrator.""" + def test_execute_with_mock_coordinator(self, xarm7_config, joint_state_zeros): + """Test execute sends trajectory to coordinator.""" module = ManipulationModule( robots=[xarm7_config], planning_timeout=10.0, @@ -368,10 +368,10 @@ def test_execute_with_mock_orchestrator(self, xarm7_config, joint_state_zeros): success = module.plan_to_joints([0.05] * 7) assert success is True - # Mock the orchestrator client + # Mock the coordinator client mock_client = MagicMock() mock_client.execute_trajectory.return_value = True - module._orchestrator_client = mock_client + module._coordinator_client = mock_client # Execute result = module.execute() @@ -379,7 +379,7 @@ def test_execute_with_mock_orchestrator(self, xarm7_config, joint_state_zeros): assert result is True assert module._state == ManipulationState.COMPLETED - # Verify orchestrator was called + # Verify coordinator was called mock_client.execute_trajectory.assert_called_once() call_args = mock_client.execute_trajectory.call_args task_name, trajectory = call_args[0] @@ -391,8 +391,8 @@ def test_execute_with_mock_orchestrator(self, xarm7_config, joint_state_zeros): finally: module.stop() - def test_execute_rejected_by_orchestrator(self, xarm7_config, joint_state_zeros): - """Test handling of orchestrator rejection.""" + def test_execute_rejected_by_coordinator(self, xarm7_config, joint_state_zeros): + """Test handling of coordinator rejection.""" module = ManipulationModule( robots=[xarm7_config], planning_timeout=10.0, @@ -408,10 +408,10 @@ def test_execute_rejected_by_orchestrator(self, xarm7_config, joint_state_zeros) # Plan a motion module.plan_to_joints([0.05] * 7) - # Mock orchestrator to reject + # Mock coordinator to reject mock_client = MagicMock() mock_client.execute_trajectory.return_value = False - module._orchestrator_client = mock_client + module._coordinator_client = mock_client # Execute result = module.execute() @@ -450,10 +450,10 @@ def test_state_transitions_during_execution(self, xarm7_config, joint_state_zero # Plan again module.plan_to_joints([0.05] * 7) - # Mock orchestrator + # Mock coordinator mock_client = MagicMock() mock_client.execute_trajectory.return_value = True - module._orchestrator_client = mock_client + module._coordinator_client = mock_client # Execute - should go to EXECUTING then COMPLETED module.execute() diff --git a/dimos/manipulation/manipulation_blueprints.py b/dimos/manipulation/manipulation_blueprints.py index 593ac94555..175533e8cb 100644 --- a/dimos/manipulation/manipulation_blueprints.py +++ b/dimos/manipulation/manipulation_blueprints.py @@ -13,11 +13,11 @@ # limitations under the License. """ -Blueprints for manipulation module integration with ControlOrchestrator. +Blueprints for manipulation module integration with ControlCoordinator. Usage: - # Start orchestrator first, then planner: - coordinator = xarm7_planner_orchestrator.build() + # Start coordinator first, then planner: + coordinator = xarm7_planner_coordinator.build() coordinator.loop() # Plan and execute via RPC client: @@ -107,7 +107,7 @@ def _make_xarm6_config( name: str = "arm", y_offset: float = 0.0, joint_prefix: str = "", - orchestrator_task: str | None = None, + coordinator_task: str | None = None, add_gripper: bool = True, ) -> RobotModelConfig: """Create XArm6 robot config. @@ -116,7 +116,7 @@ def _make_xarm6_config( name: Robot name in Drake world y_offset: Y-axis offset for base pose (for multi-arm setups) joint_prefix: Prefix for joint name mapping (e.g., "left_" or "right_") - orchestrator_task: Task name for orchestrator RPC execution + coordinator_task: Task name for coordinator RPC execution add_gripper: Whether to add the xarm gripper """ joint_names = ["joint1", "joint2", "joint3", "joint4", "joint5", "joint6"] @@ -152,7 +152,7 @@ def _make_xarm6_config( max_velocity=1.0, max_acceleration=2.0, joint_name_mapping=joint_mapping, - orchestrator_task_name=orchestrator_task, + coordinator_task_name=coordinator_task, ) @@ -160,7 +160,7 @@ def _make_xarm7_config( name: str = "arm", y_offset: float = 0.0, joint_prefix: str = "", - orchestrator_task: str | None = None, + coordinator_task: str | None = None, add_gripper: bool = False, ) -> RobotModelConfig: """Create XArm7 robot config. @@ -169,7 +169,7 @@ def _make_xarm7_config( name: Robot name in Drake world y_offset: Y-axis offset for base pose (for multi-arm setups) joint_prefix: Prefix for joint name mapping (e.g., "left_" or "right_") - orchestrator_task: Task name for orchestrator RPC execution + coordinator_task: Task name for coordinator RPC execution add_gripper: Whether to add the xarm gripper """ joint_names = ["joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "joint7"] @@ -205,7 +205,7 @@ def _make_xarm7_config( max_velocity=1.0, max_acceleration=2.0, joint_name_mapping=joint_mapping, - orchestrator_task_name=orchestrator_task, + coordinator_task_name=coordinator_task, ) @@ -213,7 +213,7 @@ def _make_piper_config( name: str = "piper", y_offset: float = 0.0, joint_prefix: str = "", - orchestrator_task: str | None = None, + coordinator_task: str | None = None, ) -> RobotModelConfig: """Create Piper robot config. @@ -221,7 +221,7 @@ def _make_piper_config( name: Robot name in Drake world y_offset: Y-axis offset for base pose (for multi-arm setups) joint_prefix: Prefix for joint name mapping (e.g., "piper_") - orchestrator_task: Task name for orchestrator RPC execution + coordinator_task: Task name for coordinator RPC execution Note: Piper has 6 revolute joints (joint1-joint6) for the arm and 2 prismatic @@ -253,7 +253,7 @@ def _make_piper_config( max_velocity=1.0, max_acceleration=2.0, joint_name_mapping=joint_mapping, - orchestrator_task_name=orchestrator_task, + coordinator_task_name=coordinator_task, ) @@ -262,7 +262,7 @@ def _make_piper_config( # ============================================================================= -# Single XArm6 planner (standalone, no orchestrator) +# Single XArm6 planner (standalone, no coordinator) xarm6_planner_only = manipulation_module( robots=[_make_xarm6_config()], planning_timeout=10.0, @@ -274,35 +274,35 @@ def _make_piper_config( ) -# Dual XArm6 planner with orchestrator integration -# Usage: Start with orchestrator_dual_mock, then plan/execute via RPC +# Dual XArm6 planner with coordinator integration +# Usage: Start with coordinator_dual_mock, then plan/execute via RPC dual_xarm6_planner = manipulation_module( robots=[ _make_xarm6_config( - "left_arm", y_offset=0.5, joint_prefix="left_", orchestrator_task="traj_left" + "left_arm", y_offset=0.5, joint_prefix="left_", coordinator_task="traj_left" ), _make_xarm6_config( - "right_arm", y_offset=-0.5, joint_prefix="right_", orchestrator_task="traj_right" + "right_arm", y_offset=-0.5, joint_prefix="right_", coordinator_task="traj_right" ), ], planning_timeout=10.0, enable_viz=True, ).transports( { - ("joint_state", JointState): LCMTransport("/orchestrator/joint_state", JointState), + ("joint_state", JointState): LCMTransport("/coordinator/joint_state", JointState), } ) -# Single XArm7 planner for orchestrator-mock -# Usage: dimos run orchestrator-mock, then dimos run xarm7-planner-orchestrator -xarm7_planner_orchestrator = manipulation_module( - robots=[_make_xarm7_config("arm", joint_prefix="arm_", orchestrator_task="traj_arm")], +# Single XArm7 planner for coordinator-mock +# Usage: dimos run coordinator-mock, then dimos run xarm7-planner-coordinator +xarm7_planner_coordinator = manipulation_module( + robots=[_make_xarm7_config("arm", joint_prefix="arm_", coordinator_task="traj_arm")], planning_timeout=10.0, enable_viz=True, ).transports( { - ("joint_state", JointState): LCMTransport("/orchestrator/joint_state", JointState), + ("joint_state", JointState): LCMTransport("/coordinator/joint_state", JointState), } ) @@ -312,5 +312,5 @@ def _make_piper_config( "XARM_GRIPPER_COLLISION_EXCLUSIONS", "dual_xarm6_planner", "xarm6_planner_only", - "xarm7_planner_orchestrator", + "xarm7_planner_coordinator", ] diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index bc5b4bd87a..e5e97feff4 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Manipulation Module - Motion planning with ControlOrchestrator execution.""" +"""Manipulation Module - Motion planning with ControlCoordinator execution.""" from __future__ import annotations @@ -90,14 +90,14 @@ class ManipulationModuleConfig(ModuleConfig): class ManipulationModule(Module): - """Motion planning module with ControlOrchestrator execution.""" + """Motion planning module with ControlCoordinator execution.""" default_config = ManipulationModuleConfig # Type annotation for the config attribute (mypy uses this) config: ManipulationModuleConfig - # Input: Joint state from orchestrator (for world sync) + # Input: Joint state from coordinator (for world sync) joint_state: In[JointState] def __init__(self, *args: object, **kwargs: object) -> None: @@ -120,8 +120,8 @@ def __init__(self, *args: object, **kwargs: object) -> None: self._planned_paths: PlannedPaths = {} self._planned_trajectories: PlannedTrajectories = {} - # Orchestrator integration (lazy initialized) - self._orchestrator_client: RPCClient | None = None + # Coordinator integration (lazy initialized) + self._coordinator_client: RPCClient | None = None logger.info("ManipulationModule initialized") @@ -507,57 +507,57 @@ def get_robot_info(self, robot_name: RobotName | None = None) -> dict[str, objec "max_velocity": config.max_velocity, "max_acceleration": config.max_acceleration, "has_joint_name_mapping": bool(config.joint_name_mapping), - "orchestrator_task_name": config.orchestrator_task_name, + "coordinator_task_name": config.coordinator_task_name, } # ========================================================================= - # Orchestrator Integration RPC Methods + # Coordinator Integration RPC Methods # ========================================================================= - def _get_orchestrator_client(self) -> RPCClient | None: - """Get or create orchestrator RPC client (lazy init).""" - if not any(c.orchestrator_task_name for _, c, _ in self._robots.values()): + def _get_coordinator_client(self) -> RPCClient | None: + """Get or create coordinator RPC client (lazy init).""" + if not any(c.coordinator_task_name for _, c, _ in self._robots.values()): return None - if self._orchestrator_client is None: - from dimos.control.orchestrator import ControlOrchestrator + if self._coordinator_client is None: + from dimos.control.coordinator import ControlCoordinator from dimos.core.rpc_client import RPCClient - self._orchestrator_client = RPCClient(None, ControlOrchestrator) - return self._orchestrator_client + self._coordinator_client = RPCClient(None, ControlCoordinator) + return self._coordinator_client - def _translate_trajectory_to_orchestrator( + def _translate_trajectory_to_coordinator( self, trajectory: JointTrajectory, robot_config: RobotModelConfig, ) -> JointTrajectory: - """Translate trajectory joint names from URDF to orchestrator namespace. + """Translate trajectory joint names from URDF to coordinator namespace. Args: trajectory: Trajectory with URDF joint names robot_config: Robot config with joint name mapping Returns: - Trajectory with orchestrator joint names + Trajectory with coordinator joint names """ if not robot_config.joint_name_mapping: return trajectory # No translation needed # Translate joint names - orchestrator_names = [ - robot_config.get_orchestrator_joint_name(j) for j in trajectory.joint_names + coordinator_names = [ + robot_config.get_coordinator_joint_name(j) for j in trajectory.joint_names ] # Create new trajectory with translated names # Note: duration is computed automatically from points in JointTrajectory.__init__ return JointTrajectory( - joint_names=orchestrator_names, + joint_names=coordinator_names, points=trajectory.points, timestamp=trajectory.timestamp, ) @rpc def execute(self, robot_name: RobotName | None = None) -> bool: - """Execute planned trajectory via ControlOrchestrator.""" + """Execute planned trajectory via ControlCoordinator.""" if (robot := self._get_robot(robot_name)) is None: return False robot_name, _, config, _ = robot @@ -565,25 +565,25 @@ def execute(self, robot_name: RobotName | None = None) -> bool: if (traj := self._planned_trajectories.get(robot_name)) is None: logger.warning("No planned trajectory") return False - if not config.orchestrator_task_name: - logger.error(f"No orchestrator_task_name for '{robot_name}'") + if not config.coordinator_task_name: + logger.error(f"No coordinator_task_name for '{robot_name}'") return False - if (client := self._get_orchestrator_client()) is None: - logger.error("No orchestrator client") + if (client := self._get_coordinator_client()) is None: + logger.error("No coordinator client") return False - translated = self._translate_trajectory_to_orchestrator(traj, config) + translated = self._translate_trajectory_to_coordinator(traj, config) logger.info( - f"Executing: task='{config.orchestrator_task_name}', {len(translated.points)} pts, {translated.duration:.2f}s" + f"Executing: task='{config.coordinator_task_name}', {len(translated.points)} pts, {translated.duration:.2f}s" ) self._state = ManipulationState.EXECUTING - if client.execute_trajectory(config.orchestrator_task_name, translated): + if client.execute_trajectory(config.coordinator_task_name, translated): logger.info("Trajectory accepted") self._state = ManipulationState.COMPLETED return True else: - return self._fail("Orchestrator rejected trajectory") + return self._fail("Coordinator rejected trajectory") @rpc def get_trajectory_status( @@ -593,9 +593,9 @@ def get_trajectory_status( if (robot := self._get_robot(robot_name)) is None: return None _, _, config, _ = robot - if not config.orchestrator_task_name or (client := self._get_orchestrator_client()) is None: + if not config.coordinator_task_name or (client := self._get_coordinator_client()) is None: return None - status = client.get_trajectory_status(config.orchestrator_task_name) + status = client.get_trajectory_status(config.coordinator_task_name) return dict(status) if status else None @property diff --git a/dimos/manipulation/planning/README.md b/dimos/manipulation/planning/README.md index 7ca4e9cd2d..eeb15e3859 100644 --- a/dimos/manipulation/planning/README.md +++ b/dimos/manipulation/planning/README.md @@ -5,11 +5,11 @@ Motion planning for robotic manipulators. Backend-agnostic design with Drake imp ## Quick Start ```bash -# Terminal 1: Mock orchestrator -dimos run orchestrator-mock +# Terminal 1: Mock coordinator +dimos run coordinator-mock # Terminal 2: Manipulation planner -dimos run xarm7-planner-orchestrator +dimos run xarm7-planner-coordinator # Terminal 3: IPython client python -m dimos.manipulation.planning.examples.manipulation_client @@ -20,7 +20,7 @@ In IPython: joints() # Get current joints plan([0.1] * 7) # Plan to target preview() # Preview in Meshcat (url() for link) -execute() # Execute via orchestrator +execute() # Execute via coordinator ``` ## Architecture @@ -67,8 +67,8 @@ config = RobotModelConfig( joint_names=["joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "joint7"], end_effector_link="link7", base_link="link_base", - joint_name_mapping={"arm_joint1": "joint1", ...}, # orchestrator <-> URDF - orchestrator_task_name="traj_arm", + joint_name_mapping={"arm_joint1": "joint1", ...}, # coordinator <-> URDF + coordinator_task_name="traj_arm", ) module = ManipulationModule( @@ -80,7 +80,7 @@ module = ManipulationModule( ) module.start() module.plan_to_joints([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]) -module.execute() # Sends to orchestrator +module.execute() # Sends to coordinator ``` ## RobotModelConfig Fields @@ -95,8 +95,8 @@ module.execute() # Sends to orchestrator | `base_link` | Base link name | | `max_velocity` | Max joint velocity (rad/s) | | `max_acceleration` | Max acceleration (rad/s²) | -| `joint_name_mapping` | Orchestrator → URDF name mapping | -| `orchestrator_task_name` | Task name for execution RPC | +| `joint_name_mapping` | Coordinator → URDF name mapping | +| `coordinator_task_name` | Task name for execution RPC | | `package_paths` | ROS package paths for meshes | | `xacro_args` | Xacro arguments (e.g., `{"dof": "7"}`) | @@ -125,8 +125,8 @@ module.execute() # Sends to orchestrator | Blueprint | Description | |-----------|-------------| -| `xarm6_planner_only` | XArm 6-DOF standalone (no orchestrator) | -| `xarm7-planner-orchestrator` | XArm 7-DOF with orchestrator | +| `xarm6_planner_only` | XArm 6-DOF standalone (no coordinator) | +| `xarm7-planner-coordinator` | XArm 7-DOF with coordinator | | `dual-xarm6-planner` | Dual XArm 6-DOF | ## Directory Structure diff --git a/dimos/manipulation/planning/examples/__init__.py b/dimos/manipulation/planning/examples/__init__.py index 6735749095..8c76a4a528 100644 --- a/dimos/manipulation/planning/examples/__init__.py +++ b/dimos/manipulation/planning/examples/__init__.py @@ -16,9 +16,9 @@ Manipulation Client - RPC client for ManipulationModule Usage: - # Start orchestrator and planner first: - dimos run orchestrator-mock - dimos run xarm7-planner-orchestrator + # Start coordinator and planner first: + dimos run coordinator-mock + dimos run xarm7-planner-coordinator # Then run the interactive client: python -m dimos.manipulation.planning.examples.manipulation_client @@ -27,7 +27,7 @@ c.joints() # Get current joint positions c.plan([0.1, ...]) # Plan to joints c.preview() # Preview in Meshcat - c.execute() # Execute via orchestrator + c.execute() # Execute via coordinator """ from dimos.manipulation.planning.examples.manipulation_client import ManipulationClient diff --git a/dimos/manipulation/planning/examples/manipulation_client.py b/dimos/manipulation/planning/examples/manipulation_client.py index aa207af948..ad5210b7e4 100644 --- a/dimos/manipulation/planning/examples/manipulation_client.py +++ b/dimos/manipulation/planning/examples/manipulation_client.py @@ -17,9 +17,9 @@ Manipulation Client - IPython interface for ManipulationModule Usage: - # Start orchestrator and planner first: - dimos run orchestrator-mock - dimos run xarm7-planner-orchestrator + # Start coordinator and planner first: + dimos run coordinator-mock + dimos run xarm7-planner-coordinator # Run interactive client: python -m dimos.manipulation.planning.examples.manipulation_client @@ -33,7 +33,7 @@ plan([0.1, ...]) # Plan to joint config plan_pose(x, y, z) # Plan to cartesian pose preview() # Preview path in Meshcat - execute() # Execute via orchestrator + execute() # Execute via coordinator box("name", x, y, z, w, h, d) # Add box obstacle sphere("name", x, y, z, r) # Add sphere obstacle @@ -172,7 +172,7 @@ def preview(self, duration: float = 3.0, robot_name: str | None = None) -> bool: return cast("bool", self._call("preview_path", duration, robot_name=robot_name)) def execute(self, robot_name: str | None = None) -> bool: - """Execute planned trajectory via orchestrator.""" + """Execute planned trajectory via coordinator.""" return cast("bool", self._call("execute", robot_name)) def has_plan(self) -> bool: @@ -312,7 +312,7 @@ def main() -> None: plan_pose(0.4, 0, 0.3) # Plan to cartesian pose (keeps orientation) plan_pose(0.4, 0, 0.3, roll=0, pitch=3.14, yaw=0) # With orientation preview() # Preview path in Meshcat - execute() # Execute via orchestrator + execute() # Execute via coordinator Obstacles: box("name", x, y, z, width, height, depth) # Add box diff --git a/dimos/manipulation/planning/monitor/world_state_monitor.py b/dimos/manipulation/planning/monitor/world_state_monitor.py index 38bb1575e1..7f2b4d36e4 100644 --- a/dimos/manipulation/planning/monitor/world_state_monitor.py +++ b/dimos/manipulation/planning/monitor/world_state_monitor.py @@ -79,7 +79,7 @@ def __init__( lock: Shared lock for thread-safe access robot_id: ID of the robot to monitor joint_names: Ordered list of joint names for this robot (URDF names) - joint_name_mapping: Maps orchestrator joint names to URDF joint names. + joint_name_mapping: Maps coordinator joint names to URDF joint names. Example: {"left_joint1": "joint1"} means messages with "left_joint1" will be mapped to URDF "joint1". If None, names must match exactly. timeout: Timeout for waiting for initial state (seconds) @@ -90,9 +90,9 @@ def __init__( self._joint_names = joint_names self._timeout = timeout - # Joint name mapping: orchestrator name -> URDF name + # Joint name mapping: coordinator name -> URDF name self._joint_name_mapping = joint_name_mapping or {} - # Build reverse mapping: URDF name -> orchestrator name + # Build reverse mapping: URDF name -> coordinator name self._reverse_mapping = {v: k for k, v in self._joint_name_mapping.items()} # Latest state @@ -185,16 +185,16 @@ def on_joint_state(self, msg: JointState) -> None: def _extract_positions(self, msg: JointState) -> NDArray[np.float64] | None: """Extract positions for our joints from JointState message. - Handles joint name translation from orchestrator namespace to URDF namespace. + Handles joint name translation from coordinator namespace to URDF namespace. If joint_name_mapping is set, message names are looked up via the reverse mapping. Args: - msg: JointState message (may use orchestrator joint names) + msg: JointState message (may use coordinator joint names) Returns: Array of joint positions or None if any joint is missing """ - # Build name->index map from message (orchestrator names) + # Build name->index map from message (coordinator names) name_to_idx = {name: i for i, name in enumerate(msg.name)} positions = [] @@ -203,7 +203,7 @@ def _extract_positions(self, msg: JointState) -> NDArray[np.float64] | None: if urdf_joint_name in name_to_idx: idx = name_to_idx[urdf_joint_name] else: - # Try reverse mapping: URDF name -> orchestrator name -> msg index + # Try reverse mapping: URDF name -> coordinator name -> msg index orch_name = self._reverse_mapping.get(urdf_joint_name) if orch_name is None or orch_name not in name_to_idx: return None # Missing joint diff --git a/dimos/manipulation/planning/spec/config.py b/dimos/manipulation/planning/spec/config.py index bb8706e6f2..600592dfa9 100644 --- a/dimos/manipulation/planning/spec/config.py +++ b/dimos/manipulation/planning/spec/config.py @@ -46,10 +46,10 @@ class RobotModelConfig: links may legitimately overlap (e.g., mimic joints). max_velocity: Maximum joint velocity for trajectory generation (rad/s) max_acceleration: Maximum joint acceleration for trajectory generation (rad/s^2) - joint_name_mapping: Maps orchestrator joint names to URDF joint names. - Example: {"left_joint1": "joint1"} means orchestrator's "left_joint1" + joint_name_mapping: Maps coordinator joint names to URDF joint names. + Example: {"left_joint1": "joint1"} means coordinator's "left_joint1" corresponds to URDF's "joint1". If empty, names are assumed to match. - orchestrator_task_name: Task name for executing trajectories via orchestrator RPC. + coordinator_task_name: Task name for executing trajectories via coordinator RPC. If set, trajectories can be executed via execute_trajectory() RPC. """ @@ -69,23 +69,23 @@ class RobotModelConfig: # Motion constraints for trajectory generation max_velocity: float = 1.0 max_acceleration: float = 2.0 - # Orchestrator integration + # Coordinator integration joint_name_mapping: dict[str, str] = field(default_factory=dict) - orchestrator_task_name: str | None = None + coordinator_task_name: str | None = None - def get_urdf_joint_name(self, orchestrator_name: str) -> str: - """Translate orchestrator joint name to URDF joint name.""" - return self.joint_name_mapping.get(orchestrator_name, orchestrator_name) + def get_urdf_joint_name(self, coordinator_name: str) -> str: + """Translate coordinator joint name to URDF joint name.""" + return self.joint_name_mapping.get(coordinator_name, coordinator_name) - def get_orchestrator_joint_name(self, urdf_name: str) -> str: - """Translate URDF joint name to orchestrator joint name.""" - for orch_name, u_name in self.joint_name_mapping.items(): + def get_coordinator_joint_name(self, urdf_name: str) -> str: + """Translate URDF joint name to coordinator joint name.""" + for coord_name, u_name in self.joint_name_mapping.items(): if u_name == urdf_name: - return orch_name + return coord_name return urdf_name - def get_orchestrator_joint_names(self) -> list[str]: - """Get joint names in orchestrator namespace.""" + def get_coordinator_joint_names(self) -> list[str]: + """Get joint names in coordinator namespace.""" if not self.joint_name_mapping: return self.joint_names - return [self.get_orchestrator_joint_name(j) for j in self.joint_names] + return [self.get_coordinator_joint_name(j) for j in self.joint_names] diff --git a/dimos/manipulation/test_manipulation_unit.py b/dimos/manipulation/test_manipulation_unit.py index b69f2cc286..719b12d8b0 100644 --- a/dimos/manipulation/test_manipulation_unit.py +++ b/dimos/manipulation/test_manipulation_unit.py @@ -46,7 +46,7 @@ def robot_config(): base_link="link_base", max_velocity=1.0, max_acceleration=2.0, - orchestrator_task_name="traj_arm", + coordinator_task_name="traj_arm", ) @@ -65,7 +65,7 @@ def robot_config_with_mapping(): "left_joint2": "joint2", "left_joint3": "joint3", }, - orchestrator_task_name="traj_left", + coordinator_task_name="traj_left", ) @@ -98,7 +98,7 @@ def _make_module(): module._world_monitor = None module._planner = None module._kinematics = None - module._orchestrator_client = None + module._coordinator_client = None return module @@ -199,25 +199,25 @@ def test_multiple_robots_require_name(self, robot_config): # ============================================================================= -# Test Joint Name Translation (for orchestrator integration) +# Test Joint Name Translation (for coordinator integration) # ============================================================================= class TestJointNameTranslation: - """Test trajectory joint name translation for orchestrator.""" + """Test trajectory joint name translation for coordinator.""" def test_no_mapping_returns_original(self, robot_config, simple_trajectory): """Without mapping, trajectory is returned unchanged.""" module = _make_module() - result = module._translate_trajectory_to_orchestrator(simple_trajectory, robot_config) + result = module._translate_trajectory_to_coordinator(simple_trajectory, robot_config) assert result is simple_trajectory # Same object def test_mapping_translates_names(self, robot_config_with_mapping, simple_trajectory): """With mapping, joint names are translated.""" module = _make_module() - result = module._translate_trajectory_to_orchestrator( + result = module._translate_trajectory_to_coordinator( simple_trajectory, robot_config_with_mapping ) assert result.joint_names == ["left_joint1", "left_joint2", "left_joint3"] @@ -230,7 +230,7 @@ def test_mapping_translates_names(self, robot_config_with_mapping, simple_trajec class TestExecute: - """Test orchestrator execution.""" + """Test coordinator execution.""" def test_execute_requires_trajectory(self, robot_config): """Execute fails without planned trajectory.""" @@ -241,7 +241,7 @@ def test_execute_requires_trajectory(self, robot_config): assert module.execute() is False def test_execute_requires_task_name(self): - """Execute fails without orchestrator_task_name.""" + """Execute fails without coordinator_task_name.""" module = _make_module() config_no_task = RobotModelConfig( name="arm", @@ -256,14 +256,14 @@ def test_execute_requires_task_name(self): assert module.execute() is False def test_execute_success(self, robot_config, simple_trajectory): - """Successful execute calls orchestrator.""" + """Successful execute calls coordinator.""" module = _make_module() module._robots = {"test_arm": ("id", robot_config, MagicMock())} module._planned_trajectories = {"test_arm": simple_trajectory} mock_client = MagicMock() mock_client.execute_trajectory.return_value = True - module._orchestrator_client = mock_client + module._coordinator_client = mock_client assert module.execute() is True assert module._state == ManipulationState.COMPLETED @@ -277,7 +277,7 @@ def test_execute_rejected(self, robot_config, simple_trajectory): mock_client = MagicMock() mock_client.execute_trajectory.return_value = False - module._orchestrator_client = mock_client + module._coordinator_client = mock_client assert module.execute() is False assert module._state == ManipulationState.FAULT @@ -292,13 +292,13 @@ class TestRobotModelConfigMapping: """Test RobotModelConfig joint name mapping helpers.""" def test_bidirectional_mapping(self, robot_config_with_mapping): - """Test URDF <-> orchestrator name translation.""" + """Test URDF <-> coordinator name translation.""" config = robot_config_with_mapping - # Orchestrator -> URDF + # Coordinator -> URDF assert config.get_urdf_joint_name("left_joint1") == "joint1" assert config.get_urdf_joint_name("unknown") == "unknown" - # URDF -> Orchestrator - assert config.get_orchestrator_joint_name("joint1") == "left_joint1" - assert config.get_orchestrator_joint_name("unknown") == "unknown" + # URDF -> Coordinator + assert config.get_coordinator_joint_name("joint1") == "left_joint1" + assert config.get_coordinator_joint_name("unknown") == "unknown" diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index e593a72472..dd18af6212 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -55,7 +55,7 @@ "unitree-go2-temporal-memory": "dimos.robot.unitree_webrtc.unitree_go2_blueprints:unitree_go2_temporal_memory", "unitree-go2-vlm-stream-test": "dimos.robot.unitree_webrtc.unitree_go2_blueprints:unitree_go2_vlm_stream_test", "xarm6-planner-only": "dimos.manipulation.manipulation_blueprints:xarm6_planner_only", - "xarm7-planner-orchestrator": "dimos.manipulation.manipulation_blueprints:xarm7_planner_orchestrator", + "xarm7-planner-coordinator": "dimos.manipulation.manipulation_blueprints:xarm7_planner_coordinator", } From 937da06bef363c1b37e25518d4f5ea9c7687fad7 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Mon, 26 Jan 2026 18:12:11 -0800 Subject: [PATCH 23/31] fixed all bluepirnts order --- dimos/robot/all_blueprints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index dd18af6212..00063be231 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -25,7 +25,6 @@ "coordinator-piper-xarm": "dimos.control.blueprints:coordinator_piper_xarm", "coordinator-xarm6": "dimos.control.blueprints:coordinator_xarm6", "coordinator-xarm7": "dimos.control.blueprints:coordinator_xarm7", - "dual-xarm6-planner": "dimos.manipulation.manipulation_blueprints:dual_xarm6_planner", "demo-camera": "dimos.hardware.sensors.camera.module:demo_camera", "demo-error-on-name-conflicts": "dimos.robot.unitree_webrtc.demo_error_on_name_conflicts:demo_error_on_name_conflicts", "demo-google-maps-skill": "dimos.agents.skills.demo_google_maps_skill:demo_google_maps_skill", @@ -33,6 +32,7 @@ "demo-object-scene-registration": "dimos.perception.demo_object_scene_registration:demo_object_scene_registration", "demo-osm": "dimos.mapping.osm.demo_osm:demo_osm", "demo-skill": "dimos.agents.skills.demo_skill:demo_skill", + "dual-xarm6-planner": "dimos.manipulation.manipulation_blueprints:dual_xarm6_planner", "unitree-g1": "dimos.robot.unitree_webrtc.unitree_g1_blueprints:unitree_g1", "unitree-g1-agentic": "dimos.robot.unitree_webrtc.unitree_g1_blueprints:unitree_g1_agentic", "unitree-g1-agentic-sim": "dimos.robot.unitree_webrtc.unitree_g1_blueprints:unitree_g1_agentic_sim", From e2bdc4baf1381aee9832f972aa2fcc4aba526733 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Mon, 26 Jan 2026 20:28:03 -0800 Subject: [PATCH 24/31] updated paths to use Path object instead of str --- dimos/e2e_tests/test_manipulation_module.py | 8 +- dimos/manipulation/manipulation_blueprints.py | 74 ++++++++++--------- dimos/manipulation/planning/README.md | 7 +- dimos/manipulation/planning/spec/config.py | 14 ++-- dimos/manipulation/planning/spec/types.py | 7 ++ .../planning/world/drake_world.py | 7 +- dimos/manipulation/test_manipulation_unit.py | 14 ++-- 7 files changed, 74 insertions(+), 57 deletions(-) diff --git a/dimos/e2e_tests/test_manipulation_module.py b/dimos/e2e_tests/test_manipulation_module.py index 7ae83d47be..c4f25cf6fe 100644 --- a/dimos/e2e_tests/test_manipulation_module.py +++ b/dimos/e2e_tests/test_manipulation_module.py @@ -33,7 +33,7 @@ ManipulationState, ) from dimos.manipulation.planning.spec import RobotModelConfig -from dimos.msgs.geometry_msgs import Pose, Quaternion, Vector3 +from dimos.msgs.geometry_msgs import Pose, PoseStamped, Quaternion, Vector3 from dimos.msgs.sensor_msgs import JointState from dimos.utils.data import get_data @@ -67,12 +67,12 @@ def _get_xarm7_config() -> RobotModelConfig: desc_path = get_data("xarm_description") return RobotModelConfig( name="test_arm", - urdf_path=str(desc_path / "urdf/xarm_device.urdf.xacro"), - base_pose=np.eye(4, dtype=np.float64), + urdf_path=desc_path / "urdf/xarm_device.urdf.xacro", + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), joint_names=["joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "joint7"], end_effector_link="link7", base_link="link_base", - package_paths={"xarm_description": str(desc_path)}, + package_paths={"xarm_description": desc_path}, xacro_args={"dof": "7", "limited": "true"}, auto_convert_meshes=True, max_velocity=1.0, diff --git a/dimos/manipulation/manipulation_blueprints.py b/dimos/manipulation/manipulation_blueprints.py index 175533e8cb..7eac1b6139 100644 --- a/dimos/manipulation/manipulation_blueprints.py +++ b/dimos/manipulation/manipulation_blueprints.py @@ -27,37 +27,63 @@ client.execute() """ -import numpy as np +from pathlib import Path from dimos.core.transport import LCMTransport from dimos.manipulation.manipulation_module import manipulation_module from dimos.manipulation.planning.spec import RobotModelConfig +from dimos.msgs.geometry_msgs import PoseStamped, Quaternion, Vector3 from dimos.msgs.sensor_msgs import JointState from dimos.utils.data import get_data +# ============================================================================= +# Pose Helpers +# ============================================================================= + + +def _make_base_pose( + x: float = 0.0, + y: float = 0.0, + z: float = 0.0, + roll: float = 0.0, + pitch: float = 0.0, + yaw: float = 0.0, +) -> PoseStamped: + """Create a base pose with optional xyz offset and rpy orientation. + + Args: + x, y, z: Position offset in meters + roll, pitch, yaw: Orientation in radians (Euler angles) + """ + return PoseStamped( + position=Vector3(x=x, y=y, z=z), + orientation=Quaternion.from_euler(Vector3(x=roll, y=pitch, z=yaw)), + ) + + # ============================================================================= # URDF Helpers # ============================================================================= -def _get_xarm_urdf_path() -> str: +def _get_xarm_urdf_path() -> Path: """Get path to xarm URDF.""" - return str(get_data("xarm_description") / "urdf/xarm_device.urdf.xacro") + return get_data("xarm_description") / "urdf/xarm_device.urdf.xacro" -def _get_xarm_package_paths() -> dict[str, str]: +def _get_xarm_package_paths() -> dict[str, Path]: """Get package paths for xarm xacro resolution.""" - return {"xarm_description": str(get_data("xarm_description"))} + return {"xarm_description": get_data("xarm_description")} -def _get_piper_urdf_path() -> str: +def _get_piper_urdf_path() -> Path: """Get path to piper URDF.""" - return str(get_data("piper_description") / "urdf/piper_description.xacro") + return get_data("piper_description") / "urdf/piper_description.xacro" -def _get_piper_package_paths() -> dict[str, str]: +def _get_piper_package_paths() -> dict[str, Path]: """Get package paths for piper xacro resolution.""" - return {"piper_description": str(get_data("piper_description"))} + return {"piper_description": get_data("piper_description")} # Piper gripper collision exclusions (parallel jaw gripper) @@ -133,15 +159,7 @@ def _make_xarm6_config( return RobotModelConfig( name=name, urdf_path=_get_xarm_urdf_path(), - base_pose=np.array( - [ - [1, 0, 0, 0], - [0, 1, 0, y_offset], - [0, 0, 1, 0], - [0, 0, 0, 1], - ], - dtype=np.float64, - ), + base_pose=_make_base_pose(y=y_offset), joint_names=joint_names, end_effector_link="link_tcp" if add_gripper else "link6", base_link="link_base", @@ -186,15 +204,7 @@ def _make_xarm7_config( return RobotModelConfig( name=name, urdf_path=_get_xarm_urdf_path(), - base_pose=np.array( - [ - [1, 0, 0, 0], - [0, 1, 0, y_offset], - [0, 0, 1, 0], - [0, 0, 0, 1], - ], - dtype=np.float64, - ), + base_pose=_make_base_pose(y=y_offset), joint_names=joint_names, end_effector_link="link_tcp" if add_gripper else "link7", base_link="link_base", @@ -234,15 +244,7 @@ def _make_piper_config( return RobotModelConfig( name=name, urdf_path=_get_piper_urdf_path(), - base_pose=np.array( - [ - [1, 0, 0, 0], - [0, 1, 0, y_offset], - [0, 0, 1, 0], - [0, 0, 0, 1], - ], - dtype=np.float64, - ), + base_pose=_make_base_pose(y=y_offset), joint_names=joint_names, end_effector_link="gripper_base", # End of arm, before gripper fingers base_link="arm_base", diff --git a/dimos/manipulation/planning/README.md b/dimos/manipulation/planning/README.md index eeb15e3859..803d8b166e 100644 --- a/dimos/manipulation/planning/README.md +++ b/dimos/manipulation/planning/README.md @@ -57,13 +57,14 @@ execute() # Execute via coordinator ## Using ManipulationModule ```python +from pathlib import Path from dimos.manipulation import ManipulationModule from dimos.manipulation.planning.spec import RobotModelConfig config = RobotModelConfig( name="xarm7", - urdf_path="/path/to/xarm7.urdf", - base_pose=np.eye(4), + urdf_path=Path("/path/to/xarm7.urdf"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), joint_names=["joint1", "joint2", "joint3", "joint4", "joint5", "joint6", "joint7"], end_effector_link="link7", base_link="link_base", @@ -89,7 +90,7 @@ module.execute() # Sends to coordinator |-------|-------------| | `name` | Robot identifier | | `urdf_path` | Path to URDF/XACRO file | -| `base_pose` | 4x4 transform for robot base | +| `base_pose` | PoseStamped for robot base in world frame | | `joint_names` | Joint names in URDF | | `end_effector_link` | EE link name | | `base_link` | Base link name | diff --git a/dimos/manipulation/planning/spec/config.py b/dimos/manipulation/planning/spec/config.py index 600592dfa9..ef8d290c43 100644 --- a/dimos/manipulation/planning/spec/config.py +++ b/dimos/manipulation/planning/spec/config.py @@ -20,9 +20,13 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: + from pathlib import Path + import numpy as np from numpy.typing import NDArray + from dimos.msgs.geometry_msgs import PoseStamped + @dataclass class RobotModelConfig: @@ -31,11 +35,11 @@ class RobotModelConfig: Attributes: name: Human-readable robot name urdf_path: Path to URDF file (can be .urdf or .xacro) - base_pose: 4x4 homogeneous transform for robot base + base_pose: Pose of robot base in world frame (position + orientation) joint_names: Ordered list of controlled joint names (in URDF namespace) end_effector_link: Name of the end-effector link for FK/IK base_link: Name of the base link (default: "base_link") - package_paths: Dict mapping package names to filesystem paths + package_paths: Dict mapping package names to filesystem Paths joint_limits_lower: Lower joint limits (radians) joint_limits_upper: Upper joint limits (radians) velocity_limits: Joint velocity limits (rad/s) @@ -54,12 +58,12 @@ class RobotModelConfig: """ name: str - urdf_path: str - base_pose: NDArray[np.float64] + urdf_path: Path + base_pose: PoseStamped joint_names: list[str] end_effector_link: str base_link: str = "base_link" - package_paths: dict[str, str] = field(default_factory=dict) + package_paths: dict[str, Path] = field(default_factory=dict) joint_limits_lower: NDArray[np.float64] | None = None joint_limits_upper: NDArray[np.float64] | None = None velocity_limits: NDArray[np.float64] | None = None diff --git a/dimos/manipulation/planning/spec/types.py b/dimos/manipulation/planning/spec/types.py index 8ce1515ae9..ec6e3cdaf3 100644 --- a/dimos/manipulation/planning/spec/types.py +++ b/dimos/manipulation/planning/spec/types.py @@ -111,6 +111,10 @@ class PlanningResult: path_length: Total path length in joint space (radians) iterations: Number of iterations/nodes expanded message: Human-readable status message + timestamps: Optional timestamps for each waypoint (seconds from start). + If provided by the planner, trajectory generator can use these directly. + velocities: Optional joint velocities at each waypoint. + If provided by the planner, trajectory generator can use these directly. """ status: PlanningStatus @@ -119,6 +123,9 @@ class PlanningResult: path_length: float = 0.0 iterations: int = 0 message: str = "" + # Optional timing fields (set by optimization-based planners) + timestamps: list[float] | None = None + velocities: list[NDArray[np.float64]] | None = None def is_success(self) -> bool: """Check if planning was successful.""" diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index 8959c66e25..d3cf4973e7 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -181,7 +181,7 @@ def add_robot(self, config: RobotModelConfig) -> WorldRobotID: def _load_urdf(self, config: RobotModelConfig) -> Any: """Load URDF/xacro and return model instance.""" - original_path = Path(config.urdf_path).resolve() + original_path = config.urdf_path.resolve() if not original_path.exists(): raise FileNotFoundError(f"URDF/xacro not found: {original_path}") @@ -197,7 +197,7 @@ def _load_urdf(self, config: RobotModelConfig) -> Any: # Register package paths if config.package_paths: for pkg_name, pkg_path in config.package_paths.items(): - self._parser.package_map().Add(pkg_name, Path(pkg_path)) + self._parser.package_map().Add(pkg_name, pkg_path) else: self._parser.package_map().Add(f"{config.name}_description", urdf_path_obj.parent) @@ -223,10 +223,11 @@ def _weld_base_if_needed(self, config: RobotModelConfig, model_instance: Any) -> pass # Weld base to world + base_transform = self._pose_to_rigid_transform(config.base_pose) self._plant.WeldFrames( self._plant.world_frame(), base_body.body_frame(), - RigidTransform(config.base_pose), + base_transform, ) def _validate_joints(self, config: RobotModelConfig, model_instance: Any) -> None: diff --git a/dimos/manipulation/test_manipulation_unit.py b/dimos/manipulation/test_manipulation_unit.py index 719b12d8b0..e81dc23c43 100644 --- a/dimos/manipulation/test_manipulation_unit.py +++ b/dimos/manipulation/test_manipulation_unit.py @@ -16,6 +16,7 @@ from __future__ import annotations +from pathlib import Path import threading from unittest.mock import MagicMock, patch @@ -27,6 +28,7 @@ ManipulationState, ) from dimos.manipulation.planning.spec import RobotModelConfig +from dimos.msgs.geometry_msgs import PoseStamped, Quaternion, Vector3 from dimos.msgs.trajectory_msgs import JointTrajectory, TrajectoryPoint # ============================================================================= @@ -39,8 +41,8 @@ def robot_config(): """Create a robot config for testing.""" return RobotModelConfig( name="test_arm", - urdf_path="/path/to/robot.urdf", - base_pose=np.eye(4, dtype=np.float64), + urdf_path=Path("/path/to/robot.urdf"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), joint_names=["joint1", "joint2", "joint3"], end_effector_link="link_tcp", base_link="link_base", @@ -55,8 +57,8 @@ def robot_config_with_mapping(): """Create a robot config with joint name mapping (dual-arm scenario).""" return RobotModelConfig( name="left_arm", - urdf_path="/path/to/robot.urdf", - base_pose=np.eye(4, dtype=np.float64), + urdf_path=Path("/path/to/robot.urdf"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), joint_names=["joint1", "joint2", "joint3"], end_effector_link="link_tcp", base_link="link_base", @@ -245,8 +247,8 @@ def test_execute_requires_task_name(self): module = _make_module() config_no_task = RobotModelConfig( name="arm", - urdf_path="/path", - base_pose=np.eye(4), + urdf_path=Path("/path"), + base_pose=PoseStamped(position=Vector3(), orientation=Quaternion()), joint_names=["j1"], end_effector_link="ee", ) From 9469a2cb01802d08c760b9abc6355fca74ea2e2a Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Mon, 26 Jan 2026 20:49:44 -0800 Subject: [PATCH 25/31] mypy type fixes --- dimos/manipulation/planning/utils/mesh_utils.py | 8 ++++---- dimos/manipulation/planning/world/drake_world.py | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/dimos/manipulation/planning/utils/mesh_utils.py b/dimos/manipulation/planning/utils/mesh_utils.py index 05bc000739..316d0e54d1 100644 --- a/dimos/manipulation/planning/utils/mesh_utils.py +++ b/dimos/manipulation/planning/utils/mesh_utils.py @@ -48,7 +48,7 @@ def prepare_urdf_for_drake( urdf_path: Path | str, - package_paths: dict[str, str] | None = None, + package_paths: dict[str, Path] | None = None, xacro_args: dict[str, str] | None = None, convert_meshes: bool = False, ) -> str: @@ -108,7 +108,7 @@ def prepare_urdf_for_drake( def _generate_cache_key( urdf_path: Path, - package_paths: dict[str, str], + package_paths: dict[str, Path], xacro_args: dict[str, str], convert_meshes: bool, ) -> str: @@ -129,7 +129,7 @@ def _generate_cache_key( def _process_xacro( xacro_path: Path, - package_paths: dict[str, str], + package_paths: dict[str, Path], xacro_args: dict[str, str], ) -> str: """Process xacro file to URDF.""" @@ -200,7 +200,7 @@ def _strip_transmission_blocks(urdf_content: str) -> str: def _resolve_package_uris( urdf_content: str, - package_paths: dict[str, str], + package_paths: dict[str, Path], output_dir: Path, ) -> str: """Resolve package:// URIs to filesystem paths.""" diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index d3cf4973e7..a056c47b26 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -197,7 +197,7 @@ def _load_urdf(self, config: RobotModelConfig) -> Any: # Register package paths if config.package_paths: for pkg_name, pkg_path in config.package_paths.items(): - self._parser.package_map().Add(pkg_name, pkg_path) + self._parser.package_map().Add(pkg_name, Path(pkg_path)) else: self._parser.package_map().Add(f"{config.name}_description", urdf_path_obj.parent) From 89006716907c7a7ee7ae5756a4bbf26333005fb1 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Thu, 29 Jan 2026 18:26:52 -0800 Subject: [PATCH 26/31] all public api boundaries use standard message types, sim specific internal storage/helpers still use NDArrays in some places --- dimos/manipulation/manipulation_module.py | 42 +++--- .../kinematics/drake_optimization_ik.py | 38 +++--- .../planning/kinematics/jacobian_ik.py | 123 +++++++++++------- .../planning/monitor/world_monitor.py | 87 +++++++------ .../monitor/world_obstacle_monitor.py | 2 +- .../planning/monitor/world_state_monitor.py | 21 +-- .../planning/planners/rrt_planner.py | 55 +++++--- dimos/manipulation/planning/spec/config.py | 9 +- dimos/manipulation/planning/spec/protocols.py | 37 +++--- dimos/manipulation/planning/spec/types.py | 22 ++-- .../manipulation/planning/utils/path_utils.py | 101 +++++++------- .../planning/world/drake_world.py | 83 +++++++----- 12 files changed, 357 insertions(+), 263 deletions(-) diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index e5e97feff4..ef92b17760 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -21,8 +21,6 @@ import threading from typing import TYPE_CHECKING, TypeAlias -import numpy as np - from dimos.core import In, Module, rpc from dimos.core.module import ModuleConfig from dimos.manipulation.planning import ( @@ -41,14 +39,11 @@ from dimos.manipulation.planning.monitor import WorldMonitor # These must be imported at runtime (not TYPE_CHECKING) for In/Out port creation -from dimos.msgs.sensor_msgs import JointState # noqa: TC001 +from dimos.msgs.sensor_msgs import JointState from dimos.msgs.trajectory_msgs import JointTrajectory from dimos.utils.logging_config import setup_logger -from dimos.utils.transform_utils import matrix_to_pose if TYPE_CHECKING: - from numpy.typing import NDArray - from dimos.core.rpc_client import RPCClient from dimos.msgs.geometry_msgs import Pose, PoseStamped @@ -259,9 +254,9 @@ def get_current_joints(self, robot_name: RobotName | None = None) -> list[float] robot_name: Robot to query (required if multiple robots configured) """ if (robot := self._get_robot(robot_name)) and self._world_monitor: - pos = self._world_monitor.get_current_positions(robot[1]) - if pos is not None: - return list(pos) + state = self._world_monitor.get_current_joint_state(robot[1]) + if state is not None: + return list(state.position) return None @rpc @@ -272,7 +267,7 @@ def get_ee_pose(self, robot_name: RobotName | None = None) -> Pose | None: robot_name: Robot to query (required if multiple robots configured) """ if (robot := self._get_robot(robot_name)) and self._world_monitor: - return matrix_to_pose(self._world_monitor.get_ee_pose(robot[1])) + return self._world_monitor.get_ee_pose(robot[1], joint_state=None) return None @rpc @@ -284,9 +279,9 @@ def is_collision_free(self, joints: list[float], robot_name: RobotName | None = robot_name: Robot to check (required if multiple robots configured) """ if (robot := self._get_robot(robot_name)) and self._world_monitor: - return self._world_monitor.is_state_valid( - robot[1], np.array(joints) - ) # robot[1] is the robot_id. + _, robot_id, config, _ = robot + joint_state = JointState(name=config.joint_names, position=joints) + return self._world_monitor.is_state_valid(robot_id, joint_state) return False # ========================================================================= @@ -333,7 +328,7 @@ def plan_to_pose(self, pose: Pose, robot_name: RobotName | None = None) -> bool: robot_name, robot_id = r assert self._world_monitor # guaranteed by _begin_planning - current = self._world_monitor.get_current_positions(robot_id) + current = self._world_monitor.get_current_joint_state(robot_id) if current is None: return self._fail("No joint state") @@ -353,11 +348,11 @@ def plan_to_pose(self, pose: Pose, robot_name: RobotName | None = None) -> bool: seed=current, check_collision=True, ) - if not ik.is_success() or ik.joint_positions is None: + if not ik.is_success() or ik.joint_state is None: return self._fail(f"IK failed: {ik.status.name}") logger.info(f"IK solved, error: {ik.position_error:.4f}m") - return self._plan_path_only(robot_name, robot_id, ik.joint_positions) + return self._plan_path_only(robot_name, robot_id, ik.joint_state) @rpc def plan_to_joints(self, joints: list[float], robot_name: RobotName | None = None) -> bool: @@ -370,23 +365,25 @@ def plan_to_joints(self, joints: list[float], robot_name: RobotName | None = Non if (r := self._begin_planning(robot_name)) is None: return False robot_name, robot_id = r + _, config, _ = self._robots[robot_name] + goal_state = JointState(name=config.joint_names, position=joints) logger.info(f"Planning to joints for {robot_name}: {[f'{j:.3f}' for j in joints]}") - return self._plan_path_only(robot_name, robot_id, np.array(joints)) + return self._plan_path_only(robot_name, robot_id, goal_state) def _plan_path_only( - self, robot_name: RobotName, robot_id: WorldRobotID, goal: NDArray[np.float64] + self, robot_name: RobotName, robot_id: WorldRobotID, goal: JointState ) -> bool: """Plan path from current position to goal, store result.""" assert self._world_monitor and self._planner # guaranteed by _begin_planning - start = self._world_monitor.get_current_positions(robot_id) + start = self._world_monitor.get_current_joint_state(robot_id) if start is None: return self._fail("No joint state") result = self._planner.plan_joint_path( world=self._world_monitor.world, robot_id=robot_id, - q_start=start, - q_goal=goal, + start=start, + goal=goal, timeout=self.config.planning_timeout, ) if not result.is_success(): @@ -396,7 +393,8 @@ def _plan_path_only( self._planned_paths[robot_name] = result.path _, _, traj_gen = self._robots[robot_name] - traj = traj_gen.generate([list(q) for q in result.path]) + # Convert JointState path to list of position lists for trajectory generator + traj = traj_gen.generate([list(state.position) for state in result.path]) self._planned_trajectories[robot_name] = traj logger.info(f"Trajectory: {traj.duration:.3f}s") diff --git a/dimos/manipulation/planning/kinematics/drake_optimization_ik.py b/dimos/manipulation/planning/kinematics/drake_optimization_ik.py index 73766f00a8..453c57138f 100644 --- a/dimos/manipulation/planning/kinematics/drake_optimization_ik.py +++ b/dimos/manipulation/planning/kinematics/drake_optimization_ik.py @@ -22,13 +22,14 @@ from dimos.manipulation.planning.spec import IKResult, IKStatus, WorldRobotID, WorldSpec from dimos.manipulation.planning.utils.kinematics_utils import compute_pose_error +from dimos.msgs.geometry_msgs import PoseStamped, Transform +from dimos.msgs.sensor_msgs import JointState from dimos.utils.logging_config import setup_logger +from dimos.utils.transform_utils import pose_to_matrix if TYPE_CHECKING: from numpy.typing import NDArray -from dimos.msgs.geometry_msgs import PoseStamped, Transform - try: from pydrake.math import RigidTransform, RotationMatrix # type: ignore[import-not-found] from pydrake.multibody.inverse_kinematics import ( # type: ignore[import-not-found] @@ -69,7 +70,7 @@ def solve( world: WorldSpec, robot_id: WorldRobotID, target_pose: PoseStamped, - seed: NDArray[np.float64] | None = None, + seed: JointState | None = None, position_tolerance: float = 0.001, orientation_tolerance: float = 0.01, check_collision: bool = True, @@ -92,7 +93,11 @@ def solve( # Get seed from current state if not provided if seed is None: with world.scratch_context() as ctx: - seed = world.get_positions(ctx, robot_id) + seed = world.get_joint_state(ctx, robot_id) + + # Extract joint names and seed positions + joint_names = seed.name + seed_positions = np.array(seed.position, dtype=np.float64) # Target transform target_transform = RigidTransform(target_matrix) @@ -101,9 +106,9 @@ def solve( best_error = float("inf") for attempt in range(max_attempts): - # Generate seed + # Generate seed positions if attempt == 0: - current_seed = seed + current_seed = seed_positions else: # Random seed within joint limits current_seed = np.random.uniform(lower_limits, upper_limits) @@ -114,19 +119,18 @@ def solve( robot_id=robot_id, target_transform=target_transform, seed=current_seed, + joint_names=joint_names, position_tolerance=position_tolerance, orientation_tolerance=orientation_tolerance, lower_limits=lower_limits, upper_limits=upper_limits, ) - if result.is_success() and result.joint_positions is not None: + if result.is_success() and result.joint_state is not None: # Check collision if requested if check_collision: - with world.scratch_context() as ctx: - world.set_positions(ctx, robot_id, result.joint_positions) - if not world.is_collision_free(ctx, robot_id): - continue # Try another seed + if not world.check_config_collision_free(robot_id, result.joint_state): + continue # Try another seed # Check error total_error = result.position_error + result.orientation_error @@ -155,6 +159,7 @@ def _solve_single( robot_id: WorldRobotID, target_transform: RigidTransform, seed: NDArray[np.float64], + joint_names: list[str], position_tolerance: float, orientation_tolerance: float, lower_limits: NDArray[np.float64], @@ -215,16 +220,18 @@ def _solve_single( joint_solution = np.clip(joint_solution, lower_limits, upper_limits) # Compute actual error using FK + solution_state = JointState(name=joint_names, position=joint_solution.tolist()) with world.scratch_context() as ctx: - world.set_positions(ctx, robot_id, joint_solution) + world.set_joint_state(ctx, robot_id, solution_state) actual_pose = world.get_ee_pose(ctx, robot_id) position_error, orientation_error = compute_pose_error( - actual_pose, + pose_to_matrix(actual_pose), target_transform.GetAsMatrix4(), ) return _create_success_result( + joint_names=joint_names, joint_positions=joint_solution, position_error=position_error, orientation_error=orientation_error, @@ -233,6 +240,7 @@ def _solve_single( def _create_success_result( + joint_names: list[str], joint_positions: NDArray[np.float64], position_error: float, orientation_error: float, @@ -240,7 +248,7 @@ def _create_success_result( ) -> IKResult: return IKResult( status=IKStatus.SUCCESS, - joint_positions=joint_positions, + joint_state=JointState(name=joint_names, position=joint_positions.tolist()), position_error=position_error, orientation_error=orientation_error, iterations=iterations, @@ -255,7 +263,7 @@ def _create_failure_result( ) -> IKResult: return IKResult( status=status, - joint_positions=None, + joint_state=None, iterations=iterations, message=message, ) diff --git a/dimos/manipulation/planning/kinematics/jacobian_ik.py b/dimos/manipulation/planning/kinematics/jacobian_ik.py index e31b16372c..5f80642058 100644 --- a/dimos/manipulation/planning/kinematics/jacobian_ik.py +++ b/dimos/manipulation/planning/kinematics/jacobian_ik.py @@ -36,11 +36,13 @@ damped_pseudoinverse, ) from dimos.utils.logging_config import setup_logger +from dimos.utils.transform_utils import pose_to_matrix if TYPE_CHECKING: from numpy.typing import NDArray -from dimos.msgs.geometry_msgs import PoseStamped, Transform +from dimos.msgs.geometry_msgs import PoseStamped, Transform, Twist, Vector3 +from dimos.msgs.sensor_msgs import JointState logger = setup_logger() @@ -91,7 +93,7 @@ def solve( world: WorldSpec, robot_id: WorldRobotID, target_pose: PoseStamped, - seed: NDArray[np.float64] | None = None, + seed: JointState | None = None, position_tolerance: float = 0.001, orientation_tolerance: float = 0.01, check_collision: bool = True, @@ -106,7 +108,7 @@ def solve( world: World for FK/collision checking robot_id: Robot to solve IK for target_pose: Target end-effector pose - seed: Initial guess (uses current position if None) + seed: Initial guess (uses current state if None) position_tolerance: Required position accuracy (meters) orientation_tolerance: Required orientation accuracy (radians) check_collision: Whether to check collision of solution @@ -115,12 +117,6 @@ def solve( Returns: IKResult with solution or failure status """ - # Convert PoseStamped to 4x4 matrix via Transform - target_matrix = Transform( - translation=target_pose.position, - rotation=target_pose.orientation, - ).to_matrix() - if not world.is_finalized: return _create_failure_result(IKStatus.NO_SOLUTION, "World must be finalized before IK") @@ -129,34 +125,38 @@ def solve( # Get seed from current state if not provided if seed is None: with world.scratch_context() as ctx: - seed = world.get_positions(ctx, robot_id) + seed = world.get_joint_state(ctx, robot_id) + + # Extract joint names for creating random seeds + joint_names = seed.name best_result: IKResult | None = None best_error = float("inf") for attempt in range(max_attempts): - # Generate seed + # Generate seed JointState if attempt == 0: current_seed = seed else: # Random seed within joint limits - current_seed = np.random.uniform(lower_limits, upper_limits) + random_positions = np.random.uniform(lower_limits, upper_limits) + current_seed = JointState(name=joint_names, position=random_positions.tolist()) # Solve iterative IK result = self.solve_iterative( world=world, robot_id=robot_id, - target_pose=target_matrix, + target_pose=target_pose, seed=current_seed, max_iterations=self._max_iterations, position_tolerance=position_tolerance, orientation_tolerance=orientation_tolerance, ) - if result.is_success() and result.joint_positions is not None: + if result.is_success() and result.joint_state is not None: # Check collision if requested if check_collision: - if not world.check_config_collision_free(robot_id, result.joint_positions): + if not world.check_config_collision_free(robot_id, result.joint_state): continue # Try another seed # Check error @@ -184,8 +184,8 @@ def solve_iterative( self, world: WorldSpec, robot_id: WorldRobotID, - target_pose: NDArray[np.float64], - seed: NDArray[np.float64], + target_pose: PoseStamped, + seed: JointState, max_iterations: int = 100, position_tolerance: float = 0.001, orientation_tolerance: float = 0.01, @@ -198,7 +198,7 @@ def solve_iterative( Args: world: World for FK/Jacobian computation robot_id: Robot to solve IK for - target_pose: Target end-effector pose (4x4 transform) + target_pose: Target end-effector pose seed: Initial joint configuration max_iterations: Maximum iterations before giving up position_tolerance: Required position accuracy (meters) @@ -207,25 +207,33 @@ def solve_iterative( Returns: IKResult with solution or failure status """ - max_iterations = max_iterations or self._max_iterations - current_joints = seed.copy() + # Convert to internal representation + target_matrix = Transform( + translation=target_pose.position, + rotation=target_pose.orientation, + ).to_matrix() + current_joints = np.array(seed.position, dtype=np.float64) + joint_names = seed.name + max_iterations = max_iterations or self._max_iterations lower_limits, upper_limits = world.get_joint_limits(robot_id) for iteration in range(max_iterations): with world.scratch_context() as ctx: - # Set current position - world.set_positions(ctx, robot_id, current_joints) + # Set current position (convert to JointState for API) + current_state = JointState(name=joint_names, position=current_joints.tolist()) + world.set_joint_state(ctx, robot_id, current_state) - # Get current pose - current_pose = world.get_ee_pose(ctx, robot_id) + # Get current pose (as matrix for error computation) + current_pose = pose_to_matrix(world.get_ee_pose(ctx, robot_id)) # Compute error - pos_error, ori_error = compute_pose_error(current_pose, target_pose) + pos_error, ori_error = compute_pose_error(current_pose, target_matrix) # Check convergence if pos_error <= position_tolerance and ori_error <= orientation_tolerance: return _create_success_result( + joint_names=joint_names, joint_positions=current_joints, position_error=pos_error, orientation_error=ori_error, @@ -233,7 +241,7 @@ def solve_iterative( ) # Compute twist to reduce error - twist = compute_error_twist(current_pose, target_pose, gain=0.5) + twist = compute_error_twist(current_pose, target_matrix, gain=0.5) # Get Jacobian J = world.get_jacobian(ctx, robot_id) @@ -262,9 +270,10 @@ def solve_iterative( # Compute final error with world.scratch_context() as ctx: - world.set_positions(ctx, robot_id, current_joints) - final_pose = world.get_ee_pose(ctx, robot_id) - pos_error, ori_error = compute_pose_error(final_pose, target_pose) + final_state = JointState(name=joint_names, position=current_joints.tolist()) + world.set_joint_state(ctx, robot_id, final_state) + final_pose = pose_to_matrix(world.get_ee_pose(ctx, robot_id)) + pos_error, ori_error = compute_pose_error(final_pose, target_matrix) return _create_failure_result( IKStatus.NO_SOLUTION, @@ -276,10 +285,10 @@ def solve_differential( self, world: WorldSpec, robot_id: WorldRobotID, - current_joints: NDArray[np.float64], - twist: NDArray[np.float64], + current_joints: JointState, + twist: Twist, dt: float, - ) -> NDArray[np.float64] | None: + ) -> JointState | None: """Single Jacobian step for velocity control. Computes joint velocities from desired end-effector twist using @@ -289,14 +298,28 @@ def solve_differential( world: World for Jacobian computation robot_id: Robot to compute for current_joints: Current joint configuration - twist: Desired 6D end-effector twist [vx, vy, vz, wx, wy, wz] + twist: Desired end-effector twist (linear + angular velocity) dt: Time step (not used, but kept for interface compatibility) Returns: - Joint velocities, or None if near singularity + JointState with velocities, or None if near singularity """ + # Convert Twist to 6D array [vx, vy, vz, wx, wy, wz] + twist_array = np.array( + [ + twist.linear.x, + twist.linear.y, + twist.linear.z, + twist.angular.x, + twist.angular.y, + twist.angular.z, + ], + dtype=np.float64, + ) + + joint_names = current_joints.name with world.scratch_context() as ctx: - world.set_positions(ctx, robot_id, current_joints) + world.set_joint_state(ctx, robot_id, current_joints) J = world.get_jacobian(ctx, robot_id) # Check for singularity @@ -308,7 +331,7 @@ def solve_differential( J_pinv = damped_pseudoinverse(J, self._damping) # Compute joint velocities - q_dot = J_pinv @ twist + q_dot = J_pinv @ twist_array # Apply velocity limits if available config = world.get_robot_config(robot_id) @@ -321,15 +344,15 @@ def solve_differential( if max_ratio > 1.0: q_dot = q_dot / max_ratio - return q_dot + return JointState(name=joint_names, velocity=q_dot.tolist()) def solve_differential_position_only( self, world: WorldSpec, robot_id: WorldRobotID, - current_joints: NDArray[np.float64], - linear_velocity: NDArray[np.float64], - ) -> NDArray[np.float64] | None: + current_joints: JointState, + linear_velocity: Vector3, + ) -> JointState | None: """Position-only differential IK using linear Jacobian. Computes joint velocities from desired linear velocity, ignoring @@ -339,13 +362,19 @@ def solve_differential_position_only( world: World for Jacobian computation robot_id: Robot to compute for current_joints: Current joint configuration - linear_velocity: Desired linear velocity [vx, vy, vz] + linear_velocity: Desired linear velocity Returns: - Joint velocities, or None if singular + JointState with velocities, or None if singular """ + # Convert Vector3 to array + vel_array = np.array( + [linear_velocity.x, linear_velocity.y, linear_velocity.z], dtype=np.float64 + ) + + joint_names = current_joints.name with world.scratch_context() as ctx: - world.set_positions(ctx, robot_id, current_joints) + world.set_joint_state(ctx, robot_id, current_joints) J = world.get_jacobian(ctx, robot_id) # Extract linear part (first 3 rows) @@ -362,13 +391,15 @@ def solve_differential_position_only( J_pinv = J_linear.T @ np.linalg.inv(JJT + self._damping**2 * I) # Compute joint velocities - return J_pinv @ linear_velocity + q_dot = J_pinv @ vel_array + return JointState(name=joint_names, velocity=q_dot.tolist()) # ============= Result Helpers ============= def _create_success_result( + joint_names: list[str], joint_positions: NDArray[np.float64], position_error: float, orientation_error: float, @@ -377,7 +408,7 @@ def _create_success_result( """Create a successful IK result.""" return IKResult( status=IKStatus.SUCCESS, - joint_positions=joint_positions, + joint_state=JointState(name=joint_names, position=joint_positions.tolist()), position_error=position_error, orientation_error=orientation_error, iterations=iterations, @@ -393,7 +424,7 @@ def _create_failure_result( """Create a failed IK result.""" return IKResult( status=status, - joint_positions=None, + joint_state=None, iterations=iterations, message=message, ) diff --git a/dimos/manipulation/planning/monitor/world_monitor.py b/dimos/manipulation/planning/monitor/world_monitor.py index f7f9e002a2..c01d9f112e 100644 --- a/dimos/manipulation/planning/monitor/world_monitor.py +++ b/dimos/manipulation/planning/monitor/world_monitor.py @@ -20,16 +20,16 @@ import threading from typing import TYPE_CHECKING, Any -import numpy as np - from dimos.manipulation.planning.factory import create_world from dimos.manipulation.planning.monitor.world_obstacle_monitor import WorldObstacleMonitor from dimos.manipulation.planning.monitor.world_state_monitor import WorldStateMonitor +from dimos.msgs.sensor_msgs import JointState from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: from collections.abc import Generator + import numpy as np from numpy.typing import NDArray from dimos.manipulation.planning.spec import ( @@ -40,7 +40,7 @@ WorldRobotID, WorldSpec, ) - from dimos.msgs.sensor_msgs import JointState + from dimos.msgs.geometry_msgs import PoseStamped from dimos.msgs.vision_msgs import Detection3D logger = setup_logger() @@ -210,23 +210,32 @@ def on_detections(self, detections: list[Detection3D]) -> None: # ============= State Access ============= - def get_current_positions(self, robot_id: WorldRobotID) -> NDArray[np.float64] | None: - """Get current joint positions. Returns None if not yet received.""" - # Try state monitor first + def get_current_joint_state(self, robot_id: WorldRobotID) -> JointState | None: + """Get current joint state. Returns None if not yet received.""" + # Try state monitor first for positions if robot_id in self._state_monitors: positions = self._state_monitors[robot_id].get_current_positions() + velocities = self._state_monitors[robot_id].get_current_velocities() if positions is not None: - return positions + 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 [], + ) # Fall back to world's live context with self._lock: ctx = self._world.get_live_context() - return self._world.get_positions(ctx, robot_id) + return self._world.get_joint_state(ctx, robot_id) - def get_current_velocities(self, robot_id: WorldRobotID) -> NDArray[np.float64] | None: - """Get current joint velocities. Returns None if not available.""" + 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: - return self._state_monitors[robot_id].get_current_velocities() + 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 None def wait_for_state(self, robot_id: WorldRobotID, timeout: float = 1.0) -> bool: @@ -255,32 +264,30 @@ def get_live_context(self) -> Any: # ============= Collision Checking ============= - def is_state_valid(self, robot_id: WorldRobotID, joint_positions: NDArray[np.float64]) -> bool: + def is_state_valid(self, robot_id: WorldRobotID, joint_state: JointState) -> bool: """Check if configuration is collision-free.""" - with self._world.scratch_context() as ctx: - self._world.set_positions(ctx, robot_id, joint_positions) - return self._world.is_collision_free(ctx, robot_id) + return self._world.check_config_collision_free(robot_id, joint_state) def is_path_valid( self, robot_id: WorldRobotID, path: JointPath, step_size: float = 0.05 ) -> bool: - """Check if path is collision-free with interpolation.""" - with self._world.scratch_context() as ctx: - for i in range(len(path) - 1): - q_start = path[i] - q_end = path[i + 1] + """Check if path is collision-free with interpolation. - # Number of interpolation steps - dist = np.linalg.norm(q_end - q_start) - num_steps = max(2, int(np.ceil(dist / step_size))) + Args: + robot_id: Robot to check + path: List of JointState waypoints + step_size: Max step size for interpolation (radians) - for j in range(num_steps): - alpha = j / (num_steps - 1) - q = q_start + alpha * (q_end - q_start) + Returns: + True if entire path is collision-free + """ + if len(path) < 2: + return len(path) == 0 or self._world.check_config_collision_free(robot_id, path[0]) - self._world.set_positions(ctx, robot_id, q) - if not self._world.is_collision_free(ctx, robot_id): - return False + # Check each edge + for i in range(len(path) - 1): + if not self._world.check_edge_collision_free(robot_id, path[i], path[i + 1], step_size): + return False return True @@ -292,25 +299,23 @@ def get_min_distance(self, robot_id: WorldRobotID) -> float: # ============= Kinematics ============= def get_ee_pose( - self, robot_id: WorldRobotID, joint_positions: NDArray[np.float64] | None = None - ) -> NDArray[np.float64]: - """Get end-effector pose as 4x4 transform. Uses current state if positions is None.""" + self, robot_id: WorldRobotID, joint_state: JointState | None = None + ) -> PoseStamped: + """Get end-effector pose. Uses current state if joint_state is None.""" with self._world.scratch_context() as ctx: - # If no positions provided, fetch current from state monitor - if joint_positions is None: - joint_positions = self.get_current_positions(robot_id) + # 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_positions is not None: - self._world.set_positions(ctx, robot_id, joint_positions) + if joint_state is not None: + self._world.set_joint_state(ctx, robot_id, joint_state) return self._world.get_ee_pose(ctx, robot_id) - def get_jacobian( - self, robot_id: WorldRobotID, joint_positions: NDArray[np.float64] - ) -> NDArray[np.float64]: + def get_jacobian(self, robot_id: WorldRobotID, joint_state: JointState) -> NDArray[np.float64]: """Get 6xN Jacobian matrix.""" with self._world.scratch_context() as ctx: - self._world.set_positions(ctx, robot_id, joint_positions) + self._world.set_joint_state(ctx, robot_id, joint_state) return self._world.get_jacobian(ctx, robot_id) # ============= Lifecycle ============= diff --git a/dimos/manipulation/planning/monitor/world_obstacle_monitor.py b/dimos/manipulation/planning/monitor/world_obstacle_monitor.py index 4401149282..cb32c16c80 100644 --- a/dimos/manipulation/planning/monitor/world_obstacle_monitor.py +++ b/dimos/manipulation/planning/monitor/world_obstacle_monitor.py @@ -90,7 +90,7 @@ def __init__( # Running state self._running = False - # Callbacks + # Callbacks: (operation, obstacle_id, obstacle) where operation is "add"/"update"/"remove" self._obstacle_callbacks: list[Callable[[str, str, Obstacle | None], None]] = [] def start(self) -> None: diff --git a/dimos/manipulation/planning/monitor/world_state_monitor.py b/dimos/manipulation/planning/monitor/world_state_monitor.py index 7f2b4d36e4..87d61bb66f 100644 --- a/dimos/manipulation/planning/monitor/world_state_monitor.py +++ b/dimos/manipulation/planning/monitor/world_state_monitor.py @@ -31,6 +31,7 @@ import numpy as np +from dimos.msgs.sensor_msgs import JointState from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: @@ -40,7 +41,6 @@ from numpy.typing import NDArray from dimos.manipulation.planning.spec import WorldSpec - from dimos.msgs.sensor_msgs import JointState logger = setup_logger() @@ -103,8 +103,8 @@ def __init__( # Running state self._running = False - # Callbacks - self._state_callbacks: list[Callable[[str, NDArray[np.float64]], None]] = [] + # Callbacks: (robot_id, joint_state) called on each state update + self._state_callbacks: list[Callable[[str, JointState], None]] = [] def start(self) -> None: """Start the state monitor.""" @@ -165,14 +165,19 @@ def on_joint_state(self, msg: JointState) -> None: # Sync to world's live context (for visualization) try: - self._world.sync_from_joint_state(self._robot_id, positions) + # Create JointState for world sync (API uses JointState) + joint_state = JointState( + name=self._joint_names, + position=positions.tolist(), + ) + self._world.sync_from_joint_state(self._robot_id, joint_state) except Exception as e: logger.error(f"Failed to sync joint state to live context: {e}") # Call registered callbacks for callback in self._state_callbacks: try: - callback(self._robot_id, positions) + callback(self._robot_id, joint_state) except Exception as e: logger.error(f"State callback error: {e}") @@ -308,18 +313,18 @@ def is_state_stale(self, max_age: float = 1.0) -> bool: def add_state_callback( self, - callback: Callable[[str, NDArray[np.float64]], None], + callback: Callable[[str, JointState], None], ) -> None: """Add callback for state updates. Args: - callback: Function called with (robot_id, positions) on each update + callback: Function called with (robot_id, joint_state) on each update """ self._state_callbacks.append(callback) def remove_state_callback( self, - callback: Callable[[str, NDArray[np.float64]], None], + callback: Callable[[str, JointState], None], ) -> None: """Remove a state callback.""" if callback in self._state_callbacks: diff --git a/dimos/manipulation/planning/planners/rrt_planner.py b/dimos/manipulation/planning/planners/rrt_planner.py index ee81b64808..e8d34d94c9 100644 --- a/dimos/manipulation/planning/planners/rrt_planner.py +++ b/dimos/manipulation/planning/planners/rrt_planner.py @@ -34,6 +34,7 @@ WorldSpec, ) from dimos.manipulation.planning.utils.path_utils import compute_path_length +from dimos.msgs.sensor_msgs import JointState from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: @@ -84,15 +85,20 @@ def plan_joint_path( self, world: WorldSpec, robot_id: WorldRobotID, - q_start: NDArray[np.float64], - q_goal: NDArray[np.float64], + start: JointState, + goal: JointState, timeout: float = 10.0, max_iterations: int = 5000, ) -> PlanningResult: """Plan collision-free path using bi-directional RRT.""" start_time = time.time() - error = self._validate_inputs(world, robot_id, q_start, q_goal) + # Extract positions as numpy arrays for internal computation + q_start = np.array(start.position, dtype=np.float64) + q_goal = np.array(goal.position, dtype=np.float64) + joint_names = start.name # Store for converting back to JointState + + error = self._validate_inputs(world, robot_id, start, goal) if error is not None: return error @@ -111,14 +117,21 @@ def plan_joint_path( ) sample = np.random.uniform(lower, upper) - extended = self._extend_tree(world, robot_id, start_tree, sample, self._step_size) + extended = self._extend_tree( + world, robot_id, start_tree, sample, self._step_size, joint_names + ) if extended is not None: connected = self._connect_tree( - world, robot_id, goal_tree, extended.config, self._connect_step_size + world, + robot_id, + goal_tree, + extended.config, + self._connect_step_size, + joint_names, ) if connected is not None: - path = self._extract_path(extended, connected) + path = self._extract_path(extended, connected, joint_names) if trees_swapped: path = list(reversed(path)) path = self._simplify_path(world, robot_id, path) @@ -142,8 +155,8 @@ def _validate_inputs( self, world: WorldSpec, robot_id: WorldRobotID, - q_start: NDArray[np.float64], - q_goal: NDArray[np.float64], + start: JointState, + goal: JointState, ) -> PlanningResult | None: """Validate planning inputs, returns error result or None if valid.""" # Check world is finalized @@ -161,21 +174,24 @@ def _validate_inputs( ) # Check start validity using context-free method - if not world.check_config_collision_free(robot_id, q_start): + if not world.check_config_collision_free(robot_id, start): return _create_failure_result( PlanningStatus.COLLISION_AT_START, "Start configuration is in collision", ) # Check goal validity using context-free method - if not world.check_config_collision_free(robot_id, q_goal): + if not world.check_config_collision_free(robot_id, goal): return _create_failure_result( PlanningStatus.COLLISION_AT_GOAL, "Goal configuration is in collision", ) - # Check limits + # Check limits (extract arrays for numpy comparison) lower, upper = world.get_joint_limits(robot_id) + q_start = np.array(start.position, dtype=np.float64) + q_goal = np.array(goal.position, dtype=np.float64) + if np.any(q_start < lower) or np.any(q_start > upper): return _create_failure_result( PlanningStatus.INVALID_START, @@ -197,6 +213,7 @@ def _extend_tree( tree: list[TreeNode], target: NDArray[np.float64], step_size: float, + joint_names: list[str], ) -> TreeNode | None: """Extend tree toward target, returns new node if successful.""" # Find nearest node @@ -212,8 +229,10 @@ def _extend_tree( new_config = nearest.config + step_size * (diff / dist) # Check validity of edge using context-free method + start_state = JointState(name=joint_names, position=nearest.config.tolist()) + end_state = JointState(name=joint_names, position=new_config.tolist()) if world.check_edge_collision_free( - robot_id, nearest.config, new_config, self._collision_step_size + robot_id, start_state, end_state, self._collision_step_size ): new_node = TreeNode(config=new_config, parent=nearest) nearest.children.append(new_node) @@ -229,11 +248,12 @@ def _connect_tree( tree: list[TreeNode], target: NDArray[np.float64], step_size: float, + joint_names: list[str], ) -> TreeNode | None: """Try to connect tree to target, returns connected node if successful.""" # Keep extending toward target while True: - result = self._extend_tree(world, robot_id, tree, target, step_size) + result = self._extend_tree(world, robot_id, tree, target, step_size, joint_names) if result is None: return None # Extension failed @@ -246,6 +266,7 @@ def _extract_path( self, start_node: TreeNode, goal_node: TreeNode, + joint_names: list[str], ) -> JointPath: """Extract path from two connected nodes.""" # Path from start node to its root (reversed to be root->node) @@ -256,9 +277,10 @@ def _extract_path( # Combine: start_root -> start_node -> goal_node -> goal_root # But we need start -> goal, so reverse the goal path - full_path = start_path + list(reversed(goal_path)) + full_path_arrays = start_path + list(reversed(goal_path)) - return full_path + # Convert to list of JointState + return [JointState(name=joint_names, position=q.tolist()) for q in full_path_arrays] def _simplify_path( self, @@ -282,6 +304,7 @@ def _simplify_path( j = np.random.randint(i + 2, len(simplified)) # Check if direct connection is valid using context-free method + # path elements are already JointState if world.check_edge_collision_free( robot_id, simplified[i], simplified[j], self._collision_step_size ): @@ -295,7 +318,7 @@ def _simplify_path( def _create_success_result( - path: list[NDArray[np.float64]], + path: JointPath, planning_time: float, iterations: int, ) -> PlanningResult: diff --git a/dimos/manipulation/planning/spec/config.py b/dimos/manipulation/planning/spec/config.py index ef8d290c43..bcaeffc9da 100644 --- a/dimos/manipulation/planning/spec/config.py +++ b/dimos/manipulation/planning/spec/config.py @@ -22,9 +22,6 @@ if TYPE_CHECKING: from pathlib import Path - import numpy as np - from numpy.typing import NDArray - from dimos.msgs.geometry_msgs import PoseStamped @@ -64,9 +61,9 @@ class RobotModelConfig: end_effector_link: str base_link: str = "base_link" package_paths: dict[str, Path] = field(default_factory=dict) - joint_limits_lower: NDArray[np.float64] | None = None - joint_limits_upper: NDArray[np.float64] | None = None - velocity_limits: NDArray[np.float64] | None = None + joint_limits_lower: list[float] | None = None + joint_limits_upper: list[float] | None = None + velocity_limits: list[float] | None = None auto_convert_meshes: bool = False xacro_args: dict[str, str] = field(default_factory=dict) collision_exclusion_pairs: list[tuple[str, str]] = field(default_factory=list) diff --git a/dimos/manipulation/planning/spec/protocols.py b/dimos/manipulation/planning/spec/protocols.py index 04322711e6..c7cc7e6053 100644 --- a/dimos/manipulation/planning/spec/protocols.py +++ b/dimos/manipulation/planning/spec/protocols.py @@ -37,6 +37,7 @@ WorldRobotID, ) from dimos.msgs.geometry_msgs import Pose, PoseStamped + from dimos.msgs.sensor_msgs import JointState @runtime_checkable @@ -72,7 +73,7 @@ def get_robot_config(self, robot_id: WorldRobotID) -> RobotModelConfig: def get_joint_limits( self, robot_id: WorldRobotID - ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: + ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: # lower limits, upper limits """Get joint limits (lower, upper) for a robot.""" ... @@ -112,19 +113,17 @@ def scratch_context(self) -> AbstractContextManager[Any]: """Get a scratch context for planning (thread-safe clone).""" ... - def sync_from_joint_state(self, robot_id: WorldRobotID, positions: NDArray[np.float64]) -> None: - """Sync live context from joint state.""" + def sync_from_joint_state(self, robot_id: WorldRobotID, joint_state: JointState) -> None: + """Sync live context from joint state message.""" ... # State Operations (require context) - def set_positions( - self, ctx: Any, robot_id: WorldRobotID, positions: NDArray[np.float64] - ) -> None: - """Set robot joint positions in a context.""" + def set_joint_state(self, ctx: Any, robot_id: WorldRobotID, joint_state: JointState) -> None: + """Set robot joint state in a context.""" ... - def get_positions(self, ctx: Any, robot_id: WorldRobotID) -> NDArray[np.float64]: - """Get robot joint positions from a context.""" + def get_joint_state(self, ctx: Any, robot_id: WorldRobotID) -> JointState: + """Get robot joint state from a context.""" ... # Collision Checking (require context) @@ -137,23 +136,23 @@ def get_min_distance(self, ctx: Any, robot_id: WorldRobotID) -> float: ... # Collision Checking (context-free, for planning) - def check_config_collision_free(self, robot_id: WorldRobotID, q: NDArray[np.float64]) -> bool: - """Check if a configuration is collision-free (manages context internally).""" + def check_config_collision_free(self, robot_id: WorldRobotID, joint_state: JointState) -> bool: + """Check if a joint state is collision-free (manages context internally).""" ... def check_edge_collision_free( self, robot_id: WorldRobotID, - q_start: NDArray[np.float64], - q_end: NDArray[np.float64], + start: JointState, + end: JointState, step_size: float = 0.05, ) -> bool: - """Check if the entire edge between two configurations is collision-free.""" + """Check if the entire edge between two joint states is collision-free.""" ... # Forward Kinematics (require context) - def get_ee_pose(self, ctx: Any, robot_id: WorldRobotID) -> NDArray[np.float64]: - """Get end-effector pose (4x4 transform).""" + def get_ee_pose(self, ctx: Any, robot_id: WorldRobotID) -> PoseStamped: + """Get end-effector pose.""" ... def get_jacobian(self, ctx: Any, robot_id: WorldRobotID) -> NDArray[np.float64]: @@ -183,7 +182,7 @@ def solve( world: WorldSpec, robot_id: WorldRobotID, target_pose: PoseStamped, - seed: NDArray[np.float64] | None = None, + seed: JointState | None = None, position_tolerance: float = 0.001, orientation_tolerance: float = 0.01, check_collision: bool = True, @@ -210,8 +209,8 @@ def plan_joint_path( self, world: WorldSpec, robot_id: WorldRobotID, - q_start: NDArray[np.float64], - q_goal: NDArray[np.float64], + start: JointState, + goal: JointState, timeout: float = 10.0, ) -> PlanningResult: """Plan a collision-free joint-space path.""" diff --git a/dimos/manipulation/planning/spec/types.py b/dimos/manipulation/planning/spec/types.py index ec6e3cdaf3..042a75878f 100644 --- a/dimos/manipulation/planning/spec/types.py +++ b/dimos/manipulation/planning/spec/types.py @@ -26,10 +26,8 @@ ) if TYPE_CHECKING: - import numpy as np - from numpy.typing import NDArray - from dimos.msgs.geometry_msgs import PoseStamped + from dimos.msgs.sensor_msgs import JointState # ============================================================================= # Semantic ID Types (documentation only, not enforced at runtime) @@ -41,8 +39,8 @@ WorldRobotID: TypeAlias = str """Internal Drake world robot ID""" -JointPath: TypeAlias = "list[NDArray[np.float64]]" -"""List of joint configurations forming a path""" +JointPath: TypeAlias = "list[JointState]" +"""List of joint states forming a path (each waypoint has names + positions)""" # ============================================================================= @@ -81,7 +79,7 @@ class IKResult: Attributes: status: Solution status - joint_positions: Solution joint positions (None if failed) + joint_state: Solution joint state with names and positions (None if failed) position_error: Cartesian position error (meters) orientation_error: Orientation error (radians) iterations: Number of iterations taken @@ -89,7 +87,7 @@ class IKResult: """ status: IKStatus - joint_positions: NDArray[np.float64] | None = None + joint_state: JointState | None = None position_error: float = 0.0 orientation_error: float = 0.0 iterations: int = 0 @@ -106,26 +104,24 @@ class PlanningResult: Attributes: status: Planning status - path: List of joint configurations (empty if failed) + path: List of joint states forming the path (empty if failed). + Each JointState contains names, positions, and optionally velocities. planning_time: Time taken to plan (seconds) path_length: Total path length in joint space (radians) iterations: Number of iterations/nodes expanded message: Human-readable status message timestamps: Optional timestamps for each waypoint (seconds from start). If provided by the planner, trajectory generator can use these directly. - velocities: Optional joint velocities at each waypoint. - If provided by the planner, trajectory generator can use these directly. """ status: PlanningStatus - path: list[NDArray[np.float64]] = field(default_factory=list) + path: list[JointState] = field(default_factory=list) planning_time: float = 0.0 path_length: float = 0.0 iterations: int = 0 message: str = "" - # Optional timing fields (set by optimization-based planners) + # Optional timing (set by optimization-based planners) timestamps: list[float] | None = None - velocities: list[NDArray[np.float64]] | None = None def is_success(self) -> bool: """Check if planning was successful.""" diff --git a/dimos/manipulation/planning/utils/path_utils.py b/dimos/manipulation/planning/utils/path_utils.py index 03eada3f34..fbf8af4032 100644 --- a/dimos/manipulation/planning/utils/path_utils.py +++ b/dimos/manipulation/planning/utils/path_utils.py @@ -32,6 +32,8 @@ import numpy as np +from dimos.msgs.sensor_msgs import JointState + if TYPE_CHECKING: from numpy.typing import NDArray @@ -48,7 +50,7 @@ def interpolate_path( between consecutive waypoints is at most `resolution`. Args: - path: Original path (list of joint configurations) + path: Original path (list of JointState waypoints) resolution: Maximum distance between waypoints (radians) Returns: @@ -56,37 +58,39 @@ def interpolate_path( Example: # After planning, interpolate for smoother execution - raw_path = planner.plan_joint_path(world, robot_id, q_start, q_goal).path + raw_path = planner.plan_joint_path(world, robot_id, start, goal).path smooth_path = interpolate_path(raw_path, resolution=0.02) """ if len(path) <= 1: - return path + return list(path) - interpolated = [path[0]] + interpolated: list[JointState] = [path[0]] + joint_names = path[0].name for i in range(len(path) - 1): - start = path[i] - end = path[i + 1] + q_start = np.array(path[i].position, dtype=np.float64) + q_end = np.array(path[i + 1].position, dtype=np.float64) - diff = end - start + diff = q_end - q_start max_diff = float(np.max(np.abs(diff))) if max_diff <= resolution: - interpolated.append(end) + interpolated.append(path[i + 1]) else: num_steps = int(np.ceil(max_diff / resolution)) for step in range(1, num_steps + 1): alpha = step / num_steps - interpolated.append(start + alpha * diff) + q_interp = q_start + alpha * diff + interpolated.append(JointState(name=joint_names, position=q_interp.tolist())) return interpolated def interpolate_segment( - start: NDArray[np.float64], - end: NDArray[np.float64], + start: JointState, + end: JointState, step_size: float, -) -> list[NDArray[np.float64]]: +) -> JointPath: """Interpolate between two configurations. Returns a list of configurations from start to end (inclusive) @@ -98,27 +102,32 @@ def interpolate_segment( step_size: Maximum step size (radians) Returns: - List of interpolated configurations [start, ..., end] + List of interpolated JointState waypoints [start, ..., end] Example: # Check collision along a segment - segment = interpolate_segment(q1, q2, step_size=0.02) - for q in segment: - if not world.is_collision_free(ctx, robot_id): + segment = interpolate_segment(start_state, end_state, step_size=0.02) + for state in segment: + if not world.check_config_collision_free(robot_id, state): return False """ - diff = end - start + q_start = np.array(start.position, dtype=np.float64) + q_end = np.array(end.position, dtype=np.float64) + joint_names = start.name + + diff = q_end - q_start distance = float(np.linalg.norm(diff)) if distance <= step_size: return [start, end] num_steps = int(np.ceil(distance / step_size)) - segment = [] + segment: JointPath = [] for i in range(num_steps + 1): alpha = i / num_steps - segment.append(start + alpha * diff) + q_interp = q_start + alpha * diff + segment.append(JointState(name=joint_names, position=q_interp.tolist())) return segment @@ -139,7 +148,7 @@ def simplify_path( Args: world: World for collision checking robot_id: Which robot - path: Original path + path: Original path (list of JointState waypoints) max_iterations: Maximum shortcutting attempts collision_step_size: Step size for collision checking along shortcuts @@ -147,11 +156,11 @@ def simplify_path( Simplified path with fewer waypoints Example: - raw_path = planner.plan_joint_path(world, robot_id, q_start, q_goal).path + raw_path = planner.plan_joint_path(world, robot_id, start, goal).path simplified = simplify_path(world, robot_id, raw_path) """ if len(path) <= 2: - return path + return list(path) simplified = list(path) @@ -163,20 +172,12 @@ def simplify_path( i = np.random.randint(0, len(simplified) - 2) j = np.random.randint(i + 2, len(simplified)) - # Check if direct connection is valid - with world.scratch_context() as ctx: - segment = interpolate_segment(simplified[i], simplified[j], collision_step_size) - valid = True - - for q in segment: - world.set_positions(ctx, robot_id, q) - if not world.is_collision_free(ctx, robot_id): - valid = False - break - - if valid: - # Remove intermediate waypoints - simplified = simplified[: i + 1] + simplified[j:] + # Check if direct connection is valid using context-free API + if world.check_edge_collision_free( + robot_id, simplified[i], simplified[j], collision_step_size + ): + # Remove intermediate waypoints + simplified = simplified[: i + 1] + simplified[j:] return simplified @@ -187,7 +188,7 @@ def compute_path_length(path: JointPath) -> float: Sums the Euclidean distances between consecutive waypoints. Args: - path: Path to measure + path: Path to measure (list of JointState waypoints) Returns: Total length in radians @@ -201,7 +202,9 @@ def compute_path_length(path: JointPath) -> float: length = 0.0 for i in range(len(path) - 1): - length += float(np.linalg.norm(path[i + 1] - path[i])) + q_curr = np.array(path[i].position, dtype=np.float64) + q_next = np.array(path[i + 1].position, dtype=np.float64) + length += float(np.linalg.norm(q_next - q_curr)) return length @@ -214,14 +217,15 @@ def is_path_within_limits( """Check if all waypoints in path are within joint limits. Args: - path: Path to check + path: Path to check (list of JointState waypoints) lower_limits: Lower joint limits (radians) upper_limits: Upper joint limits (radians) Returns: True if all waypoints are within limits """ - for q in path: + for state in path: + q = np.array(state.position, dtype=np.float64) if np.any(q < lower_limits) or np.any(q > upper_limits): return False return True @@ -235,14 +239,19 @@ def clip_path_to_limits( """Clip all waypoints in path to joint limits. Args: - path: Path to clip + path: Path to clip (list of JointState waypoints) lower_limits: Lower joint limits (radians) upper_limits: Upper joint limits (radians) Returns: Path with all waypoints clipped to limits """ - return [np.clip(q, lower_limits, upper_limits) for q in path] + clipped: list[JointState] = [] + for state in path: + q = np.array(state.position, dtype=np.float64) + q_clipped = np.clip(q, lower_limits, upper_limits) + clipped.append(JointState(name=state.name, position=q_clipped.tolist())) + return clipped def reverse_path(path: JointPath) -> JointPath: @@ -264,13 +273,13 @@ def concatenate_paths( """Concatenate multiple paths into one. Args: - *paths: Paths to concatenate + *paths: Paths to concatenate (each is a list of JointState waypoints) remove_duplicates: If True, remove duplicate waypoints at junctions Returns: Single concatenated path """ - result: list[NDArray[np.float64]] = [] + result: list[JointState] = [] for path in paths: if not path: @@ -278,7 +287,9 @@ def concatenate_paths( if remove_duplicates and result: # Check if last point matches first point (tight tolerance for joint space) - if np.allclose(result[-1], path[0], atol=1e-6, rtol=0): + q_last = np.array(result[-1].position, dtype=np.float64) + q_first = np.array(path[0].position, dtype=np.float64) + if np.allclose(q_last, q_first, atol=1e-6, rtol=0): result.extend(path[1:]) else: result.extend(path) diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index a056c47b26..4c3ac79498 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -42,6 +42,7 @@ from numpy.typing import NDArray from dimos.msgs.geometry_msgs import PoseStamped, Transform +from dimos.msgs.sensor_msgs import JointState try: from pydrake.geometry import ( # type: ignore[import-not-found] @@ -617,14 +618,17 @@ def scratch_context(self) -> Generator[Context, None, None]: yield ctx - def sync_from_joint_state(self, robot_id: WorldRobotID, positions: NDArray[np.float64]) -> None: - """Sync live context from driver's joint state. + def sync_from_joint_state(self, robot_id: WorldRobotID, joint_state: JointState) -> None: + """Sync live context from driver's joint state message. Called by StateMonitor when new JointState arrives. """ 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) + with self._lock: self._set_positions_internal(self._plant_context, robot_id, positions) @@ -634,13 +638,16 @@ def sync_from_joint_state(self, robot_id: WorldRobotID, positions: NDArray[np.fl # ============= State Operations (context-based) ============= - def set_positions( - self, ctx: Context, robot_id: WorldRobotID, positions: NDArray[np.float64] + def set_joint_state( + self, ctx: Context, robot_id: WorldRobotID, joint_state: JointState ) -> None: - """Set robot positions in given context.""" + """Set robot joint state in given context.""" 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) + # Get plant context from diagram context plant_ctx = self._diagram.GetMutableSubsystemContext(self._plant, ctx) self._set_positions_internal(plant_ctx, robot_id, positions) @@ -660,8 +667,8 @@ def _set_positions_internal( self._plant.SetPositions(plant_ctx, full_positions) - def get_positions(self, ctx: Context, robot_id: WorldRobotID) -> NDArray[np.float64]: - """Get robot positions from given context.""" + def get_joint_state(self, ctx: Context, robot_id: WorldRobotID) -> JointState: + """Get robot joint state from given context.""" if not self._finalized: raise RuntimeError("World must be finalized first") @@ -672,7 +679,8 @@ def get_positions(self, ctx: Context, robot_id: WorldRobotID) -> NDArray[np.floa plant_ctx = self._diagram.GetSubsystemContext(self._plant, ctx) full_positions = self._plant.GetPositions(plant_ctx) - return np.array([full_positions[idx] for idx in robot_data.joint_indices]) + positions = [float(full_positions[idx]) for idx in robot_data.joint_indices] + return JointState(name=robot_data.config.joint_names, position=positions) # ============= Collision Checking (context-based) ============= @@ -706,32 +714,36 @@ def get_min_distance(self, ctx: Context, robot_id: WorldRobotID) -> float: # ============= Collision Checking (context-free, for planning) ============= - def check_config_collision_free(self, robot_id: WorldRobotID, q: NDArray[np.float64]) -> bool: - """Check if a configuration is collision-free (manages context internally). + def check_config_collision_free(self, robot_id: WorldRobotID, joint_state: JointState) -> bool: + """Check if a joint state is collision-free (manages context internally). This is a convenience method for planners that don't need to manage contexts. """ with self.scratch_context() as ctx: - self.set_positions(ctx, robot_id, q) + self.set_joint_state(ctx, robot_id, joint_state) return self.is_collision_free(ctx, robot_id) def check_edge_collision_free( self, robot_id: WorldRobotID, - q_start: NDArray[np.float64], - q_end: NDArray[np.float64], + start: JointState, + end: JointState, step_size: float = 0.05, ) -> bool: - """Check if the entire edge between two configurations is collision-free. + """Check if the entire edge between two joint states is collision-free. - Interpolates between q_start and q_end at the given step_size and checks + Interpolates between start and end at the given step_size and checks each configuration for collisions. This is more efficient than checking each configuration separately as it uses a single scratch context. """ + # Extract positions as numpy arrays for interpolation + q_start = np.array(start.position, dtype=np.float64) + q_end = np.array(end.position, dtype=np.float64) + # Compute number of steps needed dist = float(np.linalg.norm(q_end - q_start)) if dist < 1e-8: - return self.check_config_collision_free(robot_id, q_start) + return self.check_config_collision_free(robot_id, start) n_steps = max(2, int(np.ceil(dist / step_size)) + 1) @@ -739,7 +751,9 @@ def check_edge_collision_free( for i in range(n_steps): t = i / (n_steps - 1) q = q_start + t * (q_end - q_start) - self.set_positions(ctx, robot_id, q) + # Create interpolated JointState + interp_state = JointState(name=start.name, position=q.tolist()) + self.set_joint_state(ctx, robot_id, interp_state) if not self.is_collision_free(ctx, robot_id): return False @@ -747,8 +761,8 @@ def check_edge_collision_free( # ============= Forward Kinematics (context-based) ============= - def get_ee_pose(self, ctx: Context, robot_id: WorldRobotID) -> NDArray[np.float64]: - """Get end-effector pose as 4x4 transform.""" + 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") @@ -761,8 +775,15 @@ def get_ee_pose(self, ctx: Context, robot_id: WorldRobotID) -> NDArray[np.float6 ee_body = robot_data.ee_frame.body() X_WE = self._plant.EvalBodyPoseInWorld(plant_ctx, ee_body) - result: NDArray[np.float64] = X_WE.GetAsMatrix4() - return result + # Extract position and quaternion from Drake transform + pos = X_WE.translation() + quat = X_WE.rotation().ToQuaternion() # Drake returns [w, x, y, z] + + return PoseStamped( + frame_id="world", + position=[float(pos[0]), float(pos[1]), float(pos[2])], + orientation=[float(quat.x()), float(quat.y()), float(quat.z()), float(quat.w())], + ) def get_link_pose( self, ctx: Context, robot_id: WorldRobotID, link_name: str @@ -859,7 +880,7 @@ def animate_path( Args: robot_id: Robot to animate - path: List of joint configurations + path: List of joint states forming the path duration: Total animation duration in seconds """ import time @@ -867,20 +888,20 @@ def animate_path( if self._meshcat is None or len(path) < 2: return - # Capture current positions of all OTHER robots so they don't snap to zero - other_robot_positions: dict[WorldRobotID, NDArray[np.float64]] = {} + # Capture current states of all OTHER robots so they don't snap to zero + other_robot_states: dict[WorldRobotID, JointState] = {} for rid, _robot_data in self._robots.items(): if rid != robot_id: - other_robot_positions[rid] = self.get_positions(self.get_live_context(), rid) + other_robot_states[rid] = self.get_joint_state(self.get_live_context(), rid) dt = duration / (len(path) - 1) - for q in path: + for joint_state in path: with self.scratch_context() as ctx: - # Restore other robots to their current positions - for rid, pos in other_robot_positions.items(): - self.set_positions(ctx, rid, pos) - # Set animated robot's position - self.set_positions(ctx, robot_id, q) + # Restore other robots to their current states + for rid, state in other_robot_states.items(): + self.set_joint_state(ctx, rid, state) + # Set animated robot's state + self.set_joint_state(ctx, robot_id, joint_state) self.publish_visualization(ctx) time.sleep(dt) From bc8954bc8c45cb1f85a7c9f5c47ed984c9ef01cc Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Thu, 29 Jan 2026 18:32:46 -0800 Subject: [PATCH 27/31] removed unused import --- dimos/e2e_tests/test_manipulation_module.py | 12 +++--------- dimos/manipulation/manipulation_module.py | 2 +- dimos/manipulation/planning/spec/protocols.py | 2 +- dimos/manipulation/planning/utils/mesh_utils.py | 1 - dimos/manipulation/planning/world/drake_world.py | 2 -- dimos/manipulation/test_manipulation_unit.py | 1 - 6 files changed, 5 insertions(+), 15 deletions(-) diff --git a/dimos/e2e_tests/test_manipulation_module.py b/dimos/e2e_tests/test_manipulation_module.py index c4f25cf6fe..858955a5c7 100644 --- a/dimos/e2e_tests/test_manipulation_module.py +++ b/dimos/e2e_tests/test_manipulation_module.py @@ -21,11 +21,10 @@ from __future__ import annotations +import importlib.util import os -import time -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock -import numpy as np import pytest from dimos.manipulation.manipulation_module import ( @@ -44,12 +43,7 @@ def _drake_available() -> bool: """Check if Drake is available.""" - try: - import pydrake - - return True - except ImportError: - return False + return importlib.util.find_spec("pydrake") is not None def _xarm_urdf_available() -> bool: diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index ef92b17760..ccb120149c 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -45,7 +45,7 @@ if TYPE_CHECKING: from dimos.core.rpc_client import RPCClient - from dimos.msgs.geometry_msgs import Pose, PoseStamped + from dimos.msgs.geometry_msgs import Pose logger = setup_logger() diff --git a/dimos/manipulation/planning/spec/protocols.py b/dimos/manipulation/planning/spec/protocols.py index c7cc7e6053..218fb89163 100644 --- a/dimos/manipulation/planning/spec/protocols.py +++ b/dimos/manipulation/planning/spec/protocols.py @@ -36,7 +36,7 @@ PlanningResult, WorldRobotID, ) - from dimos.msgs.geometry_msgs import Pose, PoseStamped + from dimos.msgs.geometry_msgs import PoseStamped from dimos.msgs.sensor_msgs import JointState diff --git a/dimos/manipulation/planning/utils/mesh_utils.py b/dimos/manipulation/planning/utils/mesh_utils.py index 316d0e54d1..46d14eb1b4 100644 --- a/dimos/manipulation/planning/utils/mesh_utils.py +++ b/dimos/manipulation/planning/utils/mesh_utils.py @@ -36,7 +36,6 @@ import re import shutil import tempfile -from typing import TYPE_CHECKING from dimos.utils.logging_config import setup_logger diff --git a/dimos/manipulation/planning/world/drake_world.py b/dimos/manipulation/planning/world/drake_world.py index 4c3ac79498..ea17684c81 100644 --- a/dimos/manipulation/planning/world/drake_world.py +++ b/dimos/manipulation/planning/world/drake_world.py @@ -17,7 +17,6 @@ from __future__ import annotations from contextlib import contextmanager -import copy from dataclasses import dataclass from pathlib import Path from threading import RLock @@ -57,7 +56,6 @@ MeshcatVisualizer, MeshcatVisualizerParams, ProximityProperties, - QueryObject, Rgba, Role, SceneGraph, diff --git a/dimos/manipulation/test_manipulation_unit.py b/dimos/manipulation/test_manipulation_unit.py index e81dc23c43..1643d2b5da 100644 --- a/dimos/manipulation/test_manipulation_unit.py +++ b/dimos/manipulation/test_manipulation_unit.py @@ -20,7 +20,6 @@ import threading from unittest.mock import MagicMock, patch -import numpy as np import pytest from dimos.manipulation.manipulation_module import ( From ffc7338b2fddf3e22ab95b50063439fdd246d468 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Sat, 31 Jan 2026 16:09:04 -0800 Subject: [PATCH 28/31] removed manipulation history reference --- dimos/manipulation/manipulation_interface.py | 64 ++++++++------------ 1 file changed, 25 insertions(+), 39 deletions(-) diff --git a/dimos/manipulation/manipulation_interface.py b/dimos/manipulation/manipulation_interface.py index eb528a3a43..f8f16cde86 100644 --- a/dimos/manipulation/manipulation_interface.py +++ b/dimos/manipulation/manipulation_interface.py @@ -13,19 +13,14 @@ # limitations under the License. """ -ManipulationInterface provides a unified interface for accessing manipulation history. +ManipulationInterface provides a unified interface for accessing manipulation data. This module defines the ManipulationInterface class, which serves as an access point -for the robot's manipulation history, agent-generated constraints, and manipulation -metadata streams. +for agent-generated constraints, manipulation tasks, and perception streams. """ -import os from typing import TYPE_CHECKING, Any -from dimos.manipulation.manipulation_history import ( # type: ignore[import-not-found,import-untyped] - ManipulationHistory, -) from dimos.types.manipulation import ( AbstractConstraint, ManipulationTask, @@ -50,28 +45,16 @@ class ManipulationInterface: def __init__( self, - output_dir: str, - new_memory: bool = False, perception_stream: Any = None, ) -> None: """ Initialize a new ManipulationInterface instance. Args: - output_dir: Directory for storing manipulation data - new_memory: If True, creates a new manipulation history from scratch perception_stream: ObjectDetectionStream instance for real-time object data """ - self.output_dir = output_dir - - # Create manipulation history directory - manipulation_dir = os.path.join(output_dir, "manipulation_history") - os.makedirs(manipulation_dir, exist_ok=True) - - # Initialize manipulation history - self.manipulation_history: ManipulationHistory = ManipulationHistory( - output_dir=manipulation_dir, new_memory=new_memory - ) + # List of manipulation tasks + self._tasks: list[ManipulationTask] = [] # List of constraints generated by the Agent via constraint generation skills self.agent_constraints: list[AbstractConstraint] = [] @@ -123,21 +106,15 @@ def get_constraint(self, constraint_id: str) -> AbstractConstraint | None: logger.warning(f"Constraint with ID {constraint_id} not found") return None - def add_manipulation_task( - self, task: ManipulationTask, manipulation_response: str | None = None - ) -> None: + def add_manipulation_task(self, task: ManipulationTask) -> None: """ - Add a manipulation task to ManipulationHistory. + Add a manipulation task. Args: task: The ManipulationTask to add - manipulation_response: Optional response from the motion planner/executor - """ - # Add task to history - self.manipulation_history.add_entry( # type: ignore[call-arg] - task=task, result=None, notes=None, manipulation_response=manipulation_response - ) + self._tasks.append(task) + logger.info(f"Added manipulation task: {task.get('id', 'unknown')}") def get_manipulation_task(self, task_id: str) -> ManipulationTask | None: """ @@ -149,7 +126,10 @@ def get_manipulation_task(self, task_id: str) -> ManipulationTask | None: Returns: The task object or None if not found """ - return self.history.get_manipulation_task(task_id) # type: ignore[attr-defined, no-any-return] + for task in self._tasks: + if task.get("id") == task_id: + return task + return None def get_all_manipulation_tasks(self) -> list[ManipulationTask]: """ @@ -158,7 +138,7 @@ def get_all_manipulation_tasks(self) -> list[ManipulationTask]: Returns: List of all manipulation tasks """ - return self.history.get_all_manipulation_tasks() # type: ignore[attr-defined, no-any-return] + return list(self._tasks) def update_task_status( self, task_id: str, status: str, result: dict[str, Any] | None = None @@ -174,7 +154,13 @@ def update_task_status( Returns: The updated task or None if task not found """ - return self.history.update_task_status(task_id, status, result) # type: ignore[attr-defined, no-any-return] + task = self.get_manipulation_task(task_id) + if task is not None: + task["status"] = status + if result is not None: + task["result"] = result + return task + return None # === Perception stream methods === @@ -260,13 +246,13 @@ def cleanup_perception_subscription(self) -> None: # === Utility methods === - def clear_history(self) -> None: + def clear(self) -> None: """ - Clear all manipulation history data and agent constraints. + Clear all manipulation tasks and agent constraints. """ - self.manipulation_history.clear() + self._tasks.clear() self.agent_constraints.clear() - logger.info("Cleared manipulation history and agent constraints") + logger.info("Cleared manipulation tasks and agent constraints") def __str__(self) -> str: """ @@ -276,7 +262,7 @@ def __str__(self) -> str: String representation with key stats """ has_stream = self.perception_stream is not None - return f"ManipulationInterface(history={self.manipulation_history}, agent_constraints={len(self.agent_constraints)}, perception_stream={has_stream}, detected_objects={len(self.latest_objects)})" + return f"ManipulationInterface(tasks={len(self._tasks)}, constraints={len(self.agent_constraints)}, perception_stream={has_stream}, detected_objects={len(self.latest_objects)})" def __del__(self) -> None: """ From 8153fde4f6ae064ddfcc031b1bd79c692dc776c8 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Sat, 31 Jan 2026 16:18:18 -0800 Subject: [PATCH 29/31] added a Jacobian typealias --- dimos/manipulation/planning/spec/__init__.py | 2 ++ dimos/manipulation/planning/spec/types.py | 10 ++++++++++ dimos/manipulation/planning/utils/__init__.py | 4 ++-- dimos/manipulation/planning/utils/kinematics_utils.py | 8 +++++--- 4 files changed, 19 insertions(+), 5 deletions(-) diff --git a/dimos/manipulation/planning/spec/__init__.py b/dimos/manipulation/planning/spec/__init__.py index 790b4ecf3f..a78fb6e5fd 100644 --- a/dimos/manipulation/planning/spec/__init__.py +++ b/dimos/manipulation/planning/spec/__init__.py @@ -24,6 +24,7 @@ from dimos.manipulation.planning.spec.types import ( CollisionObjectMessage, IKResult, + Jacobian, JointPath, Obstacle, PlanningResult, @@ -35,6 +36,7 @@ "CollisionObjectMessage", "IKResult", "IKStatus", + "Jacobian", "JointPath", "KinematicsSpec", "Obstacle", diff --git a/dimos/manipulation/planning/spec/types.py b/dimos/manipulation/planning/spec/types.py index 042a75878f..a38cc0da26 100644 --- a/dimos/manipulation/planning/spec/types.py +++ b/dimos/manipulation/planning/spec/types.py @@ -26,6 +26,9 @@ ) if TYPE_CHECKING: + import numpy as np + from numpy.typing import NDArray + from dimos.msgs.geometry_msgs import PoseStamped from dimos.msgs.sensor_msgs import JointState @@ -42,6 +45,13 @@ JointPath: TypeAlias = "list[JointState]" """List of joint states forming a path (each waypoint has names + positions)""" +# ============================================================================= +# Numeric Array Types +# ============================================================================= + +Jacobian: TypeAlias = "NDArray[np.float64]" +"""6 x n Jacobian matrix (rows: [vx, vy, vz, wx, wy, wz])""" + # ============================================================================= # Data Classes diff --git a/dimos/manipulation/planning/utils/__init__.py b/dimos/manipulation/planning/utils/__init__.py index 9073cf2203..04ec1806b5 100644 --- a/dimos/manipulation/planning/utils/__init__.py +++ b/dimos/manipulation/planning/utils/__init__.py @@ -38,14 +38,14 @@ ) __all__ = [ + # Kinematics utilities "check_singularity", "compute_error_twist", + # Path utilities "compute_path_length", "compute_pose_error", - # Kinematics utilities "damped_pseudoinverse", "get_manipulability", - # Path utilities "interpolate_path", "interpolate_segment", ] diff --git a/dimos/manipulation/planning/utils/kinematics_utils.py b/dimos/manipulation/planning/utils/kinematics_utils.py index fa6b74e61c..80536304a0 100644 --- a/dimos/manipulation/planning/utils/kinematics_utils.py +++ b/dimos/manipulation/planning/utils/kinematics_utils.py @@ -36,9 +36,11 @@ if TYPE_CHECKING: from numpy.typing import NDArray + from dimos.manipulation.planning.spec import Jacobian + def damped_pseudoinverse( - J: NDArray[np.float64], + J: Jacobian, damping: float = 0.01, ) -> NDArray[np.float64]: """Compute damped pseudoinverse of Jacobian. @@ -68,7 +70,7 @@ def damped_pseudoinverse( def check_singularity( - J: NDArray[np.float64], + J: Jacobian, threshold: float = 0.01, ) -> bool: """Check if Jacobian is near singularity. @@ -92,7 +94,7 @@ def check_singularity( return get_manipulability(J) < threshold -def get_manipulability(J: NDArray[np.float64]) -> float: +def get_manipulability(J: Jacobian) -> float: """Compute manipulability measure. The manipulability measure w = sqrt(det(J @ J^T)) represents the From 50157e000e17e3ac655391c469a9d8b935abf7e7 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Sat, 31 Jan 2026 16:42:35 -0800 Subject: [PATCH 30/31] fixed mypy errors --- dimos/manipulation/manipulation_interface.py | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/dimos/manipulation/manipulation_interface.py b/dimos/manipulation/manipulation_interface.py index f8f16cde86..524562520d 100644 --- a/dimos/manipulation/manipulation_interface.py +++ b/dimos/manipulation/manipulation_interface.py @@ -114,7 +114,7 @@ def add_manipulation_task(self, task: ManipulationTask) -> None: task: The ManipulationTask to add """ self._tasks.append(task) - logger.info(f"Added manipulation task: {task.get('id', 'unknown')}") + logger.info(f"Added manipulation task: {task.task_id or 'unknown'}") def get_manipulation_task(self, task_id: str) -> ManipulationTask | None: """ @@ -127,7 +127,7 @@ def get_manipulation_task(self, task_id: str) -> ManipulationTask | None: The task object or None if not found """ for task in self._tasks: - if task.get("id") == task_id: + if task.task_id == task_id: return task return None @@ -140,25 +140,20 @@ def get_all_manipulation_tasks(self) -> list[ManipulationTask]: """ return list(self._tasks) - def update_task_status( - self, task_id: str, status: str, result: dict[str, Any] | None = None - ) -> ManipulationTask | None: + def update_task_result(self, task_id: str, result: dict[str, Any]) -> ManipulationTask | None: """ - Update the status and result of a manipulation task. + Update the result of a manipulation task. Args: task_id: ID of the task to update - status: New status for the task (e.g., 'completed', 'failed') - result: Optional dictionary with result data + result: Result data from task execution Returns: The updated task or None if task not found """ task = self.get_manipulation_task(task_id) if task is not None: - task["status"] = status - if result is not None: - task["result"] = result + task.result = result return task return None From 6dbdf84129646cf3955ff7c3336b1d1fa6f499a3 Mon Sep 17 00:00:00 2001 From: mustafab0 Date: Sat, 31 Jan 2026 16:44:29 -0800 Subject: [PATCH 31/31] updated manipulation dependencies in pyproject.toml --- pyproject.toml | 6 ++++- uv.lock | 61 ++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 66 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index ee0adb2d6f..eeb401ecde 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -173,6 +173,9 @@ unitree = [ ] manipulation = [ + # Planning (Drake) + "drake>=1.40.0", + # Contact Graspnet Dependencies "h5py>=3.7.0", "pyrender>=0.1.45", @@ -186,8 +189,9 @@ manipulation = [ "pyyaml>=6.0", "contact-graspnet-pytorch", - # piper arm + # Hardware SDKs "piper-sdk", + "xarm-python-sdk>=1.17.0", # Visualization (Optional) "kaleido>=0.2.1", diff --git a/uv.lock b/uv.lock index ec43088319..f11421e7ce 100644 --- a/uv.lock +++ b/uv.lock @@ -1523,6 +1523,7 @@ drone = [ ] manipulation = [ { name = "contact-graspnet-pytorch" }, + { name = "drake" }, { name = "h5py" }, { name = "kaleido" }, { name = "matplotlib" }, @@ -1536,6 +1537,7 @@ manipulation = [ { name = "rtree" }, { name = "tqdm" }, { name = "trimesh" }, + { name = "xarm-python-sdk" }, ] misc = [ { name = "catkin-pkg" }, @@ -1641,6 +1643,7 @@ requires-dist = [ { name = "dimos", extras = ["agents", "web", "perception", "visualization", "sim"], marker = "extra == 'base'" }, { name = "dimos", extras = ["base"], marker = "extra == 'unitree'" }, { name = "dimos-lcm" }, + { name = "drake", marker = "extra == 'manipulation'", specifier = ">=1.40.0" }, { name = "edgetam-dimos", marker = "extra == 'misc'" }, { name = "einops", marker = "extra == 'misc'", specifier = "==0.8.1" }, { name = "empy", marker = "extra == 'misc'", specifier = "==3.3.4" }, @@ -1759,6 +1762,7 @@ requires-dist = [ { name = "unitree-webrtc-connect-leshy", marker = "extra == 'unitree'", specifier = ">=2.0.7" }, { name = "uvicorn", marker = "extra == 'web'", specifier = ">=0.34.0" }, { name = "watchdog", marker = "extra == 'dev'", specifier = ">=3.0.0" }, + { name = "xarm-python-sdk", marker = "extra == 'manipulation'", specifier = ">=1.17.0" }, { name = "xarm-python-sdk", marker = "extra == 'misc'", specifier = ">=1.17.0" }, { name = "xformers", marker = "extra == 'cuda'", specifier = ">=0.0.20" }, { name = "yapf", marker = "extra == 'misc'", specifier = "==0.40.2" }, @@ -1851,6 +1855,28 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/02/10/5da547df7a391dcde17f59520a231527b8571e6f46fc8efb02ccb370ab12/docutils-0.22.4-py3-none-any.whl", hash = "sha256:d0013f540772d1420576855455d050a2180186c91c15779301ac2ccb3eeb68de", size = 633196, upload-time = "2025-12-18T19:00:18.077Z" }, ] +[[package]] +name = "drake" +version = "1.49.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "matplotlib" }, + { name = "mosek", marker = "python_full_version < '3.15'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "pydot" }, + { name = "pyyaml" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/26/2ce3a9caf431f24e39f8b1fc7b3ebba4faafef1d61c849db3194e8d2e21d/drake-1.49.0-cp310-cp310-manylinux_2_34_x86_64.whl", hash = "sha256:6c73dbd061fcb442e82b7b5a94dadcfbf4c44949035d03394df29412114647b2", size = 41482505, upload-time = "2026-01-15T19:44:08.313Z" }, + { url = "https://files.pythonhosted.org/packages/a3/2c/b147eaeee97986d970c0618144b28049cf078c20ba73209f4db14cf9a531/drake-1.49.0-cp311-cp311-manylinux_2_34_x86_64.whl", hash = "sha256:b897f5f1516d13627ef18a8395b15f56413016d3c91c902cada76860b5cbb12c", size = 41516482, upload-time = "2026-01-15T19:44:11.342Z" }, + { url = "https://files.pythonhosted.org/packages/84/dc/c55dc5678a61e5befd3694b28e0dc5737a8422334b774a4174b517c67c22/drake-1.49.0-cp312-cp312-manylinux_2_34_x86_64.whl", hash = "sha256:b9a5b528d430764ce1670918b8679cabbb209c8daa2440824ac3a9832c686591", size = 41432263, upload-time = "2026-01-15T19:44:14.486Z" }, + { url = "https://files.pythonhosted.org/packages/54/71/4c983e130ae86479a44e8afd737fd70c12ea5cf2c71bd96491339ce74ec9/drake-1.49.0-cp313-cp313-macosx_15_0_arm64.whl", hash = "sha256:ee647a022dcce17869c7a2d3d37e1d69f1f12bca8d7e3d7e84a743cfd1230be6", size = 33962409, upload-time = "2026-01-15T19:44:17.588Z" }, + { url = "https://files.pythonhosted.org/packages/3f/a8/1a46831f5f802088df9cd92c204b888aef4e3659d9702128533aa4e5ebaa/drake-1.49.0-cp313-cp313-manylinux_2_34_x86_64.whl", hash = "sha256:0a51abf867d534cef1343381ce79883acc606d52fc56debf2dd9e306982e8910", size = 41438880, upload-time = "2026-01-15T19:44:20.265Z" }, + { url = "https://files.pythonhosted.org/packages/37/b8/053e86e0dd3209cd1ba239dd49dd3313608a63958e89914ff371113be54b/drake-1.49.0-cp314-cp314-macosx_15_0_arm64.whl", hash = "sha256:16485819958e5bdd5d7815e89c7d87ab77cf3822e19ae5f7e53dbff510a56013", size = 33998519, upload-time = "2026-01-15T19:44:22.805Z" }, + { url = "https://files.pythonhosted.org/packages/d8/60/cdbc3101bb2bd57706a6b6c5a7fc68a03270f002af1d448da875f3eff5df/drake-1.49.0-cp314-cp314-manylinux_2_34_x86_64.whl", hash = "sha256:775740e9500ab8cb2e0af0e69ab162018ac03f7553b6fe03fc6b4f03c4b01092", size = 41509337, upload-time = "2026-01-15T19:44:25.879Z" }, +] + [[package]] name = "durationpy" version = "0.10" @@ -2395,6 +2421,14 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f7/02/6e639e90f181dc9127046e00d0528f9f7ad12d428972e3a5378b9aefdb0b/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.p39.p310.p311.p312.p313-none-manylinux_2_28_x86_64.whl", hash = "sha256:7916034efa867927892635733a3b6af8cd95ceb10566fd7f1e0d2763c2ee8b12", size = 243525, upload-time = "2025-09-12T08:54:34.006Z" }, { url = "https://files.pythonhosted.org/packages/84/06/cb588ca65561defe0fc48d1df4c2ac12569b81231ae4f2b52ab37007d0bd/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.p39.p310.p311.p312.p313-none-win32.whl", hash = "sha256:6c9549da71b93e367b4d71438798daae1da2592039fd14204a80a1a2348ae127", size = 552685, upload-time = "2025-09-12T08:54:35.723Z" }, { url = "https://files.pythonhosted.org/packages/86/27/00c9c96af18ac0a5eac2ff61cbe306551a2d770d7173f396d0792ee1a59e/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.p39.p310.p311.p312.p313-none-win_amd64.whl", hash = "sha256:6292d5d6634d668cd23d337e6089491d3945a9aa4ac6e1667b0003520d7caa51", size = 559466, upload-time = "2025-09-12T08:54:37.661Z" }, + { url = "https://files.pythonhosted.org/packages/b3/87/de0b33f6f00687499ca1371f22aa73396341b85bf88f1a284f9da8842493/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-macosx_10_6_intel.whl", hash = "sha256:2aab89d2d9535635ba011fc7303390685169a1aa6731ad580d08d043524b8899", size = 105326, upload-time = "2026-01-28T05:57:56.083Z" }, + { url = "https://files.pythonhosted.org/packages/b6/a6/6ea2f73ad4474896d9e38b3ffbe6ffd5a802c738392269e99e8c6621a461/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-macosx_11_0_arm64.whl", hash = "sha256:23936202a107039b5372f0b88ae1d11080746aa1c78910a45d4a0c4cf408cfaa", size = 102180, upload-time = "2026-01-28T05:57:57.787Z" }, + { url = "https://files.pythonhosted.org/packages/58/19/d81b19e8261b9cb51b81d1402167791fef81088dfe91f0c4e9d136fdc5ca/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-manylinux2014_aarch64.whl", hash = "sha256:7be06d0838f61df67bd54cb6266a6193d54083acb3624ff3c3812a6358406fa4", size = 230038, upload-time = "2026-01-28T05:57:59.105Z" }, + { url = "https://files.pythonhosted.org/packages/e2/fa/b035636cd82198b97b51a93efe9cfc4343d6b15cefbd336a3f2be871d848/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-manylinux2014_x86_64.whl", hash = "sha256:91d36b3582a766512eff8e3b5dcc2d3ffcbf10b7cf448551085a08a10f1b8244", size = 241983, upload-time = "2026-01-28T05:58:00.352Z" }, + { url = "https://files.pythonhosted.org/packages/ff/b4/f7b6cc022dd7c68b6c702d19da5d591f978f89c958b9bd3090615db0c739/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-manylinux_2_28_aarch64.whl", hash = "sha256:27c9e9a2d5e1dc3c9e3996171d844d9df9a5a101e797cb94cce217b7afcf8fd9", size = 231053, upload-time = "2026-01-28T05:58:01.683Z" }, + { url = "https://files.pythonhosted.org/packages/5a/3f/efeb7c6801c46e11bd666a5180f0d615f74f72264212f74f39586c6fda9d/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-manylinux_2_28_x86_64.whl", hash = "sha256:ce6724bb7cb3d0543dcba17206dce909f94176e68220b8eafee72e9f92bcf542", size = 243522, upload-time = "2026-01-28T05:58:03.517Z" }, + { url = "https://files.pythonhosted.org/packages/cf/b9/b04c3aa0aad2870cfe799f32f8b59789c98e1816bbce9e83f4823c5b840b/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-win32.whl", hash = "sha256:fca724a21a372731edb290841edd28a9fb1ee490f833392752844ac807c0086a", size = 552682, upload-time = "2026-01-28T05:58:05.649Z" }, + { url = "https://files.pythonhosted.org/packages/bd/e1/6d6816b296a529ac9b897ad228b1e084eb1f92319e96371880eebdc874a6/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-win_amd64.whl", hash = "sha256:823c0bd7770977d4b10e0ed0aef2f3682276b7c88b8b65cfc540afce5951392f", size = 559464, upload-time = "2026-01-28T05:58:07.261Z" }, ] [[package]] @@ -4559,6 +4593,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl", hash = "sha256:52d4362373dcf7c52546bc4af9a86ee7c4579df9a8dc268be0a2f949d376cc9b", size = 69667, upload-time = "2025-09-02T15:23:09.635Z" }, ] +[[package]] +name = "mosek" +version = "11.1.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/04/76/895d637c7e0561fd0c630049680fe616cb279e0050597d8731b7c3da426f/mosek-11.1.2-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:c56a776188951b9028b47f9e814eb21ac0c8b44b564baf8c584277c9f61a4277", size = 8791845, upload-time = "2026-01-07T08:21:55.876Z" }, + { url = "https://files.pythonhosted.org/packages/c3/e9/253e759e6e00b9cfbb4e95e7fe079b0e971b3c81c75f059bf2c2be3216e9/mosek-11.1.2-cp39-abi3-manylinux2014_x86_64.whl", hash = "sha256:5c3566d2a603d94a1773bcd27097c8390dba1d9a1543534f3527deb56f1d0a55", size = 15359313, upload-time = "2026-01-07T08:22:00.805Z" }, + { url = "https://files.pythonhosted.org/packages/41/ea/17bb932e0d307c31de685ba817a3cba822e2757a9810e7cc516778c2baa3/mosek-11.1.2-cp39-abi3-manylinux_2_27_aarch64.whl", hash = "sha256:67c13d56a9b7adf2670e4ed6fb62aa92560ae2ff1050f6e756d0d3f82c42c19f", size = 11073007, upload-time = "2026-01-07T08:22:03.118Z" }, + { url = "https://files.pythonhosted.org/packages/f2/67/6f2b6e544cf5e284c7f0baebffbc82b55e7db5b7ed5d711b621fa965d4df/mosek-11.1.2-cp39-abi3-win_amd64.whl", hash = "sha256:ad81cfd53af508db89241c7869ddce7ceaae13ef057f7b98007d57dccbb63c92", size = 11191977, upload-time = "2026-01-07T08:22:05.845Z" }, +] + [[package]] name = "mpmath" version = "1.3.0" @@ -6716,6 +6765,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c1/60/5d4751ba3f4a40a6891f24eec885f51afd78d208498268c734e256fb13c4/pydantic_settings-2.12.0-py3-none-any.whl", hash = "sha256:fddb9fd99a5b18da837b29710391e945b1e30c135477f484084ee513adb93809", size = 51880, upload-time = "2025-11-10T14:25:45.546Z" }, ] +[[package]] +name = "pydot" +version = "4.0.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyparsing" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/50/35/b17cb89ff865484c6a20ef46bf9d95a5f07328292578de0b295f4a6beec2/pydot-4.0.1.tar.gz", hash = "sha256:c2148f681c4a33e08bf0e26a9e5f8e4099a82e0e2a068098f32ce86577364ad5", size = 162594, upload-time = "2025-06-17T20:09:56.454Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7e/32/a7125fb28c4261a627f999d5fb4afff25b523800faed2c30979949d6facd/pydot-4.0.1-py3-none-any.whl", hash = "sha256:869c0efadd2708c0be1f916eb669f3d664ca684bc57ffb7ecc08e70d5e93fee6", size = 37087, upload-time = "2025-06-17T20:09:55.25Z" }, +] + [[package]] name = "pydub" version = "0.25.1"