GraspGen generator with docker dimos module integration #1119
Conversation
Greptile OverviewGreptile SummaryThis PR introduces Docker-backed module execution to DimOS and adds GraspGen grasp generation capability. The implementation enables dependency-heavy modules (CUDA, PyTorch, EGL) to run in isolated containers while behaving like native modules from the blueprint perspective. Key Changes:
Issues Found:
Architecture Notes: Confidence Score: 3/5
Important Files Changed
Sequence DiagramsequenceDiagram
participant User
participant Agent
participant GraspingSkill
participant GraspPipeline
participant Perception
participant GraspGen
participant Docker
Note over User,Docker: Initialization Phase
User->>Agent: Start demo_perception_grasping blueprint
Agent->>Docker: Deploy GraspGenModule
Docker->>Docker: Build image with DimOS footer
Docker->>GraspGen: Start container with LCMRPC
GraspGen-->>Docker: Module ready signal
Docker-->>Agent: DockerModule handle created
Agent->>Perception: Deploy ObjectSceneRegistration
Agent->>GraspPipeline: Deploy GraspPipeline
Agent->>GraspingSkill: Deploy GraspingSkillContainer
Note over User,Docker: Grasp Generation Flow
User->>Agent: "Generate grasps for cup"
Agent->>GraspingSkill: generate_grasps("cup")
GraspingSkill->>GraspPipeline: Publish trigger ("cup")
GraspPipeline->>Perception: RPC: get_object_pointcloud_by_name("cup")
Perception-->>GraspPipeline: PointCloud2
GraspPipeline->>Perception: RPC: get_full_scene_pointcloud()
Perception-->>GraspPipeline: Scene PointCloud2
GraspPipeline->>GraspGen: RPC: generate_grasps(pc, scene_pc)
GraspGen->>GraspGen: Initialize GraspGen model (if needed)
GraspGen->>GraspGen: Run inference with CUDA
GraspGen->>GraspGen: Filter collisions (optional)
GraspGen->>GraspingSkill: Publish PoseArray via stream
GraspingSkill->>GraspingSkill: _on_grasps() callback
GraspingSkill-->>Agent: Return grasp summary
Agent-->>User: "Generated 100 grasps for 'cup'"
Note over User,Docker: Cleanup Phase
User->>Agent: Stop blueprint
Agent->>Docker: Stop all DockerModules
Docker->>GraspGen: Send SIGTERM
GraspGen->>GraspGen: Cleanup GPU memory
Docker->>Docker: Remove container
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| @@ -0,0 +1,487 @@ | |||
| # Copyright 2025-2026 Dimensional Inc. | |||
There was a problem hiding this comment.
This file appears to be a personal/test version with a typo in the filename (registreation instead of registration). Should not be committed to the main branch.
| import pickle | ||
| return pickle.dumps({"header": self.header, "poses": self.poses}) |
There was a problem hiding this comment.
Using pickle for LCM message encoding is insecure and not cross-language compatible. Consider using proper LCM encoding with .lcm schema definitions for production use.
| if not config.docker_image or not config.docker_file: | ||
| raise ValueError("docker_image and docker_file are required") |
There was a problem hiding this comment.
Both docker_image and docker_file are required but validation only checks if they're falsy. If either is an empty string, this will raise an error with the same message even though one might be set correctly.
| if not config.docker_image or not config.docker_file: | |
| raise ValueError("docker_image and docker_file are required") | |
| if not config.docker_image: | |
| raise ValueError("docker_image is required") | |
| if not config.docker_file: | |
| raise ValueError("docker_file is required") |
There was a problem hiding this comment.
The validation logic is intentionally checking that both fields are provided. The or operator ensures that if either is falsy (None, "", etc.), the error is raised - which is the correct behavior since both are required for proper Docker module operation i.e. docker_image and docker_file
There was a problem hiding this comment.
the suggested change does behaves similarly to your implementation, just has better verbose to help user identify what exactly is missing
| graspgen_path = os.environ.get("GRASPGEN_PATH", "/dimos/third_party/GraspGen") | ||
| config_name = f"graspgen_{self.config.gripper_type}.yml" | ||
|
|
||
| for subdir in ("GraspGenModels/checkpoints", "checkpoints"): | ||
| path = os.path.join(graspgen_path, subdir, config_name) | ||
| if os.path.exists(path): | ||
| return path | ||
|
|
||
| return os.path.join(graspgen_path, "checkpoints", config_name) |
There was a problem hiding this comment.
The config path search tries multiple locations but if none exist, it returns a non-existent path at line 164. This will fail later during initialization. Consider raising an error if the config file is not found.
| graspgen_path = os.environ.get("GRASPGEN_PATH", "/dimos/third_party/GraspGen") | |
| config_name = f"graspgen_{self.config.gripper_type}.yml" | |
| for subdir in ("GraspGenModels/checkpoints", "checkpoints"): | |
| path = os.path.join(graspgen_path, subdir, config_name) | |
| if os.path.exists(path): | |
| return path | |
| return os.path.join(graspgen_path, "checkpoints", config_name) | |
| graspgen_path = os.environ.get("GRASPGEN_PATH", "/dimos/third_party/GraspGen") | |
| config_name = f"graspgen_{self.config.gripper_type}.yml" | |
| for subdir in ("GraspGenModels/checkpoints", "checkpoints"): | |
| path = os.path.join(graspgen_path, subdir, config_name) | |
| if os.path.exists(path): | |
| return path | |
| raise FileNotFoundError( | |
| f"Config {config_name} not found in {graspgen_path}. " | |
| f"Searched: GraspGenModels/checkpoints, checkpoints" | |
| ) |
There was a problem hiding this comment.
Yes load_grasp_cfg() will raise an exception.
The exception is caught and handled in _initialize_graspgen() and then further caller handles it in generate_grasps()
i.e. error is logged: "Failed to initialize GraspGen: [file not found error]"
generate_grasps() returns None
Application continues running
| graspgen_path = os.environ.get("GRASPGEN_PATH", "/dimos/third_party/GraspGen") | ||
| if graspgen_path not in sys.path: | ||
| sys.path.insert(0, graspgen_path) | ||
| os.environ["PYOPENGL_PLATFORM"] = "egl" |
There was a problem hiding this comment.
The GRASPGEN_PATH environment variable defaults to /dimos/third_party/GraspGen, but from the Dockerfile the path is /app/GraspGen. This mismatch could cause initialization failures.
| graspgen_path = os.environ.get("GRASPGEN_PATH", "/dimos/third_party/GraspGen") | |
| if graspgen_path not in sys.path: | |
| sys.path.insert(0, graspgen_path) | |
| os.environ["PYOPENGL_PLATFORM"] = "egl" | |
| # Setup GraspGen path and environment | |
| graspgen_path = os.environ.get("GRASPGEN_PATH", "/app/GraspGen") | |
| if graspgen_path not in sys.path: | |
| sys.path.insert(0, graspgen_path) | |
| os.environ["PYOPENGL_PLATFORM"] = "egl" |
There was a problem hiding this comment.
os.environ.get ("GRASPGEN_PATH", "/dimos/third_party/GraspGen")
key = GRASPGEN_PATH and default is /dimos/third_party/GraspGen (only used if key not found)
Inside Docker container: GRASPGEN_PATH=/app/GraspGen is set to uses /app/GraspGen
Outside Docker (local dev): env var not set then it uses fallback /dimos/third_party/GraspGen
The default /dimos/third_party/GraspGen is just a fallback for local development - it's never used when running in Docker because the Dockerfile explicitly sets the correct path.
There was a problem hiding this comment.
But people don't have /dimos/third_party/GraspGen on their computers so why have this default value here?
There was a problem hiding this comment.
You're absolutely right. The fallback path /dimos/third_party/GraspGen doesn't exist on developers' machines and serves no practical purpose. This could be confusing and misleading.
Since this module is designed to run exclusively in Docker (as indicated by inheriting from DockerModuleConfig), the fallback should either:
- Be removed entirely and require
GRASPGEN_PATHto be set - Use a more obvious placeholder like
/path/to/graspgento make it clear it's not meant to work locally
Option 1 is cleaner since this is a Docker-only module:
| graspgen_path = os.environ.get("GRASPGEN_PATH", "/dimos/third_party/GraspGen") | |
| if graspgen_path not in sys.path: | |
| sys.path.insert(0, graspgen_path) | |
| os.environ["PYOPENGL_PLATFORM"] = "egl" | |
| graspgen_path = os.environ["GRASPGEN_PATH"] |
This way it fails fast with a clear error if someone tries to run it outside Docker without the proper environment setup.
| if not np.isfinite(points).all(): | ||
| raise ValueError("Point cloud contains NaN/Inf") | ||
| return points |
There was a problem hiding this comment.
Raising a ValueError for non-finite points will crash the grasp generation. Consider logging a warning and returning None or an empty array instead, allowing graceful degradation.
| if not np.isfinite(points).all(): | |
| raise ValueError("Point cloud contains NaN/Inf") | |
| return points | |
| points = msg.points().numpy() | |
| if not np.isfinite(points).all(): | |
| logger.warning("Point cloud contains NaN/Inf values, skipping") | |
| return np.array([]) | |
| return points |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
There was a problem hiding this comment.
The raise ValueError in _extract_points() is caught by exception in generate_grasps() which logs Grasp generation failed: Point cloud contains NaN/Inf and it returns None. So no crash
| """Temporary testing module for grasp generation pipeline. | ||
|
|
||
| Wires: GraspingSkill → GraspPipeline → Perception + GraspGen | ||
| """ |
There was a problem hiding this comment.
This is marked as a temporary testing module. Consider creating a tracking issue for replacing this with a proper production implementation.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| (echo "ERROR: Model checkpoints not found at ${CHECKPOINTS_PATH}" && exit 1) | ||
|
|
||
| # Environment variables for GraspGen | ||
| ENV GRASPGEN_PATH=/app/GraspGen |
There was a problem hiding this comment.
The GRASPGEN_PATH is set to /app/GraspGen here, but graspgen_module.py:136 defaults to /dimos/third_party/GraspGen. Ensure consistency or the module will fail to initialize.
There was a problem hiding this comment.
The Dockerfile explicitly sets ENV GRASPGEN_PATH=/app/GraspGen , which overrides the Python default. The code uses os.environ.get():graspgen_path = os.environ.get("GRASPGEN_PATH", "/dimos/third_party/GraspGen")
Inside Docker: Environment variable is set and uses /app/GraspGen
Outside Docker (local dev): Environment variable not set and uses fallback /dimos/third_party/GraspGen
Additional Comments (1)
|
|
|
||
| # Host-side Docker-backed Module handle | ||
|
|
||
| class DockerModule: |
There was a problem hiding this comment.
Why doesn't this inherit from Module such that any existing module can be run as a DockerModule
There was a problem hiding this comment.
rename to DockerHandler or something
| logger = setup_logger() | ||
|
|
||
|
|
||
| class GraspingSkillContainer(SkillModule): |
There was a problem hiding this comment.
We need to move away from using skill containers. Skills should just be in the Module itself.
| docker_file_path=Path(__file__).parent.parent / "grasping" / "docker_context" / "Dockerfile", | ||
| docker_build_context=Path(__file__).parent.parent.parent, # repo root |
There was a problem hiding this comment.
| docker_file_path=Path(__file__).parent.parent / "grasping" / "docker_context" / "Dockerfile", | |
| docker_build_context=Path(__file__).parent.parent.parent, # repo root | |
| docker_file_path=DIMOS_PROJECT_ROOT / "dimos" / "grasping" / "docker_context" / "Dockerfile", | |
| docker_build_context=DIMOS_PROJECT_ROOT |
And import from dimos.constants import DIMOS_PROJECT_ROOT
| self.camera_info.subscribe(lambda msg: setattr(self, "_camera_info", msg)) | ||
|
|
||
| aligned_frames = align_timestamped( | ||
| self.color_image.observable(), # type: ignore[no-untyped-call] | ||
| self.depth_image.observable(), # type: ignore[no-untyped-call] | ||
| buffer_size=2.0, | ||
| match_tolerance=0.1, | ||
| ) | ||
| backpressure(aligned_frames).subscribe(self._on_aligned_frames) |
There was a problem hiding this comment.
Add both subscribe() values to self._disposables
| except Exception: | ||
| scene = None |
There was a problem hiding this comment.
Why silence the error?
There was a problem hiding this comment.
Here in this scene pointcloud is optional - it's only used for collision filtering (filter_collisions defaults to False in GraspGenConfig) but the Object pointcloud is required - if that fails, the function returns early, but I agree that I can log a warning and developers can then see that why the collision filtering is off.
Thank you.
TL;DR : Adding a warning
| "GraspingSkillContainer.signal_no_grasps", | ||
| ] | ||
|
|
||
| _default_object_name: str = "object" |
There was a problem hiding this comment.
| _default_object_name: str = "object" | |
| _default_object_name: str |
There is a default in __init__. Having a default in two places means we might forget to change one.
There was a problem hiding this comment.
Yes defaults in two places is a maintenance hazard yes, I will fix it by defining the default in one place itself.
Thanks
|
|
||
| def _on_trigger(self, object_name: str) -> None: | ||
| # Use provided object name, or fall back to default | ||
| target = object_name if object_name else self._default_object_name |
There was a problem hiding this comment.
What is the purpose of _default_object_name?
object_name: str does not allow for None so the else clause would only take effect if object_name is an empty string.
Would it not be better to raise an error if an empty string is passed rather than silently replace it with "object"?
There was a problem hiding this comment.
This was added for testing convenience as this code tries to connect the perception with grasping however this grasping skill would be a minor skill in the manipulation and yes - silently falling back hides potential issues.
Thank you.
TL:DR : would update to raise the error for empty strings
| return self._latest_grasps | ||
|
|
||
| @rpc | ||
| def signal_no_grasps(self, reason: str = "unknown") -> None: |
There was a problem hiding this comment.
reason should be mandatory and not have a default. In all places where it's use it gets a value.
| def signal_no_grasps(self, reason: str = "unknown") -> None: | ||
| """Signal that grasp generation failed (e.g., no pointcloud).""" | ||
| logger.warning(f"Grasp generation failed: {reason}") | ||
| self._latest_grasps = None |
There was a problem hiding this comment.
_latest_grasps should be behind a lock
| self._latest_grasps = grasps | ||
| self._grasps_received.set() # Signal that grasps have arrived |
There was a problem hiding this comment.
It seems like you're using the event in conjunction with _latest_grasps to pass a value. If that's the case, you can use queue.Queue. They're thread safe and you don't need a lock.
| # Merge config fields into kwargs (Configurable creates config from these) | ||
| if "config" in self.kwargs: | ||
| config_dict = self.kwargs.pop("config") | ||
| # Config fields go first, extra kwargs go later | ||
| self.kwargs = {**config_dict, **self.kwargs} |
There was a problem hiding this comment.
Why is this in start()? Shouldn't it be in __init__?
| self.args = args | ||
| self.kwargs = kwargs | ||
| self._module: Module | None = None | ||
| self._shutdown = False |
There was a problem hiding this comment.
Use an Event() instead of a boolean.
There was a problem hiding this comment.
Yeah I would put up the threading.Event() .
Would use Event.wait() which blocks the thread without consuming the CPU instead of spinning up a loop every 200ms also would get the cross thread signiling with use of Event. Thanks I would always use them now. Thanks.
| self.module_path = module_path | ||
| self.args = args | ||
| self.kwargs = kwargs |
There was a problem hiding this comment.
Are these used outside StandaloneModuleRunner? If not, they should start with _.
There was a problem hiding this comment.
These are only used internally within start(). Will rename to _module_path, _args, _kwargs to match the existing_moduleand _shutdown convention. I just forgot while typing them out.
Thank you.
| elapsed = time.time() - start_time | ||
| logger.info(f"{self.remote_name} ready ({elapsed:.1f}s)") | ||
| return | ||
| except Exception: |
There was a problem hiding this comment.
I think this should be more precise. Only handle the specific error that causes you to retry.
There was a problem hiding this comment.
Yes. Catching all the exceptions might hide the real bugs and then it will silently try for 120s instead of failing fast with a useful error.
TL;DR : Will update to catch only timeout/connection errors that indicate the module isn't ready yet, and let other exceptions propagate immediately.
| # Config | ||
| config_class = getattr(module_class, "default_config", DockerModuleConfig) | ||
| config = config_class(**kwargs) | ||
| if not config.docker_image or not config.docker_file: |
There was a problem hiding this comment.
You don't need this check if you mark both as mandatory. This is needed because you've specified the type as | None. At least docker_image and docker_file should be mandatory, probably more.
| if not cfg.docker_file or not cfg.docker_image: | ||
| raise ValueError("docker_file and docker_image are required") |
There was a problem hiding this comment.
This would also not be necessary if you didn't allow them to be None.
There was a problem hiding this comment.
Due to dataclass inheritance rules, fields without defaults can't follow fields with defaults and as parent ModuleConfig has all defaulted fields, have added kw_only=True to allow required fields in the child class. It helped me to remove this check not only from this file but also in another file where we check if the docker_image or docker_path is available or no.
| content = dockerfile.read_text() | ||
|
|
||
| # Already converted? | ||
| if "DIMOS MODULE CONVERSION" in content: |
There was a problem hiding this comment.
I see you're using DIMOS MODULE CONVERSION as a sentinel value. It should be placed in a constant and used in those other strings as well. (Personally I prefer to use UUIDs as sentinel values since they're unique.)
There was a problem hiding this comment.
Created now a sentinel constant with hybrid format: "DIMOS-MODULE-CONVERSION-427593ae-c6e8-4cf1-9b2d-ee81a420a5dc"
Where I am using the UUID suffix ensures uniqueness and just have also kept prefix keeps it human-readable
Thank you.
| @@ -0,0 +1,81 @@ | |||
| # GraspGen - Grasp Pose Generation Model | |||
There was a problem hiding this comment.
This is the generated file, right? Should it have been committed?
There was a problem hiding this comment.
Yeah it is the generated file and I kept it here as an example I would take it down with the other temp files and code files like the temp_graspgen_testing.py as well. I apologize.
| self._bound_rpc_calls: dict[str, RpcCall] = {} | ||
|
|
||
| # Build image if needed | ||
| from dimos.core.docker_build import build_image, image_exists |
There was a problem hiding this comment.
Why an inline import?
There was a problem hiding this comment.
This was done defensively to avoid a potential circular import since docker_build.py imports DockerModuleConfig, however the, docker_build.py uses TYPE_CHECKING guard, so there's no actual runtime circular import. Will move this to the top-level imports for consistency.
Thank you and would keep this in mind the next time.
| # Auto-start container (must be running before transports are connected) | ||
| self.start() |
There was a problem hiding this comment.
I don't agree with starting things in __init__. This has caused us problems before. When you do work in the constructor, this causes problems like:
- hard to pass around because the object is already doing work. (This is the issue we're having with
LCM()starting the server just by initializing the object.) - hard to subclass because there's no point at which to make changes in between calling
super().__init__()and starting the module.
There was a problem hiding this comment.
Here the idea was that by the time deploy() returns the DockerModule, it should be ready to use (container running, RPC responding). This was for convenience. But your correct and yeah this would be hard to pass and, even its having inconsistent lifecycle (mostly it should have been create-> configure-> start. This does create+start ) and also would be difficult when we have unit tests without actually starting the container.
Thank you.
TL;DR : will move self.start() out of __init__ and have the caller deploy() call it explicitly
| docker_privileged: bool = False | ||
|
|
||
| # Lifecycle / overrides | ||
| docker_rm: bool = False |
There was a problem hiding this comment.
I think this shouldn't be an option and should always be done.
Is there a reason we don't want to automatically remove the container when it's done?
| self._add_entrypoint_args(cmd, cfg) | ||
| cmd.extend(cfg.docker_extra_args) | ||
|
|
||
| cmd.append(cfg.docker_image or "") |
There was a problem hiding this comment.
I'm not sure or "" makes sense. This is also a side effect of allowing None on cfg.docker_image.
| ) | ||
|
|
||
| def _add_lifecycle_args(self, cmd: list[str], cfg: DockerModuleConfig) -> None: | ||
| """Add --rm, --name, --restart args.""" |
There was a problem hiding this comment.
--restart isn't added here, it's added in _add_resource_args.
| @rpc | ||
| def start(self) -> None: | ||
| super().start() | ||
| self._disposables.add(self.trigger.subscribe(self._on_trigger)) |
There was a problem hiding this comment.
Mypy complains about this. You need:
self._disposables.add(Disposable(self.trigger.subscribe(self._on_trigger)))
because subscribe() here just returns a function.
| return None | ||
|
|
||
| try: | ||
| start_time = time.time() |
There was a problem hiding this comment.
Is start_time used for anything?
| try: | ||
| import open3d as o3d | ||
| except ImportError: | ||
| o3d = None |
There was a problem hiding this comment.
open3d is always installed so this is not needed.
| Example: docker_extra_args=["--cap-add=SYS_ADMIN", "--read-only"] | ||
| """ | ||
| # Build / image | ||
| docker_image: str | None = None |
There was a problem hiding this comment.
remove none default for both
| stream._transport = transport | ||
| return True | ||
|
|
||
| @rpc |
There was a problem hiding this comment.
@paul-nechifor I assume you had to solve this also somehow for multiprocessing? we coalesce in the follow up?
Adds GraspGen grasp generation to DimOS and introduces Docker-backed module execution, enabling dependency-heavy modules to run in containers while behaving like native modules with RPC, streaming, and autoconnect support. This skill would further be a subskill inside the manipulation stack to generate grasps.
GraspGen Module
Docker-backed Module Support