Refactor atomic_actions to typed targets + WorldState + composition by yuecideng · Pull Request #319 · DexForce/EmbodiChain

yuecideng · 2026-06-22T02:09:08Z

Description

Redesign the embodichain/lab/sim/atomic_actions/ abstraction layer from a Union+kwargs+inheritance design to a typed-target / WorldState / composition design. The full design rationale is in docs/superpowers/specs/2026-06-21-atomic-actions-redesign-design.md.

This PR carries the atomic-actions motion-primitive module (move, pick_up, move_object, place) in its redesigned form, plus the migrated tutorials, docs, and skill scaffold.

What changed in the abstraction

Typed, disjoint targets — PoseTarget / GraspTarget / HeldObjectTarget replace the five-way Union + dict/string _resolve_target DSL. Each action declares TargetType: ClassVar[type].
Single WorldState channel — replaces the four-channel held-state passing (action_context kwarg, held_object_state kwarg, self._held_object_state, get_held_object_state() + updates_held_object_state ClassVar). execute(target, state) -> ActionResult threads state explicitly; the engine is a thin sequencer.
Composition over inheritance — MoveAction, PickUpAction, MoveObjectAction, PlaceAction all inherit AtomicAction directly and compose a stateless TrajectoryBuilder. _HandCloseAction is gone; MoveAction is no longer a parent.
Name-keyed engine — AtomicActionEngine(motion_generator) + register(action) + run(steps=[(name, target), ...], state=None) -> (success, traj, final_state). Sequence order lives at the call site, not in registration order. execute_static / actions_cfg_list / SemanticAnalyzer dropped.
Affordance geometry aliasing fixed — ObjectSemantics.__post_init__ no longer mutates affordance.geometry; AntipodalAffordance takes mesh_vertices / mesh_triangles / gripper_collision_cfg / generator_cfg as direct fields.
Full-DoF trajectory return — every action returns ActionResult with a (n_envs, n_waypoints, robot.dof) trajectory; the per-call joint_ids re-indexing is gone.
validate dropped from the ABC (every prior implementation was return True # TODO).

Module layout

embodichain/lab/sim/atomic_actions/
├── affordance.py   # Affordance, AntipodalAffordance, InteractionPoints
├── core.py         # typed targets, WorldState, ActionResult, AtomicAction ABC
├── trajectory.py   # TrajectoryBuilder (composition helper)
├── actions.py      # MoveAction, PickUpAction, MoveObjectAction, PlaceAction
├── engine.py       # AtomicActionEngine + global registry
└── __init__.py

Verification

pytest tests/sim/atomic_actions/ — 60 tests pass (rewritten test suite: core, affordance, trajectory, actions, engine).
tests/sim/ collection of the 172 non-atomic tests is clean (no cross-package breakage — only the 3 tutorials + 5 test files import the package).
black --check clean across the touched surface.
Sphinx docs build succeeds.
No stale references to removed symbols (execute_static, SemanticAnalyzer, _resolve_target, MoveObjectTarget, updates_held_object_state, _HandCloseAction, custom_config on affordances) anywhere outside the historical design spec/plan under docs/superpowers/.

Notes for reviewers

End-to-end tutorial run not verified here — the 3 tutorials compile and import cleanly, but executing them requires a GPU + the full sim stack. A headless GPU smoke run (python scripts/tutorials/sim/atomic_actions.py --num_envs 1 --headless) should be done before merge.
This is a breaking change to the atomic_actions API (the only consumers are the 3 tutorial scripts and the test suite, all migrated in this PR — no task envs, RL, or agent code imports the package, so there is no deprecation window needed).
The design spec forecast larger line-count reductions than materialized; the package total dropped from ~1,872 to ~1,593 lines. The structural goals (composition, single state channel, typed targets, no dict-DSL) are all achieved — the line-count forecast was optimistic.

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

Checklist

I have run the black . command to format the code base.
I have made corresponding changes to the documentation
I have added tests that prove my fix is effective or that my feature works
Dependencies have been updated, if applicable.

🤖 Generated with Claude Code

Captures the review of embodichain/lab/sim/atomic_actions and the proposed typed-targets / WorldState / TrajectoryBuilder redesign. Co-Authored-By: Claude <noreply@anthropic.com>

13-task TDD plan implementing the design spec at docs/superpowers/specs/2026-06-21-atomic-actions-redesign-design.md. Co-Authored-By: Claude <noreply@anthropic.com>

Move Affordance, AntipodalAffordance, and InteractionPoints out of core.py into a dedicated affordance.py module. AntipodalAffordance now stores mesh_vertices, mesh_triangles, generator_cfg, gripper_collision_cfg, force_reannotate, and is_draw_grasp_xpos as typed dataclass fields rather than smuggling them through a shared geometry / custom_config dict. Remove the geometry-aliasing footgun from ObjectSemantics.__post_init__: it no longer assigns self.affordance.geometry = self.geometry, so an Affordance never silently inherits a mutable dict from its semantics parent. __post_init__ now only binds the affordance's object_label. Update actions.py, engine.py, and __init__.py imports to source the affordance types from the new module. Update test_core.py, test_engine.py, and test_actions.py to import Affordance from affordance, drop the duplicate Affordance/InteractionPoints tests now covered by test_affordance.py, and stop asserting the removed geometry alias. Co-Authored-By: Claude <noreply@anthropic.com>

…ngs) Co-Authored-By: Claude <noreply@anthropic.com>

