Update grasp pose generator and pick up action#289
Closed
matafela wants to merge 65 commits into
Closed
Conversation
Co-authored-by: chenjian <chenjian@dexforce.com>
Co-authored-by: chenjian <chenjian@dexforce.com>
Co-authored-by: chenjian <chenjian@dexforce.com>
Co-authored-by: chenjian <chenjian@dexforce.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: chenjian <chenjian@dexforce.com>
Co-authored-by: Jietao Chen <chenjietao@dexforce.com> Co-authored-by: Yueci Deng <dengyueci@qq.com>
Co-authored-by: chenjian <chenjian@dexforce.com>
Co-authored-by: chenjian <chenjian@dexforce.com>
Co-authored-by: yuanhaonan <yuanhaonan@dexforce.top>
…entation (#247) Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Copilot <copilot@github.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: chenjian <chenjian@dexforce.com> Co-authored-by: yuecideng <dengyueci@qq.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
…ration (#239) Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Chen Jian <mtfl1996@outlook.com> Co-authored-by: chenjian <chenjian@dexforce.com> Co-authored-by: Copilot <copilot@github.com>
Co-authored-by: chenjian <chenjian@dexforce.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: Yueci Deng <dengyueci@qq.com>
Co-authored-by: WaferLi <63717327+WaferLi@users.noreply.github.com> Co-authored-by: liwenfeng <liwenfeng@dexforce.top> Co-authored-by: chenjian <chenjian@dexforce.com> Co-authored-by: daojun <lookangela@qq.com> Co-authored-by: Chen Jian <mtfl1996@outlook.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
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
| @@ -698,6 +718,67 @@ def get_grasp_poses( | |||
| 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
Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 5 out of 5 changed files in this pull request and generated 10 comments.
Comments suppressed due to low confidence (1)
embodichain/toolkits/graspkit/pg_grasp/antipodal_generator.py:786
get_grasp_poses()is annotated to returnbest_open_lengthas afloat, but indexing a tensor returns a 0-dtorch.Tensor. This can lead to downstream type/device issues (e.g.,torch.tensor(open_length_list, device=...)inAntipodalAffordance.get_best_grasp_poses). Return a Python float via.item()(or update the annotation totorch.Tensor).
best_open_length = valid_open_lengths[best_idx]
if visualize_pose:
Comment on lines
+485
to
+509
| 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) | ||
| 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
470
to
471
| ) | ||
| obj_poses = semantics.entity.get_local_pose(to_matrix=True) |
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
+82
to
+83
| n_deviated_approach_directions: int = 4 | ||
| """Number of approach directions with evenly deviated angles when sampling grasp poses.""" |
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
+700
to
+705
| return ( | ||
| False, | ||
| torch.eye(4, device=self.device), | ||
| 0.0, | ||
| torch.zeros(1, device=self.device), | ||
| ) |
Comment on lines
+754
to
+757
| TODO: | ||
| 1. Support Top-k grasp poses selection. | ||
| 2. Support more selection criteria. | ||
|
|
Comment on lines
+678
to
+682
| # TODO: too slow | ||
| # # remove near grasp poses using non-maximum suppression | ||
| # nms_indices = pose_nms_indices( | ||
| # valid_grasp_poses, | ||
| # angle_th=np.pi / 18, |
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
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
+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
+700
to
+705
| return ( | ||
| False, | ||
| torch.eye(4, device=self.device), | ||
| 0.0, | ||
| torch.zeros(1, device=self.device), | ||
| ) |
Comment on lines
473
to
475
| grasp_poses_result = semantics.affordance.get_valid_grasp_poses( | ||
| obj_poses=obj_poses, approach_direction=self.approach_direction | ||
| ) |
Comment on lines
+147
to
+166
| 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 | ||
| ) | ||
| if not is_success: | ||
| logger.log_warning( | ||
| f"Failed to find valid grasp poses for {i}-th object." | ||
| ) | ||
| results.append((grasp_poses, costs)) | ||
| return results |
Comment on lines
+492
to
+499
| for i in range(n_envs): | ||
| n_pose = grasp_poses_result[i][0].shape[0] | ||
| grasp_xpos_padding[i, :n_pose] = grasp_poses_result[i][0] | ||
| grasp_cost_padding[i, :n_pose] = grasp_poses_result[i][1] | ||
| # padding with the first grasp pose, which is usually the best one, to ensure that the padded grasp poses are valid for IK computation, although they may not be optimal. | ||
| grasp_xpos_padding[i, n_pose:] = grasp_poses_result[i][0][0] | ||
| grasp_cost_padding[i, n_pose:] = grasp_poses_result[i][1][0] | ||
|
|
Comment on lines
+506
to
+513
| grasp_cost_masked = torch.where(ik_success, grasp_cost_padding, 10000.0) | ||
| best_cost, best_idx = grasp_cost_masked.min(dim=1) | ||
| is_success = best_cost < 9999.0 # usually cost < 1.0 | ||
| best_grasp_xpos = grasp_xpos_padding[ | ||
| torch.arange(n_envs, device=self.device), best_idx | ||
| ] | ||
|
|
||
| return is_success, best_grasp_xpos |
Comment on lines
+480
to
+485
| n_max_pose = 0 | ||
| for result in grasp_poses_result: | ||
| n_pose = result[0].shape[0] | ||
| if n_pose > n_max_pose: | ||
| n_max_pose = n_pose | ||
|
|
Comment on lines
1449
to
1453
| if max_velocity is not None: | ||
| drive_args["max_velocity"] = max_velocity[i].cpu().numpy() | ||
| if friction is not None: | ||
| drive_args["joint_friction"] = friction[i].cpu().numpy() | ||
| if armature is not None: | ||
| drive_args["armature"] = armature[i].cpu().numpy() | ||
| self._entities[env_idx].set_drive(**drive_args) |
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
+616
to
+626
| if self._hit_point_pairs is None: | ||
| logger.log_warning( | ||
| "No antipodal point pairs available. " | ||
| "Call generate() or annotate() first." | ||
| ) | ||
| return False, torch.eye(4, device=self.device), 0.0 | ||
| return ( | ||
| False, | ||
| torch.eye(4, device=self.device), | ||
| 0.0, | ||
| torch.zeros(1, device=self.device), | ||
| ) |
Comment on lines
+643
to
+650
| if valid_mask.sum() == 0: | ||
| logger.log_warning("No valid antipodal pairs after angle filtering.") | ||
| return False, torch.eye(4, device=self.device), 0.0 | ||
| return ( | ||
| False, | ||
| torch.eye(4, device=self.device), | ||
| 0.0, | ||
| torch.zeros(1, device=self.device), | ||
| ) |
Comment on lines
+698
to
+705
| if is_colliding.logical_not().sum() == 0: | ||
| logger.log_warning("No valid antipodal pairs after angle filtering.") | ||
| return False, torch.eye(4, device=self.device), 0.0 | ||
| logger.log_warning("No valid antipodal pairs after collision filtering.") | ||
| return ( | ||
| False, | ||
| torch.eye(4, device=self.device), | ||
| 0.0, | ||
| torch.zeros(1, device=self.device), | ||
| ) |
Comment on lines
+480
to
+488
| n_max_pose = 0 | ||
| for result in grasp_poses_result: | ||
| n_pose = result[0].shape[0] | ||
| if n_pose > n_max_pose: | ||
| n_max_pose = n_pose | ||
|
|
||
| grasp_xpos_padding = torch.zeros( | ||
| (n_envs, n_max_pose, 4, 4), dtype=torch.float32, device=self.device | ||
| ) |
Comment on lines
+492
to
+499
| for i in range(n_envs): | ||
| n_pose = grasp_poses_result[i][0].shape[0] | ||
| grasp_xpos_padding[i, :n_pose] = grasp_poses_result[i][0] | ||
| grasp_cost_padding[i, :n_pose] = grasp_poses_result[i][1] | ||
| # padding with the first grasp pose, which is usually the best one, to ensure that the padded grasp poses are valid for IK computation, although they may not be optimal. | ||
| grasp_xpos_padding[i, n_pose:] = grasp_poses_result[i][0][0] | ||
| grasp_cost_padding[i, n_pose:] = grasp_poses_result[i][1][0] | ||
|
|
Comment on lines
+22
to
+26
| import argparse | ||
| import numpy as np | ||
| import time | ||
| import torch | ||
|
|
| return interp_trajectory | ||
|
|
||
|
|
||
| def test_grasp_pose_generator(): |
Comment on lines
+678
to
+688
| # TODO: too slow | ||
| # # remove near grasp poses using non-maximum suppression | ||
| # nms_indices = pose_nms_indices( | ||
| # valid_grasp_poses, | ||
| # angle_th=np.pi / 18, | ||
| # dist_th=0.01, | ||
| # ) | ||
| # valid_grasp_poses = valid_grasp_poses[nms_indices] | ||
| # valid_open_lengths = valid_open_lengths[nms_indices] | ||
| # valid_centers = valid_centers[nms_indices] | ||
|
|
f20163f to
a912e5c
Compare
Comment on lines
+293
to
+298
| start = time.perf_counter() | ||
| is_success, grasp_poses, open_lengths, total_cost = generator.get_valid_grasp_poses( | ||
| object_pose=object_pose, | ||
| approach_direction=approach_direction, | ||
| visualize_collision=False, | ||
| ) |
| return sim | ||
|
|
||
|
|
||
| def create_robot(sim: SimulationManager, position=[0.0, 0.0, 0.0]) -> Robot: |
Comment on lines
+132
to
+137
| def create_mug(sim: SimulationManager): | ||
| mug_cfg = RigidObjectCfg( | ||
| uid="table", | ||
| shape=MeshCfg( | ||
| fpath=get_data_path("CoffeeCup/cup.ply"), | ||
| ), |
Comment on lines
+2289
to
+2293
| else: | ||
| close = ((reference_rotations @ target_rotations.T - 1.0) * 0.5).clamp_( | ||
| min=-1.0, max=1.0 | ||
| ) > rotation_cosine_th | ||
|
|
Comment on lines
+2353
to
+2354
| rotation_always_close = angle_th > math.pi | ||
| rotation_cosine_th = math.cos(float(angle_th)) if not rotation_always_close else 0.0 |
Comment on lines
+2303
to
+2310
| def pose_nms_indices( | ||
| poses: torch.Tensor, | ||
| angle_th: float = np.pi / 36, | ||
| dist_th: float = 0.003, | ||
| preserve_order: bool = False, | ||
| chunk_size: int = _POSE_NMS_CHUNK_SIZE, | ||
| ) -> torch.Tensor: | ||
| """Return pose indices after removing poses that are too close. |
Comment on lines
+22
to
+51
| import argparse | ||
| import numpy as np | ||
| import time | ||
| import torch | ||
|
|
||
| from embodichain.lab.sim import SimulationManager, SimulationManagerCfg | ||
| from embodichain.lab.sim.objects import Robot, RigidObject | ||
| from embodichain.lab.sim.utility.action_utils import interpolate_with_distance | ||
| from embodichain.lab.sim.shapes import MeshCfg | ||
| from embodichain.lab.sim.solvers import PytorchSolverCfg | ||
| from embodichain.data import get_data_path | ||
| from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser | ||
| from embodichain.utils import logger | ||
| from embodichain.lab.sim.cfg import ( | ||
| RenderCfg, | ||
| JointDrivePropertiesCfg, | ||
| RobotCfg, | ||
| LightCfg, | ||
| RigidBodyAttributesCfg, | ||
| RigidObjectCfg, | ||
| URDFCfg, | ||
| ) | ||
| from embodichain.toolkits.graspkit.pg_grasp.antipodal_generator import ( | ||
| GraspGenerator, | ||
| GraspGeneratorCfg, | ||
| AntipodalSamplerCfg, | ||
| ) | ||
| from embodichain.toolkits.graspkit.pg_grasp.gripper_collision_checker import ( | ||
| GripperCollisionCfg, | ||
| ) |
Comment on lines
+199
to
+203
| def test_grasp_pose_generator(): | ||
|
|
||
| sim = initialize_simulation() | ||
| robot = create_robot(sim, position=[0.0, 0.0, 0.0]) | ||
| mug = create_mug(sim) |
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.