Skip to content

Draft: Add grasp annotator #196

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
matafela wants to merge 7 commits into
base: main
Choose a base branch
from
Open

Draft: Add grasp annotator #196

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
5fd8dda
draft
Mar 23, 2026
1e15c77
update
Mar 23, 2026
baf731a
update comment
Mar 23, 2026
f1f043b
add viser dependence
Mar 24, 2026
1887304
Merge branch 'main' into cj/add-grasp-annotator
Mar 25, 2026
63ec5e6
update
Mar 25, 2026
73781d8
update
Mar 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
67 changes: 67 additions & 0 deletions embodichain/lab/sim/objects/rigid_object.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,11 @@
from embodichain.utils.math import convert_quat
from embodichain.utils.math import matrix_from_quat, quat_from_matrix, matrix_from_euler
from embodichain.utils import logger
from embodichain.toolkits.graspkit.pg_grasp.antipodal_annotator import (
GraspAnnotator,
GraspAnnotatorCfg,
)
Comment on lines +38 to +41
Copy link

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

RigidObject now imports GraspAnnotator (and its dependencies like viser/trimesh/open3d) at module import time. This makes the core sim objects package depend on optional UI/geometry libraries and can break environments that don’t have those extras installed, even if grasp annotation isn’t used. Consider moving these imports inside get_grasp_pose() and raising a clear, actionable error if the optional deps are missing.

Suggested change
from embodichain.toolkits.graspkit.pg_grasp.antipodal_annotator import (
GraspAnnotator,
GraspAnnotatorCfg,
)
def _load_grasp_annotator_module():
"""Lazily import the grasp annotator module.
This avoids importing optional heavy UI/geometry dependencies (e.g., viser,
trimesh, open3d) at module import time. The import is performed only when
grasp annotation functionality is actually used.
"""
try:
from embodichain.toolkits.graspkit.pg_grasp import antipodal_annotator
except ImportError as exc:
raise ImportError(
"Grasp annotator dependencies are not installed. "
"To use grasp annotation (RigidObject.get_grasp_pose and related "
"functionality), install the optional grasp/geometry extras, e.g.:\n\n"
" pip install 'embodichain[grasp]'\n\n"
"or ensure that packages like 'viser', 'trimesh', and 'open3d' are "
"available in your environment."
) from exc
return antipodal_annotator
class GraspAnnotator: # type: ignore[misc]
"""Lazy proxy for the real GraspAnnotator class.
The actual class is imported from `embodichain.toolkits.graspkit.pg_grasp`
only when this proxy is instantiated.
"""
def __new__(cls, *args, **kwargs):
module = _load_grasp_annotator_module()
real_cls = module.GraspAnnotator
return real_cls(*args, **kwargs)
class GraspAnnotatorCfg: # type: ignore[misc]
"""Lazy proxy for the real GraspAnnotatorCfg class.
The actual class is imported only when this proxy is instantiated.
"""
def __new__(cls, *args, **kwargs):
module = _load_grasp_annotator_module()
real_cls = module.GraspAnnotatorCfg
return real_cls(*args, **kwargs)

Copilot uses AI. Check for mistakes.
Comment on lines +38 to +41
Copy link

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Importing GraspAnnotator (and its dependencies like viser) at module import time makes embodichain.lab.sim.objects.rigid_object depend on the annotator stack even when users never call get_grasp_pose(). Consider moving these imports inside get_grasp_pose() (or guarding with a lazy/optional import) to reduce import-time overhead and optional-dep failures.

Suggested change
from embodichain.toolkits.graspkit.pg_grasp.antipodal_annotator import (
GraspAnnotator,
GraspAnnotatorCfg,
)
try:
from embodichain.toolkits.graspkit.pg_grasp.antipodal_annotator import (
GraspAnnotator,
GraspAnnotatorCfg,
)
except ImportError:
logger.warning(
"Optional dependency 'embodichain.toolkits.graspkit.pg_grasp.antipodal_annotator' "
"could not be imported. Grasp-related functionality may be unavailable."
)
GraspAnnotator = None # type: ignore[assignment]
GraspAnnotatorCfg = None # type: ignore[assignment]

Copilot uses AI. Check for mistakes.
import torch.nn.functional as F


@dataclass
Expand Down Expand Up @@ -1122,3 +1127,65 @@ def destroy(self) -> None:
arenas = [env]
for i, entity in enumerate(self._entities):
arenas[i].remove_actor(entity)

