Skip to content

Fix pose-sensitive MoveHeldObject prompt flow #388

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
skywhite1024 merged 1 commit into from
Jul 13, 2026
Merged

Fix pose-sensitive MoveHeldObject prompt flow #388

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
124 changes: 22 additions & 102 deletions embodichain/gen_sim/action_agent_pipeline/generation/prompt_builders.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1034,13 +1034,6 @@ def make_relative_task_prompt(
release_step_label = _relative_pose_step_label(spec, "release")
pose_sensitive = _is_pose_sensitive_placement(spec)
if pose_sensitive:
safe_high_spec = _format_high_staging_spec(active_arm, spec)
oriented_high_spec = _format_relative_pose_spec(
active_arm,
spec,
pose_kind="high",
sample_interval=45,
)
release_move_spec = _format_relative_pose_spec(
active_arm,
spec,
Expand All @@ -1049,36 +1042,28 @@ def make_relative_task_prompt(
)
place_spec = _format_release_only_place_spec(active_arm)
retreat_spec = _format_empty_hand_retreat_spec(active_arm)
edge_count = 7
release_instruction = f"""2. Move the held object to the high staging pose without changing orientation:
- {active_slot}: {safe_high_spec}
- {inactive_slot}: null

3. Apply the requested orientation at the high staging pose:
- {active_slot}: {oriented_high_spec}
- {inactive_slot}: null

4. Move the held object down to the {release_step_label} object pose:
edge_count = 5
release_instruction = f"""2. Move the held object directly to the {release_step_label} object pose while applying the requested orientation:
- {active_slot}: {release_move_spec}
- {inactive_slot}: null

5. Release the held object in-place without moving the object pose:
3. Release the held object in-place without moving the object pose:
- {active_slot}: {place_spec}
- {inactive_slot}: null

6. Retreat the now-empty end-effector upward:
4. Retreat the now-empty end-effector upward:
- {active_slot}: {retreat_spec}
- {inactive_slot}: null

7. Return the active arm to its initial pose:
5. Return the active arm to its initial pose:
- {active_slot}: {initial_spec}
- {inactive_slot}: null"""
high_instruction = release_instruction
release_rule = (
"For this pose-sensitive placement, lift with preserved orientation, "
"apply the requested orientation at the high staging pose, and only "
"then move down to the final release pose. Use the exact relative-zero "
"release-only `Place` spec shown below."
"For this pose-sensitive placement, use exactly one `MoveHeldObject` "
"to move directly to the final release object pose while applying the "
"requested orientation. Do not add staging or intermediate moves. Use "
"the exact relative-zero release-only `Place` spec shown below."
)
else:
place_spec = _format_direct_relative_place_spec(active_arm, spec)
Expand Down Expand Up @@ -1186,27 +1171,8 @@ def _single_relative_graph_steps(
edge_blocks.extend(
[
(
"Move the held object to the high staging pose without changing "
"orientation",
{
active_slot: _format_high_staging_spec(active_arm, spec),
inactive_slot: None,
},
),
(
"Apply the requested orientation at the high staging pose",
{
active_slot: _format_relative_pose_spec(
active_arm,
spec,
pose_kind="high",
sample_interval=45,
),
inactive_slot: None,
},
),
(
f"Move the held object down to the {release_step_label} object pose",
f"Move the held object directly to the {release_step_label} object "
"pose while applying the requested orientation",
{
active_slot: _format_relative_pose_spec(
active_arm,
Expand Down Expand Up @@ -1702,29 +1668,8 @@ def _dual_relative_release_edge_blocks(
if _is_pose_sensitive_placement(placement):
return [
(
f"Move `{placement.moved_runtime_uid}` to the high staging pose "
"without changing orientation",
{
active_slot: _format_high_staging_spec(active_arm, placement),
waiting_slot: waiting_value,
},
),
(
f"Apply `{placement.moved_runtime_uid}` orientation at the high "
"staging pose",
{
active_slot: _format_relative_pose_spec(
active_arm,
placement,
pose_kind="high",
sample_interval=45,
),
waiting_slot: waiting_value,
},
),
(
f"Move `{placement.moved_runtime_uid}` down to the final release "
"object pose",
f"Move `{placement.moved_runtime_uid}` directly to the final "
"release object pose while applying the requested orientation",
{
active_slot: _format_relative_pose_spec(
active_arm,
Expand Down Expand Up @@ -1769,14 +1714,15 @@ def _dual_relative_release_rule(spec: _RelativeSpecLike) -> str:
"For this dependent dual-object task, complete the first object's "
"pick-up, placement, release, retreat, and return "
"before picking up the second object. The inactive arm must remain "
"null throughout each object's sequence. Pose-sensitive objects must "
"be lifted before orientation is changed at the high staging pose."
"null throughout each object's sequence. For each pose-sensitive "
"object, use exactly one MoveHeldObject to move directly to the final "
"release object pose while applying the requested orientation."
)
if any(_is_pose_sensitive_placement(placement) for placement in spec.placements):
return (
"For pose-sensitive placements, lift while preserving orientation, "
"rotate at the high staging pose, and then move down to the final "
"release pose. The following `Place` must be the exact "
"For pose-sensitive placements, use exactly one `MoveHeldObject` to "
"move directly to the final release object pose while applying the "
"requested orientation. The following `Place` must be the exact "
"relative-zero release-only spec shown below, and then the empty hand "
"retreats upward. Any preserve placement in the same graph instead uses "
"object-aware Place directly, without MoveHeldObject."
Expand Down Expand Up @@ -1833,11 +1779,7 @@ def _relative_release_action_patterns(
if not _is_pose_sensitive_placement(placement):
return f"""- Direct orientation-preserving Place:
{_format_direct_relative_place_spec(robot_name, placement)}"""
return f"""- High staging without orientation change:
{_format_high_staging_spec(robot_name, placement)}
- Requested orientation at high staging:
{_format_relative_pose_spec(robot_name, placement, pose_kind="high", sample_interval=45)}
- Final release object pose with requested orientation:
return f"""- Direct final release object pose with requested orientation:
{_format_relative_pose_spec(robot_name, placement, pose_kind="release", sample_interval=45)}
- Release-only Place:
{_format_release_only_place_spec(robot_name)}
Expand Down Expand Up @@ -1993,8 +1935,8 @@ def _make_dual_relative_basic_background(
"Dependent objects are placed serially in dependency order."
if serial_sequence
else "Orientation-preserving placements use object-aware Place directly "
"after pickup, without MoveHeldObject. Pose-sensitive placements use "
"high staging before rotation and final placement."
"after pickup, without MoveHeldObject. Each pose-sensitive placement uses "
"exactly one direct final-pose MoveHeldObject, then release-only Place."
)
return f"""The scene comes from the exported {project_name} mesh environment.

Expand Down Expand Up @@ -2876,28 +2818,6 @@ def _relative_surface_support(
return placement.reference_runtime_uid


def _format_high_staging_spec(
robot_name: str,
placement: _RelativePlacementLike,
) -> str:
if _is_pose_sensitive_placement(placement):
return _format_relative_pose_spec(
robot_name,
placement,
pose_kind="high",
sample_interval=45,
orientation_goal="preserve",
orientation_axis="none",
align_to=None,
)
return _format_relative_pose_spec(
robot_name,
placement,
pose_kind="high",
sample_interval=45,
)


def _format_hover_move_spec(
robot_name: str,
placement: _RelativePlacementLike,
Expand Down
12 changes: 4 additions & 8 deletions tests/gen_sim/action_agent_pipeline/test_stacking_spec.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -228,7 +228,7 @@ def test_elongated_stacking_object_does_not_request_axis_alignment() -> None:
assert orientation == ("preserve", "none")


def test_pose_sensitive_relative_graph_rotates_at_high_staging() -> None:
def test_pose_sensitive_relative_graph_uses_one_final_move() -> None:
placement = _relative_step("bottle", "table", side="left")
placement = _RelativePlacementStepSpec(
**{
Expand All @@ -246,13 +246,9 @@ def test_pose_sensitive_relative_graph_rotates_at_high_staging() -> None:
and edge["left_arm_action"]["atomic_action_class"] == "MoveHeldObject"
]

assert [target["orientation_goal"] for target in move_targets] == [
"preserve",
"upright",
"upright",
]
assert move_targets[0]["offset"] == pytest.approx([0.0, 0.0, 0.3])
assert move_targets[-1]["offset"] == pytest.approx([0.0, 0.0, 0.2])
assert len(move_targets) == 1
assert move_targets[0]["orientation_goal"] == "upright"
assert move_targets[0]["offset"] == pytest.approx([0.0, 0.0, 0.2])


def _relative_step(
Expand Down
24 changes: 5 additions & 19 deletions tests/gen_sim/action_agent_pipeline/test_ur5_basket_config_generation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1958,27 +1958,14 @@ def fake_call_relative_task_llm(**kwargs):
release_offset_json = json.dumps(
summary["release_offset"], ensure_ascii=False, separators=(",", ":")
)
high_offset = list(summary["release_offset"])
high_offset[2] = round(float(high_offset[2]) + 0.25, 6)
high_offset_json = json.dumps(
high_offset, ensure_ascii=False, separators=(",", ":")
)
assert (
"Generate one deterministic nominal graph with exactly 7 nominal edges"
"Generate one deterministic nominal graph with exactly 5 nominal edges"
in task_prompt
)
for text in (task_prompt, atom_actions):
assert '"atomic_action_class":"MoveHeldObject"' in text
assert text.count('"atomic_action_class":"MoveHeldObject"') == 1
assert '"target_object_pose":{"reference":"object"' in text
assert '"obj_name":"colored_pad"' in text
assert (
f'"offset":{high_offset_json},"orientation_goal":"preserve",'
'"orientation_axis":"none"}' in text
)
assert (
f'"offset":{high_offset_json},"orientation_goal":"axis_align",'
'"orientation_axis":"long_axis","align_to":"colored_pad"' in text
)
assert f'"offset":{release_offset_json}' in text
assert '"orientation_goal":"axis_align"' in text
assert '"orientation_axis":"long_axis"' in text
Expand All @@ -1991,8 +1978,7 @@ def fake_call_relative_task_llm(**kwargs):
)
assert (
f'"atomic_action_class":"MoveEndEffector","robot_name":"{active_arm}",'
'"control":"arm","target_pose":{"reference":"relative",'
'"offset":[0.0,0.0,0.1],"frame":"world"}' in text
'"control":"arm","target_pose":{"reference":"relative"' in text
)

assert summary["orientation_goal"] == "axis_align"
Expand Down Expand Up @@ -2032,11 +2018,11 @@ def fake_call_relative_task_llm(**kwargs):
task_prompt = paths.task_prompt.read_text(encoding="utf-8")
atom_actions = paths.atom_actions.read_text(encoding="utf-8")
assert (
"Generate one deterministic nominal graph with exactly 7 nominal edges"
"Generate one deterministic nominal graph with exactly 5 nominal edges"
in task_prompt
)
for text in (task_prompt, atom_actions):
assert '"orientation_goal":"preserve","orientation_axis":"none"' in text
assert text.count('"atomic_action_class":"MoveHeldObject"') == 1
assert '"orientation_goal":"upright"' in text
assert '"orientation_axis":"none"' in text
assert '"align_to"' not in text
Expand Down