… slim ABC Co-Authored-By: Claude <noreply@anthropic.com>

…en tests) Co-Authored-By: Claude <noreply@anthropic.com>

Co-Authored-By: Claude <noreply@anthropic.com>

- MoveAction, PickUpAction, MoveObjectAction, PlaceAction all inherit AtomicAction directly. _HandCloseAction is removed. - Each holds a TrajectoryBuilder for shared helpers. - execute() takes a typed target + WorldState and returns ActionResult with a full-DoF trajectory. - Cfg classes are flat (no inheritance among Grasp/HandClose variants). Co-Authored-By: Claude <noreply@anthropic.com>

…rror types) Co-Authored-By: Claude <noreply@anthropic.com>

Drops SemanticAnalyzer, _resolve_target, _action_context, execute_static. Keeps the global register_action / unregister_action / get_registered_actions. Co-Authored-By: Claude <noreply@anthropic.com>

Co-Authored-By: Claude <noreply@anthropic.com>

Refresh the module docstring to reflect the four current primitives (move/pick_up/move_object/place) and the typed-target / WorldState design. Export list was already correct from the incremental task updates. Co-Authored-By: Claude <noreply@anthropic.com>

Co-Authored-By: Claude <noreply@anthropic.com>

Rewrite the skill to match the new API: sibling actions inheriting AtomicAction directly, TrajectoryBuilder composition, typed targets, WorldState/ActionResult contract, engine.register + run. Drops the old tuple-return / validate / _builtin_action_map / GraspActionCfg-parent guidance. Co-Authored-By: Claude <noreply@anthropic.com>

Copilot

Pull request overview

This PR refactors embodichain.lab.sim.atomic_actions into a typed-target + WorldState + composition-based design, updating the engine API, the four built-in actions, and all related tutorials/docs/tests to match the new contract.

Changes:

Replaces the old Union/dict-DSL target handling + multi-channel held-state plumbing with typed targets (PoseTarget, GraspTarget, HeldObjectTarget) and explicit WorldState/ActionResult.
Introduces TrajectoryBuilder as a shared composition helper and rewrites Move/PickUp/MoveObject/Place as sibling AtomicAction implementations.
Migrates tutorials, docs, and a rewritten pytest suite to the new AtomicActionEngine.run([(name, target), ...]) API.

Reviewed changes