def get_grasp_pose(
self,
cfg: GraspAnnotatorCfg,
approach_direction: torch.Tensor = None,
is_visual: bool = False,
Comment on lines +1131 to +1135
Copy link

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

New public method RigidObject.get_grasp_pose() adds user-facing behavior (caching + pose computation) but there is no automated test coverage for it. Since tests/sim/objects/test_rigid_object.py exists, please add at least a unit test for the deterministic math path and a smoke test for the [num_envs, 4, 4] output shape.

Copilot uses AI. Check for mistakes.
) -> torch.Tensor:
if approach_direction is None:
approach_direction = torch.tensor(
[0, 0, -1], dtype=torch.float32, device=self.device
)
approach_direction = F.normalize(approach_direction, dim=-1)
if hasattr(self, "_grasp_annotator") is False:
self._grasp_annotator = GraspAnnotator(cfg=cfg)
if hasattr(self, "_hit_point_pairs") is False or cfg.force_regenerate:
Comment on lines +1142 to +1144
Copy link

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

get_grasp_pose() caches _grasp_annotator on first call and reuses it thereafter, but it also takes a cfg parameter. If callers pass a different cfg later (e.g., different viser_port / sampling params), it will be silently ignored because the existing annotator isn’t updated/recreated. Consider recreating the annotator when cfg changes (or remove cfg from the method signature and configure via a setter/constructor).

Suggested change
if hasattr(self, "_grasp_annotator") is False:
self._grasp_annotator = GraspAnnotator(cfg=cfg)
if hasattr(self, "_hit_point_pairs") is False or cfg.force_regenerate:
# (Re)create grasp annotator if it does not exist yet or if the
# configuration has changed since the last call.
annotator_needs_update = (
not hasattr(self, "_grasp_annotator")
or not hasattr(self, "_grasp_annotator_cfg")
or self._grasp_annotator_cfg != cfg
)
if annotator_needs_update:
self._grasp_annotator = GraspAnnotator(cfg=cfg)
self._grasp_annotator_cfg = cfg
# Invalidate cached hit point pairs so they will be regenerated
# with the new annotator / configuration.
if hasattr(self, "_hit_point_pairs"):
del self._hit_point_pairs
if not hasattr(self, "_hit_point_pairs") or cfg.force_regenerate:

Copilot uses AI. Check for mistakes.
vertices = torch.tensor(
self._entities[0].get_vertices(),
dtype=torch.float32,
device=self.device,
)
triangles = torch.tensor(
self._entities[0].get_triangles(), dtype=torch.int32, device=self.device
)
scale = torch.tensor(
self._entities[0].get_body_scale(),
dtype=torch.float32,
device=self.device,
)
vertices = vertices * scale
self._hit_point_pairs = self._grasp_annotator.annotate(vertices, triangles)

poses = self.get_local_pose(to_matrix=True)
poses = torch.as_tensor(poses, dtype=torch.float32, device=self.device)
grasp_poses = []
open_lengths = []
for pose in poses:
grasp_pose, open_length = self._grasp_annotator.get_approach_grasp_poses(
self._hit_point_pairs, pose, approach_direction
)
grasp_poses.append(grasp_pose)
open_lengths.append(open_length)
grasp_poses = torch.cat(
[grasp_pose.unsqueeze(0) for grasp_pose in grasp_poses], dim=0
)

if is_visual:
vertices = self._entities[0].get_vertices()
triangles = self._entities[0].get_triangles()
scale = self._entities[0].get_body_scale()
vertices = vertices * scale
GraspAnnotator.visualize_grasp_pose(
vertices=torch.tensor(
vertices, dtype=torch.float32, device=self.device
),
triangles=torch.tensor(
triangles, dtype=torch.int32, device=self.device
),
obj_pose=poses[0],
grasp_pose=grasp_poses[0],
open_length=open_lengths[0],
)
Comment on lines +1187 to +1190
Copy link

Copilot AI Mar 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

open_lengths contains tensors from get_approach_grasp_poses(), but visualize_grasp_pose() expects open_length: float and uses it in Open3D geometry transforms. Convert to a Python float (e.g., open_length.item()) before passing it to avoid runtime type errors in visualization.

Copilot uses AI. Check for mistakes.
return grasp_poses
Loading
Loading