feat: Add atomic action abstraction layer for embodied AI motion generation

[yuecideng]

Description

This PR introduces an atomic action abstraction layer for embodied AI motion generation. The implementation provides a unified interface for atomic actions like reach, grasp, move, etc., with support for semantic object understanding and extensible custom action registration.

Key Components

1. Core Classes (core.py):

   • Affordance - Base class for affordance data (GraspPose, InteractionPoints)
   • ObjectSemantics - Semantic information about interaction targets
   • ActionCfg - Configuration class for atomic actions
   • AtomicAction - Abstract base class for all atomic actions

2. Action Implementations (actions.py):

   • ReachAction - Reach to target pose or object
   • GraspAction - Execute grasp motion
   • ReleaseAction - Release grasp
   • MoveAction - Move to target pose

3. Action Engine (engine.py):

   • AtomicActionEngine - Execution engine for atomic actions
   • Registry system for custom action registration
   • Support for action composition and chaining

Design Principles

• Leverage existing infrastructure: Uses existing MotionGenerator, PlanResult, and IK/FK solvers
• Object semantics: Actions consider semantic labels, affordance data, and geometry
• Consistent interface: Unified API across all motion primitives with predictable return types
• Extensibility: Registry-based motion primitive registration with plugin architecture

Design Document

See docs/superpowers/specs/2026-04-17-atomic-action-abstraction-design.md for detailed design specification.

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

Testing

Unit tests provided in tests/sim/atomic_actions/test_core.py covering:

• Affordance base class and subclasses (GraspPose, InteractionPoints)
• ObjectSemantics dataclass
• ActionCfg configuration
• Action registry functions

All 12 tests passing.

Checklist

• I have run the black . command to format the code base.
• I have made corresponding changes to the documentation (design spec included)
• 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

[Copilot]

Pull request overview

This PR introduces a new embodichain.lab.sim.atomic_actions module intended to provide a unified “atomic action” abstraction (reach/grasp/move/release) on top of the existing motion-planning stack, plus a small registry/engine and accompanying design spec + unit tests for core data types.

Changes:

• Added core atomic-action data models (Affordance, ObjectSemantics, ActionCfg) and an AtomicAction base class.
• Added default atomic actions (ReachAction, GraspAction, MoveAction, ReleaseAction) and an AtomicActionEngine with a global registry.
• Added a design document and unit tests for core affordance/registry helpers.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 14 comments.

Show a summary per file

File	Description
embodichain/lab/sim/atomic_actions/core.py	Defines affordance/semantics/config + base action class utilities.
embodichain/lab/sim/atomic_actions/actions.py	Implements reach/grasp/move/release actions using MotionGenerator/TOPPRA.
embodichain/lab/sim/atomic_actions/engine.py	Adds engine orchestration + global registry + placeholder semantic analyzer.
embodichain/lab/sim/atomic_actions/__init__.py	Exposes public API for atomic_actions package.
tests/sim/atomic_actions/test_core.py	Adds unit tests for affordances, ObjectSemantics, ActionCfg, registry helpers.
tests/sim/atomic_actions/__init__.py	Marks test package for atomic actions.
docs/superpowers/specs/2026-04-17-atomic-action-abstraction-design.md	Design specification for the new abstraction layer.

---

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

[Copilot]

AtomicAction.__init__ only accepts motion_generator, but other code in this PR (e.g., action implementations) calls super().__init__(motion_generator, robot, control_part, device) and later relies on self.control_part. Either update AtomicAction.__init__ to accept/store robot, control_part, and device (and keep it consistent with the concrete actions/engine), or update all actions/engine to match the existing constructor and set control_part elsewhere.

Suggested change

	):
	self.motion_generator = motion_generator
	self.robot = motion_generator.robot
	self.device = self.robot.device
	robot: Optional[Robot] = None,
	control_part: Optional[str] = None,
	device: Optional[torch.device] = None,
	):
	self.motion_generator = motion_generator
	self.robot = robot if robot is not None else motion_generator.robot
	self.control_part = (
	control_part if control_part is not None else ActionCfg.control_part
	)
	self.device = device if device is not None else self.robot.device

