Skip to content

Update grasp pose generator and pick up action#289

Closed
matafela wants to merge 65 commits into
mainfrom
cj/upgrade-grasp-generator
Closed

Update grasp pose generator and pick up action#289
matafela wants to merge 65 commits into
mainfrom
cj/upgrade-grasp-generator

Conversation

@matafela

@matafela matafela commented Jun 3, 2026

Copy link
Copy Markdown
Collaborator

Description

Upgrade grasp pose generator and pick up action.

TODO:

  • grasp pose nms
  • more grasp pose sampler from different approach direction
  • return top k but not one grasp pose.
  • pick up action choose ik valid grasp pose

Type of change

  • Enhancement (non-breaking change which improves an existing functionality)

Checklist

  • I have run the black . command to format the code base.
  • I have made corresponding changes to the documentation
  • I have added tests that prove my fix is effective or that my feature works
  • Dependencies have been updated, if applicable.

yuecideng and others added 30 commits April 12, 2026 14:58
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: 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>
Copilot AI review requested due to automatic review settings June 4, 2026 07:42

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 11 comments.

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 thread embodichain/utils/math.py
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
Comment thread embodichain/lab/sim/atomic_actions/actions.py Outdated
Comment thread embodichain/utils/math.py
Comment thread embodichain/toolkits/graspkit/pg_grasp/antipodal_generator.py
Copilot AI review requested due to automatic review settings June 9, 2026 02:46

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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 return best_open_length as a float, but indexing a tensor returns a 0-d torch.Tensor. This can lead to downstream type/device issues (e.g., torch.tensor(open_length_list, device=...) in AntipodalAffordance.get_best_grasp_poses). Return a Python float via .item() (or update the annotation to torch.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 thread embodichain/toolkits/graspkit/pg_grasp/antipodal_generator.py Outdated
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,
yuecideng and others added 3 commits June 9, 2026 18:22
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>
Copilot AI review requested due to automatic review settings June 10, 2026 08:56

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 11 comments.

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 thread embodichain/toolkits/graspkit/pg_grasp/antipodal_generator.py
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

Copilot AI review requested due to automatic review settings June 12, 2026 07:18

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 12 comments.

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]

Copilot AI review requested due to automatic review settings June 15, 2026 04:24
@yuecideng yuecideng closed this Jun 15, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 8 comments.

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 thread embodichain/utils/math.py
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 thread embodichain/utils/math.py
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 thread embodichain/utils/math.py
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

8 participants