Copilot reviewed 20 out of 21 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
`embodichain/lab/sim/atomic_actions/__init__.py`	Updates public exports for the redesigned atomic-actions API (targets/state/result/builder).
`embodichain/lab/sim/atomic_actions/core.py`	Defines typed targets, `WorldState`, `ActionResult`, and the slimmed `AtomicAction` ABC.
`embodichain/lab/sim/atomic_actions/affordance.py`	Moves affordance types into a dedicated module and reshapes `AntipodalAffordance` fields.
`embodichain/lab/sim/atomic_actions/trajectory.py`	Adds `TrajectoryBuilder` utilities for IK/FK, waypoint splitting, interpolation, etc.
`embodichain/lab/sim/atomic_actions/actions.py`	Reimplements `MoveAction`, `PickUpAction`, `MoveObjectAction`, `PlaceAction` using composition + `WorldState`.
`embodichain/lab/sim/atomic_actions/engine.py`	Replaces `execute_static` with `run(steps, state)` and removes the old target resolver/analyzer path.
`tests/sim/atomic_actions/test_affordance.py`	Adds tests for the new affordance module and anti-aliasing behavior.
`tests/sim/atomic_actions/test_trajectory.py`	Adds tests for `TrajectoryBuilder` helpers and shape/broadcasting contracts.
`tests/sim/atomic_actions/test_core.py`	Updates tests for typed targets, `WorldState`, `ActionResult`, and semantics behavior.
`tests/sim/atomic_actions/test_actions.py`	Rewrites action tests to assert full-DoF trajectories and threaded `WorldState` semantics.
`tests/sim/atomic_actions/test_engine.py`	Rewrites engine tests around name-keyed registration and `run()` sequencing semantics.
`scripts/tutorials/sim/atomic_actions.py`	Migrates tutorial to explicit action registration + `engine.run()` with typed targets.
`scripts/tutorials/atomic_action/pickup_atomic_actions.py`	New/updated tutorial showcasing `PickUpAction` using the new API.
`scripts/tutorials/atomic_action/move_object_atomic_actions.py`	New/updated tutorial showcasing `MoveObjectAction` with held-object state threading.
`skills/add-atomic-action/SKILL.md`	Updates scaffolding guidance to the new typed-target + `WorldState` + `ActionResult` contract.
`docs/source/overview/sim/atomic_actions.md`	Updates conceptual overview, supported actions table, and extension guide for the redesign.
`docs/source/tutorial/atomic_actions.rst`	Updates tutorial narrative and code snippets to the new engine/action interfaces.
`docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.rst`	Expands API reference to cover new targets/state/result/builder (but needs a small completeness tweak).
`docs/superpowers/specs/2026-06-21-atomic-actions-redesign-design.md`	Adds design spec describing the motivation and target architecture for the refactor.
`embodichain/lab/sim/solvers/pytorch_solver.py`	Adjusts default IK sampling count (tangential to atomic-actions refactor).

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

+        if qpos_seed is None:
+            qpos_seed = self.robot.get_qpos()
+        success, qpos = self.robot.compute_ik(
+            pose=target_pose.unsqueeze(0),
+            qpos_seed=qpos_seed.unsqueeze(0),
+            name=control_part,
+        )


+    def apply_local_offset(
+        self, pose: torch.Tensor, offset: torch.Tensor
+    ) -> torch.Tensor:
+        """Apply a translational offset to a batched pose."""


      Affordance
      InteractionPoints
      ObjectSemantics


    :members:
    :show-inheritance:



…ns-redesign

Copilot

Pull request overview

Copilot reviewed 18 out of 18 changed files in this pull request and generated 4 comments.

skywhite1024 · 2026-06-23T05:14:29Z

+    def resolve_pose_target(self, target: torch.Tensor, *, n_envs: int) -> torch.Tensor:
+        """Broadcast a (4, 4) pose to (n_envs, 4, 4) or validate batched shape."""
+        if not isinstance(target, torch.Tensor):
+            logger.log_error(
+                f"target must be torch.Tensor of shape (4, 4) or ({n_envs}, 4, 4)",
+                TypeError,
+            )
+        if target.shape == (4, 4):
+            target = target.unsqueeze(0).repeat(n_envs, 1, 1)
+        if target.shape != (n_envs, 4, 4):
+            logger.log_error(
+                f"target tensor must have shape (4, 4) or ({n_envs}, 4, 4), "
+                f"but got {target.shape}",
+                ValueError,
+            )
+        return target


Fixed in 5424e85. resolve_pose_target() now normalizes the target pose to the builder/robot device with float32 before broadcasting and shape validation, matching the downstream planning buffers.

skywhite1024 · 2026-06-23T05:14:39Z

+        if not (pose.dim() == 3 and pose.shape[1:] == (4, 4)):
+            logger.log_error("pose must have shape [N, 4, 4]", ValueError)
+        if offset.dim() == 1:
+            offset = offset.unsqueeze(0)
+        if not (offset.dim() == 2 and offset.shape[1] == 3):
+            logger.log_error("offset must have shape [N, 3] or [3]", ValueError)
+        result = pose.clone()
+        result[:, :3, 3] += offset
+        return result


Fixed in 5424e85. apply_local_offset() now casts offset to the pose device/dtype and explicitly validates that the offset batch size is either 1 or matches the pose batch size, so invalid shapes fail with a clear ValueError.

skywhite1024 · 2026-06-23T05:15:23Z

