feat(atomic-actions): improve upright grasp selection

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,13 +524,25 @@ 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,
             obj_poses=object_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]

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:

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)