Update atom action - #432
Open
matafela wants to merge 18 commits into
Open
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR updates the coordinated pickment workflow to use affordance-driven (auto-generated) antipodal grasp poses, and refactors atomic-action plumbing to accept a simpler grasp target.
Changes:
- Switched the coordinated pickment tutorial to use
create_antipodal_semantics()and run the action viaGraspTarget. - Added dual-arm grasp pose generation plumbing (
Affordance.get_dual_arm_valid_grasp_posesandAntipodalGenerator.get_dual_arm_valid_grasp_poses) and refactored the generator’s filtering path. - Refactored
PickUptrajectory assembly and rewiredCoordinatedPickmentto plan grasp/lift trajectories from affordance-sampled grasps.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 9 comments.
Show a summary per file
| File | Description |
|---|---|
| scripts/tutorials/atomic_action/coordinated_pickment.py | Tutorial now uses create_antipodal_semantics() + GraspTarget for coordinated pickment planning. |
| embodichain/toolkits/graspkit/pg_grasp/antipodal_generator.py | Adds dual-arm antipodal grasp selection and refactors valid-grasp filtering. |
| embodichain/lab/sim/sim_manager.py | Minor whitespace cleanup during sim manager initialization. |
| embodichain/lab/sim/atomic_actions/primitives/pick_up.py | Refactors trajectory construction into a helper; keeps behavior but reorganizes execution. |
| embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py | Changes TargetType and replumbs coordinated pickment to use affordance-driven dual-arm grasps and new planning helpers. |
| embodichain/lab/sim/atomic_actions/affordance.py | Adds get_dual_arm_valid_grasp_poses() wrapper around the antipodal generator. |
| embodichain/lab/gym/utils/gym_utils.py | Changes --num_envs default handling in CLI argument setup. |
Comments suppressed due to low confidence (2)
embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py:544
- The padding loop also assumes
grasp_poses_result[i]is a dict and does not filter out per-side failures (is_success=False). This can lead to shape errors when failed results carry placeholder poses/costs, or to selecting invalid grasps.
for i in range(n_envs):
left_result = grasp_poses_result[i]["left"]
right_result = grasp_poses_result[i]["right"]
if left_result is not None:
embodichain/toolkits/graspkit/pg_grasp/antipodal_generator.py:801
- The collision-check failure path (immediately after this
_collision_checker.query(...)call) returns placeholder values with inconsistent types/shapes ((4,4)pose, float open-length, cost=0). Downstream code assumes batched tensors and may treat the failure as a valid best-grasp due to the zero cost.
is_colliding, max_penetration = self._collision_checker.query(
object_pose,
valid_grasp_poses,
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
788
to
792
| "--num_envs", | ||
| help="The number of environments to run in parallel. " | ||
| "If not given, falls back to the gym config's `num_envs` (default 1).", | ||
| default=None, | ||
| default=1, | ||
| type=int, |
Comment on lines
+152
to
+159
| for i, obj_pose in enumerate(obj_poses): | ||
| result = self._generator.get_dual_arm_valid_grasp_poses( | ||
| object_pose=obj_pose, | ||
| approach_direction=approach_direction, | ||
| left_to_right_arm_direction=left_to_right_arm_direction, | ||
| middle_empty_ratio=middle_empty_ratio, | ||
| ) | ||
| results.append(result) |
Comment on lines
752
to
756
| return ( | ||
| False, | ||
| torch.eye(4, device=self.device), | ||
| 0.0, | ||
| torch.zeros(1, device=self.device), |
Comment on lines
390
to
394
| class CoordinatedPickment(AtomicAction): | ||
| """Pick and move a single object pinched by two hands.""" | ||
|
|
||
| TargetType: ClassVar[type] = CoordinatedPickmentTarget | ||
| TargetType: ClassVar[type] = GraspTarget | ||
|
|
Comment on lines
+668
to
+672
| def execute(self, target: GraspTarget, state: WorldState) -> ActionResult: | ||
| left_start_qpos, right_start_qpos = self._resolve_dual_arm_start(state) | ||
| left_grasp_xpos, right_grasp_xpos = self._resolve_grasp_pose( | ||
| target.semantics, left_start_qpos, right_start_qpos | ||
| ) |
Comment on lines
+504
to
+511
| for i in range(n_envs): | ||
| left_result = grasp_poses_result[i]["left"] | ||
| right_result = grasp_poses_result[i]["right"] | ||
| if left_result is None or right_result is None: | ||
| logger.log_warning( | ||
| f"Failed to find valid dual-arm grasp poses for {i}-th enviroment." | ||
| ) | ||
| continue |
Comment on lines
+704
to
+709
| full_trajectory[:, :, :] = last_qpos.unsqueeze(1) | ||
|
|
||
| hold_trajectory = torch.empty( | ||
| (self.n_envs, 0, self.robot_dof), dtype=torch.float32, device=self.device | ||
| # pading trajectory end to match the max length | ||
| full_trajectory[:, :n_left_waypoints, self.left_arm_joint_ids] = left_arm_traj | ||
| full_trajectory[:, :n_left_waypoints, self.left_hand_joint_ids] = left_hand_traj | ||
| full_trajectory[:, :n_right_waypoints, self.right_arm_joint_ids] = ( |
Comment on lines
715
to
722
| return ActionResult( | ||
| success=success_mask, | ||
| trajectory=full, | ||
| success=is_success, | ||
| trajectory=full_trajectory, | ||
| next_state=WorldState( | ||
| last_qpos=full[:, -1, :].clone(), | ||
| last_qpos=full_trajectory[:, -1, :].clone(), | ||
| held_object=None, | ||
| coordinated_held_object=coordinated_held_object, | ||
| coordinated_held_object=state.coordinated_held_object, | ||
| ), |
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 23 out of 23 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (5)
embodichain/lab/sim/atomic_actions/primitives/coordinated_pickment.py:475
- _resolve_grasp_pose() dereferences semantics.entity unconditionally, which will raise an AttributeError if the provided GraspTarget semantics has no entity. Other actions validate semantics.entity and emit a clear ValueError; CoordinatedPickment should do the same.
obj_poses = semantics.entity.get_local_pose(to_matrix=True)
embodichain/lab/sim/atomic_actions/affordance.py:162
- In get_dual_arm_valid_grasp_poses(), left_to_right_arm_direction is passed to the generator without indexing per-environment. Callers (e.g., CoordinatedPickment) provide a (n_envs, 3) tensor, but AntipodalGenerator.get_dual_arm_valid_grasp_poses expects a single (3,) direction vector; passing the batched tensor will break (e.g., .repeat(n_vert, 1) shape mismatch). Slice the direction per env (and move it to the generator device) inside the loop.
for i, obj_pose in enumerate(obj_poses):
result = self._generator.get_dual_arm_valid_grasp_poses(
object_pose=obj_pose,
approach_direction=approach_direction,
left_to_right_arm_direction=left_to_right_arm_direction,
embodichain/lab/sim/atomic_actions/primitives/hand_over.py:212
- HandOver.execute() assigns final_object_pose = self.final_object_pose and then mutates final_object_pose[:, :3, :3], which also mutates the cached tensor on the action instance. This makes subsequent executions depend on the previous object's rotation. Also, semantics.entity is dereferenced without validation.
middle_object_pose = self.middle_object_pose.clone()
final_object_pose = self.final_object_pose
# force object pose to have the same rotation as the current object pose, so that the handover is feasible.
current_object_pose = target.semantics.entity.get_local_pose(to_matrix=True)
middle_object_pose[:, :3, :3] = current_object_pose[:, :3, :3]
final_object_pose[:, :3, :3] = current_object_pose[:, :3, :3]
embodichain/lab/gym/utils/gym_utils.py:793
- --num_envs help text says that when not provided it falls back to the gym config's num_envs, but default=1 makes it always provided (so downstream can't distinguish 'unset' vs explicit 1). Either revert the default back to None or update the help + downstream fallback logic accordingly.
"--num_envs",
help="The number of environments to run in parallel. "
"If not given, falls back to the gym config's `num_envs` (default 1).",
default=1,
type=int,
)
embodichain/lab/sim/atomic_actions/primitives/pick_up.py:145
- _select_feasible_grasp_variants() computes the pre-grasp offset along each grasp pose's local Z axis (pose[:3,2]), but _get_full_pickup_trajectory() offsets along the fixed cfg.approach_direction. Since apply_local_offset adds a world-frame translation, these can diverge when the grasp pose orientation isn't aligned with approach_direction, causing execution-time IK failures for grasps that were deemed feasible during selection. Use the grasp pose's Z axis here for consistency.
pre_grasp_xpos = self.builder.apply_local_offset(
grasp_xpos, -self.approach_direction * self.cfg.pre_grasp_distance
)
Comment on lines
56
to
60
| self.d6 = 0.0819 | ||
| self.a2 = -0.24365 | ||
| self.a3 = -0.21325 | ||
| self.urdf_path = get_data_path("UniversalRobots/UR3/UR3.urdf") | ||
| elif self.ur_type == "ur3e": | ||
| self.d1 = 0.15185 |
9 tasks
Move action-specific target classes beside their primitives, add a shared object target base, and track held objects by control part. Update tutorials, benchmarks, tests, documentation, and agent context for the new contracts.
matafela
changed the base branch from
main
to
refactor/atomic-action-target-contracts
July 31, 2026 08:05
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 17 out of 17 changed files in this pull request and generated 1 comment.
Suppressed comments (2)
embodichain/lab/sim/atomic_actions/primitives/hand_over.py:212
execute()aliasesself.final_object_poseand then mutates it in-place when forcing the rotation. That permanently changes the cached pose for future calls. Also,target.semantics.entityis optional and is dereferenced without a guard, which will raise an AttributeError if the caller supplies semantics without an entity.
final_object_pose = self.final_object_pose
# force object pose to have the same rotation as the current object pose, so that the handover is feasible.
current_object_pose = target.semantics.entity.get_local_pose(to_matrix=True)
middle_object_pose[:, :3, :3] = current_object_pose[:, :3, :3]
final_object_pose[:, :3, :3] = current_object_pose[:, :3, :3]
embodichain/lab/sim/atomic_actions/affordance.py:156
left_to_right_arm_directionis forwarded to the grasp generator without being moved to the generator device/dtype. When the generator runs on CUDA this will raise a device-mismatch error insideget_dual_arm_valid_grasp_posesduring tensor math.
left_to_right_arm_direction=left_to_right_arm_direction,
Comment on lines
611
to
615
| object_pose: torch.Tensor, | ||
| approach_direction: torch.Tensor, | ||
| object_part: str = "center", | ||
| visualize_collision: bool = False, | ||
| ): |
yuecideng
approved these changes
Jul 31, 2026
added 2 commits
August 2, 2026 02:42
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 19 out of 21 changed files in this pull request and generated no new comments.
Suppressed comments (7)
embodichain/lab/sim/solvers/ur_solver.py:36
URSolverinheritsBaseSolver, which always builds apk_serial_chainfromcfg.urdf_path. WithURSolverCfg.urdf_pathno longer set, initializingURSolverwill raiseValueError("Eitherurdf_pathorchainmust be provided.")fromcreate_pk_serial_chain. Restore a default URDF path on the cfg so the base class can build FK/limits safely.
class URSolverCfg(SolverCfg):
class_type: str = "URSolver"
ur_type: str = "ur10"
end_link_name: str = "ee_link"
root_link_name: str = "base_link"
embodichain/toolkits/graspkit/pg_grasp/antipodal_generator.py:613
get_valid_grasp_posesinsertedobject_partbeforevisualize_collision, which breaks any positional 3rd-argument callers (e.g. internalget_grasp_pose(..., visualize_collision)will silently pass a bool asobject_partand ignorevisualize_collision). To preserve backward compatibility, keepvisualize_collisionas the 3rd positional argument and makeobject_partkeyword-only.
def get_valid_grasp_poses(
self,
object_pose: torch.Tensor,
approach_direction: torch.Tensor,
object_part: str = "center",
embodichain/lab/sim/solvers/ur_solver.py:58
- When
ur_typeis not the default,BaseSolverstill builds FK/joint limits fromcfg.urdf_path. After removing the per-ur_typeURDF overrides, the solver will use the UR10 model even whenur_typeisur3/ur5/..., which can lead to incorrect FK/limits. Consider restoring aur_type→URDF mapping in__post_init__sourdf_pathstays consistent with the DH parameters.
self.d6 = 0.0819
self.a2 = -0.24365
self.a3 = -0.21325
embodichain/lab/sim/atomic_actions/primitives/hand_over.py:212
execute()usestarget.semantics.entity.get_local_pose(...)without validatingsemantics.entityis set, and it mutatesself.final_object_posein-place by aliasing it asfinal_object_pose. This can produce a hard-to-diagnoseAttributeErrorwhen the entity is missing and can also leak the modified rotation into subsequent calls. Prefer validatingsemantics.entityand cloningself.final_object_posebefore modifying it.
semantics = target.semantics
transfer_object_to_eef = self._resolve_transfer_object_to_eef(state)
middle_object_pose = self.middle_object_pose.clone()
final_object_pose = self.final_object_pose
# force object pose to have the same rotation as the current object pose, so that the handover is feasible.
embodichain/lab/sim/atomic_actions/primitives/place.py:299
- The new
AssembleTarget/_resolve_assemble_place_xpos()path is not covered by unit tests. Existing tests exercisePlacewithPlaceTarget, but nothing asserts the derived posebase_pose @ assemble_to_base_pose @ object_to_eefor the error paths (missing held object / missing base entity). Adding a focused mocked unit test intests/sim/atomic_actions/test_actions.pywould help prevent regressions.
def _resolve_assemble_place_xpos(
self, target: AssembleTarget, state: WorldState
) -> torch.Tensor:
"""Derive the place EEF pose from an assembly affordance.
embodichain/lab/sim/atomic_actions/primitives/hand_over.py:117
HandOveris a new atomic action with non-trivial planning/state-transfer behavior (reading a held object from the transfer arm and writing a new one for the receiving arm), but there are no unit tests validating its success/failure paths or thatWorldState.held_objectsis updated correctly. Consider adding a mocked unit test (similar to other action tests) to cover at least: missing transfer-held-object error, successful handover updates, and trajectory shape invariants.
class HandOver(AtomicAction[GraspTarget]):
"""Hand an object from one arm to the other.
The transferring arm (already holding the object) moves it to a middle
handover pose, the receiving arm approaches and grasps a different part of
embodichain/lab/sim/atomic_actions/primitives/pick_up.py:165
_select_feasible_grasp_variantsbacks off along the grasp pose's Z axis, but_get_full_pickup_trajectorybacks off along the configuredapproach_direction. SinceTrajectoryBuilder.apply_local_offsetapplies a world-frame translation, these can differ and lead to selecting a grasp variant that is feasible for one pre-grasp definition but executing a different pre-grasp during planning. Consider using the same (grasp-Z) convention in_get_full_pickup_trajectoryfor consistency.
pre_grasp_xpos = self.builder.apply_local_offset(
grasp_xpos, -self.approach_direction * self.cfg.pre_grasp_distance
)
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 19 out of 21 changed files in this pull request and generated no new comments.
Suppressed comments (7)
embodichain/lab/sim/atomic_actions/primitives/pick_up.py:365
- _select_feasible_grasp_variants() computes the pre-grasp offset along the grasp pose's z-axis, but the actual planned trajectory in _get_full_pickup_trajectory() (and the PickUpCfg docstring) uses the world-frame approach_direction. This mismatch can cause the feasibility screening to accept a grasp variant that later fails during real trajectory planning.
Use the same approach_direction-based offset here (or update both code paths consistently).
pre_grasp_variants = grasp_variants.clone()
pre_grasp_z = pre_grasp_variants[..., :3, 2]
pre_grasp_variants[..., :3, 3] -= pre_grasp_z * self.cfg.pre_grasp_distance
embodichain/lab/sim/atomic_actions/primitives/hand_over.py:212
- final_object_pose is assigned as a reference to self.final_object_pose and then modified in-place. That mutates the cached tensor stored on the action instance, so subsequent execute() calls can observe the modified pose unexpectedly.
Also, semantics.entity is used without validation; if the target semantics has no entity this will raise an AttributeError instead of a controlled ValueError.
middle_object_pose = self.middle_object_pose.clone()
final_object_pose = self.final_object_pose
# force object pose to have the same rotation as the current object pose, so that the handover is feasible.
current_object_pose = target.semantics.entity.get_local_pose(to_matrix=True)
middle_object_pose[:, :3, :3] = current_object_pose[:, :3, :3]
embodichain/lab/sim/atomic_actions/primitives/hand_over.py:231
- HandOverCfg.pre_grasp_distance is documented as an offset along the (negative) receive_approach_direction, but execute() currently offsets along the grasp pose's local z-axis. This makes the behavior differ from the config contract and can produce unintuitive approach motions when the grasp frame z is not aligned with receive_approach_direction.
receive_grasp_z = receive_grasp_xpos[..., :3, 2]
receive_pre_grasp_eef = self.builder.apply_local_offset(
receive_grasp_xpos,
-receive_grasp_z * self.cfg.pre_grasp_distance,
)
embodichain/toolkits/graspkit/pg_grasp/antipodal_generator.py:688
- get_dual_arm_valid_grasp_poses() projects points onto left_to_right_arm_direction without normalizing it. If callers pass a non-unit vector, the min/max projections and thresholds scale with its magnitude and the masks become inconsistent across call sites.
Normalize (and validate non-zero) left_to_right_arm_direction before computing projections.
# project mesh_vert_transformed to left_to_right_arm_direction and get the min and max value
n_vert = mesh_vert_transformed.shape[0]
projected = (
mesh_vert_transformed * left_to_right_arm_direction.repeat(n_vert, 1)
).sum(dim=-1)
tests/sim/objects/test_dual_arm.py:139
- This test comment no longer matches the behavior: URSolverCfg no longer pins urdf_path in post_init, so for dual-arm robots Robot.init_solver injects the combined URDF path and link names remain prefixed.
Update the comment to reflect why prefixed link names are expected here.
# URSolverCfg pins urdf_path to the single-arm URDF in __post_init__, so
# the engine keeps link names UNPREFIXED to match that URDF (arm-local).
embodichain/lab/sim/atomic_actions/primitives/place.py:77
- Place now supports AssembleTarget, but there are no behavioral tests exercising the AssembleTarget pathway (only a TargetType membership assertion). This leaves base-pose tracking, held-object transform usage, and error cases (missing held object / missing base_object_entity) unverified.
Add unit tests that run Place.execute() with an AssembleTarget and validate the resolved EEF release pose and error handling.
@dataclass(frozen=True, slots=True, eq=False)
class AssembleTarget(ActionTarget):
"""Place a held assemble object onto a base object at a relative pose.
The base object pose is read at planning time from
embodichain/lab/sim/atomic_actions/primitives/hand_over.py:117
- HandOver is a new dual-arm primitive with non-trivial state semantics (reads held-object state from transfer arm; writes it to receive arm; plans multiple segments). There are currently no unit tests covering success paths or failure modes (missing held object, missing semantics.entity, insufficient waypoints).
Add tests to validate held-object handoff bookkeeping and that per-env failures correctly return a failure ActionResult without mutating state.
class HandOver(AtomicAction[GraspTarget]):
"""Hand an object from one arm to the other.
The transferring arm (already holding the object) moves it to a middle
handover pose, the receiving arm approaches and grasps a different part of
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
TODO:
Type of change
Checklist
black .command to format the code base.