feat(atomic-actions): improve upright grasp selection

[skywhite1024]

Description

This PR extracts the reusable Atomic Action prerequisite from GenSim PR #538. It adds object-aware upright grasp ranking before GraspKit top-k truncation, rejects top/bottom clamp poses for upright transport, and checks yaw-equivalent downstream targets without persisting grounded poses in task artifacts.

Dependencies: none. This should land before the rewritten GenSim semantic task-planning stack (#533-#538).

Refs #538

Type of change

• Bug fix (non-breaking change which fixes an issue)
• Enhancement (non-breaking change which improves an existing functionality)
• New feature (non-breaking change which adds functionality)
• Breaking change (existing functionality will not work without user modification)
• Documentation update

Validation

• black ./
• python docs/scripts/check_api_docs.py
• pytest tests/docs/test_check_api_docs.py -q --confcutdir=tests/docs
• pytest tests/sim/atomic_actions/test_affordance.py tests/sim/atomic_actions/test_primitives_helpers.py tests/sim/atomic_actions/test_actions.py tests/toolkits/test_grasp_pose_generator.py -q --disable-warnings --maxfail=1 (120 passed, 1 deselected)

Checklist

• I have run Black.
• Public API documentation coverage is aligned.
• I have added focused regression tests.
• No task-engine, controller, simulator-step, or duplicate runtime code is included.

[greptile-apps]

Too many files changed for review (164 files, 100 file limit).

Bypass the limit by tagging @greptile-apps to review.

[Copilot]

Pull request overview

Improves upright grasp selection with object-aware ranking, clamp-pose filtering, and yaw-equivalent downstream checks.

Changes:

• Adds grasp-cost callbacks before GraspKit top-k selection.
• Adds upright compatibility filtering and yaw sampling.
• Adds focused helper and affordance tests.

Reviewed changes

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

Show a summary per file

File	Review summary
tests/sim/atomic_actions/test_primitives_helpers.py	Tests yaw variants and option validation.
tests/sim/atomic_actions/test_affordance.py	Tests grasp-cost callback forwarding.
embodichain/toolkits/graspkit/pg_grasp/antipodal_generator.py	Adds custom cost adjustment before top-k selection. Open nit: add regression coverage proving callback-adjusted costs determine the selected subset (2 votes).
embodichain/lab/sim/atomic_actions/primitives/pick_up.py	Implements upright filtering, ranking, and yaw checks. Open critical issue regarding subclass callback compatibility (1 vote), moderate issue regarding mesh vertex lookup (3 votes), and nit regarding upright-compatibility regression coverage (2 votes).
embodichain/lab/sim/atomic_actions/affordance.py	Exposes grasp-cost callback forwarding.

Suppressed comments (9)

embodichain/lab/sim/atomic_actions/primitives/pick_up.py:625

• pickup_success now includes every downstream target check, so this warning can be emitted when the vertical pickup path is feasible but no downstream target is reachable. The message is therefore misleading in the new failure mode; describe the combined pickup/downstream feasibility failure instead.

 device=self.device, dtype=torch.float32

embodichain/lab/sim/atomic_actions/primitives/pick_up.py:618

• The yaw-equivalent downstream feasibility loop and its per-variant qpos selection are not exercised by the added tests. Existing PickUp tests either provide an explicit grasp or mock _resolve_grasp_pose, so they bypass _select_feasible_grasp_variants; a regression that ignores nonzero yaw targets or carries the wrong seed would therefore pass. Add a focused test with a mocked batch-IK response that succeeds only for a non-identity yaw variant and checks the resulting plan mask/seed behavior.

 grasp_success, grasp_qpos = self._compute_batch_candidate_ik(
grasp_variants, pre_grasp_qpos, manipulator
)
lift_success, lift_qpos = self._compute_batch_candidate_ik(
lift_variants, grasp_qpos, manipulator
)
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 = (
upright_compatible
& alignment_success
& pre_grasp_success
& grasp_success
& lift_success
)
downstream_success_counts: list[list[int]] = []
object_to_eef_variants = torch.matmul(
pose_inv(object_poses)[:, None, None], grasp_variants

embodichain/lab/sim/atomic_actions/primitives/pick_up.py:572

• There is no regression test for the new upright-compatibility invariant: the existing pickup tests bypass this sampled-grasp path, so a closing axis aligned with the object upright could be accepted without being detected. Add a focused test that supplies side and top/bottom clamp candidates and asserts only the side candidate remains feasible when rotate_upright is enabled.

This issue also appears on line 562 of the same file.

 self,
grasp_xpos: torch.Tensor,
start_qpos: torch.Tensor,
object_poses: torch.Tensor,
manipulator: JointPositionTarget,
options: PickUpOptions,
approach_direction: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Choose a TCP-roll variant with a feasible pickup and transport path."""
num_envs, n_pose = grasp_xpos.shape[:2]
mirrored_grasp_xpos = grasp_xpos.clone()

embodichain/lab/sim/atomic_actions/primitives/pick_up.py:172

• upright_yaw_samples is passed to range in _upright_yaw_pose_variants, but this validation only checks its numeric value. For example, PickUpOptions(upright_yaw_samples=1.5) is accepted here and then fails with a TypeError during planning instead of rejecting the invalid option at construction. Validate that the value is an integer as well as positive.

 raise ValueError("approach_direction must have shape (3,).")
if not torch.isfinite(self.approach_direction).all():

embodichain/lab/sim/atomic_actions/primitives/pick_up.py:566

• The new upright path has no action-level regression coverage: no test drives _upright_grasp_compatibility_mask to reject a top/bottom clamp or exercises the downstream yaw loop with one yaw failing and an equivalent yaw succeeding. The helper test only checks shape/translation, so these selection regressions could pass; add focused PickUp tests with mocked IK (and the object-height ranking case).

 self,
grasp_xpos: torch.Tensor,
start_qpos: torch.Tensor,
object_poses: torch.Tensor,
manipulator: JointPositionTarget,

embodichain/lab/sim/atomic_actions/primitives/pick_up.py:618

• newly_solved retains only the first successful yaw branch for each grasp candidate, and this break stops evaluating the remaining branches once one candidate per environment succeeds. With multiple downstream_object_target_poses, downstream_seed then carries that first branch's joint solution into the next target, so a different yaw-equivalent solution that is needed for the later target is discarded and a valid grasp sequence can be rejected. Preserve all reachable yaw branches (or evaluate the target sequence per branch) instead of greedily pruning here.

 object_to_eef_variants = torch.matmul(
pose_inv(object_poses)[:, None, None], grasp_variants

embodichain/lab/sim/atomic_actions/primitives/pick_up.py:478

• This new keyword is passed whenever upright sampling is enabled, but the repository already has a concrete AntipodalAffordance override in scripts/benchmark/atomic_action/common.py:575-583 whose signature does not accept grasp_cost_fn. Using PickUpOptions(rotate_upright=...) with that existing affordance therefore raises TypeError before planning. Update the override and any other subclasses to accept and forward the optional callback, or use a compatible extension point.

 }
return self.build_plan(
request,
context,
success=success_mask,

embodichain/lab/sim/atomic_actions/primitives/pick_up.py:725

• This mask is only called from _select_feasible_grasp_variants, but that selector is bypassed when GraspGoal.grasp_xpos is explicitly provided. Therefore an explicit grasp whose closing axis is aligned with the object's upright direction still passes when rotate_upright is enabled, despite the new upright-transport rule claiming to reject top/bottom clamps. Apply the same compatibility check to the explicit-pose path or centralize the screening before trajectory generation.

 ik_success = pickup_success[env_idx, pose_idx, best_variant_idx]
return selected_grasp_xpos, ik_success
def _approach_alignment_mask(
self,
grasp_poses: torch.Tensor,
options: PickUpOptions,
approach_direction: torch.Tensor,
) -> torch.Tensor:
"""Return candidates whose final TCP z-axis follows the approach direction."""
max_angle = options.approach_alignment_max_angle
if options.rotate_upright is not None or max_angle is None:
return torch.ones(
grasp_poses.shape[:3], dtype=torch.bool, device=grasp_poses.device
)
grasp_z = torch.nn.functional.normalize(grasp_poses[..., :3, 2], dim=-1)
alignment = torch.sum(grasp_z * approach_direction, dim=-1)
return alignment >= math.cos(float(max_angle))
def _compute_batch_candidate_ik(
self,
poses: torch.Tensor,
joint_seed: torch.Tensor,
manipulator: JointPositionTarget,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Solve candidate IK poses while preserving the candidate dimensions."""
num_envs, n_pose, n_variant = poses.shape[:3]

embodichain/lab/sim/atomic_actions/primitives/pick_up.py:154

• The new public upright_yaw_samples option is missing from the canonical PickUpOptions table in docs/source/overview/sim/atomic_actions/builtin_actions.md (lines 392-399), which currently documents the other selection and downstream-reachability fields. Without that entry, users cannot discover how to enable the new yaw-equivalent reachability behavior.

 obj_upright_direction: torch.Tensor | None = None
"""Optional object local direction used to choose the upright grasp rotation."""

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[skywhite1024]

Superseded by #542 after moving the head branch from the temporary codex/ namespace to ljd/gen-sim-atomic-prerequisite. The commit and code diff are unchanged. The GenSim stack has been recreated as Stack #543 with #542 as its bottom layer.