remove lower xpos in axis align - #557
Merged
Merged
Conversation
Greptile SummaryThe PR changes AxisAlign to finish at the aligned pose while retaining its grasp, removing the lowering and release stages.
Confidence Score: 4/5The PR does not appear safe to merge until AxisAlign records the held-object relation required by subsequent held-object actions. AxisAlign now ends with the gripper closed around the object but still emits an empty state delta, leaving direct follow-up MoveHeldObject or Pour operations unable to satisfy their symbolic held-object precondition. Files Needing Attention: embodichain/lab/sim/atomic_actions/primitives/axis_align.py
|
| Filename | Overview |
|---|---|
| embodichain/lab/sim/atomic_actions/primitives/axis_align.py | Removes lowering and release so AxisAlign ends grasping the object, but the previously reported symbolic held-state omission remains. |
| tests/sim/atomic_actions/test_actions.py | Updates trajectory assertions for the shorter retained-grasp AxisAlign plan. |
| docs/source/overview/sim/atomic_actions/builtin_actions.md | Documents that AxisAlign now retains the grasp after alignment. |
Reviews (3): Last reviewed commit: "update" | Re-trigger Greptile
Contributor
There was a problem hiding this comment.
Pull request overview
This PR updates the axis_align atomic action to remove the post-alignment “lowering” waypoint and adjusts the associated tutorial, docs, and tests to match the new motion profile.
Changes:
- Removed
lower_distancesupport and the lowering keyframe fromAxisAlignplanning. - Updated the AxisAlign tutorial and documentation to reflect the revised behavior.
- Updated unit tests to match the new segment list / waypoint expectations.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
embodichain/lab/sim/atomic_actions/primitives/axis_align.py |
Removes lowering/release-related logic and updates segment planning and waypoint budgeting. |
tests/sim/atomic_actions/test_actions.py |
Updates AxisAlign tests for the new segment count and final pose expectations. |
scripts/tutorials/atomic_action/axis_align.py |
Removes the lower_distance option and updates tutorial description. |
docs/source/overview/sim/atomic_actions/builtin_actions.md |
Updates AxisAlign documentation and capability matrix entry to match revised behavior. |
Suppressed comments (1)
embodichain/lab/sim/atomic_actions/primitives/axis_align.py:387
- The PR description says “Remove lower xpos in axis align”, but the implementation also removes the final
open/release segment and leaves the gripper closed at the end of the trajectory (hand joints stay atGRASP_COMMAND). This is a larger behavioral change than “remove lowering” and also means the action now likely ends with an object still physically grasped whileexpected_effectsremains empty (no held-object task-state update). Please clarify intended semantics and either (a) restore theopensegment (removing only the lowering), or (b) keep the grasp but update the symbolic task-state effects/docs/tests to reflect the new postcondition.
# Once the gripper is closed, lifting and alignment form one continuous
# held-object phase. Passing only those two semantic
# endpoints retains the required ordering without expanding the rotation
# into many CuRobo plan_pose calls. Together with the open-gripper phase,
# the action uses two MotionGenerator.generate calls and four backend
# target plans instead of n_align + 3 backend target plans.
post_close_xpos = torch.cat([lift_xpos[:, None], align_xpos], dim=1)
post_close_success, post_close_arm = self._plan_pose_phase(
post_close_xpos,
pre_close_arm[:, -1],
manipulator,
request,
n_lift + n_align,
interpolation_dt,
)
success = grasp_success & normalize_success_mask(
pre_close_success & post_close_success,
num_envs=self.num_envs,
device=self.device,
name="Axis-align trajectory success",
)
hand_close = interpolate_hand_qpos(
hand_open_qpos,
hand_grasp_qpos,
n_waypoints=options.hand_interp_steps,
)
segment_lengths = {
"approach": pre_close_arm.shape[1],
"close": hand_close.shape[1],
"manipulate": post_close_arm.shape[1],
}
full = torch.empty(
(self.num_envs, sum(segment_lengths.values()), self.robot_dof),
dtype=context.robot.qpos.dtype,
device=self.device,
)
full[:] = context.last_qpos.unsqueeze(1)
offset = pre_close_arm.shape[1]
full[:, :offset, arm_joint_ids] = pre_close_arm
full[:, :offset, hand_joint_ids] = hand_open_qpos.unsqueeze(1)
stop = offset + hand_close.shape[1]
full[:, offset:stop, arm_joint_ids] = pre_close_arm[:, -1].unsqueeze(1)
full[:, offset:stop, hand_joint_ids] = hand_close
offset = stop
stop = offset + post_close_arm.shape[1]
full[:, offset:stop, arm_joint_ids] = post_close_arm
full[:, offset:stop, hand_joint_ids] = hand_grasp_qpos.unsqueeze(1)
return self.build_plan(
request,
context,
success=success,
trajectory=TimedTrajectory.from_uniform_step(
full,
env_ids=context.env_ids,
step_dt=interpolation_dt,
),
expected_effects=StateDelta(),
segment_lengths=segment_lengths,
)
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
91
to
108
| class AxisAlignOptions(PickUpOptions): | ||
| """Per-invocation grasp-and-axis-alignment behavior.""" | ||
|
|
||
| target_axis: torch.Tensor = torch.tensor([0.0, 0.0, 1.0]) | ||
| """Desired world-frame axis, shape ``(3,)`` or ``(B, 3)``.""" | ||
|
|
||
| lower_distance: float = 0.03 | ||
| """World-Z distance (m) to lower the aligned object before release.""" | ||
|
|
||
| def __post_init__(self) -> None: | ||
| PickUpOptions.__post_init__(self) | ||
| if ( | ||
| not isinstance(self.target_axis, torch.Tensor) | ||
| or self.target_axis.dim() not in (1, 2) | ||
| or self.target_axis.shape[-1] != 3 | ||
| or not torch.isfinite(self.target_axis).all() | ||
| ): | ||
| raise ValueError("target_axis must be a finite (3,) or (B, 3) tensor.") | ||
| if torch.any(torch.linalg.vector_norm(self.target_axis, dim=-1) <= 1.0e-6): | ||
| raise ValueError("target_axis must be non-zero.") | ||
| if not math.isfinite(self.lower_distance): | ||
| raise ValueError("lower_distance must be finite.") | ||
| if self.lower_distance < 0.0: | ||
| raise ValueError("lower_distance must be non-negative.") | ||
| object.__setattr__(self, "target_axis", self.target_axis.clone()) |
| | `move_joints` | `JointPositionGoal` | `primary.motion` | named target only: command matching `target` on `primary.motion` | none | none | | ||
| | `pick_up` | `GraspGoal` | `primary.motion`, `primary.grasp` | `primary.grasp`: `open`, `grasp` | semantic object/entity | attach object to the `primary.motion` target | | ||
| | `axis_align` | `AxisAlignGoal` | `primary.motion`, `primary.grasp` | `primary.grasp`: `open`, `grasp` | unheld object with `AxisAlignAffordance` | open-loop pick, align, lower, and release | | ||
| | `axis_align` | `AxisAlignGoal` | `primary.motion`, `primary.grasp` | `primary.grasp`: `open`, `grasp` | unheld object with `AxisAlignAffordance` | open-loop pick and align while retaining the grasp | |
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.
Suppressed comments (1)
tests/gym/envs/test_official_task_layout.py:57
- This PR is described as removing the "lower xpos" step from AxisAlign, but it also removes the
rearrangementtask fromCONFIG_DEFINED_EXPERT_TASKS/TABLEWARE_CONFIG_TASKS(and related tests elsewhere). If dropping the rearrangement tableware task is part of the intended change, please update the PR title/description accordingly or split the task-removal into a separate PR so the scope is clear.
REMOVED_AGENT_ENV_IDS = {"PourWaterAgent-v3", "RearrangementAgent-v3"}
CONFIG_DEFINED_EXPERT_TASKS = {"pour_water"}
RL_SIMULATOR_ENV_IDS = {"CartPoleRL", "PushCubeRL"}
TABLEWARE_CONFIG_TASKS = {
"blocks_ranking_rgb",
"blocks_ranking_size",
"match_object_container",
"place_object_drawer",
"pour_water",
"scoop_ice",
"stack_blocks_two",
"stack_cups",
}
Comment on lines
354
to
358
| segment_lengths = { | ||
| "approach": pre_close_arm.shape[1], | ||
| "close": hand_close.shape[1], | ||
| "manipulate": post_close_arm.shape[1], | ||
| "open": hand_open.shape[1], | ||
| } |
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.
Suppressed comments (1)
embodichain/lab/sim/atomic_actions/primitives/axis_align.py:109
AxisAlignOptions.lower_distancewas removed from the dataclass, which is a breaking public-API change for any caller still passing that keyword argument. Since this PR is labeled as a non-breaking enhancement, consider retaining the field as a deprecated/ignored option (or alternatively update the PR/type/versioning to reflect a breaking change).
@dataclass(frozen=True, slots=True, eq=False)
class AxisAlignOptions(PickUpOptions):
"""Per-invocation grasp-and-axis-alignment behavior."""
target_axis: torch.Tensor = torch.tensor([0.0, 0.0, 1.0])
"""Desired world-frame axis, shape ``(3,)`` or ``(B, 3)``."""
def __post_init__(self) -> None:
PickUpOptions.__post_init__(self)
if (
not isinstance(self.target_axis, torch.Tensor)
or self.target_axis.dim() not in (1, 2)
or self.target_axis.shape[-1] != 3
or not torch.isfinite(self.target_axis).all()
):
raise ValueError("target_axis must be a finite (3,) or (B, 3) tensor.")
if torch.any(torch.linalg.vector_norm(self.target_axis, dim=-1) <= 1.0e-6):
raise ValueError("target_axis must be non-zero.")
object.__setattr__(self, "target_axis", self.target_axis.clone())
Comment on lines
68
to
71
| assert errors == [] | ||
|
|
||
|
|
||
| def test_registered_context_and_source_paths_exist() -> None: | ||
| context_map = _load_context_map() | ||
| missing_paths: list[str] = [] | ||
|
|
||
| for context_path in context_map["defaults"]["contexts"]: | ||
| resolved = _AGENT_CONTEXT_ROOT / context_path | ||
| if not resolved.exists(): | ||
| missing_paths.append(str(resolved.relative_to(_REPOSITORY_ROOT))) | ||
|
|
||
| for topic in context_map["topics"]: | ||
| for context_path in topic["paths"]: | ||
| resolved = _AGENT_CONTEXT_ROOT / context_path | ||
| if not resolved.exists(): | ||
| missing_paths.append(str(resolved.relative_to(_REPOSITORY_ROOT))) | ||
| for source_path in topic["source_of_truth"]: | ||
| resolved = _REPOSITORY_ROOT / source_path | ||
| if not resolved.exists(): | ||
| missing_paths.append(str(resolved.relative_to(_REPOSITORY_ROOT))) | ||
|
|
||
| assert missing_paths == [] | ||
|
|
||
|
|
||
| def test_related_topics_reference_registered_ids() -> None: |
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
Type of change
Checklist
black .command to format the code base.python docs/scripts/check_api_docs.py), if applicable