diff --git a/.agents/skills/add-atomic-action/SKILL.md b/.agents/skills/add-atomic-action/SKILL.md index a007497c4..d6782728a 100644 --- a/.agents/skills/add-atomic-action/SKILL.md +++ b/.agents/skills/add-atomic-action/SKILL.md @@ -29,6 +29,7 @@ Inspect only the files relevant to the requested skill: | Engine-owned planning resources | `embodichain/lab/sim/atomic_actions/runtime.py` | | Reference implementations | `embodichain/lab/sim/atomic_actions/primitives/` | | Static compiler and execution session | `engine.py`, `execution.py` | +| Controller-facing execution ports | `runner.py`, `sim_adapter.py` | The public contract is: @@ -203,8 +204,10 @@ invocation = ActionInvocation( compiled = engine.compile((invocation,)) ``` -Use `engine.start(...).tick(...)` instead when dynamic scene updates or online -error recovery are required. +For dynamic scene updates or online error recovery, create a session with +`engine.start(...)`, then connect it to observation, command, and clock ports +through `ExecutionRunner`. Use non-blocking `runner.step()` in an existing event +loop or `runner.run_until_blocked()` in a simple application. ## 5. Export and document @@ -250,4 +253,4 @@ then use the `pre-commit-check` skill before committing. | Return an arm-only tensor | Embed into full robot DoF. | | Mutate held state after planning | Declare a `StateDelta`. | | Treat `plan_success` as physical success | Verify effects during execution. | -| Step the simulator from the action | Emit plans; let the caller own execution. | +| Step the simulator from the action | Emit plans; connect execution through `ExecutionRunner`. | diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 118de8e45..a70975f96 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -438,6 +438,10 @@ topics: - ActionPlan - PlanningContext - ExecutionSession + - ExecutionRunner + - ObservationProvider + - CommandSink + - SimulationExecutionAdapter - StateDelta - held_objects - ActionBinding @@ -472,6 +476,8 @@ topics: - embodichain/lab/sim/atomic_actions/state.py - embodichain/lab/sim/atomic_actions/plans.py - embodichain/lab/sim/atomic_actions/execution.py + - embodichain/lab/sim/atomic_actions/runner.py + - embodichain/lab/sim/atomic_actions/sim_adapter.py - embodichain/lab/sim/atomic_actions/engine.py - embodichain/lab/sim/atomic_actions/trajectory.py - embodichain/lab/sim/atomic_actions/primitives/ diff --git a/agent_context/topics/atomic-actions/atomic-actions.md b/agent_context/topics/atomic-actions/atomic-actions.md index 609269d9e..55ed9fccd 100644 --- a/agent_context/topics/atomic-actions/atomic-actions.md +++ b/agent_context/topics/atomic-actions/atomic-actions.md @@ -78,13 +78,21 @@ snapshot every time the action plans. Its entity ID is recorded in ```python session = engine.start(invocations, initial_context) -tick = session.tick(latest_context, effect_success=None) +runner = ExecutionRunner( + session, + observation_provider, + command_sink, + clock=execution_clock, +) +result = runner.step(effect_success=None) ``` -An `ExecutionSession` emits at most one `JointCommand` per tick. The command's -per-environment `hold_duration` preserves `TimedTrajectory.dt` for the caller's -control loop: command `i` carries the arrival interval `dt[:, i]`, which is -normally zero for the initial waypoint. The session monitors: +`ExecutionSession` owns deterministic planning progress and recovery state. It +emits at most one `JointCommand` per tick. The command's per-environment +`hold_duration` schedules the next feedback cycle from `TimedTrajectory.dt`: +command `i` carries the arrival interval `dt[:, i + 1]` leading to the next +waypoint. The final command reuses its own interval as a settling window. The +session monitors: - joint tracking error against the previous command; - translation/rotation drift of referenced scene entities; @@ -113,6 +121,40 @@ The replacement must keep the active `skill_id` and `invocation_id`. The session resolves a new snapshot, resets that revision's recovery budgets, and replans from the latest context. +`ExecutionRunner` owns the controller-facing lifecycle around a session: + +- `ObservationProvider.observe(task_state)` supplies a fresh, monotonically + timestamped `PlanningContext` when a feedback cycle is due; +- `CommandSink.send/hold/cancel` returns a `CommandAcknowledgement` with + `accepted`, `rejected`, or `timed_out` status; +- `ExecutionClock` supplies monotonic time and backend waiting; +- non-blocking `step()` dispatches only when the current command's + `hold_duration` has elapsed; +- `run_until_blocked()` is a convenience loop that waits through the clock and + stops at a terminal state or an unhandled effect-verification boundary; the + runner remembers that boundary so a later verifier call can resume it; +- cancellation, observation/session exceptions, and negative acknowledgements + enter a best-effort cancel-then-hold path. + +`TimedTrajectory.dt[:, i]` is the interval leading to sample `i`. +`ExecutionSession` dispatches sample zero immediately, then maps each following +arrival interval to the preceding command's `JointCommand.hold_duration`. The +final sample uses its own interval again as a settling window before terminal +validation. Batched execution currently advances at a synchronized barrier +using the longest active row interval. + +`SimulationExecutionAdapter` implements observation, command, and clock ports +for a `SimulationManager`/`Robot` pair. Its `sleep()` advances an integral +number of physics steps, so simulation execution does not depend on wall time. +Stable context IDs are correlation identifiers; the adapter maps command rows +to simulation robot indices rather than using those IDs as array indices. +Real-device adapters should implement the same protocols and enforce the passed +acknowledgement timeout in their transport/controller layer. + +The latest validated session context is retained for safe hold if the first +live observation fails. Environment IDs must remain stable and ordered for the +entire session; robot and scene timestamps and scene versions must be monotonic. + ## Parameter ownership Goal dataclasses carry only semantic task intent. They do not carry robot part @@ -135,6 +177,11 @@ The module-level `register_action()` API is a process-wide extension-type discovery catalog only; it neither binds actions nor changes an engine's default built-in set. +`ExecutionRunnerCfg` is intentionally separate from action options. It +configures controller acknowledgement deadlines, scheduler cadence, and final +safe-hold behavior for one runner instance; it does not change skill planning +semantics and does not belong in `ActionInvocation` or an invocation revision. + Every `ActionBinding` value is a `RobotCfg.control_parts` key. It is not a link, TCP-frame, joint, or scene-object name. Planning services validate those names and resolve immutable `ResolvedControlPart` values containing full-robot joint @@ -177,6 +224,11 @@ tutorial may derive a simple profile from limits explicitly. | `coordinated_placement` | `CoordinatedPlacementGoal` | `placing`, `support` | | `hand_over` | `GraspGoal` | `source`, `destination` | +`GraspGoal.grasp_xpos` accepts an explicit pose tensor, a late-bound +`SceneEntityPose`, or `None` for affordance sampling. A `SceneEntityPose` +registers the referenced entity as a recovery dependency, allowing an executing +`PickUp` to replan when the grasp target moves. + ## Extension rules 1. Define a frozen action-owned goal dataclass with `goal_kind`. @@ -188,4 +240,5 @@ tutorial may derive a simple profile from limits explicitly. 7. Declare symbolic changes with `StateDelta`; do not mutate context or commit physical effects during planning. 8. Keep scene stepping, controller I/O, and task-graph/MLLM logic outside the - atomic action. + atomic action. Put execution-loop I/O behind the runner protocols rather than + calling a simulator or device from `plan()` or `ExecutionSession`. 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 29d8b335b..61684182d 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 @@ -39,9 +39,23 @@ embodichain.lab.sim.atomic_actions AtomicAction AtomicActionEngine ExecutionSession + ExecutionRunner + ExecutionRunnerCfg + RunnerStep + RunnerStatus + ObservationProvider + CommandSink + CommandAcknowledgement + CommandAckStatus + CommandDispatch + CommandOperation + ExecutionClock + SimulationExecutionAdapter ExecutionTick JointCommand ExecutionEvent + ExecutionEventKind + ExecutionStatus .. rubric:: Built-in goals and actions @@ -151,6 +165,46 @@ Engine and execution .. autoclass:: ExecutionSession :members: +.. autoclass:: ExecutionRunner + :members: + +.. autoclass:: ExecutionRunnerCfg + :members: + :exclude-members: __init__, copy, replace, to_dict + +.. autoclass:: ObservationProvider + :members: + +.. autoclass:: CommandSink + :members: + +.. autoclass:: ExecutionClock + :members: + +.. autoclass:: MonotonicExecutionClock + :members: + +.. autoclass:: SimulationExecutionAdapter + :members: + +.. autoclass:: CommandAcknowledgement + :members: + +.. autoclass:: CommandAckStatus + :members: + +.. autoclass:: CommandDispatch + :members: + +.. autoclass:: CommandOperation + :members: + +.. autoclass:: RunnerStep + :members: + +.. autoclass:: RunnerStatus + :members: + .. autoclass:: ExecutionTick :members: @@ -160,6 +214,12 @@ Engine and execution .. autoclass:: ExecutionEvent :members: +.. autoclass:: ExecutionEventKind + :members: + +.. autoclass:: ExecutionStatus + :members: + Semantic objects and helpers ---------------------------- diff --git a/docs/source/overview/sim/atomic_actions/builtin_actions.md b/docs/source/overview/sim/atomic_actions/builtin_actions.md index 0fd388c87..826d789b3 100644 --- a/docs/source/overview/sim/atomic_actions/builtin_actions.md +++ b/docs/source/overview/sim/atomic_actions/builtin_actions.md @@ -201,7 +201,8 @@ entity as a recovery dependency. | `Press.xpos` | yes | yes | | `CoordinatedPickGoal.object_target_pose` / `object_initial_pose` | yes | yes | | `CoordinatedPlacementGoal` placing/support poses | yes | yes | -| `PickUp` / `HandOver` semantic entity lookup | not through `SceneEntityPose` | no automatic scene dependency | +| `PickUp.grasp_xpos` | yes | yes | +| `PickUp` / `HandOver` `ObjectSemantics.entity` lookup | not through `SceneEntityPose` | no automatic scene dependency | | `AssembleGoal` base entity lookup | not through `SceneEntityPose` | latest pose is used when replanning, but base movement alone does not trigger it | ### Parameter ownership @@ -310,10 +311,12 @@ bound manipulator. | Effect | write `HeldObjectState` for the bound manipulator and clear overlapping coordinated attachment state | | Verification | the attachment effect must be verified during closed-loop execution | -`grasp_xpos` may be `(4, 4)` or `(B, 4, 4)`. When omitted, the action samples -valid affordance grasps, evaluates reachability, and stores the selected -`object_to_eef` transform in the expected held-object state. Later -object-centric skills reuse that transform. +`grasp_xpos` may be `(4, 4)`, `(B, 4, 4)`, or a `SceneEntityPose`. A scene +reference resolves the latest grasp pose and registers its entity as a recovery +dependency, so material target motion invalidates and replans an executing +`PickUp`. When omitted, the action samples valid affordance grasps, evaluates +reachability, and stores the selected `object_to_eef` transform in the expected +held-object state. Later object-centric skills reuse that transform. `PickUp` requires `open` and `grasp` commands on the bound end-effector profile. Important `PickUpOptions` fields: @@ -327,9 +330,9 @@ Important `PickUpOptions` fields: | `downstream_object_target_poses` | Optional future reachability constraints used in grasp selection | | `obj_upright_direction`, `rotate_upright` | Optional orientation-selection behavior | -The semantic entity pose is read when planning occurs, but it is not currently a -`SceneEntityPose` dependency. Object motion alone therefore does not trigger -automatic dynamic-goal replanning. +Reading `ObjectSemantics.entity` remains a live planning lookup rather than an +automatic dependency. Use an explicit `SceneEntityPose` in `grasp_xpos` when +object motion should trigger dynamic-goal replanning. **Example:** `scripts/tutorials/atomic_action/pickup.py` diff --git a/docs/source/overview/sim/atomic_actions/index.md b/docs/source/overview/sim/atomic_actions/index.md index 3a63ca6b9..e428388b3 100644 --- a/docs/source/overview/sim/atomic_actions/index.md +++ b/docs/source/overview/sim/atomic_actions/index.md @@ -60,6 +60,13 @@ and whole-body control are not implemented by this module yet. | | v v CompiledTrajectory JointCommand + events + | + v + ExecutionRunner + observe / schedule / dispatch + | + v + ObservationProvider + CommandSink + Clock ``` The boundary is deliberate: @@ -71,9 +78,19 @@ The boundary is deliberate: | Perception and grounding | Agent adapter or user application | Builds scene snapshots and resource bindings, or supplies already-grounded values directly | | Deterministic motion planning | Atomic action module | Produces an `ActionPlan` from an invocation and context | | Motion-generation resources | `AtomicActionEngine` | Owns one robot, motion generator, planner backend, device, trajectory builder, and control-part command profiles | -| Robot/simulator stepping | Application control loop | Consumes `JointCommand`; the session never steps the simulator itself | +| Recovery state | `ExecutionSession` | Consumes fresh contexts, emits at most one `JointCommand` per tick, and owns bounded recovery/revision state | +| Scheduling and controller lifecycle | `ExecutionRunner` | Observes only when due, dispatches timed commands, records acknowledgements, and performs safe stop | +| Robot/simulator I/O | `ObservationProvider`, `CommandSink`, and `ExecutionClock` adapters | Isolates observation, command transport, and time/physics advancement from planning and session state | | Physical-effect verification | Application observer | Verifies grasp, release, handover, and other symbolic effects | +`ExecutionRunner.step()` is non-blocking. Its convenience +`run_until_blocked()` loop waits or advances simulation through an injected +clock. Observation errors, rejected or timed-out commands, session failures, +and explicit cancellation trigger a best-effort cancel-then-hold sequence. +`SimulationExecutionAdapter` implements all three ports for a simulation robot; +real hardware integrations implement the same protocols without changing +action planning or recovery state. + ### Caller entry points The engine supports two first-class caller paths. An Action Agent emits a @@ -154,6 +171,7 @@ from leaking into an Action Agent schema. | `ActionControlOverrides` | Optional role-scoped command replacements for one invocation revision | Persistent robot configuration | | `MotionPolicy` | Motion source, sample count, timing, limits, collision option, typed planner options | Skill semantics or robot-resource names | | `RecoveryPolicy` | Replan/retry budgets, tracking and dynamic-goal thresholds, phase timeout | Controller state or mutable counters | +| `ExecutionRunnerCfg` | Runner-level acknowledgement deadlines, minimum feedback cadence, and completion hold policy | Skill behavior, planning resources, or invocation revision data | | `PlanningContext` | Measured `RobotObservation`, verified `TaskState`, versioned `SceneSnapshot`, stable environment IDs | Hypothetical simulator mutation | | `ActionPlan` | Per-environment planning result, scene-bound phases, timed trajectories, diagnostics, expected `StateDelta` | Proof that a grasp/release/contact physically succeeded | @@ -344,6 +362,9 @@ by the engine after resolving an invocation: | `AtomicAction.plan(request, context)` | Atomic-action implementer | Consumes an immutable `ResolvedActionRequest` and returns an `ActionPlan` | | `engine.plan_action(action, invocation, context)` | Extension or isolated test | Temporarily binds and plans an unregistered action instance; built-in parameter variants should use invocation `skill_options` instead | | `session.revise_current(invocation)` | Runtime orchestrator or Action Agent | Replaces the active logical call with a newer revision and replans from the latest observed context | +| `runner.step(effect_success=...)` | Non-blocking controller integration | Observes and dispatches only when the next timed command is due | +| `runner.run_until_blocked(...)` | Simple blocking application or tutorial | Advances the injected clock until terminal or external effect verification is required | +| `runner.cancel(reason)` | Explicit safe stop | Requests controller cancellation followed by an observed-position hold | Application code should start with `engine.plan()`, `engine.compile()`, or `engine.start()` unless it specifically needs one of these extension points. @@ -463,6 +484,35 @@ while session.status is ExecutionStatus.RUNNING: latest_context = observe_context() ``` +For most applications, use `ExecutionRunner` to keep scheduling and controller +acknowledgement handling outside the session: + +```python +adapter = SimulationExecutionAdapter(sim, robot, scene_supplier=read_scene) +initial_context = adapter.observe( + TaskState.empty(robot.get_qpos().shape[0], robot.device) +) +session = engine.start((moving_goal,), initial_context) +runner = ExecutionRunner(session, adapter, adapter, clock=adapter) +result = runner.run_until_blocked() +``` + +`ExecutionRunner.step()` is the non-blocking entry point for an application +that already owns its event loop. It observes only when the previous command's +`hold_duration` has elapsed, dispatches active commands through `CommandSink`, +and records accepted, rejected, or timed-out acknowledgements. Cancellation, +observation/session exceptions, and negative acknowledgements enter a +best-effort cancel-then-hold path. + +`TimedTrajectory.dt[:, i]` is the interval leading to sample `i`. +`ExecutionSession` maps each following arrival interval onto the preceding +command's post-dispatch hold, while the final sample reuses its own interval as +a settling window before terminal validation. A batched runner uses the longest +active row interval as its synchronized barrier. +`SimulationExecutionAdapter.sleep()` converts that interval to an integral +number of physics steps instead of using wall-clock sleep. Stable `env_ids` +remain correlation identifiers and are not used as simulator array indices. + On each tick, the session can detect: - joint tracking error relative to the previously emitted command; @@ -495,10 +545,11 @@ retain their runtime identity. Each emitted `JointCommand` carries a per-environment `hold_duration` derived from the plan's `TimedTrajectory.dt`. The application control loop must respect that timing after dispatching the command and before requesting the next -observation. `dt[:, i]` is the arrival interval for waypoint `i`, so the first -command normally carries zero duration. For a synchronized batch, the caller -should wait for the longest duration among active rows. A passive hold command -has zero duration. +observation. `dt[:, i]` is the arrival interval leading to waypoint `i`, so the +first waypoint is dispatched immediately and command `i` carries `dt[:, i + 1]` +until the next waypoint is due. The final command reuses `dt[:, -1]` as a +settling window. For a synchronized batch, the caller should wait for the +longest duration among active rows. A passive hold command has zero duration. Use an explicit newer revision when the application or Action Agent decides to change runtime behavior: @@ -604,4 +655,6 @@ See {doc}`builtin_actions` for the shipped skill catalog and visual demos, and - {doc}`../planners/motion_generator` — the motion generator owned by the engine - {doc}`../sim_robot` — robot control parts and kinematic configuration -- `scripts/tutorials/atomic_action/` — focused examples for every built-in skill +- {doc}`/tutorial/atomic_actions` — static, closed-loop, and recovery examples +- `scripts/tutorials/atomic_action/moving_target_recovery.py` — runnable runner + example that visibly moves a late-bound target, replans, and picks up the cube diff --git a/docs/source/tutorial/atomic_actions.rst b/docs/source/tutorial/atomic_actions.rst index fc3b306f8..94575776a 100644 --- a/docs/source/tutorial/atomic_actions.rst +++ b/docs/source/tutorial/atomic_actions.rst @@ -99,6 +99,7 @@ Focused examples live under ``scripts/tutorials/atomic_action``: * ``coordinated_pickment.py`` * ``coordinated_placement.py`` * ``hand_over.py`` +* ``moving_target_recovery.py`` The scripts are interactive by default. Add ``--auto_play`` to skip prompts; combine it with ``--headless --device cpu`` for a headless run that records @@ -260,18 +261,45 @@ must be resolved from the latest scene snapshot: ), ) - latest_context = initial_context - session = engine.start((invocation,), latest_context) - while session.status.value == "running": - tick = session.tick(latest_context) - if tick.command is not None: - send_joint_command(tick.command) - latest_context = observe_context() + from embodichain.lab.sim.atomic_actions import ( + ExecutionRunner, + SimulationExecutionAdapter, + TaskState, + ) + + adapter = SimulationExecutionAdapter(sim, robot, scene_supplier=read_scene) + task = TaskState.empty(robot.get_qpos().shape[0], robot.device) + initial_context = adapter.observe(task) + session = engine.start((invocation,), initial_context) + runner = ExecutionRunner(session, adapter, adapter, clock=adapter) + result = runner.run_until_blocked() + +The session owns planning progress and bounded recovery. The runner owns the +outer lifecycle: it requests fresh observations, schedules each command from +the :class:`~embodichain.lab.sim.atomic_actions.TimedTrajectory` time deltas, +checks controller acknowledgements, and performs cancel-then-hold on failure. +The simulation adapter advances physics instead of sleeping in wall-clock time. +``ExecutionRunnerCfg`` contains runner-level transport and scheduling settings; +it is not an atomic-action option and is not replaced by invocation revision. + +For an application that already owns its event loop, call the non-blocking +:meth:`~embodichain.lab.sim.atomic_actions.ExecutionRunner.step` method. A step +with ``is_waiting`` set has not consumed a new observation or effect result; use +its ``wait_duration`` to schedule the next call. + +The complete simulation example starts with a visible cube directly in front of +the robot, then applies a short horizontal force pulse so physics and friction +slide it sideways during one ``PickUp`` invocation whose +``GraspGoal.grasp_xpos`` is a ``SceneEntityPose``. The session observes +``dynamic_goal_changed`` and ``replanned`` events, discards the entire stale +approach/close/lift plan, and rebuilds it from the cube's new location. The +replanned action closes the gripper, verifies the physical lift, and finishes +while holding the cube. The original and regenerated goal axes remain visible +for comparison: -The session emits one command per tick. It compares observations with the last -command, detects material motion of referenced scene entities, enforces phase -timeouts, and replans from the latest observation within the recovery budget. -It does not own the simulator or controller loop. +.. code-block:: bash + + python scripts/tutorials/atomic_action/moving_target_recovery.py --headless --auto_play --device cpu Recovery replans reuse one immutable invocation-revision snapshot. If an application intentionally changes the goal, options, policy, binding, or a @@ -311,13 +339,18 @@ external per-environment verification mask: .. code-block:: python - tick = session.tick(latest_context) - if any(event.kind.value == "effect_verification_required" for event in tick.events): - verified = verify_grasp_or_release() - tick = session.tick(latest_context, effect_success=verified) + def verify_effect(context, tick): + return verify_grasp_or_release(context) + + result = runner.run_until_blocked(effect_verifier=verify_effect) This prevents a successful trajectory plan from being mistaken for a successful -physical grasp or release. +physical grasp or release. If verification is asynchronous, omit the callback; +``run_until_blocked`` returns at the verification boundary and the application +can later resume with ``runner.step(effect_success=verified)`` when the next +cycle is due, or call ``run_until_blocked(effect_verifier=...)`` again. The +runner remembers the pending boundary even though the session emits its event +only once. Adding an action ---------------- diff --git a/embodichain/lab/sim/atomic_actions/__init__.py b/embodichain/lab/sim/atomic_actions/__init__.py index 7564a52b4..5b0ad449a 100644 --- a/embodichain/lab/sim/atomic_actions/__init__.py +++ b/embodichain/lab/sim/atomic_actions/__init__.py @@ -21,8 +21,9 @@ :class:`PlanningContext` through :meth:`AtomicAction.plan`. Planning is side-effect free: it returns an :class:`ActionPlan` with timed motion, completion criteria, diagnostics, and uncommitted expected task-state effects. -:class:`AtomicActionEngine` can compile a static sequence; closed-loop execution -belongs to an execution session. +:class:`AtomicActionEngine` can compile a static sequence. For closed-loop use, +:class:`ExecutionSession` owns recovery and invocation-revision state while +:class:`ExecutionRunner` connects it to observations, commands, and time. """ from __future__ import annotations @@ -102,6 +103,23 @@ PressGoal, PressOptions, ) +from .runner import ( + CommandAcknowledgement, + CommandAckStatus, + CommandDispatch, + CommandOperation, + CommandSink, + EffectVerifier, + ExecutionClock, + ExecutionRunner, + ExecutionRunnerCfg, + MonotonicExecutionClock, + ObservationProvider, + RunnerStatus, + RunnerStep, + RunnerStepCallback, +) +from .sim_adapter import SceneSnapshotSupplier, SimulationExecutionAdapter from .state import ( CoordinatedHeldObjectState, EntityState, @@ -129,6 +147,11 @@ "AtomicActionEngine", "BUILTIN_ACTION_TYPES", "CompiledTrajectory", + "CommandAcknowledgement", + "CommandAckStatus", + "CommandDispatch", + "CommandOperation", + "CommandSink", "CompletionCondition", "CompletionConditionKind", "ControlCommand", @@ -142,8 +165,12 @@ "CoordinatedPlacementOptions", "EndEffectorPoseGoal", "EntityState", + "EffectVerifier", + "ExecutionClock", "ExecutionEvent", "ExecutionEventKind", + "ExecutionRunner", + "ExecutionRunnerCfg", "ExecutionSession", "ExecutionStatus", "ExecutionTick", @@ -158,6 +185,7 @@ "JointCommand", "JointPositionCommand", "MotionPolicy", + "MonotonicExecutionClock", "MoveEndEffector", "MoveEndEffectorOptions", "MoveHeldObject", @@ -167,6 +195,7 @@ "ObjectActionGoal", "ObjectSemantics", "OPEN_COMMAND", + "ObservationProvider", "PhaseSpec", "PickUp", "PickUpOptions", @@ -185,10 +214,15 @@ "ResolvedActionBinding", "ResolvedControlPart", "RobotObservation", + "RunnerStatus", + "RunnerStep", + "RunnerStepCallback", "SceneSnapshot", + "SceneSnapshotSupplier", "SceneEntityPose", "SkillDescriptor", "StateDelta", + "SimulationExecutionAdapter", "TaskState", "TimedTrajectory", "TrajectoryBuilder", diff --git a/embodichain/lab/sim/atomic_actions/execution.py b/embodichain/lab/sim/atomic_actions/execution.py index bccb920d8..06cd11864 100644 --- a/embodichain/lab/sim/atomic_actions/execution.py +++ b/embodichain/lab/sim/atomic_actions/execution.py @@ -91,7 +91,7 @@ class JointCommand: active_mask: torch.Tensor env_ids: torch.Tensor hold_duration: torch.Tensor - """Per-environment arrival interval to wait after dispatching this command.""" + """Per-environment delay before the next observation/command cycle.""" def __post_init__(self) -> None: if self.positions.dim() != 2: @@ -276,6 +276,11 @@ def revise_current(self, invocation: ActionInvocation) -> None: ExecutionEventKind.INVOCATION_REVISED, ) + @property + def latest_context(self) -> PlanningContext: + """Latest validated context with the session's verified task state.""" + return self._context + def tick( self, context: PlanningContext, @@ -627,7 +632,14 @@ def _command_at( ) self._last_command = positions.clone() self._last_command_mask = active_mask.clone() - hold_duration = phase.trajectory.dt[:, waypoint_index] + # ``dt[:, i]`` leads to waypoint ``i``. After dispatching waypoint + # ``i``, wait for ``dt[:, i + 1]`` before the next dispatch. Reuse the + # final arrival interval as its terminal settling window. + next_waypoint_index = min( + waypoint_index + 1, + phase.trajectory.waypoint_count - 1, + ) + hold_duration = phase.trajectory.dt[:, next_waypoint_index] return JointCommand( positions=positions, velocities=velocities, diff --git a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py index 6320f9a6f..72798560e 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py +++ b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py @@ -39,7 +39,12 @@ from ..control import GRASP_COMMAND, OPEN_COMMAND from ..core import AtomicAction, ObjectSemantics from ..effects import StateDelta -from ..goals import ObjectActionGoal, validate_pose_tensor +from ..goals import ( + ObjectActionGoal, + PoseGoalValue, + resolve_pose_goal, + validate_pose_goal, +) from ..invocation import ActionOptions, ResolvedActionRequest from ..plans import ActionPlan from ..policies import MotionPolicy @@ -52,18 +57,20 @@ class GraspGoal(ObjectActionGoal): goal_kind: ClassVar[str] = "grasp" - grasp_xpos: torch.Tensor | None = None + grasp_xpos: PoseGoalValue | None = None """Optional end-effector grasp pose. - When omitted, :class:`PickUp` selects a grasp from the target affordance. - Supplying a pose with shape ``(4, 4)`` or ``(n_envs, 4, 4)`` skips grasp - sampling. + When omitted, :class:`PickUp` selects a grasp from the target affordance. An + explicit tensor or late-bound + :class:`~embodichain.lab.sim.atomic_actions.goals.SceneEntityPose` skips + grasp sampling. Late-bound poses also declare the scene dependency used by + closed-loop execution recovery. """ def __post_init__(self) -> None: ObjectActionGoal.__post_init__(self) if self.grasp_xpos is not None: - validate_pose_tensor(self.grasp_xpos, "grasp_xpos", allow_waypoints=False) + validate_pose_goal(self.grasp_xpos, "grasp_xpos", allow_waypoints=False) @dataclass(frozen=True, slots=True, eq=False) @@ -294,7 +301,8 @@ def plan( ) else: grasp_xpos = self.builder.resolve_pose_target( - target.grasp_xpos, n_envs=self.n_envs + resolve_pose_goal(target.grasp_xpos, context, name="grasp_xpos"), + n_envs=self.n_envs, ) if options.rotate_upright is not None: grasp_xpos = self._upright_adjusted_grasp_poses( diff --git a/embodichain/lab/sim/atomic_actions/runner.py b/embodichain/lab/sim/atomic_actions/runner.py new file mode 100644 index 000000000..0c82ab6b6 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/runner.py @@ -0,0 +1,793 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Controller-independent scheduling for closed-loop atomic-action execution.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from enum import Enum +import math +import time +from typing import Protocol, runtime_checkable + +import torch + +from embodichain.utils import configclass + +from .execution import ( + ExecutionEventKind, + ExecutionSession, + ExecutionStatus, + ExecutionTick, + JointCommand, +) +from .state import PlanningContext, TaskState + + +class CommandAckStatus(str, Enum): + """Outcome reported by a command transport or controller.""" + + ACCEPTED = "accepted" + REJECTED = "rejected" + TIMED_OUT = "timed_out" + + +@dataclass(frozen=True, slots=True) +class CommandAcknowledgement: + """Synchronous acknowledgement returned by a :class:`CommandSink`.""" + + status: CommandAckStatus + """Transport/controller acknowledgement status.""" + + message: str = "" + """Human-readable diagnostic intended for logs, not policy branching.""" + + def __post_init__(self) -> None: + if not isinstance(self.status, CommandAckStatus): + raise TypeError("status must be a CommandAckStatus.") + if not isinstance(self.message, str): + raise TypeError("message must be a string.") + + @property + def accepted(self) -> bool: + """Whether the controller accepted the requested operation.""" + return self.status is CommandAckStatus.ACCEPTED + + @classmethod + def accepted_ack(cls, message: str = "") -> CommandAcknowledgement: + """Build an accepted acknowledgement. + + Args: + message: Optional controller diagnostic. + + Returns: + Accepted acknowledgement. + """ + return cls(CommandAckStatus.ACCEPTED, message) + + +class CommandOperation(str, Enum): + """Command-sink operation recorded by an execution runner.""" + + SEND = "send" + HOLD = "hold" + CANCEL = "cancel" + + +@dataclass(frozen=True, slots=True) +class CommandDispatch: + """Auditable record of one controller operation and acknowledgement.""" + + operation: CommandOperation + acknowledgement: CommandAcknowledgement + + def __post_init__(self) -> None: + if not isinstance(self.operation, CommandOperation): + raise TypeError("operation must be a CommandOperation.") + if not isinstance(self.acknowledgement, CommandAcknowledgement): + raise TypeError("acknowledgement must be a CommandAcknowledgement.") + + +@runtime_checkable +class ObservationProvider(Protocol): + """Source of fresh planning contexts for feedback-driven execution.""" + + def observe(self, task_state: TaskState) -> PlanningContext: + """Capture the latest robot and scene state. + + Args: + task_state: Runner-owned, externally verified symbolic task state. + + Returns: + Fresh context with stable, ordered environment IDs. + """ + + +@runtime_checkable +class CommandSink(Protocol): + """Controller boundary used by :class:`ExecutionRunner`.""" + + def send( + self, + command: JointCommand, + *, + timeout: float, + ) -> CommandAcknowledgement: + """Submit an active joint command and acknowledge its acceptance. + + Args: + command: Full-robot command with an explicit active mask. Inactive + rows contain hold targets and must not retain stale commands. + timeout: Maximum acknowledgement latency in seconds. + + Returns: + Transport or controller acknowledgement. + """ + + def hold( + self, + command: JointCommand, + *, + timeout: float, + ) -> CommandAcknowledgement: + """Hold the supplied observed position as a safety command. + + Args: + command: Full-robot observed-position hold command. + timeout: Maximum acknowledgement latency in seconds. + + Returns: + Transport or controller acknowledgement. + """ + + def cancel(self, *, timeout: float) -> CommandAcknowledgement: + """Cancel any controller-side command that has not completed. + + Args: + timeout: Maximum acknowledgement latency in seconds. + + Returns: + Transport or controller acknowledgement. + """ + + +@runtime_checkable +class ExecutionClock(Protocol): + """Clock abstraction used for deterministic and simulation scheduling.""" + + def now(self) -> float: + """Return a monotonic timestamp in seconds. + + Returns: + Monotonic timestamp in seconds. + """ + + def sleep(self, duration: float) -> None: + """Wait or advance the execution backend by ``duration`` seconds. + + Args: + duration: Non-negative duration in seconds. + """ + + +class MonotonicExecutionClock: + """Wall-clock implementation backed by :mod:`time`.""" + + def now(self) -> float: + """Return the current monotonic wall-clock time. + + Returns: + Monotonic wall-clock timestamp in seconds. + """ + return time.monotonic() + + def sleep(self, duration: float) -> None: + """Sleep for a non-negative wall-clock duration. + + Args: + duration: Requested duration in seconds. + """ + if not math.isfinite(duration) or duration < 0.0: + raise ValueError("duration must be finite and non-negative.") + time.sleep(duration) + + +@configclass +class ExecutionRunnerCfg: + """Transport and scheduling policy for an :class:`ExecutionRunner`.""" + + command_timeout: float = 1.0 + """Maximum time allowed for a command acknowledgement.""" + + safe_stop_timeout: float = 1.0 + """Maximum time allowed for each cancel or hold acknowledgement.""" + + minimum_cycle_time: float = 1.0e-3 + """Minimum delay between feedback cycles, including passive hold cycles.""" + + hold_on_completion: bool = True + """Whether to issue a final hold after the session completes.""" + + def __post_init__(self) -> None: + for name in ("command_timeout", "safe_stop_timeout"): + value = getattr(self, name) + if not math.isfinite(value) or value <= 0.0: + raise ValueError(f"{name} must be finite and greater than zero.") + if not math.isfinite(self.minimum_cycle_time) or self.minimum_cycle_time < 0.0: + raise ValueError("minimum_cycle_time must be finite and non-negative.") + if not isinstance(self.hold_on_completion, bool): + raise TypeError("hold_on_completion must be a bool.") + + +class RunnerStatus(str, Enum): + """Lifecycle status owned by an :class:`ExecutionRunner`.""" + + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + CANCELLED = "cancelled" + + +@dataclass(frozen=True, slots=True, eq=False) +class RunnerStep: + """Result of one non-blocking execution-runner update.""" + + status: RunnerStatus + timestamp: float + wait_duration: float + context: PlanningContext | None + tick: ExecutionTick | None + dispatches: tuple[CommandDispatch, ...] + command_count: int + message: str | None = None + """Terminal or failure diagnostic, when available.""" + + def __post_init__(self) -> None: + if not isinstance(self.status, RunnerStatus): + raise TypeError("status must be a RunnerStatus.") + if not math.isfinite(self.timestamp) or self.timestamp < 0.0: + raise ValueError("timestamp must be finite and non-negative.") + if not math.isfinite(self.wait_duration) or self.wait_duration < 0.0: + raise ValueError("wait_duration must be finite and non-negative.") + if self.command_count < 0: + raise ValueError("command_count must be non-negative.") + if self.message is not None and not isinstance(self.message, str): + raise TypeError("message must be a string or None.") + object.__setattr__(self, "dispatches", tuple(self.dispatches)) + + @property + def is_waiting(self) -> bool: + """Whether no session tick was due during this update.""" + return ( + self.status is RunnerStatus.RUNNING + and self.tick is None + and self.wait_duration > 0.0 + ) + + +EffectVerifier = Callable[[PlanningContext, ExecutionTick], torch.Tensor | None] +"""Callback that verifies a pending semantic effect for each environment.""" + +RunnerStepCallback = Callable[[RunnerStep], None] +"""Optional observer called after every blocking runner-loop iteration.""" + + +class ExecutionRunner: + """Connect an execution session to observation, controller, and time ports. + + :meth:`step` is non-blocking. It observes and advances the session only when + the next command is due according to :attr:`JointCommand.hold_duration`. + :meth:`run_until_blocked` supplies the blocking loop for tutorials and simple + applications. Controller rejection, timeout, observation failure, and + session exceptions all trigger a best-effort cancel-then-hold sequence. + + Args: + session: Stateful atomic-action execution session. + observation_provider: Source of fresh robot and scene observations. + command_sink: Controller or simulation command boundary. + clock: Optional scheduler clock. Defaults to monotonic wall time. + cfg: Optional acknowledgement, scheduling, and completion policy. + """ + + def __init__( + self, + session: ExecutionSession, + observation_provider: ObservationProvider, + command_sink: CommandSink, + *, + clock: ExecutionClock | None = None, + cfg: ExecutionRunnerCfg | None = None, + ) -> None: + if not isinstance(session, ExecutionSession): + raise TypeError("session must be an ExecutionSession.") + if not isinstance(observation_provider, ObservationProvider): + raise TypeError("observation_provider must implement ObservationProvider.") + if not isinstance(command_sink, CommandSink): + raise TypeError("command_sink must implement CommandSink.") + if clock is not None and not isinstance(clock, ExecutionClock): + raise TypeError("clock must implement ExecutionClock.") + if cfg is not None and not isinstance(cfg, ExecutionRunnerCfg): + raise TypeError("cfg must be an ExecutionRunnerCfg.") + self._session = session + self._observation_provider = observation_provider + self._command_sink = command_sink + self._clock = clock or MonotonicExecutionClock() + self.cfg = cfg or ExecutionRunnerCfg() + self._status = RunnerStatus.RUNNING + self._next_step_at = self._clock_now() + self._last_context: PlanningContext | None = session.latest_context + self._command_count = 0 + self._message: str | None = None + self._effect_verification_pending = False + self._effect_context: PlanningContext | None = None + self._effect_tick: ExecutionTick | None = None + + @property + def session(self) -> ExecutionSession: + """Execution session advanced by this runner.""" + return self._session + + @property + def status(self) -> RunnerStatus: + """Current runner lifecycle status.""" + return self._status + + @property + def command_count(self) -> int: + """Number of active commands accepted by the sink.""" + return self._command_count + + @property + def effect_verification_pending(self) -> bool: + """Whether execution is waiting for an external semantic-effect result.""" + return self._effect_verification_pending + + def step( + self, + *, + effect_success: torch.Tensor | None = None, + ) -> RunnerStep: + """Perform one due observation/session/controller update without sleeping. + + Args: + effect_success: Optional per-environment verification mask. If this + call occurs before the next cycle is due, it is not consumed and + must be supplied again on a later call. + + Returns: + Runner status, optional session tick, controller acknowledgements, + and time remaining before another update is due. + """ + now = self._clock_now() + if self._status is not RunnerStatus.RUNNING: + return self._result(timestamp=now) + wait_duration = self._remaining_wait(now) + if wait_duration > 0.0: + return self._result( + timestamp=now, + wait_duration=wait_duration, + ) + + try: + context = self._observation_provider.observe(self._session.task_state) + if not isinstance(context, PlanningContext): + raise TypeError( + "ObservationProvider.observe() must return PlanningContext." + ) + except Exception as exc: + return self._fail( + f"Observation provider failed: {type(exc).__name__}: {exc}", + context=self._last_context, + ) + self._last_context = context + + try: + tick = self._session.tick(context, effect_success=effect_success) + except Exception as exc: + return self._fail( + f"Execution session failed: {type(exc).__name__}: {exc}", + context=context, + ) + self._update_effect_boundary(context, tick, effect_success) + + dispatches: list[CommandDispatch] = [] + if tick.command is not None: + operation = ( + CommandOperation.SEND + if bool(tick.command.active_mask.any().item()) + else CommandOperation.HOLD + ) + dispatch = self._dispatch(operation, tick.command) + dispatches.append(dispatch) + if not dispatch.acknowledgement.accepted: + failure = dispatch.acknowledgement + message = ( + "Controller did not accept the requested command: " + f"{failure.status.value}." + ) + if failure.message: + message += f" {failure.message}" + return self._fail( + message, + context=context, + tick=tick, + dispatches=dispatches, + ) + if operation is CommandOperation.SEND: + self._command_count += 1 + interval = self._command_interval(tick.command) + self._next_step_at = self._clock_now() + interval + else: + self._next_step_at = self._clock_now() + + if tick.status is ExecutionStatus.COMPLETED: + if self.cfg.hold_on_completion: + hold_dispatch = self._dispatch( + CommandOperation.HOLD, + self._hold_command(context), + ) + dispatches.append(hold_dispatch) + if not hold_dispatch.acknowledgement.accepted: + failure = hold_dispatch.acknowledgement + message = ( + "Final safety hold was not accepted: " + f"{failure.status.value}." + ) + if failure.message: + message += f" {failure.message}" + return self._fail( + message, + context=context, + tick=tick, + dispatches=dispatches, + ) + self._status = RunnerStatus.COMPLETED + self._next_step_at = self._clock_now() + elif tick.status is ExecutionStatus.FAILED: + return self._fail( + "Execution session exhausted its recovery budget.", + context=context, + tick=tick, + dispatches=dispatches, + ) + + return self._result( + timestamp=self._clock_now(), + context=context, + tick=tick, + dispatches=dispatches, + wait_duration=self._remaining_wait(self._clock_now()), + ) + + def cancel(self, reason: str = "Execution cancelled by caller.") -> RunnerStep: + """Cancel controller work and hold the latest observed position. + + Args: + reason: Human-readable cancellation reason. + + Returns: + Terminal runner step. The status is ``cancelled`` only when both + cancel and hold are acknowledged; otherwise it is ``failed``. + """ + if not isinstance(reason, str) or not reason: + raise ValueError("reason must be a non-empty string.") + now = self._clock_now() + if self._status is not RunnerStatus.RUNNING: + return self._result(timestamp=now) + context = self._observe_for_stop() + dispatches = self._safe_stop(context) + if all(item.acknowledgement.accepted for item in dispatches): + self._status = RunnerStatus.CANCELLED + self._message = reason + else: + self._status = RunnerStatus.FAILED + self._message = f"{reason} Safe stop acknowledgement failed." + self._clear_effect_boundary() + self._next_step_at = self._clock_now() + return self._result( + timestamp=self._clock_now(), + context=context, + dispatches=dispatches, + ) + + def run_until_blocked( + self, + *, + effect_verifier: EffectVerifier | None = None, + on_step: RunnerStepCallback | None = None, + max_steps: int = 100_000, + ) -> RunnerStep: + """Run with clock-driven waiting until terminal or effect verification blocks. + + Args: + effect_verifier: Optional callback used after an + ``effect_verification_required`` event. Without one, the method + returns the running step so the caller can verify externally. + on_step: Optional callback for tracing or tutorial visualization. + max_steps: Hard bound on loop iterations. + + Returns: + Terminal step, or a running step blocked on external verification. + """ + if max_steps <= 0: + raise ValueError("max_steps must be greater than zero.") + pending_effect: torch.Tensor | None = None + now = self._clock_now() + last_result = self._result( + timestamp=now, + wait_duration=self._remaining_wait(now), + context=self._effect_context, + tick=self._effect_tick, + ) + if self._effect_verification_pending: + if ( + effect_verifier is None + or self._effect_context is None + or self._effect_tick is None + ): + return last_result + try: + pending_effect = effect_verifier( + self._effect_context, + self._effect_tick, + ) + except Exception as exc: + return self._fail( + f"Effect verifier failed: {type(exc).__name__}: {exc}", + context=self._effect_context, + tick=self._effect_tick, + ) + if pending_effect is None: + return last_result + for _ in range(max_steps): + result = self.step(effect_success=pending_effect) + if result.tick is not None: + pending_effect = None + if on_step is not None: + try: + on_step(result) + except Exception as exc: + return self._fail( + f"Runner step callback failed: {type(exc).__name__}: {exc}", + context=result.context or self._last_context, + tick=result.tick, + dispatches=list(result.dispatches), + ) + last_result = result + if result.status is not RunnerStatus.RUNNING: + return result + verification_required = result.tick is not None and any( + event.kind is ExecutionEventKind.EFFECT_VERIFICATION_REQUIRED + for event in result.tick.events + ) + if verification_required: + if effect_verifier is None or result.context is None: + return result + try: + pending_effect = effect_verifier(result.context, result.tick) + except Exception as exc: + return self._fail( + f"Effect verifier failed: {type(exc).__name__}: {exc}", + context=result.context, + tick=result.tick, + dispatches=list(result.dispatches), + ) + if pending_effect is None: + return result + if result.wait_duration > 0.0: + try: + self._clock.sleep(result.wait_duration) + except Exception as exc: + return self._fail( + f"Execution clock failed: {type(exc).__name__}: {exc}", + context=result.context or self._last_context, + tick=result.tick, + dispatches=list(result.dispatches), + ) + return self._fail( + f"Execution runner exceeded max_steps={max_steps}.", + context=last_result.context or self._last_context, + tick=last_result.tick, + dispatches=list(last_result.dispatches), + ) + + def _update_effect_boundary( + self, + context: PlanningContext, + tick: ExecutionTick, + effect_success: torch.Tensor | None, + ) -> None: + """Remember or clear the external effect-verification boundary.""" + verification_required = any( + event.kind is ExecutionEventKind.EFFECT_VERIFICATION_REQUIRED + for event in tick.events + ) + if verification_required: + self._effect_verification_pending = True + self._effect_context = context + self._effect_tick = tick + elif effect_success is not None and self._effect_verification_pending: + self._clear_effect_boundary() + + def _clear_effect_boundary(self) -> None: + """Clear a remembered external effect-verification boundary.""" + self._effect_verification_pending = False + self._effect_context = None + self._effect_tick = None + + def _clock_now(self) -> float: + """Read and validate the injected monotonic clock.""" + value = float(self._clock.now()) + if not math.isfinite(value) or value < 0.0: + raise ValueError("ExecutionClock.now() must be finite and non-negative.") + return value + + def _command_interval(self, command: JointCommand) -> float: + """Resolve a synchronized batch interval from per-environment durations.""" + durations = ( + command.hold_duration[command.active_mask] + if command.active_mask.any() + else command.hold_duration + ) + requested = float(durations.max().item()) if durations.numel() else 0.0 + return max(requested, self.cfg.minimum_cycle_time) + + def _remaining_wait(self, now: float) -> float: + """Return scheduled wait while absorbing float32 timing roundoff.""" + remaining = self._next_step_at - now + tolerance = max(1.0e-9, self.cfg.minimum_cycle_time * 1.0e-6) + return remaining if remaining > tolerance else 0.0 + + def _dispatch( + self, + operation: CommandOperation, + command: JointCommand | None, + ) -> CommandDispatch: + """Call one sink operation and convert exceptions to rejection acks.""" + try: + if operation is CommandOperation.SEND: + if command is None: + raise ValueError("SEND requires a JointCommand.") + acknowledgement = self._command_sink.send( + command, + timeout=self.cfg.command_timeout, + ) + elif operation is CommandOperation.HOLD: + if command is None: + raise ValueError("HOLD requires a JointCommand.") + acknowledgement = self._command_sink.hold( + command, + timeout=self.cfg.safe_stop_timeout, + ) + else: + acknowledgement = self._command_sink.cancel( + timeout=self.cfg.safe_stop_timeout + ) + if not isinstance(acknowledgement, CommandAcknowledgement): + raise TypeError( + "CommandSink methods must return CommandAcknowledgement." + ) + except Exception as exc: + acknowledgement = CommandAcknowledgement( + CommandAckStatus.REJECTED, + f"{type(exc).__name__}: {exc}", + ) + return CommandDispatch(operation, acknowledgement) + + def _observe_for_stop(self) -> PlanningContext | None: + """Best-effort observation used to build a cancellation hold command.""" + try: + context = self._observation_provider.observe(self._session.task_state) + if not isinstance(context, PlanningContext): + return self._last_context + self._last_context = context + return context + except Exception: + return self._last_context + + def _safe_stop( + self, + context: PlanningContext | None, + ) -> list[CommandDispatch]: + """Attempt controller cancellation followed by an observed-position hold.""" + dispatches = [self._dispatch(CommandOperation.CANCEL, None)] + if context is not None: + dispatches.append( + self._dispatch(CommandOperation.HOLD, self._hold_command(context)) + ) + return dispatches + + @staticmethod + def _hold_command(context: PlanningContext) -> JointCommand: + """Build an all-environment passive hold command from an observation.""" + return JointCommand( + positions=context.robot.qpos, + velocities=torch.zeros_like(context.robot.qpos), + active_mask=torch.zeros( + context.batch_size, + dtype=torch.bool, + device=context.robot.qpos.device, + ), + env_ids=context.env_ids, + hold_duration=torch.zeros( + context.batch_size, + dtype=torch.float32, + device=context.robot.qpos.device, + ), + ) + + def _fail( + self, + message: str, + *, + context: PlanningContext | None, + tick: ExecutionTick | None = None, + dispatches: list[CommandDispatch] | None = None, + ) -> RunnerStep: + """Enter failed state after a best-effort cancel-then-hold sequence.""" + records = list(dispatches or ()) + records.extend(self._safe_stop(context)) + self._status = RunnerStatus.FAILED + self._message = message + self._clear_effect_boundary() + self._next_step_at = self._clock_now() + return self._result( + timestamp=self._clock_now(), + context=context, + tick=tick, + dispatches=records, + ) + + def _result( + self, + *, + timestamp: float, + wait_duration: float = 0.0, + context: PlanningContext | None = None, + tick: ExecutionTick | None = None, + dispatches: list[CommandDispatch] | tuple[CommandDispatch, ...] = (), + ) -> RunnerStep: + """Build an immutable runner result.""" + return RunnerStep( + status=self._status, + timestamp=timestamp, + wait_duration=wait_duration, + context=context, + tick=tick, + dispatches=tuple(dispatches), + command_count=self._command_count, + message=self._message, + ) + + +__all__ = [ + "CommandAckStatus", + "CommandAcknowledgement", + "CommandDispatch", + "CommandOperation", + "CommandSink", + "EffectVerifier", + "ExecutionClock", + "ExecutionRunner", + "ExecutionRunnerCfg", + "MonotonicExecutionClock", + "ObservationProvider", + "RunnerStatus", + "RunnerStep", + "RunnerStepCallback", +] diff --git a/embodichain/lab/sim/atomic_actions/sim_adapter.py b/embodichain/lab/sim/atomic_actions/sim_adapter.py new file mode 100644 index 000000000..20f43922f --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/sim_adapter.py @@ -0,0 +1,293 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Simulation ports for :class:`~.runner.ExecutionRunner`.""" + +from __future__ import annotations + +from collections.abc import Callable +import math +from typing import TYPE_CHECKING + +import torch + +from .execution import JointCommand +from .runner import ( + CommandAcknowledgement, + CommandAckStatus, +) +from .state import PlanningContext, RobotObservation, SceneSnapshot, TaskState + +if TYPE_CHECKING: + from embodichain.lab.sim.objects import Robot + from embodichain.lab.sim.sim_manager import SimulationManager + + +SceneSnapshotSupplier = Callable[[float], SceneSnapshot] +"""Callback that returns the latest scene snapshot for a simulation timestamp.""" + + +class SimulationExecutionAdapter: + """Adapt a simulation robot to observation, command, and clock protocols. + + The adapter writes joint targets synchronously. Time advances only through + :meth:`sleep`, which converts the requested runner interval to an integral + number of physics updates. This makes :meth:`ExecutionRunner.run_until_blocked` + deterministic and avoids wall-clock sleeps in headless simulation. + + Args: + simulation: Simulation manager advanced by the execution clock. + robot: Robot observed and commanded by the adapter. + physics_dt: Optional physics period. Defaults to the simulation config. + env_ids: Optional stable correlation IDs matching every robot row. They + are not used as simulator indices; row order maps to robot instances. + scene_supplier: Optional callback for versioned scene observations. + initial_time: Initial elapsed simulation time in seconds. + """ + + def __init__( + self, + simulation: SimulationManager, + robot: Robot, + *, + physics_dt: float | None = None, + env_ids: torch.Tensor | None = None, + scene_supplier: SceneSnapshotSupplier | None = None, + initial_time: float = 0.0, + ) -> None: + if not math.isfinite(initial_time) or initial_time < 0.0: + raise ValueError("initial_time must be finite and non-negative.") + resolved_physics_dt = ( + float(simulation.sim_config.physics_dt) + if physics_dt is None + else float(physics_dt) + ) + if not math.isfinite(resolved_physics_dt) or resolved_physics_dt <= 0.0: + raise ValueError("physics_dt must be finite and greater than zero.") + qpos = robot.get_qpos() + if not isinstance(qpos, torch.Tensor) or qpos.dim() != 2: + raise ValueError("robot.get_qpos() must return shape (B, robot_dof).") + if env_ids is None: + env_ids = torch.arange(qpos.shape[0], dtype=torch.long, device=qpos.device) + if ( + not isinstance(env_ids, torch.Tensor) + or env_ids.dtype != torch.long + or env_ids.shape != (qpos.shape[0],) + ): + raise ValueError("env_ids must be int64 with one ID per robot row.") + if env_ids.device != qpos.device: + raise ValueError("env_ids and robot state must share a device.") + if torch.unique(env_ids).numel() != env_ids.numel(): + raise ValueError("env_ids must be unique.") + + self.simulation = simulation + self.robot = robot + self.physics_dt = resolved_physics_dt + self.env_ids = env_ids.clone() + self._robot_env_indices = list(range(qpos.shape[0])) + self.scene_supplier = scene_supplier + self._elapsed_time = float(initial_time) + + def now(self) -> float: + """Return elapsed simulation time in seconds. + + Returns: + Elapsed simulation time in seconds. + """ + return self._elapsed_time + + def sleep(self, duration: float) -> None: + """Advance physics by at least the requested duration. + + Args: + duration: Requested simulated duration in seconds. + """ + if not math.isfinite(duration) or duration < 0.0: + raise ValueError("duration must be finite and non-negative.") + if duration == 0.0: + return + step_count = max(1, math.ceil(duration / self.physics_dt)) + self.simulation.update(physics_dt=self.physics_dt, step=step_count) + self._elapsed_time += step_count * self.physics_dt + + def observe(self, task_state: TaskState) -> PlanningContext: + """Capture full-robot state and the latest supplied scene snapshot. + + Args: + task_state: Verified symbolic state owned by the execution session. + + Returns: + Planning context timestamped with elapsed simulation time. + """ + qpos = self.robot.get_qpos() + qvel = self._read_optional_tensor("get_qvel") + if qvel is None: + qvel = torch.zeros_like(qpos) + qeffort = self._read_optional_tensor("get_qf") + if qeffort is None: + qeffort = self._read_optional_proprioception_tensor("qf") + scene = ( + SceneSnapshot(timestamp=self._elapsed_time, version=0) + if self.scene_supplier is None + else self.scene_supplier(self._elapsed_time) + ) + if not isinstance(scene, SceneSnapshot): + raise TypeError("scene_supplier must return a SceneSnapshot.") + return PlanningContext( + robot=RobotObservation( + timestamp=self._elapsed_time, + qpos=qpos, + qvel=qvel, + qeffort=qeffort, + ), + task=task_state, + scene=scene, + env_ids=self.env_ids, + ) + + def send( + self, + command: JointCommand, + *, + timeout: float, + ) -> CommandAcknowledgement: + """Write active targets and observed-position holds as one batch. + + Args: + command: Full-robot batched command. Inactive rows already contain + observed-position holds and are written with active rows so no + environment continues tracking a stale target. + timeout: Positive acknowledgement deadline. Simulation writes are + synchronous, so this is validated but otherwise unused. + + Returns: + Accepted acknowledgement or a rejected diagnostic. + """ + self._validate_timeout(timeout) + try: + self._validate_command(command) + self.robot.set_qpos( + command.positions, + env_ids=self._robot_env_indices, + ) + if command.velocities is not None: + self.robot.set_qvel( + command.velocities, + env_ids=self._robot_env_indices, + ) + return CommandAcknowledgement.accepted_ack() + except Exception as exc: + return CommandAcknowledgement( + CommandAckStatus.REJECTED, + f"{type(exc).__name__}: {exc}", + ) + + def hold( + self, + command: JointCommand, + *, + timeout: float, + ) -> CommandAcknowledgement: + """Set every represented environment to an observed-position hold. + + Args: + command: Full-robot hold positions. ``active_mask`` is intentionally + ignored because safety hold applies to every environment row. + timeout: Positive acknowledgement deadline. + + Returns: + Accepted acknowledgement or a rejected diagnostic. + """ + self._validate_timeout(timeout) + try: + self._validate_command(command) + self.robot.set_qpos( + command.positions, + env_ids=self._robot_env_indices, + ) + if command.velocities is not None: + self.robot.set_qvel( + command.velocities, + env_ids=self._robot_env_indices, + ) + return CommandAcknowledgement.accepted_ack() + except Exception as exc: + return CommandAcknowledgement( + CommandAckStatus.REJECTED, + f"{type(exc).__name__}: {exc}", + ) + + def cancel(self, *, timeout: float) -> CommandAcknowledgement: + """Acknowledge cancellation of synchronous simulation target writes. + + Args: + timeout: Positive acknowledgement deadline. + + Returns: + Accepted acknowledgement. The following ``hold`` call installs the + actual safe target. + """ + self._validate_timeout(timeout) + return CommandAcknowledgement.accepted_ack( + "Simulation commands are synchronous; no queued command remained." + ) + + def _read_optional_tensor(self, method_name: str) -> torch.Tensor | None: + """Read an optional full-robot tensor from the robot API.""" + method = getattr(self.robot, method_name, None) + if not callable(method): + return None + try: + value = method() + except (AttributeError, NotImplementedError): + return None + return value if isinstance(value, torch.Tensor) else None + + def _read_optional_proprioception_tensor( + self, + field_name: str, + ) -> torch.Tensor | None: + """Read an optional tensor from the robot proprioception mapping.""" + method = getattr(self.robot, "get_proprioception", None) + if not callable(method): + return None + try: + value = method()[field_name] + except (AttributeError, KeyError, NotImplementedError, TypeError): + return None + return value if isinstance(value, torch.Tensor) else None + + def _validate_command(self, command: JointCommand) -> None: + """Validate command identity and shape against the attached robot.""" + if not isinstance(command, JointCommand): + raise TypeError("command must be a JointCommand.") + qpos = self.robot.get_qpos() + if command.positions.shape != qpos.shape: + raise ValueError( + "Command shape must match full robot qpos, " + f"got {tuple(command.positions.shape)} and {tuple(qpos.shape)}." + ) + if not torch.equal(command.env_ids, self.env_ids): + raise ValueError("Command env_ids must match the simulation adapter.") + + @staticmethod + def _validate_timeout(timeout: float) -> None: + """Validate an acknowledgement timeout.""" + if not math.isfinite(timeout) or timeout <= 0.0: + raise ValueError("timeout must be finite and greater than zero.") + + +__all__ = ["SceneSnapshotSupplier", "SimulationExecutionAdapter"] diff --git a/scripts/tutorials/atomic_action/moving_target_recovery.py b/scripts/tutorials/atomic_action/moving_target_recovery.py new file mode 100644 index 000000000..0fbde22b0 --- /dev/null +++ b/scripts/tutorials/atomic_action/moving_target_recovery.py @@ -0,0 +1,472 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Replan one PickUp action after its visible target moves.""" + +from __future__ import annotations + +import argparse +import math +import sys +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[3] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +import torch + +from embodichain.lab.sim import SimulationManager, VisualMaterialCfg +from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, + Affordance, + AtomicActionEngine, + ControlPartCommandProfile, + EntityState, + ExecutionEventKind, + ExecutionRunner, + ExecutionRunnerCfg, + ExecutionTick, + GraspGoal, + MotionPolicy, + ObjectSemantics, + PickUpOptions, + PlanningContext, + RecoveryPolicy, + RunnerStatus, + RunnerStep, + SceneEntityPose, + SceneSnapshot, + SimulationExecutionAdapter, + TaskState, +) +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg +from embodichain.lab.sim.objects import RigidObject +from embodichain.lab.sim.shapes import CubeCfg +from embodichain.utils import logger +from scripts.tutorials.atomic_action.tutorial_utils import ( + add_ur5_gripper_robot, + create_toppra_motion_generator, + create_tutorial_argument_parser, + create_tutorial_simulation, + draw_axis_marker, + get_hand_open_close_qpos, + initialize_pre_pick_robot_pose, + make_top_down_eef_pose, + prepare_tutorial_scene, + run_tutorial, + serve_tutorial_scene, + start_auto_play_recording, + stop_auto_play_recording, +) + +TARGET_ENTITY_ID = "moving_target" +TARGET_SIZE = (0.05, 0.05, 0.05) +INITIAL_TARGET_POSITION = (-0.42, -0.18, 0.5 * TARGET_SIZE[2]) +MOVED_TARGET_POSITION = (-0.42, 0.12, 0.5 * TARGET_SIZE[2]) +PICK_SAMPLE_COUNT = 120 +HAND_INTERP_STEPS = 12 +PICK_LIFT_HEIGHT = 0.16 +MINIMUM_LIFT_HEIGHT = 0.08 +MAXIMUM_HELD_DISTANCE = 0.10 +MOVE_AFTER_COMMAND = 20 +TARGET_MOVE_DURATION = 0.6 +TARGET_PUSH_DURATION = 0.12 +TARGET_PUSH_FORCE = 1.25 +GOAL_TRANSLATION_THRESHOLD = 0.04 +TRACKING_ERROR_THRESHOLD = 1.0 +POST_EXECUTION_UPDATES = 120 + + +class _MovingTargetScene: + """Publish a versioned target pose and physically push it exactly once.""" + + def __init__( + self, + target: RigidObject, + destination: tuple[float, float, float], + ) -> None: + self.target = target + self.destination = torch.tensor( + destination, + dtype=torch.float32, + device=target.device, + ) + self.version = 0 + self.moved = False + + def snapshot(self, timestamp: float) -> SceneSnapshot: + """Return the target pose used to ground the late-bound goal. + + Args: + timestamp: Current elapsed simulation time. + + Returns: + Versioned scene snapshot containing the target pose. + """ + return SceneSnapshot( + timestamp=timestamp, + version=self.version, + entities={ + TARGET_ENTITY_ID: EntityState( + self.target.get_local_pose(to_matrix=True) + ) + }, + ) + + def push( + self, + clock: SimulationExecutionAdapter, + *, + duration: float, + force_duration: float, + force_magnitude: float, + ) -> torch.Tensor: + """Push the visible target with a short force pulse. + + Args: + clock: Simulation adapter used to advance physics. + duration: Total time allowed for the push and natural deceleration. + force_duration: Time spent applying the horizontal force. + force_magnitude: Magnitude of the applied force in newtons. + + Returns: + Batched target pose after the physical motion. + """ + if self.moved: + return self.target.get_local_pose(to_matrix=True) + if not math.isfinite(duration) or duration <= 0.0: + raise ValueError("duration must be finite and greater than zero.") + if ( + not math.isfinite(force_duration) + or force_duration <= 0.0 + or force_duration > duration + ): + raise ValueError( + "force_duration must be finite, greater than zero, and no " + "greater than duration." + ) + if not math.isfinite(force_magnitude) or force_magnitude <= 0.0: + raise ValueError("force_magnitude must be finite and greater than zero.") + + start_pose = self.target.get_local_pose(to_matrix=True).clone() + planar_offset = self.destination - start_pose[:, :3, 3] + planar_offset[:, 2] = 0.0 + planar_distance = torch.linalg.vector_norm(planar_offset, dim=1) + if torch.any(planar_distance <= torch.finfo(planar_offset.dtype).eps): + raise ValueError("destination must differ from the current planar pose.") + force = force_magnitude * planar_offset / planar_distance.unsqueeze(-1) + + self.target.set_body_type("dynamic") + self.target.clear_dynamics() + step_count = max(1, math.ceil(duration / clock.physics_dt)) + force_step_count = min( + step_count, + max(1, math.ceil(force_duration / clock.physics_dt)), + ) + for step_index in range(step_count): + if step_index < force_step_count: + self.target.add_force_torque(force=force) + clock.sleep(clock.physics_dt) + self.target.clear_dynamics() + + pose = self.target.get_local_pose(to_matrix=True) + self.version += 1 + self.moved = True + return pose + + +def _create_moving_target(sim: SimulationManager) -> RigidObject: + """Create the bright cube, held kinematic until the physical push.""" + return sim.add_rigid_object( + cfg=RigidObjectCfg( + uid=TARGET_ENTITY_ID, + shape=CubeCfg( + size=list(TARGET_SIZE), + visual_material=VisualMaterialCfg( + uid="moving_target_blue", + base_color=[0.05, 0.30, 1.0, 1.0], + metallic=0.15, + roughness=0.3, + ), + ), + attrs=RigidBodyAttributesCfg( + mass=0.05, + dynamic_friction=0.97, + static_friction=0.99, + enable_ccd=True, + ), + body_type="kinematic", + max_convex_hull_num=16, + init_pos=INITIAL_TARGET_POSITION, + ) + ) + + +def _compose_goal_pose( + target_pose: torch.Tensor, + relative_pose: torch.Tensor, +) -> torch.Tensor: + """Compose batched target poses with the target-to-EEF transform.""" + relative_batch = relative_pose.unsqueeze(0).expand(target_pose.shape[0], -1, -1) + return torch.bmm(target_pose, relative_batch) + + +def parse_arguments() -> argparse.Namespace: + """Parse command-line arguments for the moving-target tutorial.""" + parser = create_tutorial_argument_parser( + "Demonstrate ExecutionRunner replanning after a visible target move." + ) + parser.add_argument( + "--no_target_motion", + action="store_true", + help="Keep the target fixed to run the no-replanning control case.", + ) + return parser.parse_args() + + +def main() -> None: + """Replan a late-bound PickUp request and lift the relocated cube.""" + args = parse_arguments() + sim = create_tutorial_simulation(args) + robot = add_ur5_gripper_robot(sim) + target = _create_moving_target(sim) + sim.update(step=10) + target_scene = _MovingTargetScene(target, MOVED_TARGET_POSITION) + sim_runtime = SimulationExecutionAdapter( + sim, + robot, + scene_supplier=target_scene.snapshot, + ) + motion_gen = create_toppra_motion_generator(robot) + hand_open, hand_close = get_hand_open_close_qpos(robot) + initialize_pre_pick_robot_pose(robot, target, hand_open) + if args.no_target_motion: + target.set_body_type("dynamic") + target.clear_dynamics() + + target_to_grasp = make_top_down_eef_pose( + torch.zeros(3, dtype=torch.float32, device=sim.device) + ) + initial_target_pose = target.get_local_pose(to_matrix=True) + draw_axis_marker( + sim, + "moving_target_original_goal", + _compose_goal_pose(initial_target_pose, target_to_grasp), + axis_len=0.10, + ) + + semantics = ObjectSemantics( + affordance=Affordance(), + geometry={}, + label="cube", + entity=target, + ) + binding = ActionBinding( + manipulators={"primary": "arm"}, + end_effectors={"primary": "hand"}, + ) + engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=hand_open, + grasp=hand_close, + ) + }, + ) + pick_invocation = ActionInvocation( + skill_id="pick_up", + goal=GraspGoal( + semantics, + grasp_xpos=SceneEntityPose( + TARGET_ENTITY_ID, + relative_pose=target_to_grasp, + ), + ), + binding=binding, + motion_policy=MotionPolicy( + sample_count=PICK_SAMPLE_COUNT, + control_dt=2.0 * sim_runtime.physics_dt, + ), + recovery_policy=RecoveryPolicy( + max_replans=2, + max_phase_retries=1, + tracking_error_threshold=TRACKING_ERROR_THRESHOLD, + goal_translation_threshold=GOAL_TRANSLATION_THRESHOLD, + phase_timeout=30.0, + ), + skill_options=PickUpOptions( + pre_grasp_distance=0.15, + lift_height=PICK_LIFT_HEIGHT, + hand_interp_steps=HAND_INTERP_STEPS, + ), + ) + task_state = TaskState.empty(robot.get_qpos().shape[0], robot.device) + initial_context = sim_runtime.observe(task_state) + session = engine.start((pick_invocation,), initial_context) + runner = ExecutionRunner( + session=session, + observation_provider=sim_runtime, + command_sink=sim_runtime, + clock=sim_runtime, + cfg=ExecutionRunnerCfg(minimum_cycle_time=sim_runtime.physics_dt), + ) + + wait_for_user = prepare_tutorial_scene( + sim, + args, + "Watch the blue cube, then press Enter to run recovering PickUp...", + ) + observed_events: set[ExecutionEventKind] = set() + plan_start_command = 0 + pickup_dynamics_cleared = False + + clear_after_pick_command = ( + round((PICK_SAMPLE_COUNT - HAND_INTERP_STEPS) * 0.6) + HAND_INTERP_STEPS + ) + + def on_step(step: RunnerStep) -> None: + nonlocal pickup_dynamics_cleared, plan_start_command + if ( + not args.no_target_motion + and not target_scene.moved + and step.command_count >= MOVE_AFTER_COMMAND + ): + logger.log_warning( + f"Applying a {TARGET_PUSH_FORCE:.2f} N force pulse to the blue " + "target while the robot holds its current command." + ) + moved_pose = target_scene.push( + sim_runtime, + duration=TARGET_MOVE_DURATION, + force_duration=TARGET_PUSH_DURATION, + force_magnitude=TARGET_PUSH_FORCE, + ) + draw_axis_marker( + sim, + "moving_target_replanned_goal", + _compose_goal_pose(moved_pose, target_to_grasp), + axis_len=0.10, + ) + displacement = torch.linalg.vector_norm( + moved_pose[:, :3, 3] - initial_target_pose[:, :3, 3], + dim=1, + ) + logger.log_warning( + "The force pulse moved the blue target after " + f"{step.command_count} accepted commands by " + f"{displacement.detach().cpu().tolist()} m; the original goal " + "axis remains visible." + ) + if step.tick is None: + return + for event in step.tick.events: + observed_events.add(event.kind) + if event.kind in { + ExecutionEventKind.DYNAMIC_GOAL_CHANGED, + ExecutionEventKind.REPLANNED, + ExecutionEventKind.RECOVERY_EXHAUSTED, + }: + env_ids = event.env_mask.nonzero(as_tuple=False).flatten().tolist() + logger.log_info( + f"Execution event {event.kind.value}: env rows={env_ids}; " + f"{event.message}" + ) + if event.kind is ExecutionEventKind.REPLANNED: + plan_start_command = step.command_count + logger.log_info( + "PickUp discarded the stale plan and restarted from the " + "latest cube pose.", + color="green", + ) + if ( + (args.no_target_motion or target_scene.moved) + and not pickup_dynamics_cleared + and step.command_count - plan_start_command > clear_after_pick_command + ): + target.clear_dynamics() + pickup_dynamics_cleared = True + + def verify_pickup_effect( + _context: PlanningContext, + _: ExecutionTick, + ) -> torch.Tensor: + """Verify that the cube rose with, and remains near, the end effector.""" + cube_position = target.get_local_pose(to_matrix=True)[:, :3, 3] + eef_position = robot.compute_fk( + qpos=robot.get_qpos(name="arm"), + name="arm", + to_matrix=True, + )[:, :3, 3] + lift_height = cube_position[:, 2] - 0.5 * TARGET_SIZE[2] + held_distance = torch.linalg.vector_norm(cube_position - eef_position, dim=1) + success = (lift_height >= MINIMUM_LIFT_HEIGHT) & ( + held_distance <= MAXIMUM_HELD_DISTANCE + ) + logger.log_info( + "PickUp verification: " + f"lift={lift_height.detach().cpu().tolist()} m, " + f"cube-to-EEF={held_distance.detach().cpu().tolist()} m, " + f"success={success.detach().cpu().tolist()}." + ) + return success + + recording_started = start_auto_play_recording( + sim, + args, + video_prefix="moving_target_recovery_auto_play", + look_at=( + (-1.25, -1.15, 0.95), + (-0.32, -0.02, 0.25), + (0.0, 0.0, 1.0), + ), + ) + try: + result = runner.run_until_blocked( + effect_verifier=verify_pickup_effect, + on_step=on_step, + ) + for _ in range(POST_EXECUTION_UPDATES): + sim_runtime.sleep(sim_runtime.physics_dt) + finally: + stop_auto_play_recording(sim, recording_started) + + if result.status is not RunnerStatus.COMPLETED: + raise RuntimeError(f"Closed-loop execution failed: {result.message}") + if not args.no_target_motion: + if not target_scene.moved: + raise RuntimeError("Execution completed before the target could move.") + if ExecutionEventKind.DYNAMIC_GOAL_CHANGED not in observed_events: + raise RuntimeError("The target move was not reported as a dynamic change.") + if ExecutionEventKind.REPLANNED not in observed_events: + raise RuntimeError("The target move did not trigger replanning.") + if not pickup_dynamics_cleared: + raise RuntimeError("The cube was not stabilized after gripper closure.") + logger.log_info( + f"Execution completed and lifted the cube after {result.command_count} " + "accepted commands.", + color="green", + ) + + serve_tutorial_scene(sim, args) + if wait_for_user: + input("Press Enter to exit the simulation...") + + +if __name__ == "__main__": + run_tutorial(main) diff --git a/tests/sim/atomic_actions/test_actions.py b/tests/sim/atomic_actions/test_actions.py index d15386c28..bb7c9fbc8 100644 --- a/tests/sim/atomic_actions/test_actions.py +++ b/tests/sim/atomic_actions/test_actions.py @@ -41,6 +41,8 @@ CoordinatedPlacementGoal, CoordinatedPlacementOptions, EndEffectorPoseGoal, + EntityState, + ExecutionEventKind, GraspGoal, HandOver, HandOverOptions, @@ -65,6 +67,7 @@ PressGoal, PressOptions, RobotObservation, + SceneEntityPose, SceneSnapshot, TaskState, ) @@ -198,16 +201,39 @@ def _plan_action( return action.plan(action.resolve_request(invocation), context) -def _context(task: TaskState | None = None) -> PlanningContext: +def _context( + task: TaskState | None = None, + *, + scene: SceneSnapshot | None = None, + timestamp: float = 0.0, +) -> PlanningContext: qpos = torch.zeros(NUM_ENVS, ROBOT_DOF) return PlanningContext( - robot=RobotObservation(timestamp=0.0, qpos=qpos, qvel=torch.zeros_like(qpos)), + robot=RobotObservation( + timestamp=timestamp, + qpos=qpos, + qvel=torch.zeros_like(qpos), + ), task=task or TaskState.empty(batch_size=NUM_ENVS, device="cpu"), - scene=SceneSnapshot.empty(), + scene=SceneSnapshot.empty() if scene is None else scene, env_ids=torch.arange(NUM_ENVS), ) +def _target_scene( + pose: torch.Tensor, + *, + timestamp: float, + version: int, +) -> SceneSnapshot: + """Build a versioned target snapshot for late-bound grasp tests.""" + return SceneSnapshot( + timestamp=timestamp, + version=version, + entities={"target": EntityState(pose)}, + ) + + def _binding() -> ActionBinding: return ActionBinding( manipulators={"primary": "arm"}, @@ -694,6 +720,101 @@ def test_pick_explicit_grasp_bypasses_sampling_and_records_grasp() -> None: assert torch.allclose(held.grasp_xpos, grasp) +def test_pick_resolves_late_bound_scene_grasp_and_declares_dependency() -> None: + generator = _motion_generator() + target_pose = torch.eye(4).repeat(NUM_ENVS, 1, 1) + target_pose[:, 0, 3] = torch.tensor([0.1, 0.2]) + relative_pose = torch.eye(4) + relative_pose[2, 3] = 0.05 + entity = Mock() + entity.get_local_pose.return_value = target_pose + semantics = ObjectSemantics( + affordance=Affordance(), + geometry={}, + label="late-bound-grasp-object", + entity=entity, + ) + action = _bind_action(generator, PickUp()) + context = _context(scene=_target_scene(target_pose, timestamp=0.0, version=0)) + + plan = _plan_action( + action, + _invocation( + "pick_up", + GraspGoal( + semantics=semantics, + grasp_xpos=SceneEntityPose( + "target", + relative_pose=relative_pose, + ), + ), + sample_count=20, + ), + context, + ) + projected = plan.expected_effects.apply(context.task, plan.plan_success) + expected_grasp = torch.bmm( + target_pose, + relative_pose.unsqueeze(0).expand(NUM_ENVS, -1, -1), + ) + + held = projected.get_held_object("arm") + assert held is not None + assert torch.allclose(held.grasp_xpos, expected_grasp) + assert plan.phases[0].spec.scene_dependencies == ("target",) + + +def test_pick_session_replans_when_late_bound_target_moves() -> None: + generator = _motion_generator() + initial_pose = torch.eye(4).repeat(NUM_ENVS, 1, 1) + moved_pose = initial_pose.clone() + moved_pose[:, 1, 3] = 0.3 + entity = Mock() + entity.get_local_pose.return_value = initial_pose + semantics = ObjectSemantics( + affordance=Affordance(), + geometry={}, + label="moving-grasp-object", + entity=entity, + ) + engine = AtomicActionEngine( + generator, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=torch.zeros(HAND_DOF), + grasp=torch.ones(HAND_DOF), + ) + }, + load_builtins=False, + ) + engine.register(PickUp()) + invocation = _invocation( + "pick_up", + GraspGoal( + semantics=semantics, + grasp_xpos=SceneEntityPose("target"), + ), + sample_count=20, + ) + initial_context = _context( + scene=_target_scene(initial_pose, timestamp=0.0, version=0) + ) + session = engine.start((invocation,), initial_context) + session.tick(initial_context) + entity.get_local_pose.return_value = moved_pose + + recovered = session.tick( + _context( + scene=_target_scene(moved_pose, timestamp=0.1, version=1), + timestamp=0.1, + ) + ) + + event_kinds = {event.kind for event in recovered.events} + assert ExecutionEventKind.DYNAMIC_GOAL_CHANGED in event_kinds + assert ExecutionEventKind.REPLANNED in event_kinds + + def test_pick_uses_binding_control_part_as_effect_resource() -> None: generator = _motion_generator() action = _bind_action(generator, PickUp()) diff --git a/tests/sim/atomic_actions/test_engine_per_env.py b/tests/sim/atomic_actions/test_engine_per_env.py index 588a9a5ca..4d12c3253 100644 --- a/tests/sim/atomic_actions/test_engine_per_env.py +++ b/tests/sim/atomic_actions/test_engine_per_env.py @@ -261,7 +261,7 @@ def test_session_completes_incremental_command_sequence() -> None: assert final.eligible_mask.tolist() == [True] -def test_session_commands_preserve_waypoint_arrival_intervals() -> None: +def test_session_commands_schedule_arrivals_and_final_settling() -> None: engine, _ = _engine() engine.register(NonuniformTimingAction()) session = engine.start( @@ -284,8 +284,8 @@ def test_session_commands_preserve_waypoint_arrival_intervals() -> None: ], dim=1, ) - assert torch.allclose(command_durations, torch.tensor([[0.0, 0.1, 0.3]])) - assert torch.allclose(command_durations.sum(dim=1), torch.tensor([0.4])) + assert torch.allclose(command_durations, torch.tensor([[0.1, 0.3, 0.3]])) + assert torch.allclose(command_durations[:, :-1].sum(dim=1), torch.tensor([0.4])) def test_request_snapshot_preserves_live_entity_identity() -> None: diff --git a/tests/sim/atomic_actions/test_runner.py b/tests/sim/atomic_actions/test_runner.py new file mode 100644 index 000000000..6cf1a0099 --- /dev/null +++ b/tests/sim/atomic_actions/test_runner.py @@ -0,0 +1,484 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Tests for controller-independent atomic-action execution scheduling.""" + +from __future__ import annotations + +from collections import deque +from typing import ClassVar +from unittest.mock import Mock + +import pytest +import torch + +from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, + ActionOptions, + ActionPlan, + Affordance, + AtomicAction, + AtomicActionEngine, + CommandAcknowledgement, + CommandAckStatus, + CommandOperation, + EndEffectorPoseGoal, + ExecutionEventKind, + ExecutionRunner, + ExecutionRunnerCfg, + HeldObjectState, + JointCommand, + MotionPolicy, + ObjectSemantics, + PlanningContext, + RecoveryPolicy, + ResolvedActionRequest, + RobotObservation, + RunnerStatus, + SceneSnapshot, + StateDelta, + TaskState, + TimedTrajectory, +) + +BATCH_SIZE = 1 +ROBOT_DOF = 2 +FIRST_INTERVAL = 0.1 +SECOND_INTERVAL = 0.2 +MINIMUM_CYCLE_TIME = 0.01 +TARGET_POSITION = 1.0 + + +class FakeClock: + """Deterministic clock used by non-blocking runner tests.""" + + def __init__(self) -> None: + self.time = 0.0 + self.sleeps: list[float] = [] + + def now(self) -> float: + """Return deterministic time.""" + return self.time + + def sleep(self, duration: float) -> None: + """Advance deterministic time.""" + self.sleeps.append(duration) + self.time += duration + + def advance(self, duration: float) -> None: + """Advance time outside the runner's blocking loop.""" + self.time += duration + + +class FakeObservationProvider: + """In-memory robot observation provider.""" + + def __init__(self, clock: FakeClock, batch_size: int = BATCH_SIZE) -> None: + self.clock = clock + self.qpos = torch.zeros(batch_size, ROBOT_DOF) + self.fail = False + + def observe(self, task_state: TaskState) -> PlanningContext: + """Return the current in-memory robot state.""" + if self.fail: + raise RuntimeError("observation unavailable") + return PlanningContext( + robot=RobotObservation( + timestamp=self.clock.now(), + qpos=self.qpos, + qvel=torch.zeros_like(self.qpos), + ), + task=task_state, + scene=SceneSnapshot(timestamp=self.clock.now(), version=0), + env_ids=torch.arange(self.qpos.shape[0], dtype=torch.long), + ) + + +class FakeCommandSink: + """Recording command sink with configurable acknowledgements and tracking.""" + + def __init__(self, provider: FakeObservationProvider) -> None: + self.provider = provider + self.send_statuses: deque[CommandAckStatus] = deque() + self.follow_commands: deque[bool] = deque() + self.sent: list[JointCommand] = [] + self.send_times: list[float] = [] + self.held: list[JointCommand] = [] + self.cancel_count = 0 + + def send( + self, + command: JointCommand, + *, + timeout: float, + ) -> CommandAcknowledgement: + """Record an active command and optionally update observed qpos.""" + self.sent.append(command) + self.send_times.append(self.provider.clock.now()) + status = ( + self.send_statuses.popleft() + if self.send_statuses + else CommandAckStatus.ACCEPTED + ) + follows = self.follow_commands.popleft() if self.follow_commands else True + if status is CommandAckStatus.ACCEPTED and follows: + self.provider.qpos = command.positions.clone() + return CommandAcknowledgement(status) + + def hold( + self, + command: JointCommand, + *, + timeout: float, + ) -> CommandAcknowledgement: + """Record and apply a hold command.""" + self.held.append(command) + self.provider.qpos = command.positions.clone() + return CommandAcknowledgement.accepted_ack() + + def cancel(self, *, timeout: float) -> CommandAcknowledgement: + """Record controller cancellation.""" + self.cancel_count += 1 + return CommandAcknowledgement.accepted_ack() + + +class TimedAction(AtomicAction[EndEffectorPoseGoal, ActionOptions]): + """Test action with explicit non-uniform command intervals.""" + + skill_id: ClassVar[str] = "timed" + GoalType: ClassVar[type] = EndEffectorPoseGoal + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) + + def __init__(self, *, with_effect: bool = False) -> None: + super().__init__() + self.with_effect = with_effect + self.plan_count = 0 + + def plan( + self, + request: ResolvedActionRequest[EndEffectorPoseGoal, ActionOptions], + context: PlanningContext, + ) -> ActionPlan: + """Plan three samples with intervals 0.1 s and 0.2 s.""" + goal = self.require_goal(request) + self.plan_count += 1 + assert isinstance(goal.xpos, torch.Tensor) + target_value = float(goal.xpos[0, 3]) + target = torch.full_like(context.robot.qpos, target_value) + midpoint = torch.lerp(context.robot.qpos, target, 0.5) + positions = torch.stack([context.robot.qpos, midpoint, target], dim=1) + dt = torch.tensor( + [[0.0, FIRST_INTERVAL, SECOND_INTERVAL]], + dtype=torch.float32, + ).repeat(context.batch_size, 1) + if context.batch_size > 1: + dt[1, 1:] *= 2.0 + trajectory = TimedTrajectory.from_positions( + positions, + env_ids=context.env_ids, + control_dt=request.motion_policy.control_dt, + dt=dt, + ) + effects = StateDelta() + if self.with_effect: + semantics = ObjectSemantics( + affordance=Affordance(), geometry={}, label="runner-object" + ) + held = HeldObjectState( + semantics=semantics, + object_to_eef=torch.eye(4), + grasp_xpos=torch.eye(4), + ) + effects = StateDelta(held_object_updates={"arm": held}) + return self.build_plan( + request, + context, + success=True, + trajectory=trajectory, + expected_effects=effects, + ) + + +def _make_runner( + *, + with_effect: bool = False, + batch_size: int = BATCH_SIZE, +) -> tuple[ + ExecutionRunner, + FakeClock, + FakeObservationProvider, + FakeCommandSink, + TimedAction, +]: + clock = FakeClock() + provider = FakeObservationProvider(clock, batch_size) + sink = FakeCommandSink(provider) + robot = Mock() + robot.device = torch.device("cpu") + robot.dof = ROBOT_DOF + robot.control_parts = {"arm": object()} + robot.get_qpos.return_value = torch.zeros(batch_size, ROBOT_DOF) + robot.get_joint_ids.return_value = list(range(ROBOT_DOF)) + generator = Mock() + generator.robot = robot + generator.device = torch.device("cpu") + generator.planner.cfg.planner_type = "stub" + action = TimedAction(with_effect=with_effect) + engine = AtomicActionEngine(generator) + engine.register(action) + initial_task = TaskState.empty(batch_size, "cpu") + initial_context = provider.observe(initial_task) + goal_pose = torch.eye(4) + goal_pose[0, 3] = TARGET_POSITION + invocation = ActionInvocation( + skill_id="timed", + goal=EndEffectorPoseGoal(goal_pose), + binding=ActionBinding(manipulators={"primary": "arm"}), + motion_policy=MotionPolicy(sample_count=3, control_dt=FIRST_INTERVAL), + recovery_policy=RecoveryPolicy( + max_replans=2, + tracking_error_threshold=0.05, + phase_timeout=10.0, + ), + ) + session = engine.start((invocation,), initial_context) + runner = ExecutionRunner( + session, + provider, + sink, + clock=clock, + cfg=ExecutionRunnerCfg(minimum_cycle_time=MINIMUM_CYCLE_TIME), + ) + return runner, clock, provider, sink, action + + +def test_runner_dispatches_only_when_timed_waypoint_is_due() -> None: + runner, clock, _, sink, _ = _make_runner() + + first = runner.step() + early = runner.step() + clock.advance(FIRST_INTERVAL) + second = runner.step() + clock.advance(SECOND_INTERVAL) + third = runner.step() + + assert first.command_count == 1 + assert first.wait_duration == pytest.approx(FIRST_INTERVAL) + assert early.is_waiting + assert len(sink.sent) == 3 + assert sink.send_times == pytest.approx( + [0.0, FIRST_INTERVAL, FIRST_INTERVAL + SECOND_INTERVAL] + ) + assert second.command_count == 2 + assert second.wait_duration == pytest.approx(SECOND_INTERVAL) + assert third.command_count == 3 + assert third.wait_duration == pytest.approx(SECOND_INTERVAL) + + +def test_runner_uses_the_longest_active_batch_interval_as_a_barrier() -> None: + runner, clock, _, _, _ = _make_runner(batch_size=2) + + first = runner.step() + clock.advance(2.0 * FIRST_INTERVAL) + second = runner.step() + + assert first.wait_duration == pytest.approx(2.0 * FIRST_INTERVAL) + assert second.wait_duration == pytest.approx(2.0 * SECOND_INTERVAL) + + +def test_runner_completes_and_holds_after_last_command_settles() -> None: + runner, clock, _, sink, _ = _make_runner() + + runner.step() + clock.advance(FIRST_INTERVAL) + runner.step() + clock.advance(SECOND_INTERVAL) + runner.step() + clock.advance(SECOND_INTERVAL) + completed = runner.step() + + assert completed.status is RunnerStatus.COMPLETED + assert completed.command_count == 3 + assert [item.operation for item in completed.dispatches] == [CommandOperation.HOLD] + assert len(sink.held) == 1 + + +@pytest.mark.parametrize( + "status", + [CommandAckStatus.REJECTED, CommandAckStatus.TIMED_OUT], +) +def test_runner_safely_stops_when_command_is_not_accepted( + status: CommandAckStatus, +) -> None: + runner, _, _, sink, _ = _make_runner() + sink.send_statuses.append(status) + + failed = runner.step() + + assert failed.status is RunnerStatus.FAILED + assert [item.operation for item in failed.dispatches] == [ + CommandOperation.SEND, + CommandOperation.CANCEL, + CommandOperation.HOLD, + ] + assert sink.cancel_count == 1 + assert failed.message is not None and status.value in failed.message + + +def test_runner_cancel_performs_cancel_then_hold() -> None: + runner, _, _, sink, _ = _make_runner() + + cancelled = runner.cancel("operator stop") + repeated = runner.step() + + assert cancelled.status is RunnerStatus.CANCELLED + assert [item.operation for item in cancelled.dispatches] == [ + CommandOperation.CANCEL, + CommandOperation.HOLD, + ] + assert cancelled.message == "operator stop" + assert repeated.status is RunnerStatus.CANCELLED + assert repeated.dispatches == () + assert sink.cancel_count == 1 + + +def test_runner_replans_from_observation_after_tracking_error() -> None: + runner, clock, _, sink, action = _make_runner() + sink.follow_commands.extend([True, False, True]) + + runner.step() + clock.advance(FIRST_INTERVAL) + runner.step() + clock.advance(SECOND_INTERVAL) + recovered = runner.step() + + assert action.plan_count == 2 + assert recovered.tick is not None + event_kinds = {event.kind for event in recovered.tick.events} + assert ExecutionEventKind.TRACKING_ERROR in event_kinds + assert ExecutionEventKind.REPLANNED in event_kinds + assert recovered.status is RunnerStatus.RUNNING + + +def test_runner_surfaces_explicit_invocation_revision() -> None: + runner, _, _, _, action = _make_runner() + revised_pose = torch.eye(4) + revised_pose[0, 3] = 2.0 * TARGET_POSITION + revised = ActionInvocation( + skill_id="timed", + goal=EndEffectorPoseGoal(revised_pose), + binding=ActionBinding(manipulators={"primary": "arm"}), + motion_policy=MotionPolicy(sample_count=3, control_dt=FIRST_INTERVAL), + recovery_policy=RecoveryPolicy( + max_replans=2, + tracking_error_threshold=0.05, + phase_timeout=10.0, + ), + revision=1, + ) + + runner.session.revise_current(revised) + result = runner.step() + + assert action.plan_count == 2 + assert result.tick is not None + assert any( + event.kind is ExecutionEventKind.INVOCATION_REVISED + and event.invocation_revision == 1 + for event in result.tick.events + ) + + +def test_runner_fails_safely_when_observation_provider_raises() -> None: + runner, _, provider, sink, _ = _make_runner() + provider.fail = True + + failed = runner.step() + + assert failed.status is RunnerStatus.FAILED + assert [item.operation for item in failed.dispatches] == [ + CommandOperation.CANCEL, + CommandOperation.HOLD, + ] + assert len(sink.held) == 1 + assert sink.cancel_count == 1 + assert failed.message is not None and "observation unavailable" in failed.message + + +def test_blocking_runner_uses_clock_and_completes() -> None: + runner, clock, _, _, _ = _make_runner() + + completed = runner.run_until_blocked() + + assert completed.status is RunnerStatus.COMPLETED + assert completed.command_count == 3 + assert clock.sleeps == pytest.approx( + [FIRST_INTERVAL, SECOND_INTERVAL, SECOND_INTERVAL] + ) + + +def test_blocking_runner_safely_stops_when_the_clock_fails() -> None: + runner, clock, _, sink, _ = _make_runner() + + def fail_sleep(duration: float) -> None: + raise RuntimeError("clock backend unavailable") + + clock.sleep = fail_sleep + + failed = runner.run_until_blocked() + + assert failed.status is RunnerStatus.FAILED + assert [item.operation for item in failed.dispatches[-2:]] == [ + CommandOperation.CANCEL, + CommandOperation.HOLD, + ] + assert sink.cancel_count == 1 + assert failed.message is not None and "clock backend unavailable" in failed.message + + +def test_blocking_runner_verifies_effect_before_committing_task_state() -> None: + runner, _, _, _, _ = _make_runner(with_effect=True) + + completed = runner.run_until_blocked( + effect_verifier=lambda context, tick: torch.ones( + context.batch_size, dtype=torch.bool + ) + ) + + assert completed.status is RunnerStatus.COMPLETED + assert completed.tick is not None + assert completed.tick.task_state.get_held_object("arm") is not None + + +def test_blocking_runner_resumes_a_stored_effect_verification_boundary() -> None: + runner, _, _, _, _ = _make_runner(with_effect=True) + + blocked = runner.run_until_blocked() + + assert blocked.status is RunnerStatus.RUNNING + assert runner.effect_verification_pending is True + + completed = runner.run_until_blocked( + effect_verifier=lambda context, tick: torch.ones( + context.batch_size, dtype=torch.bool + ) + ) + + assert runner.effect_verification_pending is False + assert completed.status is RunnerStatus.COMPLETED + assert completed.tick is not None + assert completed.tick.task_state.get_held_object("arm") is not None diff --git a/tests/sim/atomic_actions/test_sim_adapter.py b/tests/sim/atomic_actions/test_sim_adapter.py new file mode 100644 index 000000000..8a208f6ea --- /dev/null +++ b/tests/sim/atomic_actions/test_sim_adapter.py @@ -0,0 +1,198 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Tests for the simulation execution-runner adapter.""" + +from __future__ import annotations + +from unittest.mock import Mock + +import pytest +import torch + +from embodichain.lab.sim.atomic_actions import ( + CommandAckStatus, + JointCommand, + SceneSnapshot, + SimulationExecutionAdapter, + TaskState, +) + +BATCH_SIZE = 2 +ROBOT_DOF = 3 +PHYSICS_DT = 0.01 + + +def _simulation_and_robot() -> tuple[Mock, Mock]: + simulation = Mock() + simulation.sim_config.physics_dt = PHYSICS_DT + robot = Mock() + robot.get_qpos.return_value = torch.zeros(BATCH_SIZE, ROBOT_DOF) + robot.get_qvel.return_value = torch.full((BATCH_SIZE, ROBOT_DOF), 0.1) + robot.get_qf.return_value = torch.full((BATCH_SIZE, ROBOT_DOF), 0.2) + robot.get_proprioception.return_value = {} + return simulation, robot + + +def _command( + *, + env_ids: torch.Tensor | None = None, + active_mask: torch.Tensor | None = None, +) -> JointCommand: + return JointCommand( + positions=torch.ones(BATCH_SIZE, ROBOT_DOF), + velocities=torch.full((BATCH_SIZE, ROBOT_DOF), 0.5), + active_mask=( + torch.tensor([True, False]) if active_mask is None else active_mask + ), + env_ids=( + torch.arange(BATCH_SIZE, dtype=torch.long) if env_ids is None else env_ids + ), + hold_duration=torch.full((BATCH_SIZE,), 0.1), + ) + + +def test_simulation_adapter_observes_full_robot_state() -> None: + simulation, robot = _simulation_and_robot() + adapter = SimulationExecutionAdapter(simulation, robot) + + context = adapter.observe(TaskState.empty(BATCH_SIZE, "cpu")) + + assert context.robot.timestamp == 0.0 + assert torch.equal(context.robot.qpos, robot.get_qpos.return_value) + assert torch.equal(context.robot.qvel, robot.get_qvel.return_value) + assert torch.equal(context.robot.qeffort, robot.get_qf.return_value) + assert context.scene.version == 0 + + +@pytest.mark.parametrize("error", [AttributeError, NotImplementedError]) +def test_simulation_adapter_treats_unavailable_effort_as_optional( + error: type[Exception], +) -> None: + simulation, robot = _simulation_and_robot() + robot.get_qf.side_effect = error("effort unavailable") + adapter = SimulationExecutionAdapter(simulation, robot) + + context = adapter.observe(TaskState.empty(BATCH_SIZE, "cpu")) + + assert context.robot.qeffort is None + + +def test_simulation_adapter_falls_back_to_proprioception_effort() -> None: + simulation, robot = _simulation_and_robot() + robot.get_qf.side_effect = AttributeError("effort unavailable") + expected_qeffort = torch.full((BATCH_SIZE, ROBOT_DOF), 0.3) + robot.get_proprioception.return_value = {"qf": expected_qeffort} + adapter = SimulationExecutionAdapter(simulation, robot) + + context = adapter.observe(TaskState.empty(BATCH_SIZE, "cpu")) + + assert torch.equal(context.robot.qeffort, expected_qeffort) + + +def test_simulation_adapter_sends_active_rows_and_inactive_holds_together() -> None: + simulation, robot = _simulation_and_robot() + adapter = SimulationExecutionAdapter(simulation, robot) + command = _command() + + acknowledgement = adapter.send(command, timeout=1.0) + + assert acknowledgement.status is CommandAckStatus.ACCEPTED + sent_qpos = robot.set_qpos.call_args.args[0] + sent_qvel = robot.set_qvel.call_args.args[0] + assert torch.equal(sent_qpos, command.positions) + assert torch.equal(sent_qvel, command.velocities) + assert robot.set_qpos.call_args.kwargs["env_ids"] == [0, 1] + assert robot.set_qvel.call_args.kwargs["env_ids"] == [0, 1] + + +def test_simulation_adapter_send_writes_a_pure_hold_batch() -> None: + simulation, robot = _simulation_and_robot() + adapter = SimulationExecutionAdapter(simulation, robot) + command = _command(active_mask=torch.zeros(BATCH_SIZE, dtype=torch.bool)) + + acknowledgement = adapter.send(command, timeout=1.0) + + assert acknowledgement.status is CommandAckStatus.ACCEPTED + robot.set_qpos.assert_called_once_with(command.positions, env_ids=[0, 1]) + robot.set_qvel.assert_called_once_with(command.velocities, env_ids=[0, 1]) + + +def test_simulation_adapter_keeps_stable_ids_separate_from_robot_indices() -> None: + simulation, robot = _simulation_and_robot() + stable_ids = torch.tensor([10, 20], dtype=torch.long) + adapter = SimulationExecutionAdapter(simulation, robot, env_ids=stable_ids) + command = _command(env_ids=stable_ids) + + acknowledgement = adapter.send(command, timeout=1.0) + + assert acknowledgement.status is CommandAckStatus.ACCEPTED + assert robot.set_qpos.call_args.kwargs["env_ids"] == [0, 1] + + +def test_simulation_adapter_hold_targets_every_environment() -> None: + simulation, robot = _simulation_and_robot() + adapter = SimulationExecutionAdapter(simulation, robot) + command = _command() + + acknowledgement = adapter.hold(command, timeout=1.0) + + assert acknowledgement.status is CommandAckStatus.ACCEPTED + robot.set_qpos.assert_called_once_with(command.positions, env_ids=[0, 1]) + robot.set_qvel.assert_called_once_with(command.velocities, env_ids=[0, 1]) + + +def test_simulation_adapter_sleep_advances_integral_physics_steps() -> None: + simulation, robot = _simulation_and_robot() + adapter = SimulationExecutionAdapter(simulation, robot) + + adapter.sleep(0.025) + + simulation.update.assert_called_once_with(physics_dt=PHYSICS_DT, step=3) + assert adapter.now() == pytest.approx(0.03) + + +def test_simulation_adapter_supplies_elapsed_time_to_scene_callback() -> None: + simulation, robot = _simulation_and_robot() + timestamps: list[float] = [] + + def scene_supplier(timestamp: float) -> SceneSnapshot: + timestamps.append(timestamp) + return SceneSnapshot(timestamp=timestamp, version=3) + + adapter = SimulationExecutionAdapter( + simulation, + robot, + scene_supplier=scene_supplier, + ) + adapter.sleep(PHYSICS_DT) + + context = adapter.observe(TaskState.empty(BATCH_SIZE, "cpu")) + + assert timestamps == pytest.approx([PHYSICS_DT]) + assert context.scene.version == 3 + + +def test_simulation_adapter_rejects_changed_environment_identity() -> None: + simulation, robot = _simulation_and_robot() + adapter = SimulationExecutionAdapter(simulation, robot) + command = _command(env_ids=torch.tensor([1, 0], dtype=torch.long)) + + acknowledgement = adapter.send(command, timeout=1.0) + + assert acknowledgement.status is CommandAckStatus.REJECTED + assert "env_ids" in acknowledgement.message + robot.set_qpos.assert_not_called()