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 c597adca7..0924419ed 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -435,6 +435,10 @@ topics: - ActionPlan - PlanningContext - ExecutionSession + - ExecutionRunner + - ObservationProvider + - CommandSink + - SimulationExecutionAdapter - StateDelta - held_objects - ActionBinding @@ -464,6 +468,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 c6f69d65a..a6e686746 100644 --- a/agent_context/topics/atomic-actions/atomic-actions.md +++ b/agent_context/topics/atomic-actions/atomic-actions.md @@ -61,10 +61,17 @@ 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 and monitors: +`ExecutionSession` owns deterministic planning progress and recovery state. It +emits at most one `JointCommand` per tick and monitors: - joint tracking error against the previous command; - translation/rotation drift of referenced scene entities; @@ -88,6 +95,39 @@ 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` maps this to `JointCommand.hold_duration`; the final sample +uses its own interval 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 @@ -101,6 +141,11 @@ selection behavior. An action constructor may accept `default_options`; an invocation's `skill_options` replaces them for that call. There is no `ActionCfg` or built-in `*Cfg` layer. +`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 @@ -154,4 +199,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/index.md b/docs/source/overview/sim/atomic_actions/index.md index b6f3bef89..13439f54f 100644 --- a/docs/source/overview/sim/atomic_actions/index.md +++ b/docs/source/overview/sim/atomic_actions/index.md @@ -49,17 +49,25 @@ and whole-body control are not implemented by this module yet. | +-- MotionGenerator / planner backend | | +-- device and shared TrajectoryBuilder | | | -| registered AtomicAction.plan(...) -> ActionPlan | +| registered AtomicAction.plan(request, context) | +| -> ActionPlan | +--------------------------+----------------------------------+ | +------------+-------------+ | | v v - compile(...) start(...) / tick(...) - fixed projection observed closed loop + compile(...) start(...) + fixed projection ExecutionSession.tick(...) | | v v - CompiledTrajectory JointCommand + events + CompiledTrajectory JointCommand + recovery events + | + v + ExecutionRunner + observe / schedule / dispatch + | + v + ObservationProvider + CommandSink + Clock ``` The boundary is deliberate: @@ -71,9 +79,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 @@ -123,6 +141,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 | @@ -276,10 +295,13 @@ Both instances still borrow the same engine-owned motion generator. |---|---|---| | `AtomicAction.plan(request, context)` | Implementing a skill | Consumes an engine-resolved immutable request; application code normally calls it through the engine | | `engine.plan(invocation, context)` | Planning one registered skill | Resolves the registered action, binds shared resources, and validates its plan | -| `engine.plan_action(action, invocation, context)` | Planning an unregistered configured instance | Supports multiple configurations with one `skill_id` and one engine backend | +| `engine.plan_action(action, invocation, context)` | Planning an unregistered action instance | Supports multiple default-option variants with one `skill_id` and one engine backend | | `engine.compile(invocations, context)` | Fixed-scene/offline sequence planning | Returns one concatenated `CompiledTrajectory` and a hypothetical projected context | | `engine.start(invocations, context)` | Observed incremental execution | Returns an `ExecutionSession`; each `tick()` emits at most one command and recovery events | | `session.revise_current(invocation)` | Explicit runtime parameter/goal update | Requires a newer revision of the active logical invocation and replans from the latest 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 | `AtomicAction.plan()` is therefore not a second execution API. It is the polymorphic implementation point used by the engine. Neither it nor the engine @@ -360,6 +382,34 @@ 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`; the final +sample's interval is also its 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; @@ -482,4 +532,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/tracking_error_recovery.py` — runnable runner + example with an injected tracking disturbance diff --git a/docs/source/tutorial/atomic_actions.rst b/docs/source/tutorial/atomic_actions.rst index a5635cfc1..ef9954047 100644 --- a/docs/source/tutorial/atomic_actions.rst +++ b/docs/source/tutorial/atomic_actions.rst @@ -42,8 +42,8 @@ responsible for choosing a physically compatible pair. The engine exclusively owns the ``MotionGenerator``, shared trajectory builder, and control-part profiles. Atomic action constructors accept only optional -typed default options; ``register()`` binds each action to the engine resources. Use -``engine.plan_action(action, invocation, context)`` for an unregistered, +typed default options; ``register()`` binds each action to the engine resources. +Use ``engine.plan_action(action, invocation, context)`` for an unregistered, default-option-specific action instance. Runnable examples @@ -61,6 +61,7 @@ Focused examples live under ``scripts/tutorials/atomic_action``: * ``coordinated_pickment.py`` * ``coordinated_placement.py`` * ``hand_over.py`` +* ``tracking_error_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 @@ -170,18 +171,39 @@ 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 deliberately changes a measured joint position, +observes ``tracking_error`` and ``replanned`` events, and finishes the regenerated +trajectory: -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/tracking_error_recovery.py --headless Recovery replans reuse one immutable invocation-revision snapshot. If an application intentionally changes the goal, options, policy, binding, or a @@ -221,13 +243,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 f4ca7c71f..3dec06766 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 @@ -101,6 +102,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, @@ -127,6 +145,11 @@ "AtomicAction", "AtomicActionEngine", "CompiledTrajectory", + "CommandAcknowledgement", + "CommandAckStatus", + "CommandDispatch", + "CommandOperation", + "CommandSink", "CompletionCondition", "CompletionConditionKind", "ControlCommand", @@ -140,8 +163,12 @@ "CoordinatedPlacementOptions", "EndEffectorPoseGoal", "EntityState", + "EffectVerifier", + "ExecutionClock", "ExecutionEvent", "ExecutionEventKind", + "ExecutionRunner", + "ExecutionRunnerCfg", "ExecutionSession", "ExecutionStatus", "ExecutionTick", @@ -156,6 +183,7 @@ "JointCommand", "JointPositionCommand", "MotionPolicy", + "MonotonicExecutionClock", "MoveEndEffector", "MoveEndEffectorOptions", "MoveHeldObject", @@ -165,6 +193,7 @@ "ObjectActionGoal", "ObjectSemantics", "OPEN_COMMAND", + "ObservationProvider", "PhaseSpec", "PickUp", "PickUpOptions", @@ -183,10 +212,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 748cebed8..87e27356c 100644 --- a/embodichain/lab/sim/atomic_actions/execution.py +++ b/embodichain/lab/sim/atomic_actions/execution.py @@ -90,6 +90,8 @@ class JointCommand: velocities: torch.Tensor | None active_mask: torch.Tensor env_ids: torch.Tensor + hold_duration: torch.Tensor + """Per-environment time to hold this command before the next observation.""" def __post_init__(self) -> None: if self.positions.dim() != 2: @@ -107,15 +109,29 @@ def __post_init__(self) -> None: self.positions.shape[0], ): raise ValueError("JointCommand.env_ids must be int64 with shape (B,).") + if not isinstance(self.hold_duration, torch.Tensor): + raise TypeError("JointCommand.hold_duration must be a torch.Tensor.") + if self.hold_duration.shape != (self.positions.shape[0],): + raise ValueError("JointCommand.hold_duration must have shape (B,).") + if ( + not torch.isfinite(self.hold_duration).all() + or (self.hold_duration < 0.0).any() + ): + raise ValueError( + "JointCommand.hold_duration must contain finite non-negative values." + ) if self.active_mask.device != self.positions.device: raise ValueError("JointCommand tensors must share a device.") if self.env_ids.device != self.positions.device: raise ValueError("JointCommand tensors must share a device.") + if self.hold_duration.device != self.positions.device: + raise ValueError("JointCommand tensors must share a device.") object.__setattr__(self, "positions", self.positions.clone()) if self.velocities is not None: object.__setattr__(self, "velocities", self.velocities.clone()) object.__setattr__(self, "active_mask", self.active_mask.clone()) object.__setattr__(self, "env_ids", self.env_ids.clone()) + object.__setattr__(self, "hold_duration", self.hold_duration.clone()) @dataclass(frozen=True, slots=True, eq=False) @@ -256,6 +272,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, @@ -607,11 +628,14 @@ def _command_at( ) self._last_command = positions.clone() self._last_command_mask = active_mask.clone() + next_index = min(waypoint_index + 1, phase.trajectory.waypoint_count - 1) + hold_duration = phase.trajectory.dt[:, next_index] return JointCommand( positions=positions, velocities=velocities, active_mask=active_mask, env_ids=phase.trajectory.env_ids, + hold_duration=hold_duration, ) def _hold_command(self) -> JointCommand: @@ -621,6 +645,11 @@ def _hold_command(self) -> JointCommand: velocities=torch.zeros_like(self._context.robot.qpos), active_mask=torch.zeros_like(self._eligible), env_ids=self._context.env_ids, + hold_duration=torch.zeros( + self._context.batch_size, + dtype=torch.float32, + device=self._context.robot.qpos.device, + ), ) def _terminal_error(self, phase: PlannedPhase) -> torch.Tensor: diff --git a/embodichain/lab/sim/atomic_actions/runner.py b/embodichain/lab/sim/atomic_actions/runner.py new file mode 100644 index 000000000..9bf077ce5 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/runner.py @@ -0,0 +1,791 @@ +# ---------------------------------------------------------------------------- +# 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: + assert command is not None + acknowledgement = self._command_sink.send( + command, + timeout=self.cfg.command_timeout, + ) + elif operation is CommandOperation.HOLD: + assert command is not None + 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..871712c7e --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/sim_adapter.py @@ -0,0 +1,279 @@ +# ---------------------------------------------------------------------------- +# 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") + 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) + if not command.active_mask.any(): + return CommandAcknowledgement.accepted_ack("No active rows.") + 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 _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/tracking_error_recovery.py b/scripts/tutorials/atomic_action/tracking_error_recovery.py new file mode 100644 index 000000000..4a2c2b1d0 --- /dev/null +++ b/scripts/tutorials/atomic_action/tracking_error_recovery.py @@ -0,0 +1,235 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Demonstrate closed-loop recovery from an injected joint tracking error.""" + +from __future__ import annotations + +import argparse +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.gym.utils.gym_utils import add_env_launcher_args_to_parser +from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, + AtomicActionEngine, + ExecutionEventKind, + ExecutionRunner, + ExecutionRunnerCfg, + JointPositionGoal, + MotionPolicy, + MoveJoints, + PlanningContext, + RecoveryPolicy, + RunnerStatus, + RunnerStep, + SimulationExecutionAdapter, + TaskState, +) +from embodichain.lab.sim.objects import Robot +from embodichain.utils import logger +from scripts.tutorials.atomic_action.tutorial_utils import ( + add_ur5_gripper_robot, + create_toppra_motion_generator, + create_tutorial_simulation, + prepare_tutorial_scene, + run_tutorial, + serve_tutorial_scene, + start_auto_play_recording, + stop_auto_play_recording, +) + +SAMPLE_COUNT = 80 +INJECTION_AFTER_COMMAND = 3 +TRACKING_ERROR_OFFSET = 0.35 +TRACKING_ERROR_THRESHOLD = 0.08 +POST_EXECUTION_UPDATES = 80 + + +class _OneShotTrackingErrorInjector: + """Decorate simulation observations with one deterministic disturbance.""" + + def __init__( + self, + adapter: SimulationExecutionAdapter, + robot: Robot, + *, + joint_id: int, + offset: float, + ) -> None: + self._adapter = adapter + self._robot = robot + self._joint_id = joint_id + self._offset = offset + self._pending = False + self.injected = False + + def arm(self) -> None: + """Request a disturbance immediately before the next observation.""" + if not self.injected: + self._pending = True + + def observe(self, task_state: TaskState) -> PlanningContext: + """Inject one physical-state offset, then capture the observation. + + Args: + task_state: Session-owned verified task state. + + Returns: + Latest simulation planning context. + """ + if self._pending: + qpos = self._robot.get_qpos().clone() + qpos[:, self._joint_id] += self._offset + self._robot.set_qpos(qpos, target=False) + self._robot.set_qvel(torch.zeros_like(qpos), target=False) + self._pending = False + self.injected = True + logger.log_warning( + "Injected a joint-position disturbance before observation; " + "the session should detect tracking error and replan." + ) + return self._adapter.observe(task_state) + + +def parse_arguments() -> argparse.Namespace: + """Parse command-line arguments for the recovery tutorial.""" + parser = argparse.ArgumentParser( + description="Demonstrate ExecutionRunner tracking-error recovery." + ) + add_env_launcher_args_to_parser(parser) + parser.add_argument("--auto_play", action="store_true") + parser.add_argument( + "--no_error_injection", + action="store_true", + help="Run the closed-loop trajectory without the demonstration disturbance.", + ) + return parser.parse_args() + + +def main() -> None: + """Execute MoveJoints and recover after a one-shot state disturbance.""" + args = parse_arguments() + sim = create_tutorial_simulation(args) + robot = add_ur5_gripper_robot(sim) + motion_gen = create_toppra_motion_generator(robot) + adapter = SimulationExecutionAdapter(sim, robot) + + target = torch.tensor( + [0.35, -1.20, 1.30, -1.65, -1.57, 0.20], + dtype=torch.float32, + device=sim.device, + ) + engine = AtomicActionEngine(motion_generator=motion_gen) + engine.register(MoveJoints()) + invocation = ActionInvocation( + skill_id="move_joints", + goal=JointPositionGoal(target), + binding=ActionBinding(manipulators={"primary": "arm"}), + motion_policy=MotionPolicy( + sample_count=SAMPLE_COUNT, + control_dt=2.0 * adapter.physics_dt, + ), + recovery_policy=RecoveryPolicy( + max_replans=2, + tracking_error_threshold=TRACKING_ERROR_THRESHOLD, + phase_timeout=20.0, + ), + ) + task_state = TaskState.empty(robot.get_qpos().shape[0], robot.device) + initial_context = adapter.observe(task_state) + session = engine.start((invocation,), initial_context) + + arm_joint_id = robot.get_joint_ids(name="arm")[0] + observation_provider = _OneShotTrackingErrorInjector( + adapter, + robot, + joint_id=arm_joint_id, + offset=TRACKING_ERROR_OFFSET, + ) + runner = ExecutionRunner( + session, + observation_provider, + adapter, + clock=adapter, + cfg=ExecutionRunnerCfg(minimum_cycle_time=adapter.physics_dt), + ) + + wait_for_user = prepare_tutorial_scene( + sim, + args, + "Inspect the robot, then press Enter to start closed-loop execution...", + ) + recovery_observed = False + + def on_step(step: RunnerStep) -> None: + nonlocal recovery_observed + if ( + not args.no_error_injection + and not observation_provider.injected + and step.command_count >= INJECTION_AFTER_COMMAND + ): + observation_provider.arm() + if step.tick is None: + return + for event in step.tick.events: + if event.kind in { + ExecutionEventKind.TRACKING_ERROR, + 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}" + ) + recovery_observed |= event.kind is ExecutionEventKind.REPLANNED + + recording_started = start_auto_play_recording( + sim, + args, + video_prefix="tracking_error_recovery_auto_play", + ) + try: + result = runner.run_until_blocked(on_step=on_step) + for _ in range(POST_EXECUTION_UPDATES): + adapter.sleep(adapter.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_error_injection and not recovery_observed: + raise RuntimeError("The injected tracking error did not trigger replanning.") + logger.log_info( + f"Execution completed 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_runner.py b/tests/sim/atomic_actions/test_runner.py new file mode 100644 index 000000000..037b57c86 --- /dev/null +++ b/tests/sim/atomic_actions/test_runner.py @@ -0,0 +1,471 @@ +# ---------------------------------------------------------------------------- +# 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 +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.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) + 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=0.01), + ) + 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() + + assert first.command_count == 1 + assert first.wait_duration == pytest.approx(FIRST_INTERVAL) + assert early.is_waiting + assert len(sink.sent) == 2 + assert second.command_count == 2 + assert second.wait_duration == pytest.approx(SECOND_INTERVAL) + + +def test_runner_uses_the_longest_active_batch_interval_as_a_barrier() -> None: + runner, _, _, _, _ = _make_runner(batch_size=2) + + first = runner.step() + + assert first.wait_duration == pytest.approx(2.0 * FIRST_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..19b2988c7 --- /dev/null +++ b/tests/sim/atomic_actions/test_sim_adapter.py @@ -0,0 +1,167 @@ +# ---------------------------------------------------------------------------- +# 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) + return simulation, robot + + +def _command(*, env_ids: 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]), + 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_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_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()