[Copilot]

ReachAction.__init__ calls super().__init__(motion_generator, robot, control_part, device), but AtomicAction.__init__ currently only takes motion_generator. This mismatch will prevent instantiation of default actions in AtomicActionEngine.

Suggested change

	super().__init__(motion_generator, robot, control_part, device)
	super().__init__(motion_generator)
	self.robot = robot
	self.control_part = control_part
	self.device = device

[Copilot]

ReachAction.execute builds target_states mixing MoveType.JOINT_MOVE and MoveType.EEF_MOVE while also setting MotionGenOptions(is_interpolate=True). MotionGenerator.generate only supports pre-interpolation when all states share the same move_type, so this will error (or produce invalid stacks). Use options.start_qpos for the start state and make all target_states EEF_MOVE, or disable pre-interpolation / split planning into separate stages.

[Copilot]

ReachAction.validate passes qpos_seed= into Robot.compute_ik, but the API uses joint_seed=. As written, this will raise a TypeError when validation is called.

Suggested change

	qpos_seed=qpos_seed.unsqueeze(0),
	joint_seed=qpos_seed.unsqueeze(0),

[Copilot]

ToppraPlanOptions does not support velocity_limit / acceleration_limit keyword args (it uses a constraints dict instead). Passing these will raise a TypeError when constructing ToppraPlanOptions; map the limits into constraints={"velocity": ..., "acceleration": ...} or extend the planner options type if needed.

Suggested change

	velocity_limit=velocity_limit,
	acceleration_limit=acceleration_limit,
	constraints={
	"velocity": velocity_limit,
	"acceleration": acceleration_limit,
	},

[Copilot]

GraspAction.validate passes qpos_seed= into Robot.compute_ik, but the API uses joint_seed=. This will raise a TypeError when validation is called.

Suggested change

	qpos_seed=qpos_seed.unsqueeze(0),
	joint_seed=qpos_seed.unsqueeze(0),

[Copilot]

MoveAction.execute references kwargs.get(...) when building ToppraPlanOptions, but execute() does not accept **kwargs. This will raise NameError: name 'kwargs' is not defined. Either add **kwargs to the signature or remove the kwargs usage.

Suggested change

	sample_interval=kwargs.get("sample_interval", 30),
	sample_interval=30,

[Copilot]

SemanticAnalyzer.analyze creates tensors (torch.eye, torch.zeros) on the default device (CPU). If the rest of the action stack operates on GPU, downstream ops like object_pose @ grasp_pose can fail with device mismatch. Consider constructing these tensors on the engine/analyzer device (or explicitly moving affordance tensors to self.device).

[Copilot]

PlanResult is used in these test-only AtomicAction implementations, but it is never imported in this test module. This will raise NameError when the test class methods are executed; import PlanResult from embodichain.lab.sim.planners or reference it via the module where it is defined.

[Copilot]

ReachAction.execute calls self._get_current_qpos(), but AtomicAction in core.py does not define this helper (it exists in the design doc but not in the implementation). This will raise AttributeError at runtime unless all concrete actions re-implement it; consider adding _get_current_qpos() to AtomicAction (likely using robot.get_qpos(name=self.control_part)[0]).

[Copilot]

Pull request overview

Copilot reviewed 19 out of 19 changed files in this pull request and generated 11 comments.

---

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

[Copilot]

When the target is ObjectSemantics, _resolve_grasp_pose() returns is_success as a boolean tensor (one per env). This if not is_success: check will raise due to ambiguous tensor truthiness. Convert to a scalar (e.g., is_success.all().item()) or branch per-environment and propagate failures in a structured way.

[Copilot]

The execute() return type annotation says list[float] for joint_ids, but these IDs are used as tensor indices throughout the engine/actions and should be list[int] (or Sequence[int]). The current annotation is misleading and can hide real typing issues.

[Copilot]