+        if self._generator is None:
+            self._init_generator()
+        results = []
+        for i, obj_pose in enumerate(obj_poses):
+            is_success, grasp_poses, _, costs = self._generator.get_valid_grasp_poses(
+                obj_pose, approach_direction
+            )


Fixed in 5424e85. get_valid_grasp_poses() now normalizes approach_direction to the grasp generator device with float32 before calling GraspGenerator, avoiding CPU/CUDA mismatches.

skywhite1024 · 2026-06-23T05:15:35Z

+        if self._generator is None:
+            self._init_generator()
+        grasp_xpos_list: list[torch.Tensor] = []
+        is_success_list: list[bool] = []
+        open_length_list: list[float] = []
+        for i, obj_pose in enumerate(obj_poses):
+            is_success, grasp_xpos, open_length = self._generator.get_grasp_poses(
+                obj_pose, approach_direction
+            )


Fixed in 5424e85. get_best_grasp_poses() uses the same approach_direction normalization path before calling into the generator, so the default CPU tensor is moved to the generator device.

…/DexForce/EmbodiChain into refactor/atomic-actions-redesign

Copilot

Pull request overview

Copilot reviewed 22 out of 22 changed files in this pull request and generated 4 comments.

+        if start_qpos is None:
+            start_qpos = self.robot.get_qpos(name=control_part)
+        if start_qpos.shape == (arm_dof,):
+            start_qpos = start_qpos.unsqueeze(0).repeat(n_envs, 1)
+        if start_qpos.shape != (n_envs, arm_dof):


+        if state is None:
+            state = WorldState(last_qpos=self.robot.get_qpos().clone())



+        object_to_eef = state.held_object.object_to_eef.to(
+            device=self.device, dtype=torch.float32
+        )
+        if object_to_eef.shape == (4, 4):
+            object_to_eef = object_to_eef.unsqueeze(0).repeat(self.n_envs, 1, 1)
+        move_eef_xpos = torch.bmm(object_target_pose, object_to_eef)


+See `Academic
+Publications <docs/source/resources/publications/README.md>`__ for a
+complete list of academic papers related to EmbodiChain.


Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

skywhite1024 · 2026-06-24T03:26:44Z

This update addresses four small cleanup items for the atomic action tutorials:

Simplified tutorial script names
- Removed the _atomic_actions suffix from the focused demo scripts:
  - move_end_effector.py
  - move_joints.py
  - pickup.py
  - move_held_object.py
  - place.py
Removed the old aggregate tutorial
- Removed scripts/tutorials/atomic_action/atomic_actions.py
- Its coverage is now provided by the focused primitive demos
- Updated the tutorial and overview docs to remove stale script references
Reused duplicated recording/view logic
- Added scripts/tutorials/atomic_action/tutorial_utils.py
- Centralized recording parameters, viewer size, default camera look-at, and axis-marker drawing
- This keeps the individual tutorial scripts shorter and easier to scan
Made axis visualization semantically meaningful
- move_end_effector: shows the target EEF pose
- move_joints: shows the start EEF axis
- pickup: shows the target object pose
- place: shows the place target pose
- move_held_object: shows the final object target pose
- Also adjusted the default recording view to be closer to the robot and workspace

skywhite1024 and others added 30 commits June 16, 2026 14:05

Add upright atomic action

aad9cc3

Address upright action review feedback

d258064

Address upright action follow-up fixes

5552597

wip

0561aae

wip

74fbe43

wip

9da06f9

wip

2aaac19

Add move object atomic action

ddd1236

Keep move_object demo hold-only

dfc0800

add pickup_atomic_actions.py

bb41004

Address atomic action review feedback

9eb8fb4

Address move object atomic action review feedback

7f808b6

Minimize pickup changes for move object action

6cc298a

Fix batched pickup success handling

ef98a81

Add atomic actions redesign design spec

07dcdd5

Captures the review of embodichain/lab/sim/atomic_actions and the proposed typed-targets / WorldState / TrajectoryBuilder redesign. Co-Authored-By: Claude <noreply@anthropic.com>

Add atomic actions redesign implementation plan

fc91edc

13-task TDD plan implementing the design spec at docs/superpowers/specs/2026-06-21-atomic-actions-redesign-design.md. Co-Authored-By: Claude <noreply@anthropic.com>

Address Task 1 review feedback (drop _draw_grasp_xpos, polish docstri…

8c72bfe

…ngs) Co-Authored-By: Claude <noreply@anthropic.com>

