From 0bbeb2dac3a084d316cf12fbd447798261bdc80c Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 2 Aug 2026 11:06:46 +0000 Subject: [PATCH 1/5] feat(atomic-actions): add closed-loop execution runner --- .agents/skills/add-atomic-action/SKILL.md | 9 +- agent_context/MAP.yaml | 6 + .../topics/atomic-actions/atomic-actions.md | 58 +- .../embodichain.lab.sim.atomic_actions.rst | 60 ++ .../overview/sim/atomic_actions/index.md | 55 +- docs/source/tutorial/atomic_actions.rst | 59 +- .../lab/sim/atomic_actions/__init__.py | 38 +- .../lab/sim/atomic_actions/execution.py | 5 + embodichain/lab/sim/atomic_actions/runner.py | 791 ++++++++++++++++++ .../lab/sim/atomic_actions/sim_adapter.py | 279 ++++++ .../atomic_action/tracking_error_recovery.py | 233 ++++++ tests/sim/atomic_actions/test_runner.py | 479 +++++++++++ tests/sim/atomic_actions/test_sim_adapter.py | 167 ++++ 13 files changed, 2210 insertions(+), 29 deletions(-) create mode 100644 embodichain/lab/sim/atomic_actions/runner.py create mode 100644 embodichain/lab/sim/atomic_actions/sim_adapter.py create mode 100644 scripts/tutorials/atomic_action/tracking_error_recovery.py create mode 100644 tests/sim/atomic_actions/test_runner.py create mode 100644 tests/sim/atomic_actions/test_sim_adapter.py 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..8ec231a28 100644 --- a/agent_context/topics/atomic-actions/atomic-actions.md +++ b/agent_context/topics/atomic-actions/atomic-actions.md @@ -78,13 +78,20 @@ 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` 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: - joint tracking error against the previous command; - translation/rotation drift of referenced scene entities; @@ -113,6 +120,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 @@ -135,6 +175,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 @@ -188,4 +233,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 3a63ca6b9..d20127e0b 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,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; @@ -604,4 +653,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 fc3b306f8..83799ab09 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`` +* ``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 @@ -260,18 +261,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 @@ -311,13 +333,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..d7f115095 100644 --- a/embodichain/lab/sim/atomic_actions/execution.py +++ b/embodichain/lab/sim/atomic_actions/execution.py @@ -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, 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..bf232feca --- /dev/null +++ b/scripts/tutorials/atomic_action/tracking_error_recovery.py @@ -0,0 +1,233 @@ +# ---------------------------------------------------------------------------- +# 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, + 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) + 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..7068ceafc --- /dev/null +++ b/tests/sim/atomic_actions/test_runner.py @@ -0,0 +1,479 @@ +# ---------------------------------------------------------------------------- +# 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.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=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(MINIMUM_CYCLE_TIME) + second = runner.step() + clock.advance(FIRST_INTERVAL) + third = runner.step() + + assert first.command_count == 1 + assert first.wait_duration == pytest.approx(MINIMUM_CYCLE_TIME) + assert early.is_waiting + assert len(sink.sent) == 3 + assert second.command_count == 2 + assert second.wait_duration == pytest.approx(FIRST_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(MINIMUM_CYCLE_TIME) + second = runner.step() + + assert first.wait_duration == pytest.approx(MINIMUM_CYCLE_TIME) + assert second.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(MINIMUM_CYCLE_TIME) + runner.step() + clock.advance(FIRST_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(MINIMUM_CYCLE_TIME) + runner.step() + clock.advance(FIRST_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( + [MINIMUM_CYCLE_TIME, FIRST_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() From 3df6c535a6e0e0f25c59734b650a62cd144526fd Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 9 Aug 2026 20:57:51 +0800 Subject: [PATCH 2/5] feat(atomic-actions): demonstrate moving target recovery --- .../overview/sim/atomic_actions/index.md | 4 +- docs/source/tutorial/atomic_actions.rst | 14 +- .../atomic_action/moving_target_recovery.py | 483 ++++++++++++++++++ .../atomic_action/tracking_error_recovery.py | 233 --------- 4 files changed, 494 insertions(+), 240 deletions(-) create mode 100644 scripts/tutorials/atomic_action/moving_target_recovery.py delete mode 100644 scripts/tutorials/atomic_action/tracking_error_recovery.py diff --git a/docs/source/overview/sim/atomic_actions/index.md b/docs/source/overview/sim/atomic_actions/index.md index d20127e0b..9fc51c7b4 100644 --- a/docs/source/overview/sim/atomic_actions/index.md +++ b/docs/source/overview/sim/atomic_actions/index.md @@ -654,5 +654,5 @@ 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 - {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 +- `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 83799ab09..733d89ff6 100644 --- a/docs/source/tutorial/atomic_actions.rst +++ b/docs/source/tutorial/atomic_actions.rst @@ -99,7 +99,7 @@ Focused examples live under ``scripts/tutorials/atomic_action``: * ``coordinated_pickment.py`` * ``coordinated_placement.py`` * ``hand_over.py`` -* ``tracking_error_recovery.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 @@ -287,13 +287,17 @@ For an application that already owns its event loop, call the non-blocking 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 complete simulation example starts with a visible cube directly in front of +the robot, then slides it sideways while the robot is approaching. The session +observes ``dynamic_goal_changed`` and ``replanned`` events, discards the stale +path, and approaches the cube's new location. The same session then executes +``PickUp``, closes the gripper, verifies the physical lift, and finishes while +holding the cube. The original and regenerated goal axes remain visible for +comparison: .. code-block:: bash - python scripts/tutorials/atomic_action/tracking_error_recovery.py --headless + 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 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..7744de0b3 --- /dev/null +++ b/scripts/tutorials/atomic_action/moving_target_recovery.py @@ -0,0 +1,483 @@ +# ---------------------------------------------------------------------------- +# 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 after a visible target move, then grasp the relocated cube.""" + +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, + EndEffectorPoseGoal, + 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, + 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]) +TARGET_TO_EEF_HEIGHT = 0.30 +MOVE_SAMPLE_COUNT = 80 +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 +GOAL_TRANSLATION_THRESHOLD = 0.04 +TRACKING_ERROR_THRESHOLD = 0.25 +POST_EXECUTION_UPDATES = 120 + + +class _MovingTargetScene: + """Publish a versioned target pose and move 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 move( + self, + clock: SimulationExecutionAdapter, + *, + duration: float, + ) -> torch.Tensor: + """Animate the visible target and advance the scene version. + + Args: + clock: Simulation adapter used to advance physics between poses. + duration: Requested target-motion duration in seconds. + + Returns: + Updated batched target pose. + """ + 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.") + start_pose = self.target.get_local_pose(to_matrix=True).clone() + step_count = max(1, math.ceil(duration / clock.physics_dt)) + pose = start_pose.clone() + for step_index in range(1, step_count + 1): + alpha = step_index / step_count + pose[:, :3, 3] = torch.lerp( + start_pose[:, :3, 3], + self.destination, + alpha, + ) + self.target.set_local_pose(pose) + clock.sleep(clock.physics_dt) + self.version += 1 + self.moved = True + return pose + + +def _create_moving_target(sim: SimulationManager) -> RigidObject: + """Create the bright cube, initially kinematic for scripted relocation.""" + 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 toward a relocated cube, close the gripper, and lift it.""" + 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) + adapter = 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) + hand_open_batch = hand_open.unsqueeze(0).repeat(robot.get_qpos().shape[0], 1) + for target_value in (False, True): + robot.set_qpos(hand_open_batch, name="hand", target=target_value) + robot.clear_dynamics() + + target_to_eef = make_top_down_eef_pose( + torch.tensor( + [0.0, 0.0, TARGET_TO_EEF_HEIGHT], + 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_eef), + axis_len=0.10, + ) + + grasp_target_position = ( + INITIAL_TARGET_POSITION if args.no_target_motion else MOVED_TARGET_POSITION + ) + grasp_pose = make_top_down_eef_pose( + torch.tensor( + grasp_target_position, + dtype=torch.float32, + device=sim.device, + ) + ) + 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, + ) + }, + ) + move_invocation = ActionInvocation( + skill_id="move_end_effector", + goal=EndEffectorPoseGoal( + SceneEntityPose( + TARGET_ENTITY_ID, + relative_pose=target_to_eef, + ) + ), + binding=binding, + motion_policy=MotionPolicy( + sample_count=MOVE_SAMPLE_COUNT, + control_dt=2.0 * adapter.physics_dt, + ), + recovery_policy=RecoveryPolicy( + max_replans=2, + tracking_error_threshold=TRACKING_ERROR_THRESHOLD, + goal_translation_threshold=GOAL_TRANSLATION_THRESHOLD, + phase_timeout=20.0, + ), + ) + pick_invocation = ActionInvocation( + skill_id="pick_up", + goal=GraspGoal(semantics, grasp_xpos=grasp_pose), + binding=binding, + motion_policy=MotionPolicy( + sample_count=PICK_SAMPLE_COUNT, + control_dt=2.0 * adapter.physics_dt, + ), + recovery_policy=RecoveryPolicy( + max_phase_retries=1, + tracking_error_threshold=TRACKING_ERROR_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 = adapter.observe(task_state) + session = engine.start((move_invocation, pick_invocation), initial_context) + runner = ExecutionRunner( + session, + adapter, + adapter, + clock=adapter, + cfg=ExecutionRunnerCfg(minimum_cycle_time=adapter.physics_dt), + ) + + wait_for_user = prepare_tutorial_scene( + sim, + args, + "Watch the blue cube, then press Enter to replan and pick it up...", + ) + dynamic_change_observed = False + replan_observed = False + pickup_start_command: int | None = None + 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 dynamic_change_observed, replan_observed + nonlocal pickup_dynamics_cleared, pickup_start_command + if ( + not args.no_target_motion + and not target_scene.moved + and step.command_count >= MOVE_AFTER_COMMAND + ): + logger.log_warning( + f"Animating the blue target for {TARGET_MOVE_DURATION:.1f} s " + "while the robot holds its current command." + ) + moved_pose = target_scene.move( + adapter, + duration=TARGET_MOVE_DURATION, + ) + draw_axis_marker( + sim, + "moving_target_replanned_goal", + _compose_goal_pose(moved_pose, target_to_eef), + axis_len=0.10, + ) + displacement = torch.linalg.vector_norm( + moved_pose[:, :3, 3] - initial_target_pose[:, :3, 3], + dim=1, + ) + logger.log_warning( + "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: + 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}" + ) + dynamic_change_observed |= ( + event.kind is ExecutionEventKind.DYNAMIC_GOAL_CHANGED + ) + replan_observed |= event.kind is ExecutionEventKind.REPLANNED + if ( + event.kind is ExecutionEventKind.ACTION_PLANNED + and event.invocation_index == 1 + and pickup_start_command is None + ): + target.set_body_type("dynamic") + target.clear_dynamics() + pickup_start_command = step.command_count + logger.log_info( + "The approach completed; the cube is now dynamic and PickUp " + "is starting.", + color="green", + ) + if ( + pickup_start_command is not None + and not pickup_dynamics_cleared + and step.command_count - pickup_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): + 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_target_motion: + if not target_scene.moved: + raise RuntimeError("Execution completed before the target could move.") + if not dynamic_change_observed: + raise RuntimeError("The target move was not reported as a dynamic change.") + if not replan_observed: + raise RuntimeError("The target move did not trigger replanning.") + if pickup_start_command is None: + raise RuntimeError("PickUp did not start after the recovered approach.") + 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/scripts/tutorials/atomic_action/tracking_error_recovery.py b/scripts/tutorials/atomic_action/tracking_error_recovery.py deleted file mode 100644 index bf232feca..000000000 --- a/scripts/tutorials/atomic_action/tracking_error_recovery.py +++ /dev/null @@ -1,233 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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, - 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) - 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) From 85658aac4003c3464378c331c1420d764f350d6f Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 9 Aug 2026 21:22:41 +0800 Subject: [PATCH 3/5] fix(atomic-actions): correct runner timing and adapter state --- .../topics/atomic-actions/atomic-actions.md | 16 +++++---- .../overview/sim/atomic_actions/index.md | 24 +++++++------ .../lab/sim/atomic_actions/execution.py | 11 ++++-- embodichain/lab/sim/atomic_actions/runner.py | 6 ++-- .../lab/sim/atomic_actions/sim_adapter.py | 18 ++++++++-- .../sim/atomic_actions/test_engine_per_env.py | 6 ++-- tests/sim/atomic_actions/test_runner.py | 29 ++++++++------- tests/sim/atomic_actions/test_sim_adapter.py | 35 +++++++++++++++++-- 8 files changed, 104 insertions(+), 41 deletions(-) diff --git a/agent_context/topics/atomic-actions/atomic-actions.md b/agent_context/topics/atomic-actions/atomic-actions.md index 8ec231a28..6e209b4f7 100644 --- a/agent_context/topics/atomic-actions/atomic-actions.md +++ b/agent_context/topics/atomic-actions/atomic-actions.md @@ -89,9 +89,10 @@ result = runner.step(effect_success=None) `ExecutionSession` owns deterministic planning progress and recovery state. It 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: +`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; @@ -136,10 +137,11 @@ replans from the latest context. 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. +`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 diff --git a/docs/source/overview/sim/atomic_actions/index.md b/docs/source/overview/sim/atomic_actions/index.md index 9fc51c7b4..e428388b3 100644 --- a/docs/source/overview/sim/atomic_actions/index.md +++ b/docs/source/overview/sim/atomic_actions/index.md @@ -504,13 +504,14 @@ 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. +`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: @@ -544,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: diff --git a/embodichain/lab/sim/atomic_actions/execution.py b/embodichain/lab/sim/atomic_actions/execution.py index d7f115095..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: @@ -632,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/runner.py b/embodichain/lab/sim/atomic_actions/runner.py index 9bf077ce5..0c82ab6b6 100644 --- a/embodichain/lab/sim/atomic_actions/runner.py +++ b/embodichain/lab/sim/atomic_actions/runner.py @@ -661,13 +661,15 @@ def _dispatch( """Call one sink operation and convert exceptions to rejection acks.""" try: if operation is CommandOperation.SEND: - assert command is not None + 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: - assert command is not None + if command is None: + raise ValueError("HOLD requires a JointCommand.") acknowledgement = self._command_sink.hold( command, timeout=self.cfg.safe_stop_timeout, diff --git a/embodichain/lab/sim/atomic_actions/sim_adapter.py b/embodichain/lab/sim/atomic_actions/sim_adapter.py index 871712c7e..20f43922f 100644 --- a/embodichain/lab/sim/atomic_actions/sim_adapter.py +++ b/embodichain/lab/sim/atomic_actions/sim_adapter.py @@ -137,6 +137,8 @@ def observe(self, task_state: TaskState) -> PlanningContext: 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 @@ -177,8 +179,6 @@ def send( 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, @@ -256,6 +256,20 @@ def _read_optional_tensor(self, method_name: str) -> torch.Tensor | None: 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): 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 index 7068ceafc..6cf1a0099 100644 --- a/tests/sim/atomic_actions/test_runner.py +++ b/tests/sim/atomic_actions/test_runner.py @@ -116,6 +116,7 @@ def __init__(self, provider: FakeObservationProvider) -> None: 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 @@ -127,6 +128,7 @@ def send( ) -> 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 @@ -269,17 +271,20 @@ def test_runner_dispatches_only_when_timed_waypoint_is_due() -> None: first = runner.step() early = runner.step() - clock.advance(MINIMUM_CYCLE_TIME) - second = 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(MINIMUM_CYCLE_TIME) + 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(FIRST_INTERVAL) + assert second.wait_duration == pytest.approx(SECOND_INTERVAL) assert third.command_count == 3 assert third.wait_duration == pytest.approx(SECOND_INTERVAL) @@ -288,22 +293,22 @@ 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(MINIMUM_CYCLE_TIME) + clock.advance(2.0 * FIRST_INTERVAL) second = runner.step() - assert first.wait_duration == pytest.approx(MINIMUM_CYCLE_TIME) - assert second.wait_duration == pytest.approx(2.0 * FIRST_INTERVAL) + 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(MINIMUM_CYCLE_TIME) 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 @@ -355,10 +360,10 @@ 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(MINIMUM_CYCLE_TIME) runner.step() clock.advance(FIRST_INTERVAL) + runner.step() + clock.advance(SECOND_INTERVAL) recovered = runner.step() assert action.plan_count == 2 @@ -422,7 +427,7 @@ def test_blocking_runner_uses_clock_and_completes() -> None: assert completed.status is RunnerStatus.COMPLETED assert completed.command_count == 3 assert clock.sleeps == pytest.approx( - [MINIMUM_CYCLE_TIME, FIRST_INTERVAL, SECOND_INTERVAL] + [FIRST_INTERVAL, SECOND_INTERVAL, SECOND_INTERVAL] ) diff --git a/tests/sim/atomic_actions/test_sim_adapter.py b/tests/sim/atomic_actions/test_sim_adapter.py index 19b2988c7..8a208f6ea 100644 --- a/tests/sim/atomic_actions/test_sim_adapter.py +++ b/tests/sim/atomic_actions/test_sim_adapter.py @@ -43,14 +43,21 @@ def _simulation_and_robot() -> tuple[Mock, 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) -> JointCommand: +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]), + 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 ), @@ -84,6 +91,18 @@ def test_simulation_adapter_treats_unavailable_effort_as_optional( 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) @@ -100,6 +119,18 @@ def test_simulation_adapter_sends_active_rows_and_inactive_holds_together() -> N 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) From 9eeb7d420e2b822b99f7bc90746e4770713c557b Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 9 Aug 2026 21:51:12 +0800 Subject: [PATCH 4/5] feat(atomic-actions): support late-bound pickup goals --- .../topics/atomic-actions/atomic-actions.md | 5 + .../sim/atomic_actions/builtin_actions.md | 19 +-- docs/source/tutorial/atomic_actions.rst | 13 +- .../sim/atomic_actions/primitives/pick_up.py | 22 ++- .../atomic_action/moving_target_recovery.py | 135 +++++++----------- tests/sim/atomic_actions/test_actions.py | 127 +++++++++++++++- 6 files changed, 210 insertions(+), 111 deletions(-) diff --git a/agent_context/topics/atomic-actions/atomic-actions.md b/agent_context/topics/atomic-actions/atomic-actions.md index 6e209b4f7..55ed9fccd 100644 --- a/agent_context/topics/atomic-actions/atomic-actions.md +++ b/agent_context/topics/atomic-actions/atomic-actions.md @@ -224,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`. 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/tutorial/atomic_actions.rst b/docs/source/tutorial/atomic_actions.rst index 733d89ff6..7a82971c9 100644 --- a/docs/source/tutorial/atomic_actions.rst +++ b/docs/source/tutorial/atomic_actions.rst @@ -288,12 +288,13 @@ 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 slides it sideways while the robot is approaching. The session -observes ``dynamic_goal_changed`` and ``replanned`` events, discards the stale -path, and approaches the cube's new location. The same session then executes -``PickUp``, closes the gripper, verifies the physical lift, and finishes while -holding the cube. The original and regenerated goal axes remain visible for -comparison: +the robot, then slides 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: .. code-block:: bash 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/scripts/tutorials/atomic_action/moving_target_recovery.py b/scripts/tutorials/atomic_action/moving_target_recovery.py index 7744de0b3..8db1d12c3 100644 --- a/scripts/tutorials/atomic_action/moving_target_recovery.py +++ b/scripts/tutorials/atomic_action/moving_target_recovery.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Replan after a visible target move, then grasp the relocated cube.""" +"""Replan one PickUp action after its visible target moves.""" from __future__ import annotations @@ -36,7 +36,6 @@ Affordance, AtomicActionEngine, ControlPartCommandProfile, - EndEffectorPoseGoal, EntityState, ExecutionEventKind, ExecutionRunner, @@ -66,6 +65,7 @@ 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, @@ -78,8 +78,6 @@ 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]) -TARGET_TO_EEF_HEIGHT = 0.30 -MOVE_SAMPLE_COUNT = 80 PICK_SAMPLE_COUNT = 120 HAND_INTERP_STEPS = 12 PICK_LIFT_HEIGHT = 0.16 @@ -88,7 +86,7 @@ MOVE_AFTER_COMMAND = 20 TARGET_MOVE_DURATION = 0.6 GOAL_TRANSLATION_THRESHOLD = 0.04 -TRACKING_ERROR_THRESHOLD = 0.25 +TRACKING_ERROR_THRESHOLD = 1.0 POST_EXECUTION_UPDATES = 120 @@ -214,50 +212,36 @@ def parse_arguments() -> argparse.Namespace: def main() -> None: - """Replan toward a relocated cube, close the gripper, and lift it.""" + """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) - adapter = SimulationExecutionAdapter( + 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) - hand_open_batch = hand_open.unsqueeze(0).repeat(robot.get_qpos().shape[0], 1) - for target_value in (False, True): - robot.set_qpos(hand_open_batch, name="hand", target=target_value) - robot.clear_dynamics() - - target_to_eef = make_top_down_eef_pose( - torch.tensor( - [0.0, 0.0, TARGET_TO_EEF_HEIGHT], - dtype=torch.float32, - device=sim.device, - ) + 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_eef), + _compose_goal_pose(initial_target_pose, target_to_grasp), axis_len=0.10, ) - grasp_target_position = ( - INITIAL_TARGET_POSITION if args.no_target_motion else MOVED_TARGET_POSITION - ) - grasp_pose = make_top_down_eef_pose( - torch.tensor( - grasp_target_position, - dtype=torch.float32, - device=sim.device, - ) - ) semantics = ObjectSemantics( affordance=Affordance(), geometry={}, @@ -277,37 +261,25 @@ def main() -> None: ) }, ) - move_invocation = ActionInvocation( - skill_id="move_end_effector", - goal=EndEffectorPoseGoal( - SceneEntityPose( - TARGET_ENTITY_ID, - relative_pose=target_to_eef, - ) - ), - binding=binding, - motion_policy=MotionPolicy( - sample_count=MOVE_SAMPLE_COUNT, - control_dt=2.0 * adapter.physics_dt, - ), - recovery_policy=RecoveryPolicy( - max_replans=2, - tracking_error_threshold=TRACKING_ERROR_THRESHOLD, - goal_translation_threshold=GOAL_TRANSLATION_THRESHOLD, - phase_timeout=20.0, - ), - ) pick_invocation = ActionInvocation( skill_id="pick_up", - goal=GraspGoal(semantics, grasp_xpos=grasp_pose), + 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 * adapter.physics_dt, + 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( @@ -317,24 +289,23 @@ def main() -> None: ), ) task_state = TaskState.empty(robot.get_qpos().shape[0], robot.device) - initial_context = adapter.observe(task_state) - session = engine.start((move_invocation, pick_invocation), initial_context) + initial_context = sim_runtime.observe(task_state) + session = engine.start((pick_invocation,), initial_context) runner = ExecutionRunner( - session, - adapter, - adapter, - clock=adapter, - cfg=ExecutionRunnerCfg(minimum_cycle_time=adapter.physics_dt), + 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 replan and pick it up...", + "Watch the blue cube, then press Enter to run recovering PickUp...", ) - dynamic_change_observed = False - replan_observed = False - pickup_start_command: int | None = None + observed_events: set[ExecutionEventKind] = set() + plan_start_command = 0 pickup_dynamics_cleared = False clear_after_pick_command = ( @@ -342,8 +313,7 @@ def main() -> None: ) def on_step(step: RunnerStep) -> None: - nonlocal dynamic_change_observed, replan_observed - nonlocal pickup_dynamics_cleared, pickup_start_command + nonlocal pickup_dynamics_cleared, plan_start_command if ( not args.no_target_motion and not target_scene.moved @@ -354,13 +324,15 @@ def on_step(step: RunnerStep) -> None: "while the robot holds its current command." ) moved_pose = target_scene.move( - adapter, + sim_runtime, duration=TARGET_MOVE_DURATION, ) + target.set_body_type("dynamic") + target.clear_dynamics() draw_axis_marker( sim, "moving_target_replanned_goal", - _compose_goal_pose(moved_pose, target_to_eef), + _compose_goal_pose(moved_pose, target_to_grasp), axis_len=0.10, ) displacement = torch.linalg.vector_norm( @@ -376,6 +348,7 @@ def on_step(step: RunnerStep) -> None: 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, @@ -386,27 +359,17 @@ def on_step(step: RunnerStep) -> None: f"Execution event {event.kind.value}: env rows={env_ids}; " f"{event.message}" ) - dynamic_change_observed |= ( - event.kind is ExecutionEventKind.DYNAMIC_GOAL_CHANGED - ) - replan_observed |= event.kind is ExecutionEventKind.REPLANNED - if ( - event.kind is ExecutionEventKind.ACTION_PLANNED - and event.invocation_index == 1 - and pickup_start_command is None - ): - target.set_body_type("dynamic") - target.clear_dynamics() - pickup_start_command = step.command_count + if event.kind is ExecutionEventKind.REPLANNED: + plan_start_command = step.command_count logger.log_info( - "The approach completed; the cube is now dynamic and PickUp " - "is starting.", + "PickUp discarded the stale plan and restarted from the " + "latest cube pose.", color="green", ) if ( - pickup_start_command is not None + (args.no_target_motion or target_scene.moved) and not pickup_dynamics_cleared - and step.command_count - pickup_start_command > clear_after_pick_command + and step.command_count - plan_start_command > clear_after_pick_command ): target.clear_dynamics() pickup_dynamics_cleared = True @@ -451,7 +414,7 @@ def verify_pickup_effect( on_step=on_step, ) for _ in range(POST_EXECUTION_UPDATES): - adapter.sleep(adapter.physics_dt) + sim_runtime.sleep(sim_runtime.physics_dt) finally: stop_auto_play_recording(sim, recording_started) @@ -460,12 +423,10 @@ def verify_pickup_effect( if not args.no_target_motion: if not target_scene.moved: raise RuntimeError("Execution completed before the target could move.") - if not dynamic_change_observed: + if ExecutionEventKind.DYNAMIC_GOAL_CHANGED not in observed_events: raise RuntimeError("The target move was not reported as a dynamic change.") - if not replan_observed: + if ExecutionEventKind.REPLANNED not in observed_events: raise RuntimeError("The target move did not trigger replanning.") - if pickup_start_command is None: - raise RuntimeError("PickUp did not start after the recovered approach.") if not pickup_dynamics_cleared: raise RuntimeError("The cube was not stabilized after gripper closure.") logger.log_info( 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()) From 4a49273d859642199aa76ef44a2ba61e2b8dea6b Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 9 Aug 2026 22:01:28 +0800 Subject: [PATCH 5/5] feat(atomic-actions): push recovery target with force --- docs/source/tutorial/atomic_actions.rst | 3 +- .../atomic_action/moving_target_recovery.py | 72 +++++++++++++------ 2 files changed, 52 insertions(+), 23 deletions(-) diff --git a/docs/source/tutorial/atomic_actions.rst b/docs/source/tutorial/atomic_actions.rst index 7a82971c9..94575776a 100644 --- a/docs/source/tutorial/atomic_actions.rst +++ b/docs/source/tutorial/atomic_actions.rst @@ -288,7 +288,8 @@ 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 slides it sideways during one ``PickUp`` invocation whose +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 diff --git a/scripts/tutorials/atomic_action/moving_target_recovery.py b/scripts/tutorials/atomic_action/moving_target_recovery.py index 8db1d12c3..0fbde22b0 100644 --- a/scripts/tutorials/atomic_action/moving_target_recovery.py +++ b/scripts/tutorials/atomic_action/moving_target_recovery.py @@ -85,13 +85,15 @@ 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 move it exactly once.""" + """Publish a versioned target pose and physically push it exactly once.""" def __init__( self, @@ -126,44 +128,70 @@ def snapshot(self, timestamp: float) -> SceneSnapshot: }, ) - def move( + def push( self, clock: SimulationExecutionAdapter, *, duration: float, + force_duration: float, + force_magnitude: float, ) -> torch.Tensor: - """Animate the visible target and advance the scene version. + """Push the visible target with a short force pulse. Args: - clock: Simulation adapter used to advance physics between poses. - duration: Requested target-motion duration in seconds. + 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: - Updated batched target pose. + 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)) - pose = start_pose.clone() - for step_index in range(1, step_count + 1): - alpha = step_index / step_count - pose[:, :3, 3] = torch.lerp( - start_pose[:, :3, 3], - self.destination, - alpha, - ) - self.target.set_local_pose(pose) + 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, initially kinematic for scripted relocation.""" + """Create the bright cube, held kinematic until the physical push.""" return sim.add_rigid_object( cfg=RigidObjectCfg( uid=TARGET_ENTITY_ID, @@ -320,15 +348,15 @@ def on_step(step: RunnerStep) -> None: and step.command_count >= MOVE_AFTER_COMMAND ): logger.log_warning( - f"Animating the blue target for {TARGET_MOVE_DURATION:.1f} s " - "while the robot holds its current command." + 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.move( + moved_pose = target_scene.push( sim_runtime, duration=TARGET_MOVE_DURATION, + force_duration=TARGET_PUSH_DURATION, + force_magnitude=TARGET_PUSH_FORCE, ) - target.set_body_type("dynamic") - target.clear_dynamics() draw_axis_marker( sim, "moving_target_replanned_goal", @@ -340,7 +368,7 @@ def on_step(step: RunnerStep) -> None: dim=1, ) logger.log_warning( - "Moved the blue target after " + "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."