From a150ba8861538e92feff4c45f7e31f8d4872cf8d Mon Sep 17 00:00:00 2001 From: Jalaj Shukla Date: Tue, 10 Feb 2026 13:49:05 -0800 Subject: [PATCH 1/7] Add GraspGen integration to ManipulationModule via @rpc generate_grasps Adds a generate_grasps RPC method to ManipulationModule that lazily starts a GraspGen Docker container (GPU) on first call and returns ranked grasp poses (PoseArray) for a given point cloud. - Config: graspgen_docker_image, gripper_type, num_grasps, topk - _get_graspgen(): lazy DockerModule init with LFS checkpoint pull - Cleanup in stop() to tear down Docker container - GPU integration test with synthetic cube point cloud --- dimos/manipulation/manipulation_module.py | 63 +++++++++++++++++++- dimos/manipulation/test_manipulation_unit.py | 38 ++++++++++++ 2 files changed, 100 insertions(+), 1 deletion(-) diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index ccb120149c..8e2ebfad47 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -44,8 +44,10 @@ from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: + from dimos.core.docker_runner import DockerModule from dimos.core.rpc_client import RPCClient - from dimos.msgs.geometry_msgs import Pose + from dimos.msgs.geometry_msgs import Pose, PoseArray + from dimos.msgs.sensor_msgs import PointCloud2 logger = setup_logger() @@ -83,6 +85,12 @@ class ManipulationModuleConfig(ModuleConfig): planner_name: str = "rrt_connect" # "rrt_connect" kinematics_name: str = "jacobian" # "jacobian" or "drake_optimization" + # GraspGen Docker settings (optional) + 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 + class ManipulationModule(Module): """Motion planning module with ControlCoordinator execution.""" @@ -118,6 +126,9 @@ def __init__(self, *args: object, **kwargs: object) -> None: # Coordinator integration (lazy initialized) self._coordinator_client: RPCClient | None = None + # GraspGen Docker module (lazy initialized on first generate_grasps call) + self._graspgen: DockerModule | None = None + logger.info("ManipulationModule initialized") @rpc @@ -596,6 +607,51 @@ def get_trajectory_status( status = client.get_trajectory_status(config.coordinator_task_name) return dict(status) if status else None + # ========================================================================= + # GraspGen Integration RPC Methods + # ========================================================================= + + def _get_graspgen(self) -> DockerModule | None: + """Get or create GraspGen Docker module (lazy init).""" + if self._graspgen is not None: + return self._graspgen + + from dimos.core.docker_runner import DockerModule + from dimos.manipulation.grasping.graspgen_module import GraspGenModule + from dimos.utils.data import get_data + from dimos.utils.path_utils import get_project_root + + # Ensure GraspGen model checkpoints are pulled from LFS + get_data("models_graspgen") + + project_root = get_project_root() + docker_file = project_root / "dimos" / "manipulation" / "grasping" / "docker_context" / "Dockerfile" + self._graspgen = DockerModule( + GraspGenModule, + docker_file=docker_file, + docker_build_context=project_root, + docker_image=self.config.graspgen_docker_image, + docker_env={"CI": "1"}, # skip interactive system config prompt in container + gripper_type=self.config.graspgen_gripper_type, + num_grasps=self.config.graspgen_num_grasps, + topk_num_grasps=self.config.graspgen_topk_num_grasps, + ) + self._graspgen.start() + return self._graspgen + + @rpc + def generate_grasps( + self, + pointcloud: PointCloud2, + scene_pointcloud: PointCloud2 | None = None, + ) -> PoseArray | None: + """Generate grasp poses for the given point cloud via GraspGen (Docker/GPU). + """ + graspgen = self._get_graspgen() + if graspgen is None: + return None + return graspgen.generate_grasps(pointcloud, scene_pointcloud) + @property def world_monitor(self) -> WorldMonitor | None: """Access the world monitor for advanced obstacle/world operations.""" @@ -655,6 +711,11 @@ def stop(self) -> None: """Stop the manipulation module.""" logger.info("Stopping ManipulationModule") + # Stop GraspGen Docker container + if self._graspgen is not None: + self._graspgen.stop() + self._graspgen = None + # Stop world monitor (includes visualization thread) if self._world_monitor is not None: self._world_monitor.stop_all_monitors() diff --git a/dimos/manipulation/test_manipulation_unit.py b/dimos/manipulation/test_manipulation_unit.py index 1643d2b5da..1d7b67edc6 100644 --- a/dimos/manipulation/test_manipulation_unit.py +++ b/dimos/manipulation/test_manipulation_unit.py @@ -24,6 +24,7 @@ from dimos.manipulation.manipulation_module import ( ManipulationModule, + ManipulationModuleConfig, ManipulationState, ) from dimos.manipulation.planning.spec import RobotModelConfig @@ -100,6 +101,7 @@ def _make_module(): module._planner = None module._kinematics = None module._coordinator_client = None + module._graspgen = None return module @@ -303,3 +305,39 @@ def test_bidirectional_mapping(self, robot_config_with_mapping): # URDF -> Coordinator assert config.get_coordinator_joint_name("joint1") == "left_joint1" assert config.get_coordinator_joint_name("unknown") == "unknown" + + +# ============================================================================= +# Test GraspGen Integration +# ============================================================================= + + +@pytest.mark.gpu +@pytest.mark.cuda +class TestGraspGenIntegration: + """Integration test for GraspGen via DockerModule (requires GPU + Docker).""" + + def test_generate_grasps_cube_pointcloud(self): + """Generate grasps for a synthetic cube point cloud (500 points).""" + import numpy as np + + from dimos.msgs.sensor_msgs import PointCloud2 + + # Create a cube point cloud: 500 random points in a 5cm cube centered at (0.3, 0, 0.2) + rng = np.random.default_rng(42) + points = rng.uniform(-0.025, 0.025, size=(500, 3)).astype(np.float32) + points += np.array([0.3, 0.0, 0.2], dtype=np.float32) + pointcloud = PointCloud2.from_numpy(points, frame_id="world") + + module = _make_module() + module.config = ManipulationModuleConfig() + try: + result = module.generate_grasps(pointcloud) + + assert result is not None, "generate_grasps returned None — Docker/GPU may not be available" + assert len(result.poses) > 0, "Expected at least one grasp pose" + print(f"Generated {len(result.poses)} grasp poses for cube point cloud") + finally: + if module._graspgen is not None: + module._graspgen.stop() + module._graspgen = None From b5a5d1b8f1b9dd9b67c1fabbb9f70d49e81b1d2d Mon Sep 17 00:00:00 2001 From: Jalaj Shukla Date: Tue, 10 Feb 2026 14:29:25 -0800 Subject: [PATCH 2/7] Add GraspGen integration to ManipulationModule via @rpc generate_grasps Adds a generate_grasps RPC method to ManipulationModule that lazily starts a GraspGen Docker container (GPU) on first call and returns ranked grasp poses (PoseArray) for a given point cloud. - Config: graspgen_docker_image, gripper_type, num_grasps, topk - _get_graspgen(): lazy DockerModule init with LFS checkpoint pull - Cleanup in stop() to tear down Docker container --- dimos/manipulation/manipulation_module.py | 6 ---- dimos/manipulation/test_manipulation_unit.py | 37 -------------------- 2 files changed, 43 deletions(-) diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index 8e2ebfad47..094b271c90 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -607,10 +607,6 @@ def get_trajectory_status( status = client.get_trajectory_status(config.coordinator_task_name) return dict(status) if status else None - # ========================================================================= - # GraspGen Integration RPC Methods - # ========================================================================= - def _get_graspgen(self) -> DockerModule | None: """Get or create GraspGen Docker module (lazy init).""" if self._graspgen is not None: @@ -645,8 +641,6 @@ def generate_grasps( pointcloud: PointCloud2, scene_pointcloud: PointCloud2 | None = None, ) -> PoseArray | None: - """Generate grasp poses for the given point cloud via GraspGen (Docker/GPU). - """ graspgen = self._get_graspgen() if graspgen is None: return None diff --git a/dimos/manipulation/test_manipulation_unit.py b/dimos/manipulation/test_manipulation_unit.py index 1d7b67edc6..5783106b59 100644 --- a/dimos/manipulation/test_manipulation_unit.py +++ b/dimos/manipulation/test_manipulation_unit.py @@ -24,7 +24,6 @@ from dimos.manipulation.manipulation_module import ( ManipulationModule, - ManipulationModuleConfig, ManipulationState, ) from dimos.manipulation.planning.spec import RobotModelConfig @@ -305,39 +304,3 @@ def test_bidirectional_mapping(self, robot_config_with_mapping): # URDF -> Coordinator assert config.get_coordinator_joint_name("joint1") == "left_joint1" assert config.get_coordinator_joint_name("unknown") == "unknown" - - -# ============================================================================= -# Test GraspGen Integration -# ============================================================================= - - -@pytest.mark.gpu -@pytest.mark.cuda -class TestGraspGenIntegration: - """Integration test for GraspGen via DockerModule (requires GPU + Docker).""" - - def test_generate_grasps_cube_pointcloud(self): - """Generate grasps for a synthetic cube point cloud (500 points).""" - import numpy as np - - from dimos.msgs.sensor_msgs import PointCloud2 - - # Create a cube point cloud: 500 random points in a 5cm cube centered at (0.3, 0, 0.2) - rng = np.random.default_rng(42) - points = rng.uniform(-0.025, 0.025, size=(500, 3)).astype(np.float32) - points += np.array([0.3, 0.0, 0.2], dtype=np.float32) - pointcloud = PointCloud2.from_numpy(points, frame_id="world") - - module = _make_module() - module.config = ManipulationModuleConfig() - try: - result = module.generate_grasps(pointcloud) - - assert result is not None, "generate_grasps returned None — Docker/GPU may not be available" - assert len(result.poses) > 0, "Expected at least one grasp pose" - print(f"Generated {len(result.poses)} grasp poses for cube point cloud") - finally: - if module._graspgen is not None: - module._graspgen.stop() - module._graspgen = None From 45e5e1d4bd09bc05d05de3d5c9b0515fc3f8b23d Mon Sep 17 00:00:00 2001 From: JalajShuklaSS <144572499+JalajShuklaSS@users.noreply.github.com> Date: Tue, 10 Feb 2026 22:36:11 +0000 Subject: [PATCH 3/7] CI code cleanup --- dimos/manipulation/manipulation_module.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index 094b271c90..278451b679 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -621,7 +621,9 @@ def _get_graspgen(self) -> DockerModule | None: get_data("models_graspgen") project_root = get_project_root() - docker_file = project_root / "dimos" / "manipulation" / "grasping" / "docker_context" / "Dockerfile" + docker_file = ( + project_root / "dimos" / "manipulation" / "grasping" / "docker_context" / "Dockerfile" + ) self._graspgen = DockerModule( GraspGenModule, docker_file=docker_file, From 7b3d9e8fcc2c698d0fe596c3046dfaf44b111107 Mon Sep 17 00:00:00 2001 From: Jalaj Shukla Date: Wed, 11 Feb 2026 09:31:11 -0800 Subject: [PATCH 4/7] Add collision, visualization, and threshold config params to GraspGen integration --- dimos/manipulation/manipulation_module.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index 278451b679..3b2064ac63 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -90,6 +90,10 @@ class ManipulationModuleConfig(ModuleConfig): 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: str = "/tmp/grasp_visualization.json" class ManipulationModule(Module): @@ -633,6 +637,10 @@ def _get_graspgen(self) -> DockerModule | None: 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=self.config.graspgen_visualization_output_path, ) self._graspgen.start() return self._graspgen From 54d082cb6ee5b2fb6cd426e64212c1f40411219b Mon Sep 17 00:00:00 2001 From: Jalaj Shukla Date: Wed, 11 Feb 2026 09:55:21 -0800 Subject: [PATCH 5/7] mypy fixes --- dimos/manipulation/manipulation_module.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index 3b2064ac63..9bcf63de72 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -629,7 +629,7 @@ def _get_graspgen(self) -> DockerModule | None: project_root / "dimos" / "manipulation" / "grasping" / "docker_context" / "Dockerfile" ) self._graspgen = DockerModule( - GraspGenModule, + GraspGenModule, # type: ignore[arg-type] docker_file=docker_file, docker_build_context=project_root, docker_image=self.config.graspgen_docker_image, @@ -654,7 +654,8 @@ def generate_grasps( graspgen = self._get_graspgen() if graspgen is None: return None - return graspgen.generate_grasps(pointcloud, scene_pointcloud) + result: PoseArray | None = graspgen.generate_grasps(pointcloud, scene_pointcloud) + return result @property def world_monitor(self) -> WorldMonitor | None: From d8bfa5e81938288e4e730762c6bdcf92f48bee1d Mon Sep 17 00:00:00 2001 From: Jalaj Shukla Date: Mon, 16 Feb 2026 18:48:53 -0800 Subject: [PATCH 6/7] Address PR review comments for GraspGen integration - Thread-safe lazy init with double-checked locking - Move inline imports to top-level - Use DIMOS_PROJECT_ROOT constant - Use Path type for visualization output path --- dimos/manipulation/manipulation_module.py | 104 +++++++++++++--------- 1 file changed, 64 insertions(+), 40 deletions(-) diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index 9bcf63de72..948d1fa1c7 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -18,11 +18,15 @@ from dataclasses import dataclass, field from enum import Enum +from pathlib import Path import threading from typing import TYPE_CHECKING, TypeAlias +from dimos.constants import DIMOS_PROJECT_ROOT from dimos.core import In, Module, rpc +from dimos.core.docker_runner import DockerModule as DockerRunner from dimos.core.module import ModuleConfig +from dimos.manipulation.grasping.graspgen_module import GraspGenModule from dimos.manipulation.planning import ( JointPath, JointTrajectoryGenerator, @@ -41,10 +45,10 @@ # These must be imported at runtime (not TYPE_CHECKING) for In/Out port creation from dimos.msgs.sensor_msgs import JointState from dimos.msgs.trajectory_msgs import JointTrajectory +from dimos.utils.data import get_data from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: - from dimos.core.docker_runner import DockerModule from dimos.core.rpc_client import RPCClient from dimos.msgs.geometry_msgs import Pose, PoseArray from dimos.msgs.sensor_msgs import PointCloud2 @@ -64,6 +68,10 @@ PlannedTrajectories: TypeAlias = dict[RobotName, JointTrajectory] """Maps robot_name -> planned trajectory""" +# The host-side path (graspgen_visualization_output_path) is volume-mounted here. +_GRASPGEN_VIZ_CONTAINER_DIR = "/output/graspgen" +_GRASPGEN_VIZ_CONTAINER_PATH = f"{_GRASPGEN_VIZ_CONTAINER_DIR}/visualization.json" + class ManipulationState(Enum): """State machine for manipulation module.""" @@ -93,7 +101,9 @@ class ManipulationModuleConfig(ModuleConfig): graspgen_grasp_threshold: float = -1.0 graspgen_filter_collisions: bool = False graspgen_save_visualization_data: bool = False - graspgen_visualization_output_path: str = "/tmp/grasp_visualization.json" + graspgen_visualization_output_path: Path = field( + default_factory=lambda: Path.home() / ".dimos" / "graspgen" / "visualization.json" + ) class ManipulationModule(Module): @@ -130,8 +140,8 @@ def __init__(self, *args: object, **kwargs: object) -> None: # Coordinator integration (lazy initialized) self._coordinator_client: RPCClient | None = None - # GraspGen Docker module (lazy initialized on first generate_grasps call) - self._graspgen: DockerModule | None = None + # GraspGen Docker runner (lazy initialized on first generate_grasps call) + self._graspgen: DockerRunner | None = None logger.info("ManipulationModule initialized") @@ -611,39 +621,50 @@ def get_trajectory_status( status = client.get_trajectory_status(config.coordinator_task_name) return dict(status) if status else None - def _get_graspgen(self) -> DockerModule | None: - """Get or create GraspGen Docker module (lazy init).""" + def _get_graspgen(self) -> DockerRunner: + """Get or create GraspGen Docker module (lazy init, thread-safe).""" + # Fast path: already initialized (no lock needed for read) if self._graspgen is not None: return self._graspgen - from dimos.core.docker_runner import DockerModule - from dimos.manipulation.grasping.graspgen_module import GraspGenModule - from dimos.utils.data import get_data - from dimos.utils.path_utils import get_project_root + # Slow path: need to initialize (acquire lock to prevent race condition) + with self._lock: + # Double-check: another thread may have initialized while we waited for lock + if self._graspgen is not None: + return self._graspgen - # Ensure GraspGen model checkpoints are pulled from LFS - get_data("models_graspgen") + # Ensure GraspGen model checkpoints are pulled from LFS + get_data("models_graspgen") - project_root = get_project_root() - docker_file = ( - project_root / "dimos" / "manipulation" / "grasping" / "docker_context" / "Dockerfile" - ) - self._graspgen = DockerModule( - GraspGenModule, # type: ignore[arg-type] - docker_file=docker_file, - docker_build_context=project_root, - docker_image=self.config.graspgen_docker_image, - docker_env={"CI": "1"}, # skip interactive system config prompt in container - 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=self.config.graspgen_visualization_output_path, - ) - self._graspgen.start() - return self._graspgen + docker_file = ( + DIMOS_PROJECT_ROOT / "dimos" / "manipulation" / "grasping" / "docker_context" / "Dockerfile" + ) + + # Auto-mount host directory for visualization output when enabled. + 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"}, # skip interactive system config prompt in container + 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 # cache only after successful start + return self._graspgen @rpc def generate_grasps( @@ -651,11 +672,13 @@ def generate_grasps( pointcloud: PointCloud2, scene_pointcloud: PointCloud2 | None = None, ) -> PoseArray | None: - graspgen = self._get_graspgen() - if graspgen is None: + """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] + except Exception as e: + logger.error(f"Grasp generation failed: {e}") return None - result: PoseArray | None = graspgen.generate_grasps(pointcloud, scene_pointcloud) - return result @property def world_monitor(self) -> WorldMonitor | None: @@ -716,10 +739,11 @@ def stop(self) -> None: """Stop the manipulation module.""" logger.info("Stopping ManipulationModule") - # Stop GraspGen Docker container - if self._graspgen is not None: - self._graspgen.stop() - self._graspgen = None + # Stop GraspGen Docker container (thread-safe access to shared state) + with self._lock: + if self._graspgen is not None: + self._graspgen.stop() + self._graspgen = None # Stop world monitor (includes visualization thread) if self._world_monitor is not None: From f264b18cc59c8342299c64b9dcf7363c25fee154 Mon Sep 17 00:00:00 2001 From: JalajShuklaSS <144572499+JalajShuklaSS@users.noreply.github.com> Date: Tue, 17 Feb 2026 02:51:17 +0000 Subject: [PATCH 7/7] CI code cleanup --- dimos/manipulation/manipulation_module.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/dimos/manipulation/manipulation_module.py b/dimos/manipulation/manipulation_module.py index 948d1fa1c7..c947bcd60e 100644 --- a/dimos/manipulation/manipulation_module.py +++ b/dimos/manipulation/manipulation_module.py @@ -637,7 +637,12 @@ def _get_graspgen(self) -> DockerRunner: get_data("models_graspgen") docker_file = ( - DIMOS_PROJECT_ROOT / "dimos" / "manipulation" / "grasping" / "docker_context" / "Dockerfile" + DIMOS_PROJECT_ROOT + / "dimos" + / "manipulation" + / "grasping" + / "docker_context" + / "Dockerfile" ) # Auto-mount host directory for visualization output when enabled.