Rewrite atomic_actions.core: typed targets, WorldState, ActionResult,…

2ed3473

… slim ABC Co-Authored-By: Claude <noreply@anthropic.com>

Address Task 2 review feedback (drop quoted annotations, tighten froz…

25119e4

…en tests) Co-Authored-By: Claude <noreply@anthropic.com>

Add TrajectoryBuilder composition helper

5d5a71f

Co-Authored-By: Claude <noreply@anthropic.com>

Address Task 3 review feedback (convention conformance + rename)

c59c557

Co-Authored-By: Claude <noreply@anthropic.com>

Address Task 4 review feedback (robust allocation, drop dead state, e…

373b5f0

…rror types) Co-Authored-By: Claude <noreply@anthropic.com>

Rewrite AtomicActionEngine with name-keyed run(steps, state)

cb8a795

Drops SemanticAnalyzer, _resolve_target, _action_context, execute_static. Keeps the global register_action / unregister_action / get_registered_actions. Co-Authored-By: Claude <noreply@anthropic.com>

Address Task 5 review feedback (header, type hints, typing style)

80eec60

Co-Authored-By: Claude <noreply@anthropic.com>

Migrate atomic_actions tutorial to typed-target engine API

432f293

Co-Authored-By: Claude <noreply@anthropic.com>

Migrate move_object tutorial to new atomic_actions API

03f715e

Co-Authored-By: Claude <noreply@anthropic.com>

Migrate pickup tutorial to new atomic_actions API

a96c05d

Co-Authored-By: Claude <noreply@anthropic.com>

claude added 2 commits June 22, 2026 01:45

Update atomic_actions docs for redesign

b895342

Co-Authored-By: Claude <noreply@anthropic.com>

Copilot AI review requested due to automatic review settings June 22, 2026 02:09

yuecideng added breaking refactor motion gen Things related to motion generation for robot labels Jun 22, 2026

Copilot started reviewing on behalf of yuecideng June 22, 2026 02:09 View session

Copilot AI reviewed Jun 22, 2026

View reviewed changes

yuecideng added 2 commits June 22, 2026 23:23

Merge remote-tracking branch 'origin/main' into refactor/atomic-actio…

9e80546

…ns-redesign

wip

bd16287

Copilot AI review requested due to automatic review settings June 22, 2026 16:21

Copilot started reviewing on behalf of yuecideng June 22, 2026 16:21 View session

Copilot AI reviewed Jun 22, 2026

View reviewed changes

skywhite1024 and others added 4 commits June 23, 2026 12:20

Address atomic action device review feedback

5424e85

wip

65fca7b

Merge branch 'refactor/atomic-actions-redesign' of https://github.com…

c0c8222

…/DexForce/EmbodiChain into refactor/atomic-actions-redesign

wip

5b9008c

Copilot AI review requested due to automatic review settings June 23, 2026 07:41

Copilot started reviewing on behalf of yuecideng June 23, 2026 07:41 View session

Copilot AI reviewed Jun 23, 2026

View reviewed changes

skywhite1024 added 2 commits June 23, 2026 17:31

Add focused atomic action demos

d12b1db

Remove debug state from atomic action demos

a187be8

Copilot AI review requested due to automatic review settings June 23, 2026 10:03

Copilot AI reviewed Jun 23, 2026

skywhite1024 added 2 commits June 23, 2026 20:20

Use sugar box in atomic action demos

e7b25f9

Clean up atomic action tutorials

3cb3689

Copilot AI review requested due to automatic review settings June 24, 2026 03:26

Copilot AI reviewed Jun 24, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Refactor atomic_actions to typed targets + WorldState + composition#319

Refactor atomic_actions to typed targets + WorldState + composition#319
yuecideng wants to merge 42 commits into
mainfrom
refactor/atomic-actions-redesign

yuecideng commented Jun 22, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

skywhite1024 Jun 23, 2026

Uh oh!

skywhite1024 Jun 23, 2026

Uh oh!

skywhite1024 Jun 23, 2026

Uh oh!

skywhite1024 Jun 23, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

skywhite1024 commented Jun 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		if state is None:
		state = WorldState(last_qpos=self.robot.get_qpos().clone())

Conversation

yuecideng commented Jun 22, 2026

Description

What changed in the abstraction

Module layout

Verification

Notes for reviewers

Type of change

Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

skywhite1024 Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

skywhite1024 Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

skywhite1024 Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

skywhite1024 Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

skywhite1024 commented Jun 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants