Update grasp pose generator and pick up action#289
Open
matafela wants to merge 5 commits into
Open
Conversation
added 3 commits
June 3, 2026 16:27
Contributor
There was a problem hiding this comment.
Pull request overview
This PR enhances the grasp generation + pick-up pipeline by expanding grasp sampling (multiple approach directions), filtering near-duplicate grasp poses via NMS, and updating the pick-up action to prefer IK-feasible grasp candidates.
Changes:
- Added
pose_nmsto filter grasp pose candidates by rotational/translation proximity. - Updated antipodal grasp generation to sample additional (deviated) approach directions and apply NMS before collision filtering.
- Updated pick-up grasp resolution to choose a grasp pose that is IK-valid (instead of purely cost-based).
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 11 comments.
Show a summary per file
| File | Description |
|---|---|
| scripts/tutorials/grasp/grasp_generator.py | Tweaks tutorial sampling count and pose visualization setting. |
| embodichain/utils/math.py | Adds pose_nms utility for pose non-maximum suppression. |
| embodichain/toolkits/graspkit/pg_grasp/antipodal_generator.py | Adds multi-direction grasp pose sampling and NMS; introduces get_valid_grasp_poses. |
| embodichain/lab/sim/atomic_actions/core.py | Adds affordance helper to return valid grasp pose candidates + costs. |
| embodichain/lab/sim/atomic_actions/actions.py | Updates pick-up grasp selection to filter candidates by IK feasibility. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+485
to
+503
| n_poses = grasp_poses.shape[0] | ||
| init_qpos_repeat = init_qpos[:, None, :].repeat(1, n_poses, 1) | ||
| grasp_poses_repeat = grasp_poses[None, :, :].repeat(n_envs, 1, 1, 1) | ||
| ik_success, qpos = self.robot.compute_batch_ik( | ||
| pose=grasp_poses_repeat, | ||
| name=self.cfg.control_part, | ||
| joint_seed=init_qpos_repeat, | ||
| ) | ||
| if ik_success.sum() == 0: | ||
| is_success_list.append(False) | ||
| grasp_xpos_list.append( | ||
| torch.zeros(4, 4, dtype=torch.float32, device=self.device) | ||
| ) | ||
| continue | ||
| valid_mask = ik_success[0] | ||
| valid_poses = grasp_poses[valid_mask] | ||
| valid_costs = grasp_costs[valid_mask] | ||
| best_idx = torch.argmin(valid_costs) | ||
| best_grasp_xpos = valid_poses[best_idx] |
Comment on lines
+506
to
+509
| is_success = torch.tensor(is_success_list, device=self.device) | ||
| grasp_xpos = torch.stack(grasp_xpos_list, dim=0) | ||
|
|
||
| return is_success, grasp_xpos |
Comment on lines
+477
to
+481
| # Get best grasp pose for each object | ||
| n_envs = obj_poses.shape[0] | ||
| init_qpos = self.robot.get_qpos(name=self.cfg.control_part) | ||
| grasp_xpos_list = [] | ||
| is_success_list = [] |
Comment on lines
+618
to
+623
| return ( | ||
| False, | ||
| torch.eye(4, device=self.device), | ||
| 0.0, | ||
| torch.zeros(1, device=self.device), | ||
| ) |
Comment on lines
+642
to
+647
| return ( | ||
| False, | ||
| torch.eye(4, device=self.device), | ||
| 0.0, | ||
| torch.zeros(1, device=self.device), | ||
| ) |
Comment on lines
+675
to
+683
| # remove near grasp poses using non-maximum suppression | ||
| nms_grasp_poses, nms_indices = pose_nms( | ||
| valid_grasp_poses, | ||
| angle_th=np.pi / 30, | ||
| dist_th=0.005, | ||
| ) | ||
| valid_grasp_poses = valid_grasp_poses[nms_indices] | ||
| valid_open_lengths = valid_open_lengths[nms_indices] | ||
| valid_centers = valid_centers[nms_indices] |
Comment on lines
+82
to
+83
| n_deviated_approach_directions: int = 4 | ||
| """Number of approach directions with evenly deviated angles when sampling grasp poses.""" |
Comment on lines
+147
to
+160
| def get_valid_grasp_poses( | ||
| self, | ||
| obj_poses: torch.Tensor, | ||
| approach_direction: torch.Tensor = torch.tensor( | ||
| [0, 0, -1], dtype=torch.float32 | ||
| ), | ||
| ) -> list[tuple[torch.Tensor, torch.Tensor]]: | ||
| if self.generator is None: | ||
| self._init_generator() | ||
| results = [] | ||
| for i, obj_pose in enumerate(obj_poses): | ||
| is_success, grasp_poses, _, costs = self.generator.get_valid_grasp_poses( | ||
| obj_pose, approach_direction | ||
| ) |
Comment on lines
+2270
to
+2274
| def pose_nms( | ||
| poses: torch.Tensor, | ||
| angle_th: float = np.pi / 36, | ||
| dist_th: float = 0.003, | ||
| ) -> tuple[torch.Tensor, torch.Tensor]: |
Comment on lines
+2300
to
+2304
| keep_indices_list: list[int] = [] | ||
|
|
||
| for pose_idx in range(poses.shape[0]): | ||
| if not keep_indices_list: | ||
| keep_indices_list.append(pose_idx) |
added 2 commits
June 4, 2026 11:51
Comment on lines
+485
to
+505
| n_poses = grasp_poses.shape[0] | ||
| init_qpos_repeat = init_qpos[:, None, :].repeat(1, n_poses, 1) | ||
| grasp_poses_repeat = grasp_poses[None, :, :].repeat(n_envs, 1, 1, 1) | ||
| ik_success, qpos = self.robot.compute_batch_ik( | ||
| pose=grasp_poses_repeat, | ||
| name=self.cfg.control_part, | ||
| joint_seed=init_qpos_repeat, | ||
| ) | ||
| if ik_success.sum() == 0: | ||
| is_success_list.append(False) | ||
| grasp_xpos_list.append( | ||
| torch.zeros(4, 4, dtype=torch.float32, device=self.device) | ||
| ) | ||
| continue | ||
| valid_mask = ik_success[0] | ||
| valid_poses = grasp_poses[valid_mask] | ||
| valid_costs = grasp_costs[valid_mask] | ||
| best_idx = torch.argmin(valid_costs) | ||
| best_grasp_xpos = valid_poses[best_idx] | ||
| is_success_list.append(True) | ||
| grasp_xpos_list.append(best_grasp_xpos) |
Comment on lines
+147
to
+156
| def get_valid_grasp_poses( | ||
| self, | ||
| obj_poses: torch.Tensor, | ||
| approach_direction: torch.Tensor = torch.tensor( | ||
| [0, 0, -1], dtype=torch.float32 | ||
| ), | ||
| ) -> list[tuple[torch.Tensor, torch.Tensor]]: | ||
| if self.generator is None: | ||
| self._init_generator() | ||
| results = [] |
Comment on lines
+621
to
+626
| return ( | ||
| False, | ||
| torch.eye(4, device=self.device), | ||
| 0.0, | ||
| torch.zeros(1, device=self.device), | ||
| ) |
Comment on lines
+645
to
+650
| return ( | ||
| False, | ||
| torch.eye(4, device=self.device), | ||
| 0.0, | ||
| torch.zeros(1, device=self.device), | ||
| ) |
Comment on lines
+699
to
+704
| return ( | ||
| False, | ||
| torch.eye(4, device=self.device), | ||
| 0.0, | ||
| torch.zeros(1, device=self.device), | ||
| ) |
Comment on lines
717
to
720
| center_distance = torch.norm(valid_centers - mesh_center, dim=-1) | ||
| center_cost = center_distance / center_distance.max() | ||
| length_cost = 1 - valid_open_lengths / valid_open_lengths.max() | ||
| total_cost = 0.3 * angle_cost + 0.3 * length_cost + 0.4 * center_cost |
Comment on lines
+82
to
+83
| n_deviated_approach_directions: int = 4 | ||
| """Number of approach directions with evenly deviated angles when sampling grasp poses.""" |
Comment on lines
+2270
to
+2276
| def pose_nms_indices( | ||
| poses: torch.Tensor, | ||
| angle_th: float = np.pi / 36, | ||
| dist_th: float = 0.003, | ||
| preserve_order: bool = False, | ||
| ) -> torch.Tensor: | ||
| """Return pose indices after removing poses that are too close. |
Comment on lines
270
to
275
| is_success, grasp_pose, open_length = grasp_generator.get_grasp_poses( | ||
| obj_pose, | ||
| approach_direction, | ||
| visualize_collision=False, | ||
| visualize_pose=False, | ||
| visualize_pose=True, | ||
| ) |
| is_success = torch.tensor(is_success_list, device=self.device) | ||
| grasp_xpos = torch.stack(grasp_xpos_list, dim=0) | ||
|
|
||
| return is_success, grasp_xpos |
yuecideng
requested changes
Jun 4, 2026
| init_qpos = self.robot.get_qpos(name=self.cfg.control_part) | ||
| grasp_xpos_list = [] | ||
| is_success_list = [] | ||
| for i in range(n_envs): |
Contributor
There was a problem hiding this comment.
This also can be vectorized.
| return offset_pose | ||
|
|
||
|
|
||
| def pose_nms_indices( |
Contributor
There was a problem hiding this comment.
Add unitest for this function. Can we combine the two functions?
| number of sampled surface points, ray perturbation angle, and gripper jaw | ||
| distance limits. See :class:`AntipodalSamplerCfg` for details.""" | ||
|
|
||
| max_deviation_angle: float = np.pi / 6 |
Contributor
There was a problem hiding this comment.
Add a benchmark for the updated grasp generation algorithm. Two dimensions, 1. number of vertices. 2. number of grasp samples.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Upgrade grasp pose generator and pick up action.
TODO:
Type of change
Checklist
black .command to format the code base.