diff --git a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py index bfefa7656..f29f73f47 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py +++ b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py @@ -81,6 +81,33 @@ make_manipulation_slot, ) +_UPRIGHT_SIDE_GRASP_MAX_AXIS_ALIGNMENT = 0.65 +_UPRIGHT_SIDE_GRASP_MIN_AXIS_FRACTION = 0.35 +_UPRIGHT_SIDE_GRASP_MAX_AXIS_FRACTION = 0.75 +_UPRIGHT_SIDE_GRASP_HEIGHT_COST_WEIGHT = 2.0 + + +def _upright_yaw_pose_variants( + target_pose: torch.Tensor, + sample_count: int, +) -> torch.Tensor: + """Return object poses with evenly sampled world-Z yaw rotations.""" + signed_steps = [0] + for step in range(1, (sample_count + 1) // 2): + signed_steps.extend((step, -step)) + if sample_count % 2 == 0: + signed_steps.append(sample_count // 2) + angles = target_pose.new_tensor(signed_steps) * (2.0 * math.pi / sample_count) + yaw = target_pose.new_zeros((sample_count, 3, 3)) + yaw[:, 0, 0] = torch.cos(angles) + yaw[:, 0, 1] = -torch.sin(angles) + yaw[:, 1, 0] = torch.sin(angles) + yaw[:, 1, 1] = torch.cos(angles) + yaw[:, 2, 2] = 1.0 + variants = target_pose[:, None].repeat(1, sample_count, 1, 1) + variants[:, :, :3, :3] = torch.matmul(yaw[None], target_pose[:, None, :3, :3]) + return variants + @dataclass(frozen=True, slots=True, eq=False) class GraspGoal(ObjectActionGoal): @@ -127,6 +154,9 @@ class PickUpOptions(ActionOptions): downstream_object_target_poses: tuple[PoseGoalValue, ...] = () """Future object poses that must be reachable with the selected grasp.""" + upright_yaw_samples: int = 1 + """Equivalent world-yaw samples for semantically upright downstream targets.""" + obj_upright_direction: torch.Tensor | None = None """Optional object local direction used to choose the upright grasp rotation.""" @@ -142,6 +172,8 @@ def __post_init__(self) -> None: raise ValueError("lift_height must be non-negative.") if self.pre_grasp_distance < 0.0: raise ValueError("pre_grasp_distance must be non-negative.") + if self.upright_yaw_samples < 1: + raise ValueError("upright_yaw_samples must be positive.") if self.approach_direction.shape != (3,): raise ValueError("approach_direction must have shape (3,).") if not torch.isfinite(self.approach_direction).all(): @@ -492,6 +524,17 @@ def _resolve_grasp_pose( device=self.device, ) is_positive_part = options.pick_object_part == "top" + grasp_cost_fn = None + if options.rotate_upright is not None: + grasp_cost_fn = lambda candidate_object_pose, grasp_poses, costs: ( + self._upright_grasp_costs( + semantics, + candidate_object_pose, + grasp_poses, + costs, + options, + ) + ) grasp_poses_result = generator.get_valid_grasp_poses( mesh_vertices=affordance.mesh_vertices, mesh_triangles=affordance.mesh_triangles, @@ -499,6 +542,7 @@ def _resolve_grasp_pose( approach_direction=approach_direction, obj_longest_axis=obj_longest_axis, is_positive_part=is_positive_part, + pose_cost_fn=grasp_cost_fn, ) num_envs = object_pose.shape[0] n_max_pose = max(r[0].shape[0] for r in grasp_poses_result) @@ -582,8 +626,17 @@ def _select_feasible_grasp_variants( alignment_success = self._approach_alignment_mask( grasp_variants, options, approach_direction ) + upright_compatible = self._upright_grasp_compatibility_mask( + grasp_variants, + object_poses, + options, + ) pickup_success = ( - alignment_success & pre_grasp_success & grasp_success & lift_success + upright_compatible + & alignment_success + & pre_grasp_success + & grasp_success + & lift_success ) downstream_success_counts: list[list[int]] = [] object_to_eef_variants = torch.matmul( @@ -606,17 +659,37 @@ def _select_feasible_grasp_variants( f"(4, 4) or ({num_envs}, 4, 4), but got " f"{object_target_pose.shape}." ) - downstream_eef_variants = torch.matmul( - object_target_pose[:, None, None], object_to_eef_variants - ) - downstream_success, downstream_seed = self._compute_batch_candidate_ik( - downstream_eef_variants, downstream_seed, manipulator + object_target_variants = _upright_yaw_pose_variants( + object_target_pose, + options.upright_yaw_samples, ) + downstream_success = torch.zeros_like(pickup_success) + selected_qpos = downstream_seed + for yaw_target in object_target_variants.unbind(dim=1): + downstream_eef_variants = torch.matmul( + yaw_target[:, None, None], object_to_eef_variants + ) + yaw_success, yaw_qpos = self._compute_batch_candidate_ik( + downstream_eef_variants, + downstream_seed, + manipulator, + ) + newly_solved = ~downstream_success & yaw_success + selected_qpos = torch.where( + newly_solved[..., None], + yaw_qpos, + selected_qpos, + ) + downstream_success |= yaw_success + if bool((pickup_success & downstream_success).any(dim=(1, 2)).all()): + break + downstream_seed = selected_qpos pickup_success &= downstream_success downstream_success_counts.append(pickup_success.sum(dim=(1, 2)).tolist()) if not pickup_success.any(dim=(1, 2)).all(): logger.log_warning( "PickUp found no candidate with a feasible vertical pickup path: " + f"upright_compatible={upright_compatible.sum(dim=(1, 2)).tolist()}, " f"aligned={alignment_success.sum(dim=(1, 2)).tolist()}, " f"pre_grasp={pre_grasp_success.sum(dim=(1, 2)).tolist()}, " f"grasp={(pre_grasp_success & grasp_success).sum(dim=(1, 2)).tolist()}, " @@ -690,6 +763,119 @@ def _compute_batch_candidate_ik( qpos.reshape(num_envs, n_pose, n_variant, manipulator_dof), ) + def _upright_grasp_compatibility_mask( + self, + grasp_xpos: torch.Tensor, + object_poses: torch.Tensor, + options: PickUpOptions, + ) -> torch.Tensor: + """Reject upright grasps that clamp the object's support and top faces.""" + shape = grasp_xpos.shape[:3] + if options.rotate_upright is None: + return torch.ones(shape, dtype=torch.bool, device=grasp_xpos.device) + local_upright = self._normalized_obj_upright_direction(options).to( + device=grasp_xpos.device, + dtype=grasp_xpos.dtype, + ) + object_poses = object_poses.to( + device=grasp_xpos.device, + dtype=grasp_xpos.dtype, + ) + world_upright = torch.matmul(object_poses[:, :3, :3], local_upright) + closing_axes = torch.nn.functional.normalize( + grasp_xpos[..., :3, 0], + dim=-1, + ) + axis_alignment = torch.abs( + torch.sum(closing_axes * world_upright[:, None, None, :], dim=-1) + ) + return axis_alignment <= _UPRIGHT_SIDE_GRASP_MAX_AXIS_ALIGNMENT + + def _upright_grasp_costs( + self, + semantics: ObjectSemantics, + object_pose: torch.Tensor, + grasp_poses: torch.Tensor, + costs: torch.Tensor, + options: PickUpOptions, + ) -> torch.Tensor: + """Rank side grasps before generator top-k truncation.""" + local_upright = self._normalized_obj_upright_direction(options).to( + device=grasp_poses.device, + dtype=grasp_poses.dtype, + ) + object_pose = object_pose.to( + device=grasp_poses.device, + dtype=grasp_poses.dtype, + ) + world_upright = torch.matmul(object_pose[:3, :3], local_upright) + closing_axes = torch.nn.functional.normalize( + grasp_poses[:, :3, 0], + dim=-1, + ) + axis_alignment = torch.abs( + torch.sum(closing_axes * world_upright[None, :], dim=-1) + ) + adjusted = torch.where( + axis_alignment <= _UPRIGHT_SIDE_GRASP_MAX_AXIS_ALIGNMENT, + costs, + torch.full_like(costs, torch.inf), + ) + + affordance = semantics.affordance + if not isinstance(affordance, AntipodalAffordance): + return adjusted + vertices = affordance.mesh_vertices + if vertices is None: + return adjusted + vertices = vertices.to( + dtype=grasp_poses.dtype, + device=grasp_poses.device, + ) + if vertices.ndim != 2 or vertices.shape[-1] != 3 or vertices.numel() == 0: + return adjusted + vertex_axis_positions = torch.matmul(vertices, local_upright) + axis_min = vertex_axis_positions.min() + axis_extent = vertex_axis_positions.max() - axis_min + if float(axis_extent) <= 1.0e-6: + return adjusted + + relative_centers = grasp_poses[:, :3, 3] - object_pose[None, :3, 3] + center_axis_positions = torch.sum( + relative_centers * world_upright[None, :], + dim=-1, + ) + center_fractions = (center_axis_positions - axis_min) / axis_extent + interval = ( + _UPRIGHT_SIDE_GRASP_MAX_AXIS_FRACTION + - _UPRIGHT_SIDE_GRASP_MIN_AXIS_FRACTION + ) + height_penalty = ( + torch.clamp( + _UPRIGHT_SIDE_GRASP_MIN_AXIS_FRACTION - center_fractions, + min=0.0, + ) + + torch.clamp( + center_fractions - _UPRIGHT_SIDE_GRASP_MAX_AXIS_FRACTION, + min=0.0, + ) + ) / interval + return adjusted + _UPRIGHT_SIDE_GRASP_HEIGHT_COST_WEIGHT * height_penalty + + def _normalized_obj_upright_direction( + self, + options: PickUpOptions, + ) -> torch.Tensor: + """Return the configured non-zero object-local upright direction.""" + direction = options.obj_upright_direction + if direction is None: + direction = torch.tensor([0, 0, 1], dtype=torch.float32) + direction = direction.to(device=self.device, dtype=torch.float32) + norm = torch.linalg.vector_norm(direction) + if norm <= 1.0e-6: + logger.log_error("obj_upright_direction must be non-zero.", ValueError) + return direction / norm + def _upright_adjusted_grasp_poses( self, grasp_xpos: torch.Tensor, @@ -700,14 +886,14 @@ def _upright_adjusted_grasp_poses( if options.rotate_upright is None: return grasp_xpos - if options.obj_upright_direction is None: - upright_direction = torch.tensor( - [0, 0, 1], dtype=torch.float32, device=self.device - ) - else: - upright_direction = options.obj_upright_direction.to( - device=self.device, dtype=torch.float32 - ) + upright_direction = self._normalized_obj_upright_direction(options).to( + device=grasp_xpos.device, + dtype=grasp_xpos.dtype, + ) + object_pose = object_pose.to( + device=grasp_xpos.device, + dtype=grasp_xpos.dtype, + ) obj_upright = torch.matmul(object_pose[:, :3, :3], upright_direction) adjusted_grasp_xpos = grasp_xpos.clone() grasp_ry = adjusted_grasp_xpos[..., :3, 1] diff --git a/embodichain/toolkits/graspkit/pg_grasp/_antipodal_backend.py b/embodichain/toolkits/graspkit/pg_grasp/_antipodal_backend.py index 62760de10..1d5a69aa4 100644 --- a/embodichain/toolkits/graspkit/pg_grasp/_antipodal_backend.py +++ b/embodichain/toolkits/graspkit/pg_grasp/_antipodal_backend.py @@ -29,6 +29,7 @@ import viser import viser.transforms as tf +from collections.abc import Callable from pathlib import Path from typing import Any, cast @@ -491,6 +492,9 @@ def get_valid_grasp_poses( obj_longest_axis: torch.Tensor | None = None, is_positive_part: bool = True, visualize_collision: bool = False, + pose_cost_fn: ( + Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None + ) = None, ): """Filter valid grasps, optionally to one projected half of the object. @@ -502,6 +506,8 @@ def get_valid_grasp_poses( is_positive_part: When an axis is supplied, select the positive projected half if true and the negative half otherwise. visualize_collision: Whether to visualize collision checks. + pose_cost_fn: Optional callback used to rerank collision-free poses + before the top-k candidate limit is applied. Returns: Success, grasp poses, opening lengths, and grasp costs. @@ -753,6 +759,16 @@ def _filter_valid_grasp_poses( center_cost = center_distance / center_distance.max() length_cost = 1 - valid_open_lengths / valid_open_lengths.max() total_cost = 0.25 * angle_cost + 0.25 * length_cost + 0.5 * center_cost + if pose_cost_fn is not None: + adjusted_cost = pose_cost_fn(valid_grasp_poses, total_cost) + if not isinstance(adjusted_cost, torch.Tensor): + raise TypeError("pose_cost_fn must return a torch.Tensor.") + if adjusted_cost.shape != total_cost.shape: + raise ValueError("pose_cost_fn must preserve the grasp cost shape.") + total_cost = adjusted_cost.to( + device=total_cost.device, + dtype=total_cost.dtype, + ) n_valid = valid_grasp_poses.shape[0] if n_valid == 0: diff --git a/embodichain/toolkits/graspkit/pg_grasp/pose_generator.py b/embodichain/toolkits/graspkit/pg_grasp/pose_generator.py index 101475c9f..1ea4bdff6 100644 --- a/embodichain/toolkits/graspkit/pg_grasp/pose_generator.py +++ b/embodichain/toolkits/graspkit/pg_grasp/pose_generator.py @@ -18,6 +18,7 @@ from __future__ import annotations +from collections.abc import Callable from copy import deepcopy import math from typing import Literal @@ -411,8 +412,13 @@ def get_valid_grasp_poses( approach_direction: torch.Tensor, obj_longest_axis: torch.Tensor | None = None, is_positive_part: bool | torch.Tensor = True, + pose_cost_fn: ( + Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor] | None + ) = None, ) -> list[tuple[torch.Tensor, torch.Tensor]]: """Return ranked candidates, optionally from one projected axis end.""" + if pose_cost_fn is not None and not callable(pose_cost_fn): + raise TypeError("pose_cost_fn must be callable or None.") backend = self._backend(mesh_vertices, mesh_triangles) poses = self._object_poses(obj_poses, device=backend.device) directions = self._approach_directions( @@ -464,6 +470,15 @@ def get_valid_grasp_poses( approach_direction=directions[index], obj_longest_axis=None if axes is None else axes[index], is_positive_part=bool(positive_parts[index].item()), + pose_cost_fn=( + None + if pose_cost_fn is None + else lambda grasp_poses, costs: pose_cost_fn( + object_pose, + grasp_poses, + costs, + ) + ), ) if grasp_poses.shape == (4, 4): grasp_poses = grasp_poses.unsqueeze(0) diff --git a/embodichain/toolkits/graspkit/pose_generator.py b/embodichain/toolkits/graspkit/pose_generator.py index d6be0ee01..1d2786f0d 100644 --- a/embodichain/toolkits/graspkit/pose_generator.py +++ b/embodichain/toolkits/graspkit/pose_generator.py @@ -19,6 +19,7 @@ from __future__ import annotations from abc import ABC, abstractmethod +from collections.abc import Callable from copy import deepcopy import math @@ -124,8 +125,26 @@ def get_valid_grasp_poses( approach_direction: torch.Tensor, obj_longest_axis: torch.Tensor | None = None, is_positive_part: bool | torch.Tensor = True, + pose_cost_fn: ( + Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor] | None + ) = None, ) -> list[tuple[torch.Tensor, torch.Tensor]]: - """Return candidates, optionally restricted to one projected axis end.""" + """Return candidates, optionally reranked before candidate truncation. + + Args: + mesh_vertices: Target-local mesh vertices. + mesh_triangles: Target-local mesh triangle indices. + obj_poses: Batched world poses for the target object. + approach_direction: Shared or batched world-frame approach direction. + obj_longest_axis: Optional shared or batched object-selection axis. + is_positive_part: Whether to use the positive end of the selection axis. + pose_cost_fn: Optional callback receiving one object pose, its candidate + grasp poses, and their current costs. It must return costs with the + same shape and is applied before the implementation's top-k limit. + + Returns: + Per-object candidate grasp poses and their ranked costs. + """ @abstractmethod def get_best_grasp_poses( diff --git a/tests/sim/atomic_actions/test_actions.py b/tests/sim/atomic_actions/test_actions.py index 9ed24021f..907dc805e 100644 --- a/tests/sim/atomic_actions/test_actions.py +++ b/tests/sim/atomic_actions/test_actions.py @@ -19,6 +19,7 @@ from __future__ import annotations import math +from collections.abc import Callable from dataclasses import replace from typing import Literal, TypeVar from unittest.mock import Mock @@ -136,6 +137,9 @@ def get_valid_grasp_poses( approach_direction: torch.Tensor, obj_longest_axis: torch.Tensor | None = None, is_positive_part: bool | torch.Tensor = True, + pose_cost_fn: ( + Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor] | None + ) = None, ) -> list[tuple[torch.Tensor, torch.Tensor]]: del ( mesh_vertices, @@ -144,13 +148,18 @@ def get_valid_grasp_poses( obj_longest_axis, is_positive_part, ) - return [ - ( - torch.eye(4, dtype=torch.float32, device=obj_poses.device).unsqueeze(0), - torch.zeros(1, dtype=torch.float32, device=obj_poses.device), - ) - for _ in range(obj_poses.shape[0]) - ] + results = [] + for object_pose in obj_poses: + grasp_poses = torch.eye( + 4, + dtype=torch.float32, + device=obj_poses.device, + ).unsqueeze(0) + costs = torch.zeros(1, dtype=torch.float32, device=obj_poses.device) + if pose_cost_fn is not None: + costs = pose_cost_fn(object_pose, grasp_poses, costs) + results.append((grasp_poses, costs)) + return results def get_best_grasp_poses( self, diff --git a/tests/sim/atomic_actions/test_primitives_helpers.py b/tests/sim/atomic_actions/test_primitives_helpers.py index 2ff754fe0..f5f297738 100644 --- a/tests/sim/atomic_actions/test_primitives_helpers.py +++ b/tests/sim/atomic_actions/test_primitives_helpers.py @@ -27,6 +27,10 @@ resolve_batched_pose, resolve_object_target, ) +from embodichain.lab.sim.atomic_actions.primitives.pick_up import ( + PickUpOptions, + _upright_yaw_pose_variants, +) BATCH_SIZE = 2 ROBOT_DOF = 6 @@ -83,3 +87,19 @@ def test_resolve_object_target_uses_custom_name_in_shape_error() -> None: device=torch.device("cpu"), name="placing_object_target_pose", ) + + +def test_upright_yaw_pose_variants_preserve_translation() -> None: + pose = torch.eye(4).repeat(2, 1, 1) + pose[:, :3, 3] = torch.tensor([[0.2, -0.1, 0.8], [-0.3, 0.4, 0.7]]) + + variants = _upright_yaw_pose_variants(pose, 4) + + assert variants.shape == (2, 4, 4, 4) + assert torch.allclose(variants[:, :, :3, 3], pose[:, None, :3, 3].expand(-1, 4, -1)) + assert torch.allclose(variants[:, 0], pose) + + +def test_upright_yaw_samples_must_be_positive() -> None: + with pytest.raises(ValueError, match="upright_yaw_samples"): + PickUpOptions(upright_yaw_samples=0) diff --git a/tests/toolkits/test_parallel_jaw_grasp_pose_generator.py b/tests/toolkits/test_parallel_jaw_grasp_pose_generator.py index aac94b82c..9fd00278f 100644 --- a/tests/toolkits/test_parallel_jaw_grasp_pose_generator.py +++ b/tests/toolkits/test_parallel_jaw_grasp_pose_generator.py @@ -255,6 +255,40 @@ def test_failed_candidates_receive_infinite_cost_and_refresh_is_service_owned( assert backend.instances[0].annotate_calls == 1 +def test_valid_candidates_apply_object_aware_cost_before_return( + backend: type[_Backend], + monkeypatch: pytest.MonkeyPatch, +) -> None: + vertices, triangles = _geometry() + generator = _generator() + generator.prepare_mesh(mesh_vertices=vertices, mesh_triangles=triangles) + object_pose = torch.eye(4) + object_pose[0, 3] = 0.25 + grasp_poses = torch.eye(4).repeat(2, 1, 1) + + def get_valid_grasp_poses(**kwargs: object) -> tuple[object, ...]: + callback = kwargs["pose_cost_fn"] + assert callable(callback) + costs = callback(grasp_poses, torch.tensor([0.2, 0.4])) + return True, grasp_poses, torch.tensor([0.02, 0.03]), costs + + monkeypatch.setattr( + backend.instances[0], + "get_valid_grasp_poses", + get_valid_grasp_poses, + ) + + results = generator.get_valid_grasp_poses( + mesh_vertices=vertices, + mesh_triangles=triangles, + obj_poses=object_pose.unsqueeze(0), + approach_direction=torch.tensor([0, 0, -1]), + pose_cost_fn=lambda obj, _grasps, current: current + obj[0, 3], + ) + + assert torch.allclose(results[0][1], torch.tensor([0.45, 0.65])) + + def test_rejects_collision_margin_outside_gripper_opening() -> None: with pytest.raises(ValueError, match="opening_margin"): AntipodalGraspPoseGenerator(