Add grasp annotator

docs/source/tutorial/grasp_generator.rst

@@ -0,0 +1,77 @@
+Generating and Executing Robot Grasps
+======================================
+
+.. currentmodule:: embodichain.lab.sim
+
+This tutorial demonstrates how to generate antipodal grasp poses for a target object and execute a full grasp trajectory with a robot arm. It covers scene initialization, robot and object creation, interactive grasp region annotation, grasp pose computation, and trajectory execution in the simulation loop.
+
+The Code
+~~~~~~~~
+
+The tutorial corresponds to the ``grasp_generator.py`` script in the ``scripts/tutorials/grasp`` directory.
+
+.. dropdown:: Code for grasp_generator.py
+   :icon: code
+
+   .. literalinclude:: ../../../scripts/tutorials/grasp/grasp_generator.py
+      :language: python
+      :linenos:
+
+
+The Code Explained
+~~~~~~~~~~~~~~~~~~
+
+Configuring the simulation
+--------------------------
+
+Command-line arguments are parsed with ``argparse`` to select the number of parallel environments, the compute device, and optional rendering features such as ray tracing and headless mode.
+
+.. literalinclude:: ../../../scripts/tutorials/grasp/grasp_generator.py
+   :language: python
+   :start-at: def parse_arguments():
+   :end-at: return parser.parse_args()
+
+The parsed arguments are passed to ``initialize_simulation``, which builds a :class:`SimulationManagerCfg` and creates the :class:`SimulationManager` instance. When ray tracing is enabled a directional :class:`cfg.LightCfg` is also added to the scene.
+
+.. literalinclude:: ../../../scripts/tutorials/grasp/grasp_generator.py
+   :language: python
+   :start-at: def initialize_simulation(args) -> SimulationManager:
+   :end-at: return sim
+
+Annotating and computing grasp poses
+-------------------------------------
+
+Grasp generation is performed by :meth:`objects.RigidObject.get_grasp_pose`, which internally runs an antipodal sampler on the object mesh. A :class:`toolkits.graspkit.pg_grasp.GraspAnnotatorCfg` controls sampler parameters (sample count, gripper jaw limits) and the interactive annotation workflow:
+
+1. Open the visualization in a browser at the reported port (e.g. ``http://localhost:11801``).
+2. Use *Rect Select Region* to highlight the area of the object that should be grasped.
+3. Click *Confirm Selection* to finalize the region.
+
+The function returns a batch of ``(N_envs, 4, 4)`` homogeneous transformation matrices representing candidate grasp frames in the world coordinate system.
+
+For each grasp pose, gripper approach direction in world coordinate is required to compute the antipodal grasp. In this tutorial, we use a fixed approach direction (straight down in world frame) for simplicity, but it can be customized based on the task or object geometry.
+
+.. literalinclude:: ../../../scripts/tutorials/grasp/grasp_generator.py
+   :language: python
+   :start-at: # get mug grasp pose
+   :end-at: logger.log_info(f"Get grasp pose cost time: {cost_time:.2f} seconds")
+
+
+The Code Execution
+~~~~~~~~~~~~~~~~~~
+
+To run the script, execute the following command from the project root:
+
+.. code-block:: bash
+
+   python scripts/tutorials/grasp/grasp_generator.py
+
+A simulation window will open showing the robot and the mug. A browser-based visualizer will also launch (default port ``11801``) for interactive grasp region annotation.
+
+You can customize the run with additional arguments:
+
+.. code-block:: bash
+
+   python scripts/tutorials/grasp/grasp_generator.py --num_envs <n> --device <cuda/cpu> --enable_rt --headless
+
+After confirming the grasp region in the browser, the script will compute a grasp pose, print the elapsed time, and then wait for you to press **Enter** before executing the full grasp trajectory in the simulation. Press **Enter** again to exit once the motion is complete.

docs/source/tutorial/index.rst

@@ -14,6 +14,7 @@ Tutorials
    sensor
    motion_gen
    gizmo
+   grasp_generator
    basic_env
    modular_env
    rl

embodichain/lab/sim/objects/rigid_object.py

@@ -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,
+)
+import torch.nn.functional as F
 
 
 @dataclass
@@ -1122,3 +1127,63 @@ 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,
+    ) -> 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:
+            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._grasp_annotator = GraspAnnotator(
+                vertices=vertices, triangles=triangles, cfg=cfg
+            )
+
+        # Annotate antipodal point pairs
+        if hasattr(self, "_hit_point_pairs") is False or cfg.force_regenerate:
+            self._hit_point_pairs = self._grasp_annotator.annotate()
+
+        poses = self.get_local_pose(to_matrix=True)
+        poses = torch.as_tensor(poses, dtype=torch.float32, device=self.device)
+        grasp_poses: tuple[torch.Tensor] = []
+        open_lengths: tuple[torch.Tensor] = []
+        for pose in poses:
+            grasp_pose, open_length = self._grasp_annotator.get_approach_grasp_poses(
+                self._hit_point_pairs, pose, approach_direction, is_visual=False
+            )
+            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
+            self._grasp_annotator.visualize_grasp_pose(
+                obj_pose=poses[0],
+                grasp_pose=grasp_poses[0],
+                open_length=open_lengths[0].item(),
+            )
+        return grasp_poses