diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 51f32e838..ff8aa5405 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -606,6 +606,7 @@ topics: - SceneSnapshotSupplier - SceneProvider - RigidObjectSceneProvider + - create_simulation_atomic_action_engine - SceneRegistry - RegistrySceneProvider - SceneEntityRef @@ -885,7 +886,6 @@ topics: - open drawer - hand over - pour water - - rearrangement - 配置生成环境 - 运行时环境注册 paths: @@ -902,7 +902,6 @@ topics: - embodichain_tasks/configs/tasks/manipulation/open_drawer/ - embodichain_tasks/configs/tasks/manipulation/hand_over/ - embodichain_tasks/configs/tasks/manipulation/tableware/pour_water/ - - embodichain_tasks/configs/tasks/manipulation/tableware/rearrangement/ related_topics: - atomic-actions - env-framework diff --git a/agent_context/topics/atomic-actions/atomic-actions.md b/agent_context/topics/atomic-actions/atomic-actions.md index 610be703c..6d181c993 100644 --- a/agent_context/topics/atomic-actions/atomic-actions.md +++ b/agent_context/topics/atomic-actions/atomic-actions.md @@ -82,6 +82,23 @@ instance of every type in `BUILTIN_ACTION_TYPES`; use `load_builtins=False` only for isolated tests or a fully custom action set. A bound action cannot be reused by another engine. +An engine may additionally borrow a default `SceneProvider`. In that case, +`engine.initial_context()` captures a `SceneSnapshot` from the provider using +the robot observation timestamp and generated environment IDs. An explicitly +supplied `scene=` snapshot takes precedence, and engines without either source +retain the empty-scene behavior. The provider is only an initial-context +convenience for direct-core planning; execution observations and scene revision +advancement remain owned by the runtime's `ObservationProvider`. + +Direct simulation callers that only need selected rigid-object poses should use +`create_simulation_atomic_action_engine(..., scene_entities=(...))`. The factory +derives canonical direct-core IDs from the supplied objects' `uid` values and +installs the default provider; it never scans `SimulationManager`. Actions then +select the entries they consume through their goal and semantic entity IDs. +Articulation/link observations, aliases, collision roles, dynamic execution, +and external perception remain explicit `SceneProvider` or `SceneRegistry` +integration paths. + ## Engine entry points Choose the public engine entry point by lifecycle, not by skill type: diff --git a/agent_context/topics/expert-programs/expert-programs.md b/agent_context/topics/expert-programs/expert-programs.md index 24ef61191..457eedcd7 100644 --- a/agent_context/topics/expert-programs/expert-programs.md +++ b/agent_context/topics/expert-programs/expert-programs.md @@ -238,19 +238,14 @@ evidence fails closed. Resource disjointness alone is never sufficient. `simulation.pour` calls after built-in Pick. The configured lowerers preserve verified held-object state and keep task motion values in the selected robot policy preset. -- `Rearrangement-v3` uses registered `simulation.push_object` calls over - predeclared utensil-to-slot routes. Each utensil is pushed, settled, and - corrected from the latest measured pose; the action short-circuits inside its - configured completion tolerance, while segment validators remain the final - physical acceptance boundary. Do not infer physical qualification from the two trajectory-only examples or from unit/fake-port tests. Physical acceptance belongs in dedicated validation integrations with measured evidence, controlled multi-seed/randomization runs, and task validators. `EmbodiedEnv` is the common execution environment for all -five references. Their Gym configs own the complete supported composition +four references. Their Gym configs own the complete supported composition roots, including the allowlisted services needed by Open Drawer, Hand Over, -Pour Water, and Rearrangement; none has a task environment module or subclass. +and Pour Water; none has a task environment module or subclass. ## Recommended change sites @@ -265,7 +260,7 @@ Pour Water, and Rearrangement; none has a task environment module or subclass. | Lazy Gym action/segment lifecycle | `bridge.py` and `embodichain/lab/gym/envs/demo.py` | | Environment adapter binding, episode program selection, success/reset | `embodichain/lab/gym/envs/embodied_env.py` and `embodichain/lab/gym/utils/registration.py` | | Config-created simple runtime and dynamic ID binding | `embodichain/lab/gym/envs/expert_program/configured_runtime.py` and `embodichain/lab/gym/utils/gym_utils.py` | -| Reference scene/profile/runtime values | `embodichain_tasks/configs/tasks/manipulation/{repeated_pick_place,open_drawer,hand_over}/env.json` and `embodichain_tasks/configs/tasks/manipulation/tableware/{pour_water,rearrangement}/env.json` | +| Reference scene/profile/runtime values | `embodichain_tasks/configs/tasks/manipulation/{repeated_pick_place,open_drawer,hand_over}/env.json` and `embodichain_tasks/configs/tasks/manipulation/tableware/pour_water/env.json` | | Configured live-service implementations | `embodichain/lab/gym/envs/expert_program/_configured_runtime_services.py` | Prefer changing the narrow owner. Do not add task-local motion generators, diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.rst index 3f33b5e59..8a32abb9a 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.atomic_actions.rst @@ -72,6 +72,7 @@ embodichain.lab.sim.atomic_actions CommandDispatch CommandOperation ExecutionClock + create_simulation_atomic_action_engine SimulationExecutionAdapter ExecutionTick EffectVerificationRequest @@ -303,6 +304,8 @@ Engine and execution .. autoclass:: MonotonicExecutionClock :members: +.. autofunction:: create_simulation_atomic_action_engine + .. autoclass:: SimulationExecutionAdapter :members: diff --git a/docs/source/api_reference/public_api.rst b/docs/source/api_reference/public_api.rst index 9b1f941b1..949cbdf40 100644 --- a/docs/source/api_reference/public_api.rst +++ b/docs/source/api_reference/public_api.rst @@ -766,6 +766,7 @@ embodichain.lab.sim.atomic_actions.sim_adapter .. autosummary:: + create_simulation_atomic_action_engine RigidObjectSceneProvider RigidObjectSceneProviderCfg SceneSnapshotSupplier diff --git a/embodichain/lab/sim/atomic_actions/__init__.py b/embodichain/lab/sim/atomic_actions/__init__.py index eac6b110d..268e50222 100644 --- a/embodichain/lab/sim/atomic_actions/__init__.py +++ b/embodichain/lab/sim/atomic_actions/__init__.py @@ -224,6 +224,7 @@ ) from .scene import SceneProvider from .sim_adapter import ( + create_simulation_atomic_action_engine, RigidObjectSceneProvider, RigidObjectSceneProviderCfg, SceneSnapshotSupplier, @@ -270,6 +271,7 @@ "CommandSink", "ControlCommand", "ControlPartCommandProfile", + "create_simulation_atomic_action_engine", "CoordinatedHeldObjectState", "CoordinatedPickGoal", "CoordinatedPickment", diff --git a/embodichain/lab/sim/atomic_actions/engine.py b/embodichain/lab/sim/atomic_actions/engine.py index 68dc04037..f38229e95 100644 --- a/embodichain/lab/sim/atomic_actions/engine.py +++ b/embodichain/lab/sim/atomic_actions/engine.py @@ -30,8 +30,9 @@ from .plans import ActionPlan, CompiledTrajectory, TimedTrajectory from .policies import MotionPolicy, RecoveryPolicy from .runtime import ActionPlanningServices +from .scene import SceneProvider from .state import PlanningContext, RobotObservation, SceneSnapshot, TaskState -from .tracking import TrackingRuntime +from .tracking import TrackingPolicy, TrackingRuntime if TYPE_CHECKING: from embodichain.lab.sim.objects import Robot @@ -52,6 +53,7 @@ def __init__( *, load_builtins: bool = True, tracking_runtime: TrackingRuntime | None = None, + scene_provider: SceneProvider | None = None, ) -> None: """Initialize one engine and bind its built-in action implementations. @@ -66,13 +68,19 @@ def __init__( tracking_runtime: Optional exact-version feedback, projector, and metric registries. Built-in joint tracking is installed when omitted. + scene_provider: Optional default scene-observation source used by + :meth:`initial_context` when the caller does not supply an + explicit scene snapshot. The provider is borrowed by reference. """ + if scene_provider is not None and not isinstance(scene_provider, SceneProvider): + raise TypeError("scene_provider must implement SceneProvider.") self._planning_services = ActionPlanningServices( motion_generator, control_profiles=control_profiles, tracking_runtime=tracking_runtime, grasp_pose_generators=grasp_pose_generators, ) + self._scene_provider = scene_provider self._actions: dict[str, AtomicAction] = {} self._skill_catalog_revision = 0 if load_builtins: @@ -194,6 +202,7 @@ def make_invocation( *, control_parts: Mapping[str, Mapping[str, str]] | None = None, motion_policy: MotionPolicy | None = None, + tracking_policy: TrackingPolicy | None = None, recovery_policy: RecoveryPolicy | None = None, skill_options: OptionsT | None = None, control_overrides: ActionControlOverrides | None = None, @@ -211,6 +220,7 @@ def make_invocation( goal: Action-specific typed goal. control_parts: Direct ``slot -> endpoint -> control_part`` mapping. motion_policy: Optional invocation motion policy. + tracking_policy: Optional typed tracking and terminal-acceptance policy. recovery_policy: Optional invocation recovery policy. skill_options: Optional action-specific invocation options. control_overrides: Optional endpoint-scoped command overrides. @@ -235,6 +245,11 @@ def make_invocation( goal=goal, binding=binding, motion_policy=MotionPolicy() if motion_policy is None else motion_policy, + tracking_policy=( + TrackingPolicy.joint_position() + if tracking_policy is None + else tracking_policy + ), recovery_policy=( RecoveryPolicy() if recovery_policy is None else recovery_policy ), @@ -394,7 +409,9 @@ def initial_context( Args: task: Optional symbolic task state; an empty state is used otherwise. - scene: Optional scene snapshot; an empty snapshot is used otherwise. + scene: Optional explicit scene snapshot. It overrides the engine's + configured scene provider; an empty snapshot is used when both + are absent. timestamp: Timestamp assigned to the captured robot observation. control_dt: Explicit command period for action-owned interpolation. @@ -410,15 +427,24 @@ def initial_context( qvel_value = candidate.to(self.device) qvel = torch.zeros_like(qpos) if qvel_value is None else qvel_value batch_size = int(qpos.shape[0]) + env_ids = torch.arange(batch_size, dtype=torch.long, device=self.device) if task is None: task = TaskState.empty(batch_size=batch_size, device=self.device) if scene is None: - scene = SceneSnapshot.empty() + if self._scene_provider is None: + scene = SceneSnapshot.empty() + else: + scene = self._scene_provider.snapshot( + timestamp=timestamp, + env_ids=env_ids.clone(), + ) + if not isinstance(scene, SceneSnapshot): + raise TypeError("scene_provider must return a SceneSnapshot.") return PlanningContext( robot=RobotObservation(timestamp=timestamp, qpos=qpos, qvel=qvel), task=task, scene=scene, - env_ids=torch.arange(batch_size, dtype=torch.long, device=self.device), + env_ids=env_ids, control_dt=control_dt, ) diff --git a/embodichain/lab/sim/atomic_actions/sim_adapter.py b/embodichain/lab/sim/atomic_actions/sim_adapter.py index 7500ed2ef..50933d614 100644 --- a/embodichain/lab/sim/atomic_actions/sim_adapter.py +++ b/embodichain/lab/sim/atomic_actions/sim_adapter.py @@ -27,6 +27,8 @@ from embodichain.utils import configclass from .bindings import JointPositionTarget, RuntimeEndpointTarget +from .control import ControlPartCommandProfile +from .engine import AtomicActionEngine from .runner import ( CommandAcknowledgement, CommandAckStatus, @@ -40,10 +42,13 @@ SceneSnapshot, TaskState, ) +from .tracking import TrackingRuntime if TYPE_CHECKING: from embodichain.lab.sim.objects import RigidObject, Robot + from embodichain.lab.sim.planners import MotionGenerator from embodichain.lab.sim.sim_manager import SimulationManager + from embodichain.toolkits.graspkit import GraspPoseGenerator @configclass @@ -239,6 +244,63 @@ def _pose_change_mask( ) +def create_simulation_atomic_action_engine( + motion_generator: MotionGenerator, + scene_entities: Sequence[RigidObject], + control_profiles: Mapping[str, ControlPartCommandProfile] | None = None, + grasp_pose_generators: Mapping[str, GraspPoseGenerator] | None = None, + *, + load_builtins: bool = True, + tracking_runtime: TrackingRuntime | None = None, +) -> AtomicActionEngine: + """Create an engine whose initial context observes selected rigid objects. + + This is the direct-simulation convenience path for offline planning. Entity + IDs are derived from each rigid object's stable ``uid``; only explicitly + supplied objects are observed. Advanced integrations that need aliases, + articulation/link state, collision roles, or an external perception source + should construct :class:`AtomicActionEngine` with their own + :class:`SceneProvider` instead. + + Args: + motion_generator: Motion-generation backend owned by the engine. + scene_entities: Non-empty sequence of simulation rigid objects to expose + in automatically captured initial scene snapshots. + control_profiles: Semantic commands keyed by robot control-part name. + grasp_pose_generators: Grasp-pose services keyed by grasp endpoint target. + load_builtins: Whether to install all built-in atomic actions. + tracking_runtime: Optional typed tracking runtime shared by action plans. + + Returns: + Engine configured with a rigid-object scene provider. + + Raises: + TypeError: If ``scene_entities`` is not a sequence. + ValueError: If an entity lacks a stable UID or UIDs are duplicated. + """ + if isinstance(scene_entities, (str, bytes)) or not isinstance( + scene_entities, Sequence + ): + raise TypeError("scene_entities must be a sequence of rigid objects.") + entities_by_id: dict[str, RigidObject] = {} + for entity in scene_entities: + entity_id = getattr(entity, "uid", None) + if not isinstance(entity_id, str) or not entity_id.strip(): + raise ValueError("Every scene entity must have a non-empty string uid.") + if entity_id in entities_by_id: + raise ValueError(f"Duplicate scene entity uid {entity_id!r}.") + entities_by_id[entity_id] = entity + + return AtomicActionEngine( + motion_generator, + control_profiles=control_profiles, + grasp_pose_generators=grasp_pose_generators, + load_builtins=load_builtins, + tracking_runtime=tracking_runtime, + scene_provider=RigidObjectSceneProvider(entities_by_id), + ) + + SceneSnapshotSupplier = Callable[[float], SceneSnapshot] """Callback that returns the latest scene snapshot for a simulation timestamp.""" @@ -617,6 +679,7 @@ def _validate_timeout(timeout: float) -> None: __all__ = [ + "create_simulation_atomic_action_engine", "RigidObjectSceneProvider", "RigidObjectSceneProviderCfg", "SceneSnapshotSupplier", diff --git a/scripts/tutorials/atomic_action/assemble.py b/scripts/tutorials/atomic_action/assemble.py index 80cf39f56..c664c90fe 100644 --- a/scripts/tutorials/atomic_action/assemble.py +++ b/scripts/tutorials/atomic_action/assemble.py @@ -40,15 +40,13 @@ from embodichain.lab.sim.atomic_actions import ( AssembleAffordance, AssembleGoal, - AtomicActionEngine, ControlPartCommandProfile, - EntityState, + create_simulation_atomic_action_engine, GraspGoal, PickUpOptions, PlaceOptions, MotionPolicy, SceneEntityPose, - SceneSnapshot, ) from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg from embodichain.data import get_data_path @@ -295,8 +293,9 @@ def run_assemble_demo( lift_height=PLACE_LIFT_HEIGHT, hand_interp_steps=PLACE_HAND_INTERP_STEPS, ) - engine = AtomicActionEngine( + engine = create_simulation_atomic_action_engine( motion_generator=motion_gen, + scene_entities=(can, cube), control_profiles={ "left_hand": ControlPartCommandProfile.joint_positions( open=left_open, @@ -320,14 +319,6 @@ def run_assemble_demo( assemble_affordance = AssembleAffordance( assemble_to_base_pose=assemble_to_base, ) - scene = SceneSnapshot( - timestamp=0.0, - version=0, - entities={ - can.uid: EntityState(can.get_local_pose(to_matrix=True)), - cube.uid: EntityState(cube.get_local_pose(to_matrix=True)), - }, - ) endpoint_mapping = {"primary": {"motion": "left_arm", "grasp": "left_hand"}} compiled = engine.compile( ( @@ -355,10 +346,7 @@ def run_assemble_demo( skill_options=place_options, ), ), - engine.initial_context( - scene=scene, - control_dt=sim.sim_config.physics_dt, - ), + engine.initial_context(control_dt=sim.sim_config.physics_dt), ) success = compiled.plan_success traj = compiled.trajectory.positions diff --git a/scripts/tutorials/atomic_action/axis_align.py b/scripts/tutorials/atomic_action/axis_align.py index f0fb8e3be..8685f3f6c 100644 --- a/scripts/tutorials/atomic_action/axis_align.py +++ b/scripts/tutorials/atomic_action/axis_align.py @@ -29,15 +29,13 @@ import torch from typing import Sequence from embodichain.lab.sim.atomic_actions import ( - AtomicActionEngine, AxisAlignAffordance, AxisAlignGoal, AxisAlignOptions, ControlPartCommandProfile, - EntityState, + create_simulation_atomic_action_engine, MotionPolicy, ObjectSemantics, - SceneSnapshot, ) from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg from embodichain.lab.sim.objects import RigidObject @@ -154,8 +152,9 @@ def main() -> None: initialize_pre_pick_robot_pose(robot, obj, hand_open) motion_gen = create_toppra_motion_generator(robot) - engine = AtomicActionEngine( + engine = create_simulation_atomic_action_engine( motion_generator=motion_gen, + scene_entities=(obj,), control_profiles={ "hand": ControlPartCommandProfile.joint_positions( open=hand_open, @@ -215,14 +214,7 @@ def main() -> None: ), ), ), - engine.initial_context( - scene=SceneSnapshot( - timestamp=0.0, - version=0, - entities={obj.uid: EntityState(obj.get_local_pose(to_matrix=True))}, - ), - control_dt=sim.sim_config.physics_dt, - ), + engine.initial_context(control_dt=sim.sim_config.physics_dt), ) if not compiled.plan_success.all(): logger.log_warning("Failed to plan AxisAlign demo trajectory.") diff --git a/scripts/tutorials/atomic_action/coordinated_pickment.py b/scripts/tutorials/atomic_action/coordinated_pickment.py index eb358f45f..b3ea01d1f 100644 --- a/scripts/tutorials/atomic_action/coordinated_pickment.py +++ b/scripts/tutorials/atomic_action/coordinated_pickment.py @@ -41,9 +41,7 @@ ControlPartCommandProfile, CoordinatedPickGoal, CoordinatedPickmentOptions, - EntityState, MotionPolicy, - SceneSnapshot, ) from embodichain.lab.sim.cfg import ( RigidBodyAttributesCfg, @@ -468,14 +466,7 @@ def run_coordinated_pickment_demo( skill_options=pickment_options, ), ), - engine.initial_context( - scene=SceneSnapshot( - timestamp=0.0, - version=0, - entities={obj.uid: EntityState(object_pose_batch)}, - ), - control_dt=sim.sim_config.physics_dt, - ), + engine.initial_context(control_dt=sim.sim_config.physics_dt), ) success = compiled.plan_success traj = compiled.trajectory.positions diff --git a/scripts/tutorials/atomic_action/coordinated_placement.py b/scripts/tutorials/atomic_action/coordinated_placement.py index dae6f6815..e78034e5b 100644 --- a/scripts/tutorials/atomic_action/coordinated_placement.py +++ b/scripts/tutorials/atomic_action/coordinated_placement.py @@ -38,8 +38,8 @@ from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.atomic_actions import ( - AtomicActionEngine, ControlPartCommandProfile, + create_simulation_atomic_action_engine, CoordinatedPlacementOptions, CoordinatedPlacementGoal, GraspGoal, @@ -581,8 +581,9 @@ def run_coordinated_placement_demo( hold_steps=6, retreat_steps=18, ) - engine = AtomicActionEngine( + engine = create_simulation_atomic_action_engine( motion_generator=motion_gen, + scene_entities=(bread, pan), control_profiles={ "left_hand": ControlPartCommandProfile.joint_positions( open=left_open, diff --git a/scripts/tutorials/atomic_action/hand_over.py b/scripts/tutorials/atomic_action/hand_over.py index 940c5f237..86f25354e 100644 --- a/scripts/tutorials/atomic_action/hand_over.py +++ b/scripts/tutorials/atomic_action/hand_over.py @@ -30,13 +30,11 @@ from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.atomic_actions import ( - AtomicActionEngine, ControlPartCommandProfile, - EntityState, + create_simulation_atomic_action_engine, HandOverGoal, HandOverOptions, MotionPolicy, - SceneSnapshot, ) from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg from embodichain.data import get_data_path @@ -213,8 +211,9 @@ def run_handover_demo( n_sample=10_000, force_refresh=False, ) - engine = AtomicActionEngine( + engine = create_simulation_atomic_action_engine( motion_generator=motion_gen, + scene_entities=(obj,), control_profiles={ "left_hand": ControlPartCommandProfile.joint_positions( open=left_open, @@ -258,14 +257,7 @@ def run_handover_demo( skill_options=handover_options, ), ), - engine.initial_context( - scene=SceneSnapshot( - timestamp=0.0, - version=0, - entities={obj.uid: EntityState(obj.get_local_pose(to_matrix=True))}, - ), - control_dt=sim.sim_config.physics_dt, - ), + engine.initial_context(control_dt=sim.sim_config.physics_dt), ) success = compiled.plan_success traj = compiled.trajectory.positions diff --git a/scripts/tutorials/atomic_action/move_held_object.py b/scripts/tutorials/atomic_action/move_held_object.py index d439c936b..2a135ba53 100644 --- a/scripts/tutorials/atomic_action/move_held_object.py +++ b/scripts/tutorials/atomic_action/move_held_object.py @@ -30,8 +30,8 @@ from embodichain.data import get_data_path from embodichain.lab.sim.atomic_actions import ( - AtomicActionEngine, ControlPartCommandProfile, + create_simulation_atomic_action_engine, EndEffectorPoseGoal, GraspGoal, HeldObjectPoseGoal, @@ -121,8 +121,9 @@ def main() -> None: motion_gen = create_curobo_motion_generator(robot) hand_open, hand_close = get_hand_open_close_qpos(robot) - engine = AtomicActionEngine( + engine = create_simulation_atomic_action_engine( motion_generator=motion_gen, + scene_entities=(obj,), control_profiles={ "hand": ControlPartCommandProfile.joint_positions( open=hand_open, diff --git a/scripts/tutorials/atomic_action/open_door.py b/scripts/tutorials/atomic_action/open_door.py index 4677a2e51..cab6cfebb 100644 --- a/scripts/tutorials/atomic_action/open_door.py +++ b/scripts/tutorials/atomic_action/open_door.py @@ -32,14 +32,12 @@ from embodichain.lab.sim.atomic_actions import ( AtomicActionEngine, ControlPartCommandProfile, - EntityState, MotionPolicy, ObjectSemantics, ObservedArticulationJointState, OpenDoorAffordance, OpenDoorGoal, OpenDoorOptions, - SceneEntityPose, SceneSnapshot, ) from embodichain.lab.sim.cfg import ( @@ -188,7 +186,7 @@ def main() -> None: "open_door", OpenDoorGoal( semantics, - SceneEntityPose(HANDLE_SCENE_ENTITY_ID), + handle_pose, open_fraction=open_fraction, ), control_parts={"primary": {"motion": "arm", "grasp": "hand"}}, @@ -205,7 +203,6 @@ def main() -> None: scene=SceneSnapshot( timestamp=0.0, version=0, - entities={HANDLE_SCENE_ENTITY_ID: EntityState(handle_pose)}, articulation_joints={ ( MICROWAVE_SCENE_ENTITY_ID, diff --git a/scripts/tutorials/atomic_action/pickup.py b/scripts/tutorials/atomic_action/pickup.py index cab98b423..ff511f978 100644 --- a/scripts/tutorials/atomic_action/pickup.py +++ b/scripts/tutorials/atomic_action/pickup.py @@ -29,13 +29,11 @@ import torch from embodichain.lab.sim.atomic_actions import ( - AtomicActionEngine, ControlPartCommandProfile, - EntityState, + create_simulation_atomic_action_engine, GraspGoal, PickUpOptions, MotionPolicy, - SceneSnapshot, ) from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg from embodichain.lab.sim.objects import RigidObject @@ -133,8 +131,9 @@ def main() -> None: initialize_pre_pick_robot_pose(robot, obj, hand_open) motion_gen = create_curobo_motion_generator(robot) - engine = AtomicActionEngine( + engine = create_simulation_atomic_action_engine( motion_generator=motion_gen, + scene_entities=(obj,), control_profiles={ "hand": ControlPartCommandProfile.joint_positions( open=hand_open, @@ -176,14 +175,7 @@ def main() -> None: ), ), ), - engine.initial_context( - scene=SceneSnapshot( - timestamp=0.0, - version=0, - entities={obj.uid: EntityState(obj.get_local_pose(to_matrix=True))}, - ), - control_dt=sim.sim_config.physics_dt, - ), + engine.initial_context(control_dt=sim.sim_config.physics_dt), ) if not compiled.plan_success.all(): logger.log_warning("Failed to plan PickUp demo trajectory.") diff --git a/scripts/tutorials/atomic_action/place.py b/scripts/tutorials/atomic_action/place.py index a7129595e..7aeb58584 100644 --- a/scripts/tutorials/atomic_action/place.py +++ b/scripts/tutorials/atomic_action/place.py @@ -29,8 +29,8 @@ import torch from embodichain.lab.sim.atomic_actions import ( - AtomicActionEngine, ControlPartCommandProfile, + create_simulation_atomic_action_engine, GraspGoal, PickUpOptions, PlaceGoal, @@ -129,8 +129,9 @@ def main() -> None: hand_open, hand_close = get_hand_open_close_qpos(robot) initialize_pre_pick_robot_pose(robot, obj, hand_open) - engine = AtomicActionEngine( + engine = create_simulation_atomic_action_engine( motion_generator=motion_gen, + scene_entities=(obj,), control_profiles={ "hand": ControlPartCommandProfile.joint_positions( open=hand_open, diff --git a/scripts/tutorials/atomic_action/pour.py b/scripts/tutorials/atomic_action/pour.py index 8c7c8721d..a57e53160 100644 --- a/scripts/tutorials/atomic_action/pour.py +++ b/scripts/tutorials/atomic_action/pour.py @@ -30,8 +30,8 @@ import torch from embodichain.lab.sim.atomic_actions import ( - AtomicActionEngine, ControlPartCommandProfile, + create_simulation_atomic_action_engine, GraspGoal, MotionPolicy, PickUpOptions, @@ -100,8 +100,9 @@ def main() -> None: initialize_pre_pick_robot_pose(robot, obj, hand_open) motion_gen = create_toppra_motion_generator(robot) - engine = AtomicActionEngine( + engine = create_simulation_atomic_action_engine( motion_generator=motion_gen, + scene_entities=(obj,), control_profiles={ "hand": ControlPartCommandProfile.joint_positions( open=hand_open, diff --git a/scripts/tutorials/atomic_action/press.py b/scripts/tutorials/atomic_action/press.py index 384b836ad..8e0400776 100644 --- a/scripts/tutorials/atomic_action/press.py +++ b/scripts/tutorials/atomic_action/press.py @@ -32,14 +32,11 @@ from embodichain.lab.sim.atomic_actions import ( AtomicActionEngine, ControlPartCommandProfile, - EntityState, MotionPolicy, ObjectSemantics, PressAffordance, PressGoal, PressOptions, - SceneEntityPose, - SceneSnapshot, ) from embodichain.lab.sim.cfg import ( ArticulationCfg, @@ -212,7 +209,7 @@ def main() -> None: "press", PressGoal( semantics, - SceneEntityPose(BUTTON_SCENE_ENTITY_ID), + target_pose, ), control_parts={"primary": {"motion": "arm", "grasp": "hand"}}, motion_policy=MotionPolicy(sample_count=PRESS_SAMPLE_INTERVAL), @@ -228,14 +225,7 @@ def main() -> None: ), ), ), - context=engine.initial_context( - scene=SceneSnapshot( - timestamp=0.0, - version=0, - entities={BUTTON_SCENE_ENTITY_ID: EntityState(target_pose)}, - ), - control_dt=sim.sim_config.physics_dt, - ), + context=engine.initial_context(control_dt=sim.sim_config.physics_dt), ) if not compiled.plan_success.all(): logger.log_warning("Failed to plan the Press demo trajectory.") diff --git a/scripts/tutorials/atomic_action/slide.py b/scripts/tutorials/atomic_action/slide.py index 5b6de48ba..9780d77a9 100644 --- a/scripts/tutorials/atomic_action/slide.py +++ b/scripts/tutorials/atomic_action/slide.py @@ -35,14 +35,11 @@ ActionInvocation, AtomicActionEngine, ControlPartCommandProfile, - EntityState, MotionPolicy, ObjectSemantics, SlideAffordance, SlideGoal, SlideOptions, - SceneEntityPose, - SceneSnapshot, ) from embodichain.lab.sim.cfg import ( ArticulationCfg, @@ -137,6 +134,7 @@ def create_drawer_semantics(drawer: Articulation) -> ObjectSemantics: def create_invocation( engine: AtomicActionEngine, semantics: ObjectSemantics, + target_pose: torch.Tensor, *, direction: Literal["pull", "push"], approach_distance: float, @@ -147,6 +145,7 @@ def create_invocation( Args: engine: Engine used to resolve the slide control-part binding. semantics: Drawer-handle semantics shared by both operations. + target_pose: Latest observed world pose of the drawer handle. direction: Whether this invocation pulls open or pushes closed. approach_distance: Pre-grasp offset opposite the approach axis. translation_distance: Drawer travel distance for this operation. @@ -158,7 +157,7 @@ def create_invocation( "slide", SlideGoal( semantics, - SceneEntityPose(HANDLE_SCENE_ENTITY_ID), + target_pose, ), control_parts={"primary": {"motion": "arm", "grasp": "hand"}}, motion_policy=MotionPolicy(sample_count=TRAJECTORY_SAMPLE_COUNT), @@ -219,7 +218,7 @@ def main() -> None: "Inspect the closed drawer, then press Enter to plan the pull...", ) - for scene_version, direction in enumerate(("pull", "push")): + for direction in ("pull", "push"): if direction == "push" and wait_for_user: input( "Pull replay finished. Press Enter to read the moved handle " @@ -232,21 +231,13 @@ def main() -> None: create_invocation( engine, semantics, + handle_pose, direction=direction, approach_distance=args.approach_distance, translation_distance=args.translation_distance, ), ), - context=engine.initial_context( - scene=SceneSnapshot( - timestamp=float(scene_version), - version=scene_version, - entities={ - HANDLE_SCENE_ENTITY_ID: EntityState(handle_pose), - }, - ), - control_dt=sim.sim_config.physics_dt, - ), + context=engine.initial_context(control_dt=sim.sim_config.physics_dt), ) if not compiled.plan_success.all(): logger.log_warning(f"Failed to plan the Slide {direction} trajectory.") diff --git a/scripts/tutorials/atomic_action/twist.py b/scripts/tutorials/atomic_action/twist.py index 44f6ebec6..88cccab65 100644 --- a/scripts/tutorials/atomic_action/twist.py +++ b/scripts/tutorials/atomic_action/twist.py @@ -32,14 +32,11 @@ from embodichain.lab.sim.atomic_actions import ( AtomicActionEngine, ControlPartCommandProfile, - EntityState, MotionPolicy, ObjectSemantics, TwistAffordance, TwistGoal, TwistOptions, - SceneEntityPose, - SceneSnapshot, ) from embodichain.lab.sim.cfg import ( ArticulationCfg, @@ -193,7 +190,7 @@ def main() -> None: "twist", TwistGoal( semantics, - SceneEntityPose(KNOB_SCENE_ENTITY_ID), + target_pose, ), control_parts={"primary": {"motion": "arm", "grasp": "hand"}}, motion_policy=MotionPolicy(sample_count=TWIST_SAMPLE_INTERVAL), @@ -204,14 +201,7 @@ def main() -> None: ), ), ), - context=engine.initial_context( - scene=SceneSnapshot( - timestamp=0.0, - version=0, - entities={KNOB_SCENE_ENTITY_ID: EntityState(target_pose)}, - ), - control_dt=sim.sim_config.physics_dt, - ), + context=engine.initial_context(control_dt=sim.sim_config.physics_dt), ) if not compiled.plan_success.all(): logger.log_warning("Failed to plan the Twist demo trajectory.") diff --git a/tests/sim/atomic_actions/test_engine.py b/tests/sim/atomic_actions/test_engine.py index b1fb5fbfd..0622cc7f5 100644 --- a/tests/sim/atomic_actions/test_engine.py +++ b/tests/sim/atomic_actions/test_engine.py @@ -35,6 +35,7 @@ AtomicActionEngine, BUILTIN_ACTION_TYPES, ControlPartCommandProfile, + EntityState, JointPositionCommand, JointPositionGoal, JointPositionTarget, @@ -46,10 +47,12 @@ PressGoal, PressOptions, ResolvedActionRequest, + SceneSnapshot, SkillBindingContract, SkillEndpointRequirement, SkillResourceSlot, TimedTrajectory, + TrackingPolicy, ) ACTION_DT = 0.02 @@ -109,6 +112,24 @@ class OtherStubAction(StubAction): skill_id: ClassVar[str] = "other_stub" +class StubSceneProvider: + """Record scene captures and return one configured snapshot.""" + + def __init__(self, scene: SceneSnapshot) -> None: + self.scene = scene + self.calls: list[tuple[float, torch.Tensor]] = [] + + def snapshot( + self, + *, + timestamp: float, + env_ids: torch.Tensor, + ) -> SceneSnapshot: + """Return the configured scene and retain owned call arguments.""" + self.calls.append((timestamp, env_ids.clone())) + return self.scene + + def _motion_generator( batch_size: int = 2, robot_dof: int = 3, @@ -195,6 +216,52 @@ def test_engine_can_disable_builtin_loading() -> None: assert _engine(load_builtins=False).actions == {} +def test_initial_context_uses_configured_scene_provider() -> None: + scene = SceneSnapshot( + timestamp=1.5, + version=3, + entities={"target": EntityState(torch.eye(4))}, + ) + provider = StubSceneProvider(scene) + engine = AtomicActionEngine( + _motion_generator(), + load_builtins=False, + scene_provider=provider, + ) + + context = engine.initial_context(timestamp=1.5) + + assert context.scene is scene + assert len(provider.calls) == 1 + timestamp, env_ids = provider.calls[0] + assert timestamp == pytest.approx(1.5) + assert torch.equal(env_ids, torch.tensor([0, 1])) + + +def test_initial_context_explicit_scene_overrides_configured_provider() -> None: + provider = StubSceneProvider(SceneSnapshot(timestamp=0.0, version=1)) + explicit_scene = SceneSnapshot(timestamp=0.0, version=2) + engine = AtomicActionEngine( + _motion_generator(), + load_builtins=False, + scene_provider=provider, + ) + + context = engine.initial_context(scene=explicit_scene) + + assert context.scene is explicit_scene + assert provider.calls == [] + + +def test_engine_rejects_invalid_scene_provider() -> None: + with pytest.raises(TypeError, match="scene_provider must implement SceneProvider"): + AtomicActionEngine( + _motion_generator(), + load_builtins=False, + scene_provider=object(), # type: ignore[arg-type] + ) + + def test_auto_registered_builtin_accepts_per_invocation_options() -> None: generator = _motion_generator(robot_dof=3) generator.robot.control_parts = {"arm": object(), "hand": object()} @@ -357,12 +424,14 @@ def test_engine_make_invocation_binds_direct_control_parts() -> None: engine.register(StubAction()) goal = JointPositionGoal(torch.ones(2, 3)) motion_policy = MotionPolicy(sample_count=2) + tracking_policy = TrackingPolicy.timed() invocation = engine.make_invocation( "stub", goal, control_parts={"primary": {"motion": "all"}}, motion_policy=motion_policy, + tracking_policy=tracking_policy, invocation_id="direct-call", revision=1, ) @@ -373,6 +442,7 @@ def test_engine_make_invocation_binds_direct_control_parts() -> None: assert invocation.skill_id == "stub" assert invocation.goal is goal assert invocation.motion_policy is motion_policy + assert invocation.tracking_policy is tracking_policy assert invocation.invocation_id == "direct-call" assert invocation.revision == 1 assert target.control_part == "all" diff --git a/tests/sim/atomic_actions/test_sim_adapter.py b/tests/sim/atomic_actions/test_sim_adapter.py index 5b29109c0..4365b9469 100644 --- a/tests/sim/atomic_actions/test_sim_adapter.py +++ b/tests/sim/atomic_actions/test_sim_adapter.py @@ -18,13 +18,14 @@ from __future__ import annotations -from unittest.mock import Mock +from unittest.mock import Mock, patch import pytest import torch from embodichain.lab.sim.atomic_actions import ( CommandAckStatus, + create_simulation_atomic_action_engine, EndpointCommand, EndpointCommandTransport, JointPositionPayload, @@ -434,6 +435,49 @@ def test_rigid_object_scene_provider_tracks_per_environment_collision_revision() assert torch.equal(changed.entities["obstacle"].pose, moved_pose) +def test_simulation_engine_factory_registers_selected_entity_uids() -> None: + cube = Mock(uid="cube") + cube_pose = torch.eye(4).repeat(BATCH_SIZE, 1, 1) + cube.get_local_pose.return_value = cube_pose + motion_generator = Mock() + control_profiles = {"hand": Mock()} + grasp_pose_generators = {"hand": Mock()} + tracking_runtime = Mock() + + with patch( + "embodichain.lab.sim.atomic_actions.sim_adapter.AtomicActionEngine" + ) as engine_type: + engine = create_simulation_atomic_action_engine( + motion_generator, + (cube,), + control_profiles, + grasp_pose_generators, + load_builtins=False, + tracking_runtime=tracking_runtime, + ) + + assert engine is engine_type.return_value + kwargs = engine_type.call_args.kwargs + assert kwargs["control_profiles"] is control_profiles + assert kwargs["grasp_pose_generators"] is grasp_pose_generators + assert kwargs["load_builtins"] is False + assert kwargs["tracking_runtime"] is tracking_runtime + scene = kwargs["scene_provider"].snapshot( + timestamp=0.0, + env_ids=torch.arange(BATCH_SIZE), + ) + assert tuple(scene.entities) == ("cube",) + assert torch.equal(scene.entities["cube"].pose, cube_pose) + + +def test_simulation_engine_factory_rejects_duplicate_entity_uids() -> None: + with pytest.raises(ValueError, match="Duplicate scene entity uid 'cube'"): + create_simulation_atomic_action_engine( + Mock(), + (Mock(uid="cube"), Mock(uid="cube")), + ) + + @pytest.mark.parametrize( "value", [-1.0, float("nan"), float("inf"), float("-inf")], diff --git a/tests/sim/atomic_actions/test_tutorial_utils.py b/tests/sim/atomic_actions/test_tutorial_utils.py index b2bb7c432..69b9b2bb6 100644 --- a/tests/sim/atomic_actions/test_tutorial_utils.py +++ b/tests/sim/atomic_actions/test_tutorial_utils.py @@ -19,6 +19,7 @@ from __future__ import annotations import importlib +import inspect import math import re import xml.etree.ElementTree as ET @@ -81,6 +82,7 @@ ATOMIC_ACTION_TUTORIAL_MODULES = ( "assemble", "axis_align", + "control_dt", "coordinated_pickment", "coordinated_placement", "dynamic_obstacle_recovery", @@ -89,9 +91,39 @@ "move_held_object", "move_joints", "moving_target_recovery", + "open_door", "pickup", "place", + "pour", "press", + "slide", + "twist", +) +RIGID_SCENE_TUTORIAL_MODULES = ( + "assemble", + "axis_align", + "coordinated_placement", + "hand_over", + "move_held_object", + "pickup", + "place", + "pour", +) +STATIC_POSE_TUTORIAL_MODULES = ( + "coordinated_pickment", + "press", + "slide", + "twist", +) +EXPLICIT_SCENE_LIFECYCLE_TUTORIAL_MODULES = ( + "dynamic_obstacle_recovery", + "moving_target_recovery", + "open_door", +) +SCENE_FREE_TUTORIAL_MODULES = ( + "control_dt", + "move_end_effector", + "move_joints", ) @@ -497,6 +529,90 @@ def test_all_atomic_action_tutorials_accept_both_robot_choices( assert franka_args.robot == "franka" +def test_place_tutorial_registers_pick_object_with_simulation_engine_factory() -> None: + module = importlib.import_module("scripts.tutorials.atomic_action.place") + args = Namespace( + auto_play=True, + force_reannotate=False, + n_sample=1, + no_vis_eef_axis=True, + robot="ur5", + ) + sim = MagicMock() + sim.device = torch.device("cpu") + sim.sim_config.physics_dt = PHYSICS_DT + robot = MagicMock() + robot.get_qpos.return_value = torch.zeros(1, 8) + obj = MagicMock() + obj.uid = "cube" + obj.get_vertices.return_value = [torch.zeros(1, 3)] + obj.get_triangles.return_value = [torch.zeros(1, 3, dtype=torch.long)] + engine = MagicMock() + engine.compile.return_value = SimpleNamespace(plan_success=torch.tensor([False])) + + with ( + patch.object(module, "parse_arguments", return_value=args), + patch.object(module, "create_tutorial_simulation", return_value=sim), + patch.object(module, "add_tutorial_robot", return_value=robot), + patch.object(module, "create_pick_object", return_value=obj), + patch.object(module, "create_curobo_motion_generator"), + patch.object( + module, + "get_hand_open_close_qpos", + return_value=( + torch.zeros(1), + torch.ones(1), + ), + ), + patch.object(module, "initialize_pre_pick_robot_pose"), + patch.object( + module, + "create_simulation_atomic_action_engine", + return_value=engine, + ) as engine_factory, + patch.object(module, "create_parallel_jaw_grasp_pose_generator"), + patch.object(module, "prepare_tutorial_scene", return_value=False), + ): + module.main() + + assert engine_factory.call_args.kwargs["scene_entities"] == (obj,) + engine.initial_context.assert_called_once_with(control_dt=PHYSICS_DT) + + +def test_atomic_action_tutorial_scene_strategies_cover_every_entry_point() -> None: + classified = ( + set(RIGID_SCENE_TUTORIAL_MODULES) + | set(STATIC_POSE_TUTORIAL_MODULES) + | set(EXPLICIT_SCENE_LIFECYCLE_TUTORIAL_MODULES) + | set(SCENE_FREE_TUTORIAL_MODULES) + ) + + assert classified == set(ATOMIC_ACTION_TUTORIAL_MODULES) + + +@pytest.mark.parametrize("module_name", RIGID_SCENE_TUTORIAL_MODULES) +def test_rigid_scene_tutorials_use_simulation_engine_factory( + module_name: str, +) -> None: + module = importlib.import_module(f"scripts.tutorials.atomic_action.{module_name}") + source = inspect.getsource(module) + + assert "create_simulation_atomic_action_engine(" in source + assert "RigidObjectSceneProvider" not in source + assert "SceneSnapshot" not in source + + +@pytest.mark.parametrize("module_name", STATIC_POSE_TUTORIAL_MODULES) +def test_static_pose_tutorials_do_not_construct_scene_snapshots( + module_name: str, +) -> None: + module = importlib.import_module(f"scripts.tutorials.atomic_action.{module_name}") + source = inspect.getsource(module) + + assert "SceneSnapshot" not in source + assert "SceneEntityPose" not in source + + def test_axis_align_tutorial_exposes_upright_and_horizontal_modes() -> None: module = importlib.import_module("scripts.tutorials.atomic_action.axis_align")