_apply_offset says it applies an offset "in local frame", but the implementation adds the offset directly to the translation component (world-frame offset). Either update the implementation to rotate the offset by the pose rotation (local-frame semantics) or update the docstring/parameter naming to reflect world-frame behavior.

Suggested change

	result = pose.clone()
	result[:, :3, 3] += offset
	if offset.shape[0] == 1 and pose.shape[0] != 1:
	offset = offset.expand(pose.shape[0], -1)
	elif offset.shape[0] != pose.shape[0]:
	logger.log_error("offset batch dimension must be 1 or match pose batch dimension")
	result = pose.clone()
	world_offset = torch.bmm(pose[:, :3, :3], offset.unsqueeze(-1)).squeeze(-1)
	result[:, :3, 3] += world_offset

[Copilot]

create_mug() sets uid="table" even though it loads a cup mesh and the rest of the test treats it as a mug. This is confusing and can cause downstream code (logging, lookups by uid/label) to be inconsistent. Use a uid consistent with the object (e.g., "mug").

Suggested change

	uid="table",
	uid="mug",

[Copilot]

AtomicAction.init uses cfg: ActionCfg = ActionCfg() as a default argument. This creates a single shared config instance at import time, which can lead to cross-instance state leakage if the config is mutated. Prefer cfg: ActionCfg | None = None and instantiate ActionCfg() inside the constructor when None.

[Copilot]

__all__ includes "GraspPose", but this module does not define or import a GraspPose symbol. This will make from embodichain.lab.sim.atomic_actions import GraspPose fail and can also confuse docs/autosummary generation. Either implement/export GraspPose or remove it from __all__. Also consider adding AntipodalAffordance to __all__ since it is part of the public API used by the tutorial.

Suggested change

	"GraspPose",
	"AntipodalAffordance",

[Copilot]

This test hard-requires CUDA (sim_device="cuda") and calls sim.init_gpu_physics(). In CI or dev environments without a GPU this will fail. Consider switching the test to CPU execution (as many other sim tests do) or marking it with pytest.mark.skipif(not torch.cuda.is_available(), ...) / a heavier integration-test marker so it doesn't gate unit test runs.

[Copilot]

PickUpAction.init overwrites the config set by the base class with self.cfg = cfg. When cfg is None (the common/default case), this sets self.cfg to None and the next line (self.cfg.approach_direction) will raise an AttributeError. Keep the config created in the parent (super().__init__(...)) or set self.cfg to the resolved default config object instead of the raw argument.

Suggested change

self.cfg = cfg

[Copilot]

PlaceAction.init has the same issue as PickUpAction: self.cfg = cfg can set self.cfg to None, but later code assumes it is a PlaceActionCfg instance (e.g., accessing hand_open_qpos/hand_close_qpos). Avoid overwriting self.cfg with the raw constructor argument; keep the resolved config from the base constructor or explicitly default it.

[Copilot]

robot.compute_ik() returns a boolean tensor (shape (n_envs,)), but this code uses if not is_success: which is invalid for multi-element tensors (and will raise "Boolean value of Tensor is ambiguous" even for shape (1,)). Use if not is_success.all().item(): (or handle per-env failures explicitly) before proceeding.

[Copilot]

Pull request overview

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

---

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

[Copilot]

The AtomicAction class docstring says actions “use PlanResult … as the return type for execute()”, but the abstract execute() signature actually returns a 3-tuple (bool, Tensor, joint_ids). This mismatch will confuse implementers and users of the abstraction layer. Update either the documentation or the signature so they agree (and also consider typing joint_ids as list[int] rather than list[float]).

[Copilot]

_resolve_target mutates and extends the returned ObjectSemantics instance (semantics.properties.update(...), semantics.uid = ...). If the semantic analyzer returns a cached object (the default path), this will leak per-call state into the cache and affect later resolutions. Also, ObjectSemantics doesn’t define a uid field, so this becomes an ad-hoc attribute. Consider copying the semantics before applying overrides and/or adding an explicit uid: str | None field to ObjectSemantics.

[Copilot]

AtomicAction._ik_solve calls Robot.compute_ik with a non-existent keyword argument qpos_seed and also assumes qpos_seed is 1D (it defaults to robot.get_qpos(), which is typically batched). This will raise at runtime and/or produce an incorrectly-shaped seed. Use the Robot API’s joint_seed parameter and normalize qpos_seed to the expected shape before calling compute_ik (e.g., accept (dof,) or (n_envs,dof) and pass through without adding extra batch dims).

Suggested change

	target_pose: Target pose [4, 4]
	qpos_seed: Seed configuration [DOF]
	Returns:
	Joint configuration [DOF]
	Raises:
	RuntimeError: If IK fails to find a solution
	"""
	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=self.control_part,
	)
	if not success.all():
	raise RuntimeError(f"IK failed for target pose: {target_pose}")
	return qpos.squeeze(0)
	target_pose: Target pose [4, 4] or [B, 4, 4]
	qpos_seed: Seed configuration [DOF] or [B, DOF]
	Returns:
	Joint configuration [DOF] or [B, DOF]
	Raises:
	RuntimeError: If IK fails to find a solution
	"""
	single_pose = target_pose.dim() == 2
	if single_pose:
	pose = target_pose.unsqueeze(0)
	elif target_pose.dim() == 3 and target_pose.shape[1:] == (4, 4):
	pose = target_pose
	else:
	raise RuntimeError(
	f"target_pose must have shape [4, 4] or [B, 4, 4], got {tuple(target_pose.shape)}"
	)
	if qpos_seed is None:
	qpos_seed = self.robot.get_qpos()
	if qpos_seed.dim() == 1:
	joint_seed = qpos_seed.unsqueeze(0)
	elif qpos_seed.dim() == 2:
	joint_seed = qpos_seed
	else:
	raise RuntimeError(
	f"qpos_seed must have shape [DOF] or [B, DOF], got {tuple(qpos_seed.shape)}"
	)
	if joint_seed.shape[0] != pose.shape[0]:
	if joint_seed.shape[0] == 1:
	joint_seed = joint_seed.expand(pose.shape[0], -1)
	else:
	raise RuntimeError(
	"Batch size mismatch between target_pose and qpos_seed: "
	f"{pose.shape[0]} != {joint_seed.shape[0]}"
	)
	success, qpos = self.robot.compute_ik(
	pose=pose,
	joint_seed=joint_seed,
	name=self.control_part,
	)
	if not success.all():
	raise RuntimeError(f"IK failed for target pose: {target_pose}")
	return qpos.squeeze(0) if single_pose else qpos

[Copilot]

AntipodalAffordance._draw_grasp_xpos references SimulationManager and MarkerCfg but neither is imported in this module, so enabling is_draw_grasp_xpos will raise a NameError. Import these (preferably inside the method to avoid circular deps, consistent with other sim modules) before calling SimulationManager.get_instance() / MarkerCfg(...).

[Copilot]

Pull request overview

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

---

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

[Copilot]

robot.compute_ik returns a boolean tensor (shape (n_envs,)). The check if not is_success: will raise "Boolean value of Tensor is ambiguous" at runtime. Use an aggregate check like if not is_success.all(): (and decide whether partial-env failure should be allowed) and handle the per-env failures accordingly.

Suggested change

	if not is_success:
	logger.log_warning(
	f"Failed to compute IK for target state {j} in some environments. "
	"The resulting trajectory may be invalid."
	if not bool(is_success.all().item()):
	failed_env_ids = (~is_success).nonzero(as_tuple=False).flatten().tolist()
	logger.log_warning(
	f"Failed to compute IK for target state {j} in environments "
	f"{failed_env_ids}. The resulting trajectory may be invalid."

[Copilot]

get_best_grasp_poses has a default approach_direction=torch.tensor(...) created on CPU at import time. If the grasp generator runs on CUDA, passing this default will cause a device mismatch inside GraspGenerator computations. Consider making the default None and creating/moving the tensor to obj_poses.device (or self.generator.device) inside the method.

Suggested change

	approach_direction: torch.Tensor = torch.tensor(
	[0, 0, -1], dtype=torch.float32
	),
	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
	if self.generator is None:
	self._init_generator()
	approach_direction: Optional[torch.Tensor] = None,
	) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
	if self.generator is None:
	self._init_generator()
	if approach_direction is None:
	approach_direction = torch.tensor(
	[0, 0, -1], dtype=torch.float32, device=obj_poses.device
	)
	else:
	approach_direction = approach_direction.to(obj_poses.device)

[Copilot]

__all__ includes "GraspPose", but there is no GraspPose symbol defined or imported in this package. This will break from embodichain.lab.sim.atomic_actions import * and can confuse doc generation. Remove it or export the correct symbol (e.g., AntipodalAffordance if that’s what was intended).

[Copilot]

The AtomicAction class docstring states that actions use PlanResult as the return type for execute(), but the abstract method is actually defined to return a raw tuple. Please update the docstring (or the method signature) so the documented return type matches the real API.

Suggested change

	All atomic actions use PlanResult from embodichain.lab.sim.planners
	as the return type for execute() method, ensuring consistency with
	the existing motion planning infrastructure.
	Atomic actions expose an ``execute()`` method that returns a tuple of
	``(is_success, trajectory, joint_ids)``, where ``is_success`` indicates
	whether planning succeeded, ``trajectory`` is the planned joint-space
	trajectory tensor, and ``joint_ids`` identifies the joints corresponding
	to that trajectory.

[Copilot]

execute_static logs an error when len(target_list) doesn’t match the number of configured actions, but it still proceeds. Because zip(...) truncates, this can silently skip actions or ignore extra targets while still returning success. This should return early (False / raise) after logging so callers don’t get a partially planned trajectory.

Suggested change

	if len(target_list) != len(action_names):
	logger.log_error(
	f"Length of target_list ({len(target_list)}) must match number of actions ({len(action_names)})."
	)
	start_qpos = self.motion_generator.robot.get_qpos()
	n_envs = start_qpos.shape[0]
	all_dof = self.motion_generator.robot.dof
	start_qpos = self.motion_generator.robot.get_qpos()
	n_envs = start_qpos.shape[0]
	all_dof = self.motion_generator.robot.dof
	if len(target_list) != len(action_names):
	logger.log_error(
	f"Length of target_list ({len(target_list)}) must match number of actions ({len(action_names)})."
	)
	return (
	False,
	torch.empty(
	size=(n_envs, 0, all_dof),
	dtype=torch.float32,
	device=self.device,
	),
	)

[Copilot]

_init_actions stores actions in a dict keyed by cfg.name. This drops duplicates and can reorder/overwrite actions if actions_cfg_list contains repeated names (e.g., two move actions), which conflicts with the documented positional sequencing behavior. Consider storing actions as an ordered list (or list of (name, instance) pairs) so duplicates and order are preserved.

[Copilot]

In _resolve_target, when the input is a dict you mutate the returned ObjectSemantics (semantics.properties.update(...) and semantics.uid = uid). If the analyzer returned a cached semantics instance, these mutations will persist across future calls for the same label. Also, ObjectSemantics doesn’t define a uid field, so this adds an ad-hoc attribute. Prefer returning a new semantics object (or a copy) whenever overrides like properties/uid are provided, and add an explicit uid field to ObjectSemantics if it’s part of the API.

[Copilot]

create_robot uses a mutable default argument (position=[0.0, 0.0, 0.0]). If the list is ever mutated, the default will be shared across calls. Use position: Sequence[float] | None = None and assign a new list inside the function.

[Copilot]

execute() is annotated to return tuple[bool, torch.Tensor, list[float]], but the third item is used as joint_ids (indices) throughout the engine/actions. This should be a list of ints (or a sequence of ints) to match usage and avoid type confusion for custom action implementers.