diff --git a/dimos/core/rpc_client.py b/dimos/core/rpc_client.py index 84de18d671..13748e9666 100644 --- a/dimos/core/rpc_client.py +++ b/dimos/core/rpc_client.py @@ -62,12 +62,16 @@ def __call__(self, *args, **kwargs): # type: ignore[no-untyped-def] # For stop, use call_nowait to avoid deadlock # (the remote side stops its RPC service before responding) if self._name == "stop": - self._rpc.call_nowait(f"{self._remote_name}/{self._name}", (args, kwargs)) # type: ignore[arg-type] + self._rpc.call_nowait(f"{self._remote_name}/{self._name}", (list(args), kwargs)) # type: ignore[arg-type] if self._stop_rpc_client: self._stop_rpc_client() return None - result, unsub_fn = self._rpc.call_sync(f"{self._remote_name}/{self._name}", (args, kwargs)) # type: ignore[arg-type] + timeout = kwargs.pop("rpc_timeout", None) + sync_args: dict[str, float] = {"rpc_timeout": timeout} if timeout is not None else {} + result, unsub_fn = self._rpc.call_sync( + f"{self._remote_name}/{self._name}", (list(args), kwargs), **sync_args + ) self._unsub_fns.append(unsub_fn) return result diff --git a/dimos/hardware/manipulators/xarm/adapter.py b/dimos/hardware/manipulators/xarm/adapter.py index 3e24c530d1..013048c8cb 100644 --- a/dimos/hardware/manipulators/xarm/adapter.py +++ b/dimos/hardware/manipulators/xarm/adapter.py @@ -20,11 +20,14 @@ from __future__ import annotations +import logging import math from typing import TYPE_CHECKING from xarm.wrapper import XArmAPI +logger = logging.getLogger(__name__) + if TYPE_CHECKING: from dimos.hardware.manipulators.registry import AdapterRegistry @@ -74,6 +77,10 @@ def connect(self) -> bool: print(f"ERROR: XArm at {self._ip} not reachable (connected=False)") return False + # Clear any stale errors and enable motion before setting mode + self._arm.clean_error() + self._arm.motion_enable(enable=True) + # Initialize to servo mode for high-frequency control self._arm.set_mode(_XARM_MODE_SERVO_CARTESIAN) # Mode 1 = servo mode self._arm.set_state(0) @@ -221,6 +228,12 @@ def write_joint_positions( # Use set_servo_angle_j for high-frequency servo control (100Hz+) # This only executes the last instruction, suitable for real-time control code: int = self._arm.set_servo_angle_j(angles, speed=100, mvacc=500) + if code != 0: + logger.error( + f"[XArm] set_servo_angle_j failed: code={code}, " + f"error_code={self._arm.error_code}, warn_code={self._arm.warn_code}, " + f"state={self._arm.state}, mode={self._arm.mode}" + ) return code == 0 def write_joint_velocities(self, velocities: list[float]) -> bool: diff --git a/dimos/manipulation/blueprints.py b/dimos/manipulation/blueprints.py index a9fb0fb44b..1117630fe1 100644 --- a/dimos/manipulation/blueprints.py +++ b/dimos/manipulation/blueprints.py @@ -37,6 +37,8 @@ from dimos.core.blueprints import autoconnect from dimos.core.transport import LCMTransport from dimos.hardware.sensors.camera.realsense.camera import RealSenseCamera +from dimos.manipulation.bt.bt_manipulation_module import bt_manipulation_module +from dimos.manipulation.bt.pick_place_module import PickPlaceModule from dimos.manipulation.manipulation_module import ManipulationModule from dimos.manipulation.pick_and_place_module import PickAndPlaceModule from dimos.manipulation.planning.spec.config import RobotModelConfig @@ -45,6 +47,7 @@ from dimos.msgs.geometry_msgs.Transform import Transform from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.perception.detection.detectors.yoloe import YoloePromptMode from dimos.perception.object_scene_registration import ObjectSceneRegistrationModule from dimos.robot.foxglove_bridge import FoxgloveBridge # TODO: migrate to rerun from dimos.utils.data import get_data @@ -210,6 +213,13 @@ def _make_xarm7_config( if add_gripper: xacro_args["add_gripper"] = "true" + # xArm7 joint limits from URDF (xarm7.urdf.xacro with limited=true) + # Without these, the planner defaults to +- pi which rejects valid joint + # positions beyond 3.14 rad (joints 1,3,5,7 go to +- 2Pi). + _TWO_PI = 2.0 * math.pi + joint_limits_lower = [-_TWO_PI, -2.059, -_TWO_PI, -0.19198, -_TWO_PI, -1.69297, -_TWO_PI] + joint_limits_upper = [_TWO_PI, 2.0944, _TWO_PI, 3.927, _TWO_PI, math.pi, _TWO_PI] + return RobotModelConfig( name=name, urdf_path=_get_xarm_urdf_path(), @@ -227,6 +237,8 @@ def _make_xarm7_config( coordinator_task_name=coordinator_task, gripper_hardware_id=gripper_hardware_id, tf_extra_links=tf_extra_links or [], + joint_limits_lower=joint_limits_lower, + joint_limits_upper=joint_limits_upper, # Home configuration: arm extended forward, elbow up (safe observe pose) home_joints=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], ) @@ -403,12 +415,16 @@ def _make_piper_config( ], planning_timeout=10.0, enable_viz=True, + graspgen_topk_num_grasps=400, + graspgen_save_visualization_data=True, ), RealSenseCamera.blueprint( base_frame_id="link7", base_transform=_XARM_PERCEPTION_CAMERA_TRANSFORM, ), - ObjectSceneRegistrationModule.blueprint(target_frame="world"), + ObjectSceneRegistrationModule.blueprint( + target_frame="world", prompt_mode=YoloePromptMode.PROMPT + ), FoxgloveBridge.blueprint(), # TODO: migrate to rerun ) .transports( @@ -458,10 +474,81 @@ def _make_piper_config( McpClient.blueprint(system_prompt=_MANIPULATION_AGENT_SYSTEM_PROMPT), ) +_BT_AGENT_SYSTEM_PROMPT = """\ +You are a robotic manipulation assistant controlling an xArm7 robot arm. + +Available skills: +- pick: Pick up an object by name. Uses DL-based grasp generation with retry and verification. \ +The BT automatically scans and detects the object. +- place: Place a held object at x, y, z position (meters, world frame). +- place_back: ONLY when user says "put it back" or "return the object". \ +Releases the object at its original pick position. +- pick_and_place: Pick an object and place it at target location. +- go_home: Move arm to home/safe position. Use for "come back", "go back", \ +"return home", "come here". Keeps held object — does NOT release it. +- stop: Emergency stop — cancels motion and opens gripper. + +IMPORTANT: "come back", "go back", "return" = go_home (keeps object). \ +"put it back", "return the object" = place_back (releases object). \ +When in doubt, use go_home. + +Workflow: Just call pick with the object name — the BT handles scanning automatically. \ +If the user mentions a named location (e.g. "the box"), ask for or remember its \ +coordinates from conversation and use place(x, y, z) directly. \ +Do NOT sequence low-level operations. Use pick/place for all manipulation. +For robot_name parameters, always omit or pass None (single-arm setup). +""" + +# BT Agent mode — uses BTManipulationModule (no skills) + BT PickPlaceModule (skills only) +# DL grasps via GraspGen Docker (requires GPU + Docker build) +_bt_xarm_perception = ( + autoconnect( + bt_manipulation_module( + robots=[ + _make_xarm7_config( + "arm", + pitch=math.radians(45), + joint_prefix="arm_", + coordinator_task="traj_arm", + add_gripper=True, + gripper_hardware_id="arm", + tf_extra_links=["link7"], + ), + ], + planning_timeout=10.0, + enable_viz=True, + graspgen_topk_num_grasps=400, + graspgen_save_visualization_data=True, + ), + RealSenseCamera.blueprint( + base_frame_id="link7", + base_transform=_XARM_PERCEPTION_CAMERA_TRANSFORM, + ), + ObjectSceneRegistrationModule.blueprint( + target_frame="world", prompt_mode=YoloePromptMode.PROMPT + ), + FoxgloveBridge.blueprint(), # TODO: migrate to rerun + ) + .transports( + { + ("joint_state", JointState): LCMTransport("/coordinator/joint_state", JointState), + } + ) + .global_config(viewer="foxglove") +) + +bt_pick_place_agent = autoconnect( + _bt_xarm_perception, + PickPlaceModule.blueprint(), + McpServer.blueprint(), + McpClient.blueprint(system_prompt=_BT_AGENT_SYSTEM_PROMPT), +) + __all__ = [ "PIPER_GRIPPER_COLLISION_EXCLUSIONS", "XARM_GRIPPER_COLLISION_EXCLUSIONS", + "bt_pick_place_agent", "dual_xarm6_planner", "xarm6_planner_only", "xarm7_planner_coordinator", diff --git a/dimos/manipulation/bt/actions.py b/dimos/manipulation/bt/actions.py new file mode 100644 index 0000000000..b1b11a2b07 --- /dev/null +++ b/dimos/manipulation/bt/actions.py @@ -0,0 +1,845 @@ +# 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. + +"""BT action leaf nodes for pick-and-place orchestration. + +Each action wraps one or more BTManipulationModule / OSR / GraspGen RPC calls, +following the py_trees lifecycle: initialise() → update() → terminate(). + +Long-running operations (trajectory execution, gripper settle) return RUNNING +and poll for completion on subsequent ticks. +""" + +from __future__ import annotations + +import math +import threading +import time +from typing import TYPE_CHECKING, Any + +import py_trees +from py_trees.common import Status + +from dimos.manipulation.grasping.gripper_adapter import GripperAdapter +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState +from dimos.utils.logging_config import setup_logger +from dimos.utils.transform_utils import offset_distance + +if TYPE_CHECKING: + from dimos.manipulation.bt.pick_place_module import PickPlaceModule + +logger = setup_logger() + + +class ManipulationAction(py_trees.behaviour.Behaviour): + """Base class for all BT action nodes. + + Holds a reference to the PickPlaceModule so leaf nodes can call RPCs + and read configuration. + """ + + def __init__(self, name: str, module: PickPlaceModule) -> None: + super().__init__(name=name) + self.module = module + self.bb = self.attach_blackboard_client(name=self.name) + # Register common keys for error/result messages + self.bb.register_key(key="error_message", access=py_trees.common.Access.WRITE) + self.bb.register_key(key="result_message", access=py_trees.common.Access.WRITE) + + +class ScanObjects(ManipulationAction): + """Set detection prompts, refresh obstacles, and populate blackboard detections.""" + + def __init__(self, name: str, module: PickPlaceModule, min_duration: float = 1.0) -> None: + super().__init__(name, module) + self.min_duration = min_duration + self.bb.register_key(key="detections", access=py_trees.common.Access.WRITE) + self.bb.register_key(key="object_name", access=py_trees.common.Access.READ) + + def update(self) -> Status: + try: + object_name: str | None = getattr(self.bb, "object_name", None) + if object_name: + self.module.get_rpc_calls("ObjectSceneRegistrationModule.set_prompts")( + text=[object_name] + ) + time.sleep(self.module.config.prompt_settle_time) + + self.module.get_rpc_calls("BTManipulationModule.refresh_obstacles")(self.min_duration) + + detections = ( + self.module.get_rpc_calls("BTManipulationModule.list_cached_detections")() or [] + ) + + self.bb.detections = detections + logger.info(f"[ScanObjects] Found {len(detections)} detection(s)") + return Status.SUCCESS if detections else Status.FAILURE + except Exception as e: + logger.error(f"[ScanObjects] Failed: {e}") + self.bb.error_message = f"Error: Scan failed — {e}" + return Status.FAILURE + + +class FindObject(ManipulationAction): + """Find target object in detection list by name or object_id.""" + + def __init__(self, name: str, module: PickPlaceModule) -> None: + super().__init__(name, module) + self.bb.register_key(key="detections", access=py_trees.common.Access.READ) + self.bb.register_key(key="object_name", access=py_trees.common.Access.READ) + self.bb.register_key(key="object_id", access=py_trees.common.Access.READ) + self.bb.register_key(key="target_object", access=py_trees.common.Access.WRITE) + + def update(self) -> Status: + detections: list[dict[str, Any]] = self.bb.detections + object_name: str = self.bb.object_name + object_id: str | None = getattr(self.bb, "object_id", None) + + for det in detections: + if object_id and det.get("object_id") == object_id: + self.bb.target_object = det + logger.info(f"[FindObject] Found by ID: {det.get('name')}") + return Status.SUCCESS + det_name = det.get("name", "") + if object_name.lower() in det_name.lower() or det_name.lower() in object_name.lower(): + self.bb.target_object = det + logger.info(f"[FindObject] Found '{det_name}' matching '{object_name}'") + return Status.SUCCESS + + available = [d.get("name", "?") for d in detections] + msg = f"Error: Object '{object_name}' not found. Available: {available}" + logger.warning(f"[FindObject] {msg}") + self.bb.error_message = msg + return Status.FAILURE + + +class GetObjectPointcloud(ManipulationAction): + """Fetch object pointcloud from ObjectSceneRegistrationModule.""" + + def __init__(self, name: str, module: PickPlaceModule) -> None: + super().__init__(name, module) + self.bb.register_key(key="target_object", access=py_trees.common.Access.READ) + self.bb.register_key(key="object_pointcloud", access=py_trees.common.Access.WRITE) + + def update(self) -> Status: + target: dict[str, Any] = self.bb.target_object + obj_id = target.get("object_id") + obj_name = target.get("name", "object") + + try: + if obj_id: + pc = self.module.get_rpc_calls( + "ObjectSceneRegistrationModule.get_object_pointcloud_by_object_id" + )(obj_id) + else: + pc = self.module.get_rpc_calls( + "ObjectSceneRegistrationModule.get_object_pointcloud_by_name" + )(obj_name) + + if pc is None: + self.bb.error_message = f"Error: No pointcloud for '{obj_name}'" + return Status.FAILURE + + self.bb.object_pointcloud = pc + logger.info(f"[GetObjectPointcloud] Got pointcloud for '{obj_name}'") + return Status.SUCCESS + except Exception as e: + logger.error(f"[GetObjectPointcloud] Failed: {e}") + self.bb.error_message = f"Error: Pointcloud fetch failed — {e}" + return Status.FAILURE + + +class GetScenePointcloud(ManipulationAction): + """Fetch full scene pointcloud for collision filtering.""" + + def __init__(self, name: str, module: PickPlaceModule) -> None: + super().__init__(name, module) + self.bb.register_key(key="target_object", access=py_trees.common.Access.READ) + self.bb.register_key(key="scene_pointcloud", access=py_trees.common.Access.WRITE) + + def update(self) -> Status: + try: + target: dict[str, Any] = self.bb.target_object + exclude_id = target.get("object_id") + pc = self.module.get_rpc_calls( + "ObjectSceneRegistrationModule.get_full_scene_pointcloud" + )(exclude_object_id=exclude_id) + self.bb.scene_pointcloud = pc + return Status.SUCCESS + except Exception as e: + logger.warning(f"[GetScenePointcloud] Could not get scene PC: {e}") + self.bb.scene_pointcloud = None + return Status.SUCCESS # Non-fatal — grasps work without scene PC + + +class GenerateGrasps(ManipulationAction): + """Generate DL-based grasps via GraspGen Docker module. + + Runs the RPC in a background thread so Docker startup doesn't block the + BT tick loop. Returns RUNNING while in-flight, falls back to heuristic on FAILURE. + """ + + def __init__(self, name: str, module: PickPlaceModule, rpc_timeout: float = 1800.0) -> None: + super().__init__(name, module) + self.rpc_timeout = rpc_timeout + self.bb.register_key(key="object_pointcloud", access=py_trees.common.Access.READ) + self.bb.register_key(key="scene_pointcloud", access=py_trees.common.Access.READ) + self.bb.register_key(key="grasp_candidates", access=py_trees.common.Access.WRITE) + self.bb.register_key(key="grasp_index", access=py_trees.common.Access.WRITE) + self._thread: threading.Thread | None = None + self._result: Any = None + self._error: Exception | None = None + self._run_id: object = object() + + def initialise(self) -> None: + self._run_id = object() # fresh token each run + current_id = self._run_id + self._result = None + self._error = None + obj_pc = self.bb.object_pointcloud + scene_pc = getattr(self.bb, "scene_pointcloud", None) + + def _run() -> None: + try: + result = self.module.get_rpc_calls("BTManipulationModule.generate_grasps")( + obj_pc, scene_pc, rpc_timeout=self.rpc_timeout + ) + if current_id is self._run_id: + self._result = result + except Exception as e: + if current_id is self._run_id: + self._error = e + + self._thread = threading.Thread(target=_run, daemon=True, name="GenerateGrasps-RPC") + self._thread.start() + logger.info("[GenerateGrasps] Started grasp generation (may include Docker startup)") + + def update(self) -> Status: + # Thread still running — Docker may be initializing + if self._thread is not None and self._thread.is_alive(): + return Status.RUNNING + + # Thread completed — check results + if self._error is not None: + logger.warning(f"[GenerateGrasps] GraspGen RPC failed: {self._error}") + self.bb.error_message = f"Error: GraspGen failed — {self._error}" + return Status.FAILURE + + result = self._result + if result is not None and result.poses: + self.bb.grasp_candidates = list(result.poses) + self.bb.grasp_index = 0 + logger.info(f"[GenerateGrasps] Generated {len(result.poses)} DL grasps") + return Status.SUCCESS + + logger.warning("[GenerateGrasps] GraspGen returned no usable grasps") + self.bb.error_message = "Error: GraspGen returned no grasps" + return Status.FAILURE + + def terminate(self, new_status: Status) -> None: + if new_status == Status.INVALID: + # BT interrupted — thread may still be running in background. + # We can't cancel a blocking RPC, but we detach our reference + # so results are discarded on next initialise(). + self._thread = None + self._result = None + self._error = None + + +class GenerateHeuristicGrasps(ManipulationAction): + """Fallback: single top-down grasp from the detection center (pitch=pi).""" + + def __init__(self, name: str, module: PickPlaceModule) -> None: + super().__init__(name, module) + self.bb.register_key(key="target_object", access=py_trees.common.Access.READ) + self.bb.register_key(key="grasp_candidates", access=py_trees.common.Access.WRITE) + self.bb.register_key(key="grasp_index", access=py_trees.common.Access.WRITE) + + def update(self) -> Status: + target: dict[str, Any] = self.bb.target_object + center: list[float] | None = target.get("center") + if center is None: + self.bb.error_message = "Error: No detection center for heuristic grasp" + return Status.FAILURE + + cx, cy, cz = float(center[0]), float(center[1]), float(center[2]) + grasp_pose = Pose( + Vector3(cx, cy, cz), + Quaternion.from_euler(Vector3(0.0, math.pi, 0.0)), + ) + self.bb.grasp_candidates = [grasp_pose] + self.bb.grasp_index = 0 + logger.info(f"[GenerateHeuristicGrasps] Top-down grasp at ({cx:.3f}, {cy:.3f}, {cz:.3f})") + return Status.SUCCESS + + +class VisualizeGrasps(ManipulationAction): + """Render grasp candidates in MeshCat (best-effort, always SUCCESS).""" + + def __init__(self, name: str, module: PickPlaceModule) -> None: + super().__init__(name, module) + self.bb.register_key(key="grasp_candidates", access=py_trees.common.Access.READ) + + def update(self) -> Status: + candidates = self.bb.grasp_candidates + if not candidates: + return Status.SUCCESS + try: + self.module.get_rpc_calls("BTManipulationModule.visualize_grasps")(candidates) + except Exception as e: + logger.warning(f"[VisualizeGrasps] Visualization failed (non-fatal): {e}") + return Status.SUCCESS + + +class AdaptGrasps(ManipulationAction): + """Adapt grasps from source gripper frame to target gripper frame. Always SUCCESS.""" + + def __init__( + self, + name: str, + module: PickPlaceModule, + source_gripper: str = "robotiq_2f_140", + target_gripper: str = "ufactory_xarm", + ) -> None: + super().__init__(name, module) + self.bb.register_key(key="grasp_candidates", access=py_trees.common.Access.READ) + self.bb.register_key(key="grasp_candidates", access=py_trees.common.Access.WRITE) + self._adapter = GripperAdapter(source=source_gripper, target=target_gripper) + + def update(self) -> Status: + candidates = self.bb.grasp_candidates + if not candidates: + return Status.SUCCESS + + adapted = self._adapter.adapt_grasps(candidates) + self.bb.grasp_candidates = adapted + return Status.SUCCESS + + +class FilterGraspWorkspace(ManipulationAction): + """Filter grasps by min-Z, max-distance, and approach angle. Sorts by quality.""" + + def __init__(self, name: str, module: PickPlaceModule) -> None: + super().__init__(name, module) + self.bb.register_key(key="grasp_candidates", access=py_trees.common.Access.READ) + self.bb.register_key(key="grasp_candidates", access=py_trees.common.Access.WRITE) + self.bb.register_key(key="grasp_index", access=py_trees.common.Access.WRITE) + + def update(self) -> Status: + cfg = self.module.config + candidates: list[Pose] = list(self.bb.grasp_candidates) + total = len(candidates) + if total == 0: + return Status.SUCCESS + + cos_threshold = math.cos(cfg.max_approach_angle) + passed: list[tuple[float, Pose]] = [] + rej_z = rej_dist = rej_angle = 0 + + for pose in candidates: + if pose.position.z < cfg.min_grasp_z: + rej_z += 1 + continue + if pose.position.magnitude() > cfg.max_grasp_distance: + rej_dist += 1 + continue + # Approach direction Z-component from quaternion: rot[2,2] = 1 - 2(qx² + qy²) + qx, qy = pose.orientation.x, pose.orientation.y + approach_z = 1.0 - 2.0 * (qx * qx + qy * qy) + if approach_z > -cos_threshold: + rej_angle += 1 + continue + passed.append((approach_z, pose)) + + # Sort by approach quality — most top-down first (lowest approach_z) + passed.sort(key=lambda t: t[0]) + filtered = [p for _, p in passed] + + logger.info( + f"[FilterGraspWorkspace] {len(filtered)}/{total} passed " + f"(rejected: {rej_z} below-Z, {rej_dist} beyond-reach, " + f"{rej_angle} steep-angle)" + ) + + if not filtered: + self.bb.error_message = ( + f"Error: No grasps in workspace — {rej_z} below table, " + f"{rej_dist} beyond reach, {rej_angle} steep approach" + ) + return Status.FAILURE + + self.bb.grasp_candidates = filtered + self.bb.grasp_index = 0 + return Status.SUCCESS + + +class SelectNextGrasp(ManipulationAction): + """Select the next grasp candidate from the ranked list.""" + + def __init__(self, name: str, module: PickPlaceModule) -> None: + super().__init__(name, module) + self.bb.register_key(key="grasp_candidates", access=py_trees.common.Access.READ) + self.bb.register_key(key="grasp_index", access=py_trees.common.Access.READ) + self.bb.register_key(key="grasp_index", access=py_trees.common.Access.WRITE) + self.bb.register_key(key="current_grasp", access=py_trees.common.Access.WRITE) + self._logged_exhausted = False + + def update(self) -> Status: + candidates: list[Pose] = self.bb.grasp_candidates + idx: int = self.bb.grasp_index + + if idx >= len(candidates): + self.bb.error_message = "Error: All grasp candidates exhausted" + if not self._logged_exhausted: + logger.warning(f"[SelectNextGrasp] Exhausted {len(candidates)} candidates") + self._logged_exhausted = True + return Status.FAILURE + + self.bb.current_grasp = candidates[idx] + self.bb.grasp_index = idx + 1 + logger.info(f"[SelectNextGrasp] Selected candidate {idx + 1}/{len(candidates)}") + return Status.SUCCESS + + +class ComputePreGrasp(ManipulationAction): + """Compute pre-grasp pose offset along approach direction.""" + + def __init__(self, name: str, module: PickPlaceModule) -> None: + super().__init__(name, module) + self.bb.register_key(key="current_grasp", access=py_trees.common.Access.READ) + self.bb.register_key(key="pre_grasp_pose", access=py_trees.common.Access.WRITE) + + def update(self) -> Status: + grasp_pose: Pose = self.bb.current_grasp + offset = self.module.config.pre_grasp_offset + self.bb.pre_grasp_pose = offset_distance(grasp_pose, offset) + return Status.SUCCESS + + +class PlanToPose(ManipulationAction): + """Plan collision-free path to a pose from the blackboard. + + The target pose is read from ``bb.{pose_key}``. + """ + + def __init__(self, name: str, module: PickPlaceModule, pose_key: str) -> None: + super().__init__(name, module) + self.pose_key = pose_key + self.bb.register_key(key=pose_key, access=py_trees.common.Access.READ) + + def update(self) -> Status: + try: + pose = getattr(self.bb, self.pose_key) + logger.info( + f"[PlanToPose] Planning to {self.pose_key} " + f"({pose.position.x:.3f}, {pose.position.y:.3f}, {pose.position.z:.3f})" + ) + self.module.get_rpc_calls("BTManipulationModule.reset")() + result = self.module.get_rpc_calls("BTManipulationModule.plan_to_pose")(pose) + if result: + return Status.SUCCESS + self.bb.error_message = f"Error: Planning to {self.pose_key} failed" + return Status.FAILURE + except Exception as e: + self.bb.error_message = f"Error: Planning exception — {e}" + logger.error(f"[PlanToPose] {e}") + return Status.FAILURE + + +class PlanToJoints(ManipulationAction): + """Plan collision-free path to joint configuration from the blackboard.""" + + def __init__(self, name: str, module: PickPlaceModule, joints_key: str) -> None: + super().__init__(name, module) + self.joints_key = joints_key + self.bb.register_key(key=joints_key, access=py_trees.common.Access.READ) + + def update(self) -> Status: + try: + joints = getattr(self.bb, self.joints_key) + if isinstance(joints, list): + joints = JointState(position=joints) + logger.info(f"[PlanToJoints] Planning to {self.joints_key}") + self.module.get_rpc_calls("BTManipulationModule.reset")() + result = self.module.get_rpc_calls("BTManipulationModule.plan_to_joints")(joints) + if result: + return Status.SUCCESS + self.bb.error_message = f"Error: Joint planning to {self.joints_key} failed" + return Status.FAILURE + except Exception as e: + self.bb.error_message = f"Error: Joint planning exception — {e}" + logger.error(f"[PlanToJoints] {e}") + return Status.FAILURE + + +class ExecuteTrajectory(ManipulationAction): + """Execute planned trajectory. Returns RUNNING while executing, cancels on interrupt.""" + + def __init__( + self, + name: str, + module: PickPlaceModule, + timeout: float = 60.0, + ) -> None: + super().__init__(name, module) + self.timeout = timeout + self._start_time: float = 0.0 + self._execute_ok: bool = False + self._seen_executing: bool = False + + def initialise(self) -> None: + self._seen_executing = False + try: + result = self.module.get_rpc_calls("BTManipulationModule.execute")() + self._execute_ok = bool(result) + except Exception as e: + logger.error(f"[ExecuteTrajectory] Execute call failed: {e}") + self._execute_ok = False + self._start_time = time.time() + + def update(self) -> Status: + if not self._execute_ok: + self.bb.error_message = "Error: Trajectory execution send failed" + return Status.FAILURE + + try: + status = self.module.get_rpc_calls("BTManipulationModule.get_trajectory_status")() + except Exception as e: + logger.warning(f"[ExecuteTrajectory] Status poll failed: {e}") + status = None + + if status is not None: + state_val = int(status.get("state", -1)) if isinstance(status, dict) else int(status) + + # Track whether we've seen EXECUTING at least once to handle + # the race where IDLE is reported before the coordinator + # transitions to EXECUTING after execute(). + if state_val == TrajectoryState.EXECUTING: + self._seen_executing = True + return Status.RUNNING + if state_val == TrajectoryState.COMPLETED: + return Status.SUCCESS + if state_val == TrajectoryState.IDLE: + if self._seen_executing: + return Status.SUCCESS + return Status.RUNNING + if state_val in (TrajectoryState.ABORTED, TrajectoryState.FAULT): + self.bb.error_message = ( + f"Error: Trajectory execution failed (state={TrajectoryState(state_val).name})" + ) + return Status.FAILURE + + if time.time() - self._start_time > self.timeout: + self.bb.error_message = "Error: Trajectory execution timed out" + return Status.FAILURE + + return Status.RUNNING + + def terminate(self, new_status: Status) -> None: + if new_status == Status.INVALID: + try: + self.module.get_rpc_calls("BTManipulationModule.cancel")() + logger.info("[ExecuteTrajectory] Cancelled trajectory on interrupt") + except Exception as e: + logger.warning(f"[ExecuteTrajectory] Cancel on interrupt failed (best-effort): {e}") + + +class SetGripper(ManipulationAction): + """Set gripper to target position. Returns RUNNING during settle, then SUCCESS.""" + + def __init__( + self, + name: str, + module: PickPlaceModule, + position: float, + settle_time: float = 0.5, + ) -> None: + super().__init__(name, module) + self.position = position + self.settle_time = settle_time + self._start_time: float = 0.0 + self._command_sent: bool = False + + def initialise(self) -> None: + try: + self.module.get_rpc_calls("BTManipulationModule.set_gripper")(self.position) + self._command_sent = True + except Exception as e: + logger.error(f"[SetGripper] Command failed: {e}") + self._command_sent = False + self._start_time = time.time() + + def update(self) -> Status: + if not self._command_sent: + self.bb.error_message = "Error: Gripper command failed" + return Status.FAILURE + + if time.time() - self._start_time >= self.settle_time: + logger.info(f"[SetGripper] Gripper set to {self.position:.2f}m") + return Status.SUCCESS + return Status.RUNNING + + +class StorePickPosition(ManipulationAction): + """Store the current grasp position for place_back().""" + + def __init__(self, name: str, module: PickPlaceModule) -> None: + super().__init__(name, module) + self.bb.register_key(key="current_grasp", access=py_trees.common.Access.READ) + + def update(self) -> Status: + grasp: Pose = self.bb.current_grasp + self.module._last_pick_position = grasp.position + logger.info( + f"[StorePickPosition] Stored ({grasp.position.x:.3f}, " + f"{grasp.position.y:.3f}, {grasp.position.z:.3f})" + ) + return Status.SUCCESS + + +class ComputePlacePose(ManipulationAction): + """Compute top-down place pose and pre-place offset.""" + + def __init__(self, name: str, module: PickPlaceModule, x: float, y: float, z: float) -> None: + super().__init__(name, module) + self.x = x + self.y = y + self.z = z + self.bb.register_key(key="place_pose", access=py_trees.common.Access.WRITE) + self.bb.register_key(key="pre_place_pose", access=py_trees.common.Access.WRITE) + + def update(self) -> Status: + place_pose = Pose( + Vector3(self.x, self.y, self.z), + Quaternion.from_euler(Vector3(0.0, math.pi, 0.0)), + ) + pre_place_pose = offset_distance(place_pose, self.module.config.pre_grasp_offset) + self.bb.place_pose = place_pose + self.bb.pre_place_pose = pre_place_pose + logger.info(f"[ComputePlacePose] Place at ({self.x:.3f}, {self.y:.3f}, {self.z:.3f})") + return Status.SUCCESS + + +class SetResultMessage(ManipulationAction): + """Set a result message on the blackboard.""" + + def __init__(self, name: str, module: PickPlaceModule, message: str) -> None: + super().__init__(name, module) + self.message = message + + def update(self) -> Status: + self.bb.result_message = self.message + return Status.SUCCESS + + +class ResetRobot(ManipulationAction): + """Call BTManipulationModule.reset() to clear fault/abort state.""" + + def update(self) -> Status: + try: + result = self.module.get_rpc_calls("BTManipulationModule.reset")() + if result: + logger.info("[ResetRobot] Reset succeeded") + return Status.SUCCESS + logger.warning("[ResetRobot] Reset returned False") + self.bb.error_message = "Error: Robot reset failed" + return Status.FAILURE + except Exception as e: + logger.error(f"[ResetRobot] Reset RPC failed: {e}") + self.bb.error_message = f"Error: Robot reset exception — {e}" + return Status.FAILURE + + +class CancelMotion(ManipulationAction): + """Cancel active motion and wait for settle. Always SUCCESS.""" + + def __init__(self, name: str, module: PickPlaceModule, settle_time: float = 0.5) -> None: + super().__init__(name, module) + self.settle_time = settle_time + self._start_time: float = 0.0 + self._cancel_sent: bool = False + + def initialise(self) -> None: + self._cancel_sent = False + self._start_time = time.time() + + def update(self) -> Status: + if not self._cancel_sent: + try: + self.module.get_rpc_calls("BTManipulationModule.cancel")() + logger.info("[CancelMotion] Cancel sent, waiting for settle") + except Exception as e: + logger.warning(f"[CancelMotion] Cancel failed (best-effort): {e}") + self._cancel_sent = True + self._start_time = time.time() + + if time.time() - self._start_time >= self.settle_time: + return Status.SUCCESS + return Status.RUNNING + + +class ClearGraspState(ManipulationAction): + """Reset all grasp-related blackboard keys for a fresh rescan. Always SUCCESS.""" + + def __init__(self, name: str, module: PickPlaceModule) -> None: + super().__init__(name, module) + self.bb.register_key(key="detections", access=py_trees.common.Access.WRITE) + self.bb.register_key(key="target_object", access=py_trees.common.Access.WRITE) + self.bb.register_key(key="object_pointcloud", access=py_trees.common.Access.WRITE) + self.bb.register_key(key="scene_pointcloud", access=py_trees.common.Access.WRITE) + self.bb.register_key(key="grasp_candidates", access=py_trees.common.Access.WRITE) + self.bb.register_key(key="grasp_index", access=py_trees.common.Access.WRITE) + self.bb.register_key(key="current_grasp", access=py_trees.common.Access.WRITE) + self.bb.register_key(key="pre_grasp_pose", access=py_trees.common.Access.WRITE) + self.bb.register_key(key="has_object", access=py_trees.common.Access.WRITE) + + def update(self) -> Status: + self.bb.detections = [] + self.bb.target_object = None + self.bb.object_pointcloud = None + self.bb.scene_pointcloud = None + self.bb.grasp_candidates = [] + self.bb.grasp_index = 0 + self.bb.current_grasp = None + self.bb.pre_grasp_pose = None + self.bb.has_object = False + return Status.SUCCESS + + +class SetHasObject(ManipulationAction): + """Write a fixed value to ``bb.has_object``. Always SUCCESS.""" + + def __init__(self, name: str, module: PickPlaceModule, value: bool) -> None: + super().__init__(name, module) + self.value = value + self.bb.register_key(key="has_object", access=py_trees.common.Access.WRITE) + + def update(self) -> Status: + self.bb.has_object = self.value + return Status.SUCCESS + + +class ProbeGripperState(ManipulationAction): + """Query gripper and set ``bb.has_object``. Always SUCCESS.""" + + def __init__(self, name: str, module: PickPlaceModule) -> None: + super().__init__(name, module) + self.bb.register_key(key="has_object", access=py_trees.common.Access.WRITE) + + def update(self) -> Status: + try: + pos = float(self.module.get_rpc_calls("BTManipulationModule.get_gripper")()) + threshold = self.module.config.gripper_grasp_threshold + open_pos = self.module.config.gripper_open_position + # Upper bound at 90% of open prevents false positive when gripper is wide open + holding = threshold < pos < open_pos * 0.9 + self.bb.has_object = holding + logger.info(f"[ProbeGripperState] pos={pos:.4f}m, holding={holding}") + except Exception as e: + logger.warning(f"[ProbeGripperState] get_gripper failed: {e}") + self.bb.has_object = False + return Status.SUCCESS + + +class ExhaustRetriesIfHolding(ManipulationAction): + """If holding an object, clear grasp candidates to stop further retries. Always SUCCESS.""" + + def __init__(self, name: str, module: PickPlaceModule) -> None: + super().__init__(name, module) + self.bb.register_key(key="has_object", access=py_trees.common.Access.READ) + self.bb.register_key(key="grasp_candidates", access=py_trees.common.Access.WRITE) + + def update(self) -> Status: + if self.bb.has_object: + self.bb.grasp_candidates = [] + self.bb.error_message = "Error: Holding object but lift failed — retries exhausted" + logger.warning("[ExhaustRetriesIfHolding] Object held — aborting retries") + return Status.SUCCESS + + +class ComputeLiftPose(ManipulationAction): + """Compute lift pose: current EE + lift_height in Z. Falls back to bb.current_grasp.""" + + def __init__( + self, name: str, module: PickPlaceModule, lift_height: float | None = None + ) -> None: + super().__init__(name, module) + self._lift_height = lift_height + self.bb.register_key(key="current_grasp", access=py_trees.common.Access.READ) + self.bb.register_key(key="lift_pose", access=py_trees.common.Access.WRITE) + + def update(self) -> Status: + lift_h = ( + self._lift_height if self._lift_height is not None else self.module.config.lift_height + ) + + try: + base_pose = self.module.get_rpc_calls("BTManipulationModule.get_ee_pose")() + except Exception as e: + logger.warning( + f"[ComputeLiftPose] get_ee_pose failed, falling back to current_grasp: {e}" + ) + base_pose = None + if base_pose is None: + try: + base_pose = self.bb.current_grasp + except (KeyError, AttributeError): + base_pose = None + if base_pose is None: + self.bb.error_message = "Error: No EE pose available for lift" + return Status.FAILURE + + self.bb.lift_pose = Pose( + Vector3(base_pose.position.x, base_pose.position.y, base_pose.position.z + lift_h), + base_pose.orientation, + ) + return Status.SUCCESS + + +class ComputeLocalRetreatPose(ManipulationAction): + """Compute retreat pose: current EE + retreat_height in Z. FAILURE if no EE pose.""" + + def __init__( + self, name: str, module: PickPlaceModule, retreat_height: float | None = None + ) -> None: + super().__init__(name, module) + self._retreat_height = retreat_height + self.bb.register_key(key="retreat_pose", access=py_trees.common.Access.WRITE) + + def update(self) -> Status: + height = ( + self._retreat_height + if self._retreat_height is not None + else self.module.config.lift_height + ) + + try: + base_pose = self.module.get_rpc_calls("BTManipulationModule.get_ee_pose")() + except Exception as e: + logger.warning(f"[ComputeLocalRetreatPose] get_ee_pose failed: {e}") + base_pose = None + + if base_pose is None: + self.bb.error_message = "Error: Cannot determine EE position for retreat" + return Status.FAILURE + + self.bb.retreat_pose = Pose( + Vector3(base_pose.position.x, base_pose.position.y, base_pose.position.z + height), + base_pose.orientation, + ) + return Status.SUCCESS diff --git a/dimos/manipulation/bt/bt_manipulation_module.py b/dimos/manipulation/bt/bt_manipulation_module.py new file mode 100644 index 0000000000..605504e95d --- /dev/null +++ b/dimos/manipulation/bt/bt_manipulation_module.py @@ -0,0 +1,264 @@ +# 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. + +"""BT-specific ManipulationModule with perception + GraspGen integration. + +Extends ManipulationModule with the RPC methods the BT PickPlaceModule needs: +- Perception: objects stream, obstacle monitor, refresh_obstacles, list_cached_detections +- Grasping: generate_grasps (GraspGen Docker), visualize_grasps (MeshCat) + +Unlike PickAndPlaceModule, this has NO @skill methods — only the BT PickPlaceModule +exposes skills to the agent, keeping a clean separation of concerns. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from dimos.constants import DIMOS_PROJECT_ROOT +from dimos.core.core import rpc +from dimos.core.docker_runner import DockerModule as DockerRunner +from dimos.core.stream import In +from dimos.manipulation.grasping.graspgen_module import GraspGenModule +from dimos.manipulation.manipulation_module import ( + ManipulationModule, + ManipulationModuleConfig, +) +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.perception.detection.type.detection3d.object import ( + Object as DetObject, +) +from dimos.utils.data import get_data +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from dimos.msgs.geometry_msgs.PoseArray import PoseArray + from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + +logger = setup_logger() + +_GRASPGEN_VIZ_CONTAINER_DIR = "/output/graspgen" +_GRASPGEN_VIZ_CONTAINER_PATH = f"{_GRASPGEN_VIZ_CONTAINER_DIR}/visualization.json" + + +class BTManipulationModuleConfig(ManipulationModuleConfig): + """ManipulationModule config with GraspGen settings for BT workflow.""" + + graspgen_docker_image: str = "dimos-graspgen:latest" + graspgen_gripper_type: str = "robotiq_2f_140" + graspgen_num_grasps: int = 400 + graspgen_topk_num_grasps: int = 100 + graspgen_grasp_threshold: float = -1.0 + graspgen_filter_collisions: bool = False + graspgen_save_visualization_data: bool = False + graspgen_visualization_output_path: Path = ( + Path.home() / ".dimos" / "graspgen" / "visualization.json" + ) + + +class BTManipulationModule(ManipulationModule): + """ManipulationModule extended with perception and GraspGen for BT workflow. + + Provides the RPC methods that BT PickPlaceModule calls: + - refresh_obstacles, list_cached_detections (perception) + - generate_grasps, visualize_grasps (grasping) + + No @skill methods — only BT PickPlaceModule registers skills with the agent. + """ + + default_config = BTManipulationModuleConfig + config: BTManipulationModuleConfig + + # Perception input from ObjectSceneRegistrationModule + objects: In[list[DetObject]] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._graspgen: DockerRunner | None = None + self._detection_snapshot: list[DetObject] = [] + + @rpc + def start(self) -> None: + super().start() + + # Subscribe to objects port for perception obstacle integration + if self.objects is not None: + self.objects.observable().subscribe(self._on_objects) # type: ignore[no-untyped-call] + logger.info("Subscribed to objects port (async)") + + # Start obstacle monitor for perception integration + if self._world_monitor is not None: + self._world_monitor.start_obstacle_monitor() + + logger.info("BTManipulationModule started") + + def _on_objects(self, objects: list[DetObject]) -> None: + try: + if self._world_monitor is not None: + self._world_monitor.on_objects(objects) + except Exception as e: + logger.error(f"Exception in _on_objects: {e}") + + @rpc + def stop(self) -> None: + logger.info("Stopping BTManipulationModule") + with self._lock: + if self._graspgen is not None: + self._graspgen.stop() + self._graspgen = None + super().stop() + + # Override parent's @skill go_home — remove skill registration so only + # the BT PickPlaceModule's go_home is exposed to the agent. + def go_home(self, robot_name: str | None = None) -> str: + """Move to home position (RPC only, not registered as agent skill).""" + return super().go_home(robot_name) + + @rpc + def refresh_obstacles(self, min_duration: float = 0.0) -> list[dict[str, Any]]: + """Refresh perception obstacles and snapshot detections.""" + if self._world_monitor is None: + return [] + result = self._world_monitor.refresh_obstacles(min_duration) + self._detection_snapshot = self._world_monitor.get_cached_objects() + logger.info(f"Detection snapshot: {[d.name for d in self._detection_snapshot]}") + return result + + @rpc + def clear_perception_obstacles(self) -> str: + """Clear all perception obstacles from the planning world.""" + if self._world_monitor is None: + return "No world monitor available" + count = self._world_monitor.clear_perception_obstacles() + self._detection_snapshot = [] + return f"Cleared {count} perception obstacle(s) from planning world" + + @rpc + def list_cached_detections(self) -> list[dict[str, Any]]: + """List cached detections from perception.""" + if self._world_monitor is None: + return [] + return self._world_monitor.list_cached_detections() + + def _get_graspgen(self) -> DockerRunner: + """Get or create GraspGen Docker module (lazy init, thread-safe).""" + if self._graspgen is not None: + return self._graspgen + + with self._lock: + if self._graspgen is not None: + return self._graspgen + + get_data("models_graspgen") + + docker_file = ( + DIMOS_PROJECT_ROOT + / "dimos" + / "manipulation" + / "grasping" + / "docker_context" + / "Dockerfile" + ) + + docker_volumes: list[tuple[str, str, str]] = [] + if self.config.graspgen_save_visualization_data: + host_dir = self.config.graspgen_visualization_output_path.parent + host_dir.mkdir(parents=True, exist_ok=True) + docker_volumes.append((str(host_dir), _GRASPGEN_VIZ_CONTAINER_DIR, "rw")) + + graspgen = DockerRunner( + GraspGenModule, # type: ignore[arg-type] + docker_file=docker_file, + docker_build_context=DIMOS_PROJECT_ROOT, + docker_image=self.config.graspgen_docker_image, + docker_env={"CI": "1"}, + docker_volumes=docker_volumes, + gripper_type=self.config.graspgen_gripper_type, + num_grasps=self.config.graspgen_num_grasps, + topk_num_grasps=self.config.graspgen_topk_num_grasps, + grasp_threshold=self.config.graspgen_grasp_threshold, + filter_collisions=self.config.graspgen_filter_collisions, + save_visualization_data=self.config.graspgen_save_visualization_data, + visualization_output_path=_GRASPGEN_VIZ_CONTAINER_PATH, + ) + graspgen.start() + self._graspgen = graspgen + return self._graspgen + + @rpc + def generate_grasps( + self, + pointcloud: PointCloud2, + scene_pointcloud: PointCloud2 | None = None, + ) -> PoseArray | None: + """Generate grasp poses via GraspGen Docker module.""" + try: + graspgen = self._get_graspgen() + return graspgen.generate_grasps(pointcloud, scene_pointcloud, rpc_timeout=300) # type: ignore[no-any-return] + except Exception as e: + logger.error(f"Grasp generation failed: {e}") + return None + + @rpc + def visualize_grasps( + self, + poses: list[Pose], + max_grasps: int = 100, + ) -> bool: + """Render grasp candidates as gripper wireframes in MeshCat.""" + import numpy as np + from pydrake.geometry import Box, Rgba + from pydrake.math import RigidTransform + + from dimos.manipulation.grasping import visualize_grasps as viz_grasps + from dimos.msgs.geometry_msgs.Transform import Transform + + assert self._world_monitor is not None, "WorldMonitor is not initialized" + meshcat = self._world_monitor.world._meshcat # type: ignore[attr-defined] + + W = viz_grasps.GRIPPER_WIDTH / 2.0 + FL = viz_grasps.FINGER_LENGTH + PD = viz_grasps.PALM_DEPTH + TUBE = 0.004 + + parts = [ + ("palm", Box(W * 2, TUBE, TUBE), [0, 0, -FL]), + ("left", Box(TUBE, TUBE, FL * 1.25), [-W, 0, -0.375 * FL]), + ("right", Box(TUBE, TUBE, FL * 1.25), [W, 0, -0.375 * FL]), + ("approach", Box(TUBE, TUBE, PD), [0, 0, -FL - PD / 2]), + ] + + meshcat.Delete("grasps") + + num = min(len(poses), max_grasps) + for i in range(num): + t = i / max(num - 1, 1) + rgba = Rgba(min(1.0, 2 * t), max(0.0, 1.0 - t), 0.0, 0.8) + grasp_mat = Transform( + translation=poses[i].position, + rotation=poses[i].orientation, + ).to_matrix() + for name, shape, xyz in parts: + local = np.eye(4) + local[:3, 3] = xyz + path = f"grasps/grasp_{i}/{name}" + meshcat.SetObject(path, shape, rgba) + meshcat.SetTransform(path, RigidTransform(grasp_mat @ local)) + + logger.info(f"Visualized {num}/{len(poses)} grasp poses in MeshCat") + return True + + +bt_manipulation_module = BTManipulationModule.blueprint diff --git a/dimos/manipulation/bt/conditions.py b/dimos/manipulation/bt/conditions.py new file mode 100644 index 0000000000..d2e876ae4c --- /dev/null +++ b/dimos/manipulation/bt/conditions.py @@ -0,0 +1,254 @@ +# 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. + +"""BT condition leaf nodes for pick-and-place orchestration. + +Condition nodes are instant checks (never return RUNNING). They test +predicates on blackboard state or robot hardware and return SUCCESS/FAILURE. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import py_trees +from py_trees.common import Status + +from dimos.msgs.trajectory_msgs.TrajectoryStatus import TrajectoryState +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from dimos.manipulation.bt.pick_place_module import PickPlaceModule + from dimos.msgs.geometry_msgs.Pose import Pose + +logger = setup_logger() + + +class ManipulationCondition(py_trees.behaviour.Behaviour): + """Base class for all BT condition nodes.""" + + def __init__(self, name: str, module: PickPlaceModule) -> None: + super().__init__(name=name) + self.module = module + self.bb = self.attach_blackboard_client(name=self.name) + + +class HasDetections(ManipulationCondition): + """Check whether the blackboard has a non-empty detection list.""" + + def __init__(self, name: str, module: PickPlaceModule) -> None: + super().__init__(name, module) + self.bb.register_key(key="detections", access=py_trees.common.Access.READ) + + def update(self) -> Status: + try: + detections = self.bb.detections + except KeyError: + return Status.FAILURE + + if detections: + return Status.SUCCESS + return Status.FAILURE + + +class GripperHasObject(ManipulationCondition): + """Verify grasp by checking gripper position against threshold. + + Position above threshold = object present (fingers stopped on object). + Sets ``bb.has_object`` accordingly. + """ + + def __init__(self, name: str, module: PickPlaceModule) -> None: + super().__init__(name, module) + self.bb.register_key(key="error_message", access=py_trees.common.Access.WRITE) + self.bb.register_key(key="has_object", access=py_trees.common.Access.WRITE) + + def update(self) -> Status: + try: + position = self.module.get_rpc_calls("BTManipulationModule.get_gripper")() + pos = float(position) + except Exception as e: + logger.error(f"[GripperHasObject] get_gripper RPC failed: {e}") + self.bb.error_message = f"Error: Gripper query failed — {e}" + self.bb.has_object = False + return Status.FAILURE + + threshold = self.module.config.gripper_grasp_threshold + # No upper bound needed — called right after CloseGripper, + # unlike ProbeGripperState which guards against open gripper. + has_object = pos > threshold + + if has_object: + logger.info( + f"[GripperHasObject] Object detected (position={pos:.4f}m > " + f"threshold={threshold:.4f}m)" + ) + self.bb.has_object = True + return Status.SUCCESS + + logger.warning( + f"[GripperHasObject] No object (position={pos:.4f}m <= threshold={threshold:.4f}m)" + ) + self.bb.has_object = False + self.bb.error_message = "Error: Grasp verification failed — gripper empty" + return Status.FAILURE + + +class RobotIsHealthy(ManipulationCondition): + """Check that the robot is ready for new commands. + + FAIL if trajectory state is EXECUTING, ABORTED, or FAULT. + FAIL if ``get_robot_info()`` returns None. + """ + + def update(self) -> Status: + try: + status = self.module.get_rpc_calls("BTManipulationModule.get_trajectory_status")() + except Exception as e: + logger.error(f"[RobotIsHealthy] get_trajectory_status failed: {e}") + return Status.FAILURE + + if status is not None: + state_val = int(status.get("state", -1)) if isinstance(status, dict) else int(status) + if state_val in ( + TrajectoryState.EXECUTING, + TrajectoryState.ABORTED, + TrajectoryState.FAULT, + ): + logger.warning(f"[RobotIsHealthy] Bad state: {TrajectoryState(state_val).name}") + return Status.FAILURE + + try: + info = self.module.get_rpc_calls("BTManipulationModule.get_robot_info")() + except Exception as e: + logger.error(f"[RobotIsHealthy] get_robot_info failed: {e}") + return Status.FAILURE + + if info is None: + logger.warning("[RobotIsHealthy] get_robot_info returned None") + return Status.FAILURE + + return Status.SUCCESS + + +class HasObject(ManipulationCondition): + """Check ``bb.has_object`` — SUCCESS if True, FAILURE otherwise.""" + + def __init__(self, name: str, module: PickPlaceModule) -> None: + super().__init__(name, module) + self.bb.register_key(key="has_object", access=py_trees.common.Access.READ) + + def update(self) -> Status: + try: + if self.bb.has_object: + return Status.SUCCESS + except (KeyError, AttributeError): + pass + return Status.FAILURE + + +class VerifyReachedPose(ManipulationCondition): + """Verify the EE has reached a target pose within tolerance.""" + + def __init__( + self, + name: str, + module: PickPlaceModule, + pose_key: str, + pos_tol: float | None = None, + rot_tol: float | None = None, + ) -> None: + super().__init__(name, module) + self.pose_key = pose_key + self.pos_tol = pos_tol + self.rot_tol = rot_tol + self.bb.register_key(key=pose_key, access=py_trees.common.Access.READ) + self.bb.register_key(key="error_message", access=py_trees.common.Access.WRITE) + + def update(self) -> Status: + pos_tol = ( + self.pos_tol if self.pos_tol is not None else self.module.config.pose_position_tolerance + ) + rot_tol = ( + self.rot_tol + if self.rot_tol is not None + else self.module.config.pose_orientation_tolerance + ) + + try: + target: Pose = getattr(self.bb, self.pose_key) + except (KeyError, AttributeError): + self.bb.error_message = f"Error: No target pose '{self.pose_key}' on blackboard" + return Status.FAILURE + + try: + actual: Pose | None = self.module.get_rpc_calls("BTManipulationModule.get_ee_pose")() + except Exception as e: + self.bb.error_message = f"Error: get_ee_pose failed — {e}" + return Status.FAILURE + + if actual is None: + self.bb.error_message = "Error: get_ee_pose returned None" + return Status.FAILURE + + pos_err = target.position.distance(actual.position) + if pos_err > pos_tol: + self.bb.error_message = ( + f"Error: Position error {pos_err:.4f}m > tolerance {pos_tol:.4f}m" + ) + logger.warning(f"[VerifyReachedPose:{self.pose_key}] {self.bb.error_message}") + return Status.FAILURE + + angle_err = target.orientation.angular_distance(actual.orientation) + if angle_err > rot_tol: + self.bb.error_message = ( + f"Error: Orientation error {angle_err:.4f}rad > tolerance {rot_tol:.4f}rad" + ) + logger.warning(f"[VerifyReachedPose:{self.pose_key}] {self.bb.error_message}") + return Status.FAILURE + + logger.info( + f"[VerifyReachedPose:{self.pose_key}] OK — " + f"pos_err={pos_err:.4f}m, angle_err={angle_err:.4f}rad" + ) + return Status.SUCCESS + + +class VerifyHoldAfterLift(ManipulationCondition): + """Re-check gripper after lift to confirm object wasn't dropped.""" + + def __init__(self, name: str, module: PickPlaceModule) -> None: + super().__init__(name, module) + self.bb.register_key(key="error_message", access=py_trees.common.Access.WRITE) + self.bb.register_key(key="has_object", access=py_trees.common.Access.WRITE) + + def update(self) -> Status: + try: + position = self.module.get_rpc_calls("BTManipulationModule.get_gripper")() + pos = float(position) + except Exception as e: + logger.error(f"[VerifyHoldAfterLift] get_gripper failed: {e}") + self.bb.error_message = f"Error: Post-lift gripper check failed — {e}" + self.bb.has_object = False + return Status.FAILURE + + threshold = self.module.config.gripper_grasp_threshold + if pos > threshold: + logger.info("[VerifyHoldAfterLift] Object still held after lift") + return Status.SUCCESS + + logger.warning("[VerifyHoldAfterLift] Object dropped during lift") + self.bb.has_object = False + self.bb.error_message = "Error: Object dropped during lift" + return Status.FAILURE diff --git a/dimos/manipulation/bt/pick_place_module.py b/dimos/manipulation/bt/pick_place_module.py new file mode 100644 index 0000000000..38bc8298b0 --- /dev/null +++ b/dimos/manipulation/bt/pick_place_module.py @@ -0,0 +1,455 @@ +# 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. + +"""BT-driven PickPlaceModule for robust pick-and-place orchestration.""" + +from __future__ import annotations + +import atexit +import subprocess +import threading +import time +from typing import TYPE_CHECKING, Any + +import py_trees +from pydantic import Field + +from dimos.agents.annotation import skill +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.transport import pLCMTransport +from dimos.manipulation.bt.trees import build_go_home_tree, build_pick_tree, build_place_tree +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from dimos.msgs.geometry_msgs.Vector3 import Vector3 + +logger = setup_logger() + + +def _cleanup_graspgen_containers() -> None: + """Stop all running GraspGen Docker containers at exit.""" + try: + result = subprocess.run( + ["docker", "ps", "-q", "--filter", "name=dimos_graspgenmodule"], + capture_output=True, + text=True, + timeout=5, + ) + container_ids = [cid for cid in result.stdout.strip().split("\n") if cid] + if container_ids: + subprocess.run(["docker", "stop", *container_ids], capture_output=True, timeout=30) + logger.info(f"[PickPlaceModule] Cleaned up {len(container_ids)} GraspGen container(s)") + except Exception: + pass + + +atexit.register(_cleanup_graspgen_containers) + + +class PickPlaceModuleConfig(ModuleConfig): + """Configuration for BT-based PickPlaceModule.""" + + # Grasp strategy — False skips DL grasps (no Docker/OSR), heuristic only + use_dl_grasps: bool = True + tick_rate_hz: float = 20.0 + max_pick_attempts: int = 3 + max_rescan_attempts: int = 3 # outer retry: rescan + regenerate grasps + + # Gripper parameters (meters) — override per-robot + gripper_open_position: float = 0.85 + gripper_close_position: float = 0.0 + gripper_grasp_threshold: float = 0.005 # above = object present + gripper_settle_time: float = 1.5 + + # Timing + execution_timeout: float | None = None # optional hard upper bound for full BT run + scan_duration: float = 1.0 + prompt_settle_time: float = 3.0 # seconds to wait after setting detection prompts + + # Gripper adaptation (GraspGen source → physical robot target) + adapt_grasps: bool = True + source_gripper: str = "robotiq_2f_140" + target_gripper: str = "ufactory_xarm" + + # Approach / lift + pre_grasp_offset: float = 0.10 + lift_height: float = 0.1 + + # Pose verification tolerances + pose_position_tolerance: float = 0.02 + pose_orientation_tolerance: float = 0.15 # radians (~8.6 deg) + grasp_position_tolerance: float = 0.05 # more lenient than pre-grasp + + # Geometric workspace filter + min_grasp_z: float = 0.05 + max_grasp_distance: float = 0.9 + max_approach_angle: float = 1.05 # radians (~60 deg) + + # Home joint configuration — resolved at runtime via get_init_joints RPC if empty + home_joints: list[float] = Field(default_factory=list) + + +class PickPlaceModule(Module): + """BT-orchestrated pick-and-place module. + + Exposes @skill methods to the agent (pick, place, stop) that build and tick + py_trees Behavior Trees. Each BT calls BTManipulationModule / OSR / GraspGen + RPCs. This module orchestrates BTManipulationModule, not replaces it. + """ + + default_config = PickPlaceModuleConfig + config: PickPlaceModuleConfig + + rpc_calls: list[str] = [ + "BTManipulationModule.plan_to_pose", + "BTManipulationModule.plan_to_joints", + "BTManipulationModule.execute", + "BTManipulationModule.get_trajectory_status", + "BTManipulationModule.cancel", + "BTManipulationModule.reset", + "BTManipulationModule.refresh_obstacles", + "BTManipulationModule.list_cached_detections", + "BTManipulationModule.get_ee_pose", + "BTManipulationModule.get_robot_info", + "BTManipulationModule.get_init_joints", + "BTManipulationModule.set_gripper", + "BTManipulationModule.get_gripper", + "BTManipulationModule.generate_grasps", + "BTManipulationModule.visualize_grasps", + "ObjectSceneRegistrationModule.set_prompts", + "ObjectSceneRegistrationModule.get_object_pointcloud_by_name", + "ObjectSceneRegistrationModule.get_object_pointcloud_by_object_id", + "ObjectSceneRegistrationModule.get_full_scene_pointcloud", + ] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._stop_event = threading.Event() + self._lock = threading.Lock() + self._last_pick_position: Vector3 | None = None + + # Direct stop channel — CLI publishes here to bypass blocked agent thread + self._stop_transport: pLCMTransport[bool] = pLCMTransport("/bt_stop") + self._stop_transport.subscribe(self._on_stop_signal) + + # Result notification — publishes BT completion to agent via /human_input + self._result_transport: pLCMTransport[str] = pLCMTransport("/human_input") + + @rpc + def start(self) -> None: + super().start() + + def _on_stop_signal(self, _msg: object) -> None: + """Handle direct stop signal from CLI (bypasses blocked agent).""" + if self._lock.locked(): + self._stop_event.set() + logger.warning("[PickPlaceModule] Direct stop signal received from CLI") + + def _start_async(self, fn: Any, *args: object, **kwargs: object) -> str | None: + """Run *fn* in a background thread; return None (started) or error string (busy). + + When *fn* finishes, the result is published to the agent via /human_input. + """ + if not self._lock.acquire(blocking=False): + return "Error: A pick/place operation is already running" + self._stop_event.clear() + + def _worker() -> None: + try: + result = fn(*args, **kwargs) + except Exception as e: + logger.error(f"[PickPlaceModule] Async operation failed: {e}") + result = f"Error: {e}" + try: + self._notify_agent(str(result)) + finally: + self._lock.release() + + threading.Thread(target=_worker, daemon=True, name="PickPlaceBT").start() + return None + + def _notify_agent(self, result: str) -> None: + """Publish BT completion result to the agent via /human_input.""" + try: + self._result_transport.publish(f"[Operation result] {result}") + except Exception as e: + logger.error(f"[PickPlaceModule] Failed to notify agent: {e}") + + def _resolve_home_joints(self) -> list[float]: + """Return home joints: config if non-empty, else robot's init joints via RPC. + + Raises RuntimeError if neither source provides joints (fail loudly + rather than silently sending wrong joint count to a different robot). + """ + cfg_joints = self.config.home_joints + if cfg_joints and any(j != 0.0 for j in cfg_joints): + return cfg_joints + + try: + init_js = self.get_rpc_calls("BTManipulationModule.get_init_joints")() + if init_js is not None and init_js.position: + return list(init_js.position) + except Exception as e: + logger.warning(f"[PickPlaceModule] Could not fetch init joints: {e}") + + if not cfg_joints: + raise RuntimeError( + "[PickPlaceModule] Cannot resolve home joints — config.home_joints is empty " + "and get_init_joints RPC failed. Set home_joints in PickPlaceModuleConfig." + ) + return cfg_joints + + def _abort_tree(self, root: py_trees.behaviour.Behaviour, open_gripper: bool = False) -> None: + """Halt the BT and cancel any in-flight motion.""" + root.stop(py_trees.common.Status.INVALID) + try: + self.get_rpc_calls("BTManipulationModule.cancel")() + except Exception: + pass + if open_gripper: + try: + self.get_rpc_calls("BTManipulationModule.set_gripper")( + self.config.gripper_open_position + ) + except Exception: + pass + + @staticmethod + def _read_blackboard(key: str, fallback: str) -> str: + """Read a single blackboard key, returning *fallback* on any error.""" + try: + bb = py_trees.blackboard.Client(name="BBReader") + bb.register_key(key=key, access=py_trees.common.Access.READ) + return str(getattr(bb, key)) + except (KeyError, AttributeError): + return fallback + + def _tick_tree(self, root: py_trees.behaviour.Behaviour) -> str: + """Run a BT tick loop until SUCCESS, FAILURE, timeout, or stop.""" + # Reset blackboard to prevent stale data from a previous run + bb = py_trees.blackboard.Client(name="TreeReset") + defaults: dict[str, Any] = { + "detections": [], + "target_object": None, + "object_pointcloud": None, + "scene_pointcloud": None, + "grasp_candidates": [], + "grasp_index": 0, + "current_grasp": None, + "pre_grasp_pose": None, + "place_pose": None, + "pre_place_pose": None, + "lift_pose": None, + "retreat_pose": None, + "has_object": False, + "error_message": "", + "result_message": "", + } + for key, default in defaults.items(): + bb.register_key(key=key, access=py_trees.common.Access.WRITE) + setattr(bb, key, default) + + tree = py_trees.trees.BehaviourTree(root=root) + tree.setup() + + start = time.time() + period = 1.0 / self.config.tick_rate_hz + + while True: + if self._stop_event.is_set(): + self._abort_tree(root, open_gripper=True) + return "Operation stopped by user" + + tree.tick() + + if root.status == py_trees.common.Status.SUCCESS: + return self._read_blackboard("result_message", "Operation completed successfully") + + if root.status == py_trees.common.Status.FAILURE: + return self._read_blackboard("error_message", "Error: Operation failed") + + if ( + self.config.execution_timeout is not None + and time.time() - start > self.config.execution_timeout + ): + self._abort_tree(root) + return "Error: Operation timed out" + + time.sleep(period) + + def _exec_pick( + self, + object_name: str, + object_id: str | None = None, + max_attempts: int | None = None, + ) -> str: + """Synchronous pick implementation (runs in background thread).""" + attempts = max_attempts or self.config.max_pick_attempts + strategy = "DL+heuristic" if self.config.use_dl_grasps else "heuristic-only" + logger.info( + f"[PickPlaceModule] pick('{object_name}', id={object_id}, " + f"attempts={attempts}, strategy={strategy})" + ) + + root = build_pick_tree( + module=self, + object_name=object_name, + object_id=object_id, + max_attempts=attempts, + home_joints_override=self._resolve_home_joints(), + ) + return self._tick_tree(root) + + @skill + def pick( + self, + object_name: str, + object_id: str | None = None, + max_attempts: int | None = None, + ) -> str | None: + """Pick up an object using BT-orchestrated grasp with DL-based grasp generation. + + Runs asynchronously — returns immediately while the BT executes in the + background. Type ``stop`` to cancel. + """ + return self._start_async(self._exec_pick, object_name, object_id, max_attempts) + + def _exec_place(self, x: float, y: float, z: float) -> str: + """Synchronous place implementation (runs in background thread).""" + logger.info(f"[PickPlaceModule] place({x:.3f}, {y:.3f}, {z:.3f})") + root = build_place_tree(module=self, x=x, y=y, z=z) + return self._tick_tree(root) + + @skill + def place(self, x: float, y: float, z: float) -> str | None: + """Place the held object at the specified position (meters, world frame). + + Runs asynchronously. Type ``stop`` to cancel. + """ + return self._start_async(self._exec_place, x, y, z) + + @skill + def place_back(self) -> str | None: + """Place the held object back at its original pick position. + + ONLY call when the user explicitly asks to return/put back the object. + For "go home" or "return home", use go_home instead (keeps the object). + Runs asynchronously. Type ``stop`` to cancel. + """ + if self._last_pick_position is None: + return "Error: No stored pick position — pick an object first" + pos = self._last_pick_position + return self._start_async(self._exec_place, pos.x, pos.y, pos.z) + + @skill + def pick_and_place( + self, + object_name: str, + place_x: float, + place_y: float, + place_z: float, + object_id: str | None = None, + max_attempts: int | None = None, + ) -> str | None: + """Pick an object and place it at the target location. + + Sequentially runs pick then place in the background. + Type ``stop`` to cancel at any time. + """ + + def _run() -> str: + pick_result = self._exec_pick(object_name, object_id, max_attempts) + if self._stop_event.is_set() or pick_result.startswith("Error"): + return pick_result + return self._exec_place(place_x, place_y, place_z) + + return self._start_async(_run) + + @skill + def go_home(self) -> str | None: + """Move to the home/safe position while keeping any held object. + + Does NOT open the gripper if holding — only opens when nothing is held. + Use for "go back", "come home", "return", etc. Runs asynchronously. + """ + + def _run() -> str: + logger.info("[PickPlaceModule] go_home()") + root = build_go_home_tree(module=self, home_joints_override=self._resolve_home_joints()) + return self._tick_tree(root) + + return self._start_async(_run) + + @skill + def detect(self, prompts: list[str]) -> str: + """Detect objects matching the given text prompts. + + Sets detection prompts on the perception system, waits for detections, + and returns a list of detected objects with their 3D positions. + + Args: + prompts: Text descriptions of objects to detect (e.g., ["cup", "bottle"]). + """ + if not prompts: + return "No prompts provided." + + try: + self.get_rpc_calls("ObjectSceneRegistrationModule.set_prompts")(text=prompts) + except Exception as e: + return f"Error setting prompts: {e}" + + import time + + time.sleep(self.config.prompt_settle_time) + + try: + self.get_rpc_calls("BTManipulationModule.refresh_obstacles")(self.config.scan_duration) + except Exception as e: + return f"Error refreshing obstacles: {e}" + + try: + detections = self.get_rpc_calls("BTManipulationModule.list_cached_detections")() or [] + except Exception as e: + return f"Error listing detections: {e}" + + if not detections: + return "No objects detected." + + lines = [f"Detected {len(detections)} object(s):"] + for det in detections: + name = det.get("name", "unknown") + center = det.get("center", [0, 0, 0]) + x, y, z = center[0], center[1], center[2] + lines.append(f" - {name}: ({x:.3f}, {y:.3f}, {z:.3f})") + return "\n".join(lines) + + @skill + def stop(self) -> None: + """Emergency stop — cancel all motion and open gripper. Safe to call at any time.""" + if self._lock.locked(): + self._stop_event.set() + logger.warning("[PickPlaceModule] Stop requested — halting BT") + return + + try: + self.get_rpc_calls("BTManipulationModule.cancel")() + except Exception: + pass + try: + self.get_rpc_calls("BTManipulationModule.set_gripper")( + self.config.gripper_open_position + ) + except Exception: + pass diff --git a/dimos/manipulation/bt/trees.py b/dimos/manipulation/bt/trees.py new file mode 100644 index 0000000000..ebeb2b7638 --- /dev/null +++ b/dimos/manipulation/bt/trees.py @@ -0,0 +1,474 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""BT tree builders and RetryOnFailure decorator for pick-and-place.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import py_trees +from py_trees.common import Status + +from dimos.manipulation.bt import actions, conditions +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from dimos.manipulation.bt.pick_place_module import PickPlaceModule + +logger = setup_logger() + + +class RetryOnFailure(py_trees.decorators.Decorator): + """Retry the child subtree on FAILURE, up to *max_attempts*. + + On each FAILURE the child is stopped (reset) so the next tick + re-runs it from scratch. SUCCESS and RUNNING pass through. + """ + + def __init__(self, name: str, child: py_trees.behaviour.Behaviour, max_attempts: int) -> None: + super().__init__(name=name, child=child) + self.max_attempts = max_attempts + self._attempt = 0 + + def initialise(self) -> None: + self._attempt = 0 + + def update(self) -> Status: + if self.decorated.status == Status.SUCCESS: + return Status.SUCCESS + + if self.decorated.status == Status.RUNNING: + return Status.RUNNING + + self._attempt += 1 + logger.info( + f"[RetryOnFailure:{self.name}] Attempt {self._attempt}/{self.max_attempts} failed" + ) + + if self._attempt >= self.max_attempts: + logger.warning( + f"[RetryOnFailure:{self.name}] All {self.max_attempts} attempts exhausted" + ) + return Status.FAILURE + + self.decorated.stop(Status.INVALID) + return Status.RUNNING + + def terminate(self, new_status: Status) -> None: + if new_status == Status.INVALID and self.decorated.status == Status.RUNNING: + self.decorated.stop(Status.INVALID) + + +def build_pick_tree( + module: PickPlaceModule, + object_name: str, + object_id: str | None = None, + max_attempts: int = 5, + home_joints_override: list[float] | None = None, +) -> py_trees.behaviour.Behaviour: + """Build a complete pick Behavior Tree with two-level retry. + + Inner retry (GraspWithRetry) sweeps grasp candidates. + Outer retry (PickWithRescan) re-scans and regenerates grasps. + + Tree structure:: + + Sequence(memory=True) "Pick" + ├── SetHasObject(False) + ├── Selector "EnsureReady" + │ ├── RobotIsHealthy + │ └── Sequence [ResetRobot, RobotIsHealthy] + ├── RetryOnFailure "PickWithRescan" + │ └── Sequence(memory=True) "ScanAndGrasp" + │ ├── ClearGraspState + │ ├── Sequence "ReturnHome" + │ ├── Selector "EnsureScanned" [HasDetections, ScanObjects] + │ ├── FindObject + │ ├── Selector "GenerateGraspCandidates" [MLGrasps, Heuristic] + │ ├── FilterGraspWorkspace + │ ├── VisualizeGrasps + │ └── RetryOnFailure "GraspWithRetry" + │ └── Selector(memory=True) "AttemptOrRecover" + │ ├── Sequence(memory=True) "GraspAttempt" + │ │ [select, pre-grasp, approach, grasp, verify, lift] + │ └── Sequence(memory=True) "RecoverThenFail" + │ [cancel, reset, recovery(local/home), Failure] + ├── StorePickPosition + └── SetResultMessage + """ + cfg = module.config + + init_state = actions.SetHasObject("InitHasObject", module, value=False) + + ensure_ready = py_trees.composites.Selector( + "EnsureReady", + memory=False, + children=[ + conditions.RobotIsHealthy("RobotIsHealthy?", module), + py_trees.composites.Sequence( + "ResetAndVerify", + memory=True, + children=[ + actions.ResetRobot("ResetRobot", module), + conditions.RobotIsHealthy("RobotIsHealthy?2", module), + ], + ), + ], + ) + + clear_state = actions.ClearGraspState("ClearGraspState", module) + + ensure_scanned = py_trees.composites.Selector( + "EnsureScanned", + memory=False, + children=[ + conditions.HasDetections("HasDetections?", module), + actions.ScanObjects("ScanObjects", module, min_duration=cfg.scan_duration), + ], + ) + + find_object = actions.FindObject("FindObject", module) + + # ML grasps need gripper adaptation (base to contact frame); heuristic don't. + adapt_grasps_nodes: list[py_trees.behaviour.Behaviour] = [] + if cfg.adapt_grasps: + adapt_grasps_nodes.append( + actions.AdaptGrasps( + "AdaptGrasps", + module, + source_gripper=cfg.source_gripper, + target_gripper=cfg.target_gripper, + ) + ) + + if cfg.use_dl_grasps: + dl_grasps = py_trees.composites.Sequence( + "DLGrasps", + memory=True, + children=[ + actions.GetObjectPointcloud("GetObjectPointcloud", module), + actions.GetScenePointcloud("GetScenePointcloud", module), + actions.GenerateGrasps("GenerateGrasps", module), + *adapt_grasps_nodes, + ], + ) + generate_grasp_candidates: py_trees.behaviour.Behaviour = py_trees.composites.Selector( + "GenerateGraspCandidates", + memory=True, + children=[ + dl_grasps, + actions.GenerateHeuristicGrasps("HeuristicGrasps", module), + ], + ) + else: + generate_grasp_candidates = actions.GenerateHeuristicGrasps("HeuristicGrasps", module) + + grasp_attempt = py_trees.composites.Sequence( + "GraspAttempt", + memory=True, + children=[ + actions.SetHasObject("ResetHasObject", module, value=False), + actions.SelectNextGrasp("SelectNextGrasp", module), + actions.ComputePreGrasp("ComputePreGrasp", module), + actions.SetGripper( + "OpenGripper", module, position=cfg.gripper_open_position, settle_time=0.5 + ), + actions.PlanToPose("PlanToPreGrasp", module, pose_key="pre_grasp_pose"), + actions.ExecuteTrajectory("ExecuteApproach", module), + conditions.VerifyReachedPose("VerifyPreGrasp", module, pose_key="pre_grasp_pose"), + actions.PlanToPose("PlanToGrasp", module, pose_key="current_grasp"), + actions.ExecuteTrajectory("ExecuteGrasp", module), + conditions.VerifyReachedPose( + "VerifyGraspPose", + module, + pose_key="current_grasp", + pos_tol=cfg.grasp_position_tolerance, + ), + actions.SetGripper( + "CloseGripper", + module, + position=cfg.gripper_close_position, + settle_time=cfg.gripper_settle_time, + ), + conditions.GripperHasObject("VerifyGrasp", module), + actions.ComputeLiftPose("ComputeLiftPose", module), + actions.PlanToPose("PlanLift", module, pose_key="lift_pose"), + actions.ExecuteTrajectory("ExecuteLift", module), + conditions.VerifyHoldAfterLift("VerifyHoldAfterLift", module), + ], + ) + + # Recovery: conditional gripper open (skip if holding). + # Two instances needed — py_trees requires unique node objects. + maybe_open_gripper = py_trees.composites.Selector( + "MaybeOpenGripper", + memory=False, + children=[ + conditions.HasObject("HasObject?", module), + actions.SetGripper( + "OpenGripperRecovery", module, position=cfg.gripper_open_position, settle_time=0.5 + ), + ], + ) + + maybe_open_gripper2 = py_trees.composites.Selector( + "MaybeOpenGripper2", + memory=False, + children=[ + conditions.HasObject("HasObject?2", module), + actions.SetGripper( + "OpenGripperRecovery2", module, position=cfg.gripper_open_position, settle_time=0.5 + ), + ], + ) + + try_local_retreat = py_trees.composites.Sequence( + "TryLocalRetreat", + memory=True, + children=[ + maybe_open_gripper, + actions.ComputeLocalRetreatPose("ComputeRetreatPose", module), + actions.PlanToPose("PlanRetreat", module, pose_key="retreat_pose"), + actions.ExecuteTrajectory("ExecuteRetreat", module), + ], + ) + + go_home_retreat = py_trees.composites.Sequence( + "GoHomeRetreat", + memory=True, + children=[ + maybe_open_gripper2, + actions.ResetRobot("ResetForHome", module), + actions.PlanToJoints("PlanToHome", module, joints_key="home_joints"), + actions.ExecuteTrajectory("ExecuteGoHome", module), + ], + ) + + recover_then_fail = py_trees.composites.Sequence( + "RecoverThenFail", + memory=True, + children=[ + actions.CancelMotion("CancelMotion", module), + actions.ResetRobot("ResetAfterFailure", module), + actions.ExhaustRetriesIfHolding("ExhaustRetriesIfHolding", module), + py_trees.composites.Selector( + "Recovery", + memory=True, + children=[try_local_retreat, go_home_retreat], + ), + py_trees.behaviours.Failure("RecoveryComplete"), + ], + ) + + # memory=True: recovery must complete before RetryOnFailure starts the next attempt + attempt_or_recover = py_trees.composites.Selector( + "AttemptOrRecover", + memory=True, + children=[grasp_attempt, recover_then_fail], + ) + + grasp_with_retry = RetryOnFailure( + "GraspWithRetry", + child=attempt_or_recover, + max_attempts=max_attempts, + ) + + # Return home before each rescan so the arm is out of the camera view + return_home = py_trees.composites.Sequence( + "ReturnHome", + memory=True, + children=[ + actions.CancelMotion("CancelBeforeHome", module, settle_time=0.3), + actions.ResetRobot("ResetBeforeHome", module), + actions.SetGripper( + "OpenGripperBeforeHome", + module, + position=cfg.gripper_open_position, + settle_time=0.5, + ), + actions.PlanToJoints("PlanToHomeForRescan", module, joints_key="home_joints"), + actions.ExecuteTrajectory("ExecuteHomeForRescan", module), + ], + ) + + scan_and_grasp = py_trees.composites.Sequence( + "ScanAndGrasp", + memory=True, + children=[ + clear_state, + return_home, + ensure_scanned, + find_object, + generate_grasp_candidates, + actions.FilterGraspWorkspace("FilterGraspWorkspace", module), + actions.VisualizeGrasps("VisualizeGrasps", module), + grasp_with_retry, + ], + ) + + pick_with_rescan = RetryOnFailure( + "PickWithRescan", + child=scan_and_grasp, + max_attempts=cfg.max_rescan_attempts, + ) + + root = py_trees.composites.Sequence( + "Pick", + memory=True, + children=[ + init_state, + ensure_ready, + pick_with_rescan, + actions.StorePickPosition("StorePickPosition", module), + actions.SetResultMessage( + "SetResult", + module, + message=f"Pick complete — grasped '{object_name}' successfully", + ), + ], + ) + + bb = py_trees.blackboard.Client(name="PickTreeSetup") + bb.register_key(key="object_name", access=py_trees.common.Access.WRITE) + bb.register_key(key="object_id", access=py_trees.common.Access.WRITE) + bb.register_key(key="home_joints", access=py_trees.common.Access.WRITE) + bb.object_name = object_name + bb.object_id = object_id + bb.home_joints = home_joints_override if home_joints_override is not None else cfg.home_joints + + return root + + +def build_place_tree( + module: PickPlaceModule, + x: float, + y: float, + z: float, +) -> py_trees.behaviour.Behaviour: + """Build a place Behavior Tree. + + Tree structure:: + + Sequence(memory=True) "Place" + ├── ComputePlacePose + ├── PlanToPose("pre_place_pose") → ExecuteTrajectory + ├── PlanToPose("place_pose") → ExecuteTrajectory + ├── SetGripper(open) + ├── PlanToPose("pre_place_pose") → ExecuteTrajectory ← retract + └── SetResultMessage + """ + cfg = module.config + + return py_trees.composites.Sequence( + "Place", + memory=True, + children=[ + actions.ComputePlacePose("ComputePlacePose", module, x=x, y=y, z=z), + actions.PlanToPose("PlanToPrePlace", module, pose_key="pre_place_pose"), + actions.ExecuteTrajectory("ExecutePrePlace", module), + actions.PlanToPose("PlanToPlace", module, pose_key="place_pose"), + actions.ExecuteTrajectory("ExecutePlace", module), + actions.SetGripper( + "OpenGripper", + module, + position=cfg.gripper_open_position, + settle_time=0.5, + ), + actions.PlanToPose("PlanRetract", module, pose_key="pre_place_pose"), + actions.ExecuteTrajectory("ExecuteRetract", module), + actions.SetResultMessage( + "SetResult", + module, + message=f"Place complete — object released at ({x:.3f}, {y:.3f}, {z:.3f})", + ), + ], + ) + + +def build_go_home_tree( + module: PickPlaceModule, + home_joints_override: list[float] | None = None, +) -> py_trees.behaviour.Behaviour: + """Build a go-home Behavior Tree. + + Tree structure:: + + Sequence(memory=True) "GoHome" + ├── ProbeGripperState + ├── Selector "MaybeOpenGripper" [HasObject?, SetGripper(open)] + ├── FailureIsSuccess "BestEffortLift" + │ └── Selector "MaybeLiftFirst" + │ ├── Inverter(HasObject?) + │ └── Sequence [ComputeLiftPose, PlanToPose, Execute] + ├── PlanToJoints("home_joints") → ExecuteTrajectory + └── SetResultMessage + """ + cfg = module.config + + root = py_trees.composites.Sequence( + "GoHome", + memory=True, + children=[ + actions.ProbeGripperState("ProbeGripper", module), + py_trees.composites.Selector( + "MaybeOpenGripper", + memory=False, + children=[ + conditions.HasObject("HasObject?", module), + actions.SetGripper( + "OpenGripper", + module, + position=cfg.gripper_open_position, + settle_time=0.5, + ), + ], + ), + # Best-effort lift if holding (avoids dragging); FailureIsSuccess + # so a lift IK failure doesn't abort the joint-space go-home. + py_trees.decorators.FailureIsSuccess( + name="BestEffortLift", + child=py_trees.composites.Selector( + "MaybeLiftFirst", + memory=False, + children=[ + py_trees.decorators.Inverter( + name="NotHolding?", + child=conditions.HasObject("HasObject?Lift", module), + ), + py_trees.composites.Sequence( + "LiftBeforeHome", + memory=True, + children=[ + actions.ComputeLiftPose( + "ComputeSafeLift", module, lift_height=cfg.lift_height + ), + actions.PlanToPose("PlanSafeLift", module, pose_key="lift_pose"), + actions.ExecuteTrajectory("ExecuteSafeLift", module), + ], + ), + ], + ), + ), + actions.PlanToJoints("PlanToHome", module, joints_key="home_joints"), + actions.ExecuteTrajectory("ExecuteGoHome", module), + actions.SetResultMessage("SetResult", module, message="Moved to home position"), + ], + ) + + bb = py_trees.blackboard.Client(name="GoHomeSetup") + bb.register_key(key="home_joints", access=py_trees.common.Access.WRITE) + bb.home_joints = home_joints_override if home_joints_override is not None else cfg.home_joints + + return root diff --git a/dimos/manipulation/grasping/docker_context/Dockerfile b/dimos/manipulation/grasping/docker_context/Dockerfile index d10b3cac76..d256fbaaad 100644 --- a/dimos/manipulation/grasping/docker_context/Dockerfile +++ b/dimos/manipulation/grasping/docker_context/Dockerfile @@ -25,9 +25,9 @@ RUN conda init bash && \ echo 'yes' | conda tos accept --override-channels --channel defaults 2>/dev/null || true && \ conda create -n app python=3.10 -y && conda clean -afy -# Clone GraspGen repository +# Clone GraspGen repository (pin to a known-good commit for reproducibility) WORKDIR /app -RUN git clone https://github.com/NVlabs/GraspGen.git && cd GraspGen && git checkout main +RUN git clone https://github.com/NVlabs/GraspGen.git && cd GraspGen && git checkout 0c8066e478ece8920a8b624dc2d0544a07ed9f10 # Install PyTorch with CUDA 12.8 RUN conda run -n app pip install --no-cache-dir \ @@ -35,7 +35,7 @@ RUN conda run -n app pip install --no-cache-dir \ # Install GraspGen package WORKDIR /app/GraspGen -RUN conda run -n app pip install --no-cache-dir -e . +RUN conda run -n app pip install --no-cache-dir --no-deps -e . # Build CUDA extensions (pointnet2_ops) RUN conda run -n app bash -c "\ @@ -47,9 +47,18 @@ RUN conda run -n app bash -c "\ # Install torch-scatter and torch-cluster (require CUDA compilation) RUN conda run -n app pip install --no-cache-dir --no-build-isolation torch-scatter torch-cluster -# Additional dependencies +# Install ALL GraspGen runtime dependencies. +# Split into groups: CUDA-compiled packages first, then pure-Python packages. +# spconv-cu120 supports CUDA 12.x (works with 12.8 at runtime). +RUN conda run -n app pip install --no-cache-dir spconv-cu120 + RUN conda run -n app pip install --no-cache-dir \ - numpy trimesh pillow pyrender imageio scipy + numpy trimesh pillow pyrender imageio scipy \ + omegaconf hydra-core "diffusers==0.11.1" "timm==1.0.15" \ + h5py scikit-learn tensorboard transformers tensordict \ + "huggingface-hub==0.25.2" addict "yapf==0.40.1" tensorboardx \ + matplotlib webdataset "yourdfpy==0.0.56" "PyOpenGL==3.1.5" \ + torch-geometric sharedarray meshcat # Model checkpoints from LFS archive COPY data/.lfs/models_graspgen.tar.gz /tmp/ diff --git a/dimos/manipulation/grasping/graspgen_module.py b/dimos/manipulation/grasping/graspgen_module.py index ae2d59512a..cdcc6ea6f4 100644 --- a/dimos/manipulation/grasping/graspgen_module.py +++ b/dimos/manipulation/grasping/graspgen_module.py @@ -92,7 +92,14 @@ def generate_grasps( """Generate grasp poses for the given pointcloud.""" try: points = self._extract_points(pointcloud) + logger.info( + f"[GraspGenModule] Input pointcloud: {len(points)} points, " + f"frame_id='{pointcloud.frame_id}'" + ) if len(points) < 10: + logger.warning( + f"[GraspGenModule] Pointcloud too small ({len(points)} < 10) — skipping inference" + ) return None # Run inference (with optional collision filtering) @@ -101,17 +108,24 @@ def generate_grasps( scene_points = self._extract_points(scene_pointcloud) grasps, scores = self._run_inference(points, scene_points) if len(grasps) == 0: + logger.warning( + f"[GraspGenModule] Inference returned 0 grasps from {len(points)} points" + ) return None # Convert and publish results pose_array = self._grasps_to_pose_array(grasps, scores, pointcloud.frame_id) self.grasps.publish(pose_array) + logger.info( + f"[GraspGenModule] Generated {len(pose_array.poses)} grasps " + f"(from {len(points)} points)" + ) if self.config.save_visualization_data: self._save_visualization_data(points, grasps, scores, pointcloud.frame_id) return pose_array except Exception as e: - logger.error(f"Grasp generation failed: {e}") + logger.error(f"[GraspGenModule] Grasp generation failed: {e}") return None def _initialize_graspgen(self) -> bool: @@ -165,6 +179,7 @@ def _run_inference( self, object_pc: np.ndarray[Any, Any], scene_pc: np.ndarray[Any, Any] | None = None ) -> tuple[np.ndarray[Any, Any], np.ndarray[Any, Any]]: if self._sampler is None: + logger.warning("[GraspGenModule] Sampler not initialized — cannot run inference") return np.array([]), np.array([]) from grasp_gen.grasp_server import GraspGenSampler # type: ignore[import-not-found] @@ -180,14 +195,29 @@ def _run_inference( if len(object_pc) > OUTLIER_REMOVAL_THRESHOLD: pc_filtered, _ = point_cloud_outlier_removal(pc_torch) object_pc_filtered = pc_filtered.numpy() + logger.info( + f"[GraspGenModule] Outlier removal: {len(object_pc)} → {len(object_pc_filtered)} points" + ) if len(object_pc_filtered) < MIN_POINTS_FOR_INFERENCE: + logger.warning( + f"[GraspGenModule] After outlier removal only {len(object_pc_filtered)} points " + f"(< {MIN_POINTS_FOR_INFERENCE}), reverting to original {len(object_pc)} points" + ) object_pc_filtered = object_pc else: object_pc_filtered = object_pc if len(object_pc_filtered) < MIN_POINTS_FOR_INFERENCE: + logger.warning( + f"[GraspGenModule] Too few points for inference: " + f"{len(object_pc_filtered)} < {MIN_POINTS_FOR_INFERENCE}" + ) return np.array([]), np.array([]) + logger.info( + f"[GraspGenModule] Running inference with {len(object_pc_filtered)} points " + f"(num_grasps={self.config.num_grasps}, topk={self.config.topk_num_grasps})" + ) grasps, scores = GraspGenSampler.run_inference( object_pc_filtered, self._sampler, @@ -198,6 +228,7 @@ def _run_inference( ) if len(grasps) == 0: + logger.warning("[GraspGenModule] Model returned 0 grasps after inference") return np.array([]), np.array([]) grasps_np = grasps.cpu().numpy() diff --git a/dimos/manipulation/grasping/gripper_adapter.py b/dimos/manipulation/grasping/gripper_adapter.py new file mode 100644 index 0000000000..9fed36af1d --- /dev/null +++ b/dimos/manipulation/grasping/gripper_adapter.py @@ -0,0 +1,76 @@ +# 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. + +"""Adapts grasp poses between gripper geometries (base frame -> TCP frame).""" + +from __future__ import annotations + +from dataclasses import dataclass + +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + + +@dataclass(frozen=True) +class GripperGeometry: + """Parallel-jaw gripper dimensions.""" + + name: str + tcp_offset: float # Base-to-TCP distance along approach axis (m) + grasp_depth_correction: float = 0.0 # Extra approach shift for contact (m) + + +GRIPPER_GEOMETRIES: dict[str, GripperGeometry] = { + "robotiq_2f_140": GripperGeometry(name="robotiq_2f_140", tcp_offset=0.175), + "ufactory_xarm": GripperGeometry(name="ufactory_xarm", tcp_offset=0.172), + "franka_panda": GripperGeometry(name="franka_panda", tcp_offset=0.103), +} + + +class GripperAdapter: + """Shifts grasp poses along the approach axis to convert base frame -> TCP frame.""" + + def __init__(self, source: str = "robotiq_2f_140", target: str = "ufactory_xarm") -> None: + self.source = GRIPPER_GEOMETRIES[source] + self.target = GRIPPER_GEOMETRIES[target] + + def adapt_grasps(self, poses: list[Pose]) -> list[Pose]: + """Shift each pose along its +Z (approach) axis by the TCP offset.""" + shift = self.source.tcp_offset + self.target.grasp_depth_correction + adapted = [self._shift(p, shift) for p in poses] + logger.info( + f"[GripperAdapter] Adapted {len(adapted)}/{len(poses)} grasps " + f"({self.source.name} -> {self.target.name}, shift={shift:.3f}m)" + ) + return adapted + + @staticmethod + def _shift(pose: Pose, distance: float) -> Pose: + approach = Transform( + translation=pose.position, + rotation=pose.orientation, + ).to_matrix()[:3, 2] + offset = approach * distance + return Pose( + position=Vector3( + x=pose.position.x + offset[0], + y=pose.position.y + offset[1], + z=pose.position.z + offset[2], + ), + orientation=pose.orientation, + ) diff --git a/dimos/manipulation/grasping/visualize_grasps.py b/dimos/manipulation/grasping/visualize_grasps.py index 53edc1bb16..3ac5f13239 100644 --- a/dimos/manipulation/grasping/visualize_grasps.py +++ b/dimos/manipulation/grasping/visualize_grasps.py @@ -26,7 +26,7 @@ FINGER_LENGTH = 0.052 PALM_DEPTH = 0.04 MAX_GRASPS = 100 -VISUALIZATION_FILE = "/tmp/grasp_visualization.json" +VISUALIZATION_FILE = str(Path.home() / ".dimos" / "graspgen" / "visualization.json") def create_gripper_geometry(transform: np.ndarray[Any, Any], color: list[float]) -> list[Any]: diff --git a/dimos/manipulation/pick_and_place_module.py b/dimos/manipulation/pick_and_place_module.py index 81e7bcf2d3..7e1fb244ce 100644 --- a/dimos/manipulation/pick_and_place_module.py +++ b/dimos/manipulation/pick_and_place_module.py @@ -235,11 +235,60 @@ def generate_grasps( """Generate grasp poses for the given point cloud via GraspGen Docker module.""" try: graspgen = self._get_graspgen() - return graspgen.generate_grasps(pointcloud, scene_pointcloud) # type: ignore[no-any-return] + return graspgen.generate_grasps(pointcloud, scene_pointcloud, rpc_timeout=300) # type: ignore[no-any-return] except Exception as e: logger.error(f"Grasp generation failed: {e}") return None + @rpc + def visualize_grasps( + self, + poses: list[Pose], + max_grasps: int = 100, + ) -> bool: + """Render grasp candidates as gripper wireframes in MeshCat.""" + import numpy as np + from pydrake.geometry import Box, Rgba + from pydrake.math import RigidTransform + + import dimos.manipulation.grasping.visualize_grasps as viz_grasps + from dimos.msgs.geometry_msgs.Transform import Transform + + assert self._world_monitor is not None, "WorldMonitor is not initialized" + meshcat = self._world_monitor.world._meshcat # type: ignore[attr-defined] + + W = viz_grasps.GRIPPER_WIDTH / 2.0 + FL = viz_grasps.FINGER_LENGTH + PD = viz_grasps.PALM_DEPTH + TUBE = 0.004 + + parts = [ + ("palm", Box(W * 2, TUBE, TUBE), [0, 0, -FL]), + ("left", Box(TUBE, TUBE, FL * 1.25), [-W, 0, -0.375 * FL]), + ("right", Box(TUBE, TUBE, FL * 1.25), [W, 0, -0.375 * FL]), + ("approach", Box(TUBE, TUBE, PD), [0, 0, -FL - PD / 2]), + ] + + meshcat.Delete("grasps") + + num = min(len(poses), max_grasps) + for i in range(num): + t = i / max(num - 1, 1) + rgba = Rgba(min(1.0, 2 * t), max(0.0, 1.0 - t), 0.0, 0.8) + grasp_mat = Transform( + translation=poses[i].position, + rotation=poses[i].orientation, + ).to_matrix() + for name, shape, xyz in parts: + local = np.eye(4) + local[:3, 3] = xyz + path = f"grasps/grasp_{i}/{name}" + meshcat.SetObject(path, shape, rgba) + meshcat.SetTransform(path, RigidTransform(grasp_mat @ local)) + + logger.info(f"Visualized {num}/{len(poses)} grasp poses in MeshCat") + return True + def _compute_pre_grasp_pose(self, grasp_pose: Pose, offset: float = 0.10) -> Pose: """Compute a pre-grasp pose offset along the approach direction (local -Z). diff --git a/dimos/manipulation/planning/spec/config.py b/dimos/manipulation/planning/spec/config.py index 80cf248f08..f8dcb8b548 100644 --- a/dimos/manipulation/planning/spec/config.py +++ b/dimos/manipulation/planning/spec/config.py @@ -16,7 +16,6 @@ from __future__ import annotations -from collections.abc import Iterable, Sequence from pathlib import Path from pydantic import Field @@ -65,7 +64,7 @@ class RobotModelConfig(ModuleConfig): velocity_limits: list[float] | None = None auto_convert_meshes: bool = False xacro_args: dict[str, str] = Field(default_factory=dict) - collision_exclusion_pairs: Iterable[tuple[str, str]] = () + 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 @@ -74,9 +73,9 @@ class RobotModelConfig(ModuleConfig): coordinator_task_name: str | None = None gripper_hardware_id: str | None = None # TF publishing for extra links (e.g., camera mount) - tf_extra_links: Sequence[str] = () + tf_extra_links: list[str] = Field(default_factory=list) # Home/observe joint configuration for go_home skill - home_joints: Iterable[float] | None = None + home_joints: list[float] | None = None # Pre-grasp offset distance in meters (along approach direction) pre_grasp_offset: float = 0.10 diff --git a/dimos/msgs/geometry_msgs/Quaternion.py b/dimos/msgs/geometry_msgs/Quaternion.py index 8e83a32e50..a78c69f194 100644 --- a/dimos/msgs/geometry_msgs/Quaternion.py +++ b/dimos/msgs/geometry_msgs/Quaternion.py @@ -16,6 +16,7 @@ from collections.abc import Sequence from io import BytesIO +import math import struct from typing import TYPE_CHECKING, BinaryIO, TypeAlias @@ -246,6 +247,16 @@ def normalize(self) -> Quaternion: raise ZeroDivisionError("Cannot normalize zero quaternion") return Quaternion(self.x / norm, self.y / norm, self.z / norm, self.w / norm) + def angular_distance(self, other: Quaternion) -> float: + """Return the angular distance (in radians) between two unit quaternions. + + Uses the formula: angle = 2 * arccos(|q1 · q2|), which gives the + minimum rotation angle between the two orientations (range [0, pi]). + """ + dot = abs(self.x * other.x + self.y * other.y + self.z * other.z + self.w * other.w) + dot = min(dot, 1.0) # clamp for numerical safety + return 2.0 * math.acos(dot) + def rotate_vector(self, vector: Vector3) -> Vector3: """Rotate a 3D vector by this quaternion. diff --git a/dimos/msgs/sensor_msgs/PointCloud2.py b/dimos/msgs/sensor_msgs/PointCloud2.py index 67af1c5ac3..90a06c62fc 100644 --- a/dimos/msgs/sensor_msgs/PointCloud2.py +++ b/dimos/msgs/sensor_msgs/PointCloud2.py @@ -100,6 +100,13 @@ def __getstate__(self) -> dict[str, object]: # Remove non-picklable objects del state["_pcd_tensor"] state["_pcd_legacy_cache"] = None + # Remove cached_property values that hold unpicklable Open3D objects + for key in ( + "axis_aligned_bounding_box", + "oriented_bounding_box", + "bounding_box_dimensions", + ): + state.pop(key, None) return state def __setstate__(self, state: dict[str, object]) -> None: diff --git a/dimos/perception/object_scene_registration.py b/dimos/perception/object_scene_registration.py index 3be2db4b47..7a58f9f7c0 100644 --- a/dimos/perception/object_scene_registration.py +++ b/dimos/perception/object_scene_registration.py @@ -70,8 +70,9 @@ def __init__( self, target_frame: str = "map", prompt_mode: YoloePromptMode = YoloePromptMode.LRPC, + **kwargs: Any, ) -> None: - super().__init__() + super().__init__(**kwargs) self._target_frame = target_frame self._prompt_mode = prompt_mode self._object_db = ObjectDB() diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 5910093d61..03a22ee5fb 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -16,6 +16,7 @@ # Run `pytest dimos/robot/test_all_blueprints_generation.py` to regenerate. all_blueprints = { + "bt-pick-place-agent": "dimos.manipulation.blueprints:bt_pick_place_agent", "coordinator-basic": "dimos.control.blueprints.basic:coordinator_basic", "coordinator-cartesian-ik-mock": "dimos.control.blueprints.teleop:coordinator_cartesian_ik_mock", "coordinator-cartesian-ik-piper": "dimos.control.blueprints.teleop:coordinator_cartesian_ik_piper", @@ -101,6 +102,7 @@ "arm-teleop-module": "dimos.teleop.quest.quest_extensions", "b-box-navigation-module": "dimos.navigation.bbox_navigation", "b1-connection-module": "dimos.robot.unitree.b1.connection", + "bt-manipulation-module": "dimos.manipulation.bt.bt_manipulation_module", "camera-module": "dimos.hardware.sensors.camera.module", "cartesian-motion-controller": "dimos.manipulation.control.servo_control.cartesian_motion_controller", "control-coordinator": "dimos.control.coordinator", @@ -151,6 +153,7 @@ "person-tracker": "dimos.perception.detection.person_tracker", "phone-teleop-module": "dimos.teleop.phone.phone_teleop_module", "pick-and-place-module": "dimos.manipulation.pick_and_place_module", + "pick-place-module": "dimos.manipulation.bt.pick_place_module", "quest-teleop-module": "dimos.teleop.quest.quest_teleop_module", "real-sense-camera": "dimos.hardware.sensors.camera.realsense.camera", "receiver-module": "dimos.utils.demo_image_encoding", diff --git a/pyproject.toml b/pyproject.toml index 7e2f38546e..f666713957 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -84,6 +84,8 @@ dependencies = [ "rerun-sdk>=0.20.0", "dimos-viewer>=0.30.0a2", "toolz>=1.1.0", + # Behavior Trees + "py-trees>=2.2,<3", "protobuf>=6.33.5,<7", "psutil>=7.0.0", "sqlite-vec>=0.1.6", @@ -327,6 +329,7 @@ docker = [ "sortedcontainers", "PyTurboJPEG", "rerun-sdk", + "langchain-core", "open3d-unofficial-arm; platform_system == 'Linux' and platform_machine == 'aarch64'", "open3d>=0.18.0; platform_system != 'Linux' or platform_machine != 'aarch64'", ] diff --git a/uv.lock b/uv.lock index 529842294b..c31455318c 100644 --- a/uv.lock +++ b/uv.lock @@ -1698,6 +1698,7 @@ dependencies = [ { name = "plum-dispatch" }, { name = "protobuf" }, { name = "psutil" }, + { name = "py-trees" }, { name = "pydantic" }, { name = "pydantic-settings" }, { name = "python-dotenv" }, @@ -1859,6 +1860,7 @@ dev = [ ] docker = [ { name = "dimos-lcm" }, + { name = "langchain-core" }, { name = "lcm" }, { 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'" }, @@ -2020,6 +2022,7 @@ requires-dist = [ { name = "langchain", marker = "extra == 'agents'", specifier = "==1.2.3" }, { name = "langchain-chroma", marker = "extra == 'agents'", specifier = ">=1,<2" }, { name = "langchain-core", marker = "extra == 'agents'", specifier = "==1.2.3" }, + { name = "langchain-core", marker = "extra == 'docker'" }, { name = "langchain-huggingface", marker = "extra == 'agents'", specifier = ">=1,<2" }, { name = "langchain-ollama", marker = "extra == 'agents'", specifier = ">=1,<2" }, { name = "langchain-openai", marker = "extra == 'agents'", specifier = ">=1,<2" }, @@ -2070,6 +2073,7 @@ requires-dist = [ { name = "psutil", specifier = ">=7.0.0" }, { name = "psycopg2-binary", marker = "extra == 'psql'", specifier = ">=2.9.11" }, { name = "py-spy", marker = "extra == 'dev'" }, + { name = "py-trees", specifier = ">=2.2,<3" }, { name = "pydantic" }, { name = "pydantic", marker = "extra == 'docker'" }, { name = "pydantic-settings", specifier = ">=2.11.0,<3" }, @@ -7119,6 +7123,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e1/da/fcc9a9fcd4ca946ff402cff20348e838b051d69f50f5d1f5dca4cd3c5eb8/py_spy-0.4.1-py2.py3-none-win_amd64.whl", hash = "sha256:d92e522bd40e9bf7d87c204033ce5bb5c828fca45fa28d970f58d71128069fdc", size = 1818784, upload-time = "2025-07-31T19:33:23.802Z" }, ] +[[package]] +name = "py-trees" +version = "2.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pydot" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/79/9c/b8c39a86ce887e3547b415e6bd61b70e1004af0a8ba1b44e75233a4406a2/py_trees-2.4.0.tar.gz", hash = "sha256:cd8738a80e7422aec2b133270b6aa92e3e83926ded2ba9b3614028dba99cdb41", size = 85359, upload-time = "2025-11-13T22:46:01.41Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7d/0a/2e0c30342586dabb0b4b7946393a7aeb725fca7b4549fb02e128a1a0fe87/py_trees-2.4.0-py3-none-any.whl", hash = "sha256:c26a2d23f70871ce093abd5f8425998529bfe7b97cc563316e9c3c9cec20c312", size = 114099, upload-time = "2025-11-13T22:45:58.884Z" }, +] + [[package]] name = "pyarrow" version = "23.0.0" @@ -7582,7 +7598,7 @@ name = "pydot" version = "4.0.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pyparsing", marker = "platform_machine != 'aarch64' or sys_platform != 'linux'" }, + { 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 = [