diff --git a/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.expert_program.rst b/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.expert_program.rst new file mode 100644 index 000000000..c4d7d1f4d --- /dev/null +++ b/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.expert_program.rst @@ -0,0 +1,95 @@ +embodichain.lab.gym.envs.expert_program +======================================= + +.. automodule:: embodichain.lab.gym.envs.expert_program + + .. autosummary:: + + ExpertProgramCfg + ExpertProgramIntegrationCfg + ExpertProgramCompiler + CompiledProgram + load_expert_program + loads_expert_program_json + parse_expert_program_json + decode_expert_program + ExpertProgramEnvironmentMixin + ExpertProgramEnvironmentAdapter + SimulationSceneBinding + SimulationResourceEndpointBinding + SimulationRobotResourceBinding + RobotResourceBinding + ControlPartEndpointBinding + ControlPartResourceBinding + SimulationRobotSkillProfileBinding + SimulationExpertProgramFactory + SimulationSegmentPolicyPort + ControlCommandStateEvidenceTracker + +.. currentmodule:: embodichain.lab.gym.envs.expert_program + +Schema and loading +------------------ + +The public decoders and file loaders support Expert Program schema versions 1 +and 2. Version 2 adds deterministic parallel blocks with explicit barriers. + +.. autoclass:: ExpertProgramCfg + :members: + +.. autoclass:: ExpertProgramIntegrationCfg + :members: + +.. autofunction:: load_expert_program + +.. autofunction:: loads_expert_program_json + +.. autofunction:: parse_expert_program_json + +.. autofunction:: decode_expert_program + +Compilation and environment integration +--------------------------------------- + +.. autoclass:: ExpertProgramCompiler + :members: + +.. autoclass:: CompiledProgram + :members: + +.. autoclass:: ExpertProgramEnvironmentMixin + :members: + +.. autoclass:: ExpertProgramEnvironmentAdapter + :members: + +Simulation integration +---------------------- + +.. autoclass:: SimulationSceneBinding + :members: + +.. autoclass:: SimulationResourceEndpointBinding + +.. autoclass:: SimulationRobotResourceBinding + +.. autoclass:: RobotResourceBinding + :members: + +.. autoclass:: ControlPartEndpointBinding + :members: + +.. autoclass:: ControlPartResourceBinding + :members: + +.. autoclass:: SimulationRobotSkillProfileBinding + :members: + +.. autoclass:: SimulationExpertProgramFactory + :members: + +.. autoclass:: SimulationSegmentPolicyPort + :members: + +.. autoclass:: ControlCommandStateEvidenceTracker + :members: diff --git a/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.rst b/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.rst index f5617a955..6c6c5dd91 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.rst @@ -21,9 +21,15 @@ through :func:`~embodichain.lab.gym.utils.registration.make`. .. autosummary:: demo + expert_program managers wrapper +.. toctree:: + :hidden: + + embodichain.lab.gym.envs.expert_program + .. currentmodule:: embodichain.lab.gym.envs Environment Classes @@ -60,6 +66,9 @@ segment spans. .. autoclass:: DemoSegment :members: +.. autoclass:: ProcessedEnvAction + :members: + .. autoclass:: DemoSegmentResult :members: diff --git a/docs/source/api_reference/embodichain/embodichain.utils.rst b/docs/source/api_reference/embodichain/embodichain.utils.rst index 36aa780f0..18b6021c8 100644 --- a/docs/source/api_reference/embodichain/embodichain.utils.rst +++ b/docs/source/api_reference/embodichain/embodichain.utils.rst @@ -19,6 +19,7 @@ and image processing. warp cfg configclass + config_paths device_utils file img_utils @@ -46,6 +47,12 @@ Configuration Classes :undoc-members: :show-inheritance: +Configuration Paths +------------------- + +.. automodule:: embodichain.utils.config_paths + :members: + Configuration Nodes ------------------- diff --git a/docs/source/overview/sim/atomic_actions/expert_programs.md b/docs/source/overview/sim/atomic_actions/expert_programs.md new file mode 100644 index 000000000..532138d90 --- /dev/null +++ b/docs/source/overview/sim/atomic_actions/expert_programs.md @@ -0,0 +1,255 @@ +(expert-programs)= + +# Declarative Expert Programs + +Expert Programs let a task describe semantic intent without implementing a +task-local motion generator. A program names registered scene entities, robot +profiles, runtime presets, semantic calls, post-policies, and validators. The +shared compiler lowers every call just in time through the same +`SemanticSkillCompiler`, `AtomicActionEngine`, and `SkillRuntime` used by the +Python semantic API. + +Use an Expert Program when later motion depends on the physical result of an +earlier call. Each call receives a fresh scene observation, owns one +`ExecutionSession`, verifies its physical effect, and commits verified symbolic +state before the next call is grounded. + +## Author a program + +Schema version 1 supports bounded `sequence`, `repeat`, `segment`, and `invoke` +nodes. Schema version 2 additionally supports deterministic `parallel` blocks +and explicit `barrier` nodes. Unknown fields, unsupported discriminators, +unbounded structures, executable values, and dotted environment traversal are +rejected before physical execution or command emission. + +`RegisteredSemanticCall` is an opaque extension boundary in these schema +versions. An extension with a physical effect must also register its typed +compiler/effect contract; a serialized call ID alone cannot manufacture effect +verification semantics. + +The repeated-cube task is configured entirely as semantic calls: + +```yaml +schema_version: 1 +program_id: repeated_cube_pick_place +integration: + robot_profile: ur5_parallel_gripper_v1 + scene_registry: multi_segments_cube_v1 + runtime_preset: safe +targets: + drop_pose: + kind: cyclic_pose + values: + - position: [-0.40, 0.48, 0.10] + quaternion_wxyz: [1.0, 0.0, 0.0, 0.0] +program: + kind: repeat + count: 3 + body: + kind: segment + name: move_cube + steps: + kind: sequence + items: + - kind: invoke + call: {kind: pick, object: cube} + - kind: invoke + call: + kind: place + object: cube + at: {kind: target_ref, target: drop_pose} + post: + - {kind: wait_stable, entity: cube, preset: rigid_object} + validators: + - kind: object_near_target + object: cube + target: drop_pose + position_tolerance: 0.12 +``` + +The top-level Gym configuration selects the file with a path relative to that +configuration file: + +```json +{ + "expert_program_path": "../../expert_program/my_task.yaml" +} +``` + +`run_env` can override it explicitly: + +```bash +python -m embodichain.lab.scripts.run_env \ + --gym_config path/to/gym.json \ + --expert-program path/to/program.yaml +``` + +## Accept untrusted model output + +Model-generated programs use the same decoder and compiler, but enter through +the narrower MLLM frontend. The trusted host owns the scene, robot profile, and +runtime preset; the model response must omit `integration` entirely: + +```python +from embodichain.agents.mllm import compile_mllm_expert_program +from embodichain.lab.gym.envs.expert_program import ExpertProgramIntegrationCfg + +compiled = compile_mllm_expert_program( + model_response, + adapter=adapter, + integration=ExpertProgramIntegrationCfg( + robot_profile="my_robot_v1", + scene_registry="my_scene_v1", + runtime_preset="safe", + ), +) +``` + +This entry point accepts exactly one bounded JSON document. It rejects duplicate +keys, non-finite or overflowing numeric values, invalid Unicode, Markdown +fences, trailing text, and every normal schema violation. Its initial policy is +deliberately smaller than the file format: only schema version 1 and curated +`pick`, `place`, `hand_over`, and `operate_articulation` calls are admitted. +The model cannot select `resources`, a hand-over `receiver`, a runtime preset, +or an explicit articulation position/displacement; articulation operations must +use a host-declared named target. Registered calls and parallel nodes remain +host-authored extensions. + +`compile_mllm_expert_program` delegates to the existing +`ExpertProgramEnvironmentAdapter.compile` method. It neither creates a second +compiler nor assembles a runtime while validating model output. + +## Integrate a simulation task + +Task code supplies typed scene and robot integration declarations once, while +the external Expert Program configuration owns task sequence and targets. The +task then delegates runtime assembly to the shared factory; it does not +construct approach, grasp, pull, or placement trajectories: + +```python +class MyTaskEnv(ExpertProgramEnvironmentMixin, EmbodiedEnv): + def __init__(self, cfg, **kwargs): + super().__init__(cfg, **kwargs) + self._expert_program_adapter = create_simulation_expert_program_adapter( + self, + scene_binding=create_my_scene_binding(), + robot_profile_binding=create_my_robot_profile_binding(), + ) + + @property + def expert_program_adapter(self): + return self._expert_program_adapter +``` + +The scene binding is authoritative for semantic identity, live pose sources, +geometry, affordances, and collision roles. The robot profile owns reusable +resources, endpoint capabilities, semantic commands, policy presets, and effect +monitor selection. `SimulationRobotSkillProfileBinding` accepts generic +`RobotResourceBinding` declarations containing arbitrary typed +`ResourceEndpoint` values; `ControlPartResourceBinding` is its stricter +joint-backed convenience. Endpoint adapters and runtime transports are the +extension boundary for mobile-base, whole-body, or non-joint controllers and +are accepted by the standard simulation helper. Task programs keep the same +semantic calls and do not gain controller-shaped fields. + +Relation and rendezvous semantics are also explicit integration capabilities. +`Place(on=...)` and `Place(inside=...)` require an exact typed/versioned +`RelationTargetGrounder` for the selected affordance payload. `HandOver` +requires the profile-selected `HandOverPoseProvider`. A direct `Place(at=...)` +does not require a relation grounder. Missing, ambiguous, or stale providers +fail during provider-aware program preflight, before the first physical action. + +The same preflight rejects a reachable `safe` preset for a dynamic scene before +the first observation when the active motion generator cannot provide the +required dynamic collision world. + +## Execution and physical effects + +`AtomicDemoBridge` yields lazy `DemoSegment` actions. Every command and settling +hold is consumed by normal `env.step()`, so action managers, recorders, rewards, +timing, and dataset boundaries remain authoritative. `BaseEnv.step_dt` is the +only control cadence; a command duration that is not representable on that grid +fails instead of being silently resampled. + +Creating a bridge materializes the bounded segment stream and performs +provider-aware semantic analysis before any command can be emitted. Sequential +stretches retain downstream object-target look-ahead across segment boundaries; +an explicit parallel block is a conservative look-ahead barrier. Runtime still +re-observes and grounds each call just in time after prior verified effects. + +The standard simulation integration verifies grasp and release with two pieces +of evidence: + +- the last exact open/grasp command accepted by the buffered Gym command sink, + tracked independently for every stable environment ID; and +- the live object-to-endpoint pose relation from the shared scene snapshot. + +The command-state update is transactional: encoder, buffer, cancellation, or +safe-stop failures invalidate it. An integration with contact, constraint, +force, or wrench sensing can install typed evidence callbacks without changing +the semantic call or program. + +Program/demo-segment metadata records runtime call results, named trajectory +segments, effect decisions, recovery events, scene and collision revisions, +settling outcomes, and validator results in deterministic JSON-safe values. +Trajectory segments are trace ranges inside one atomic plan; they do not create +independent recovery or timeout boundaries. + +Schema-version-2 parallel blocks additionally require an authoritative +`ParallelCommandSafetyValidator`. Resource-claim disjointness is necessary but +is not treated as proof of physical safety. If no validator is installed, the +parallel block refuses to start; the standard simulation adapter intentionally +does not invent one from resource names. Every parallel frame must occupy +exactly one `BaseEnv.step_dt`; shorter lanes repeat their last safe target as +hold padding, while fractional frames are rejected rather than resampled. +Version 2 also uses strict symbolic key-level conflict detection at the barrier: +two branches may not commit the same task-state key, even when their physical +changes occurred in disjoint environment rows. + +## Python semantic calls + +Standalone applications can use the same compiler and runtime through +`AtomicSkills`: + +```python +skills = AtomicSkills.from_env(runtime_provider, preset="safe") +cube = skills.scene.object("cube") +tray = skills.scene.object("tray") +result = skills.run(Pick(object=cube), Place(object=cube, on=tray)) +``` + +In this example, `runtime_provider` owns the typed relation grounder for the +tray's placement affordance. Applications without such a provider can use a +direct `SemanticPose` through `Place(at=...)`. + +`from_env` requires an explicit `SkillRuntimeProvider`; it never scans arbitrary +environment attributes. Gym demonstration environments intentionally use the +lazy bridge instead, because a synchronous runtime would bypass the required +`env.step()` handshake. Advanced applications may use +`AtomicSkills.from_components(...)` with explicit observation, command, +evidence, and clock ports. + +For the lower-level planning and execution contracts, see {doc}`index`. For +robot resource and endpoint declarations, see {doc}`robot_skill_profiles`. + +## Capability status + +| Surface | Shared contract | Standard simulation integration | +| --- | --- | --- | +| `Pick` | Compiler, runtime, effect verification | Antipodal grasp binding plus motion/grasp resources | +| `Place(at=...)` | Object-centric lowering with verified held state | Direct semantic pose target | +| `Place(on=...)` / `Place(inside=...)` | Exact typed relation dispatch | Integration must install the matching `RelationTargetGrounder` | +| `HandOver` | Coordinated call, state flow, and effect contract | Embodiment must install its named `HandOverPoseProvider` and evidence sources | +| `OperateArticulation` | Named/absolute/displacement target and joint effect | Link, joint, operation-affordance, and interaction endpoint bindings | +| Registered calls | Typed call catalog and explicit lowerer | Physical extensions must add an explicit effect contract | +| Mobile/whole-body extensions | Generic resources, claims, endpoint targets, command frames, and routing | Requires a reusable semantic skill/lowerer plus matching adapter, payload, transport, and effect integration; no curated navigation or whole-body skill is installed today | +| Parallel blocks | Shared-clock coordinator and strict barrier merge | Requires an authoritative `ParallelCommandSafetyValidator`; none is inferred by default | + +The table separates implemented reusable contracts from embodiment-specific +providers. It is not a claim that every row has completed task-level physical +simulation acceptance. + +The Open Drawer vertical slice has completed its supported-simulation physical +run and reached the configured drawer joint target. Repeated cube pick/place has +completed one physical Pick/Place/settle/validator cycle; its full three-cycle +run remains in threshold calibration. diff --git a/embodichain/lab/gym/envs/__init__.py b/embodichain/lab/gym/envs/__init__.py index 14c7e98bf..19dc53837 100644 --- a/embodichain/lab/gym/envs/__init__.py +++ b/embodichain/lab/gym/envs/__init__.py @@ -21,6 +21,7 @@ from .base_env import * from .demo import * from .embodied_env import * +from .settling import * from .wrapper import * # Official task environments live in the bundled ``embodichain_tasks`` import diff --git a/embodichain/lab/gym/envs/demo.py b/embodichain/lab/gym/envs/demo.py index 695500a8d..dc3be1ef6 100644 --- a/embodichain/lab/gym/envs/demo.py +++ b/embodichain/lab/gym/envs/demo.py @@ -20,9 +20,14 @@ from collections.abc import Callable, Iterable, Mapping from dataclasses import dataclass, field, replace -from typing import Any +import math +from types import MappingProxyType +from typing import Any, Literal import torch +from tensordict import TensorDict + +from embodichain.lab.sim.types import EnvAction __all__ = [ "DEMO_ANNOTATION_KEYS", @@ -30,6 +35,7 @@ "DemoEpisodeResult", "DemoSegment", "DemoSegmentResult", + "ProcessedEnvAction", "execute_demo_episode", "resolve_demo_segments", ] @@ -50,6 +56,69 @@ """Per-frame annotation keys stored in expert rollout buffers.""" +def _json_safe_copy(value: Any, *, field_name: str) -> Any: + """Return an owned JSON value without implicit type coercion.""" + if value is None or type(value) in {bool, int, str}: + return value + if type(value) is float: + if not math.isfinite(value): + raise ValueError(f"{field_name} contains a non-finite float.") + return value + if isinstance(value, Mapping): + result: dict[str, Any] = {} + for key, item in value.items(): + if type(key) is not str or not key or key != key.strip(): + raise ValueError( + f"{field_name} mapping keys must be non-empty strings " + "without outer whitespace." + ) + result[key] = _json_safe_copy( + item, + field_name=f"{field_name}.{key}", + ) + return result + if isinstance(value, (list, tuple)): + return [ + _json_safe_copy(item, field_name=f"{field_name}[{index}]") + for index, item in enumerate(value) + ] + raise TypeError(f"{field_name} contains non-JSON value {type(value).__name__}.") + + +@dataclass(frozen=True, slots=True, eq=False) +class ProcessedEnvAction: + """Owned controller-ready action that must still pass through ``env.step``. + + Semantic runtimes and demonstration bridges may already have produced the + action-manager output (for example, a full joint-position command assembled + from typed runtime endpoints). Wrapping it prevents the environment from + applying the pre-action transform a second time while retaining the normal + simulation, manager, recorder, reward, and dataset step lifecycle. + + Args: + value: Controller-ready tensor or ``TensorDict``. + metadata: JSON-compatible provenance attached by the producer. The + environment does not interpret this mapping. + """ + + value: EnvAction + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not isinstance(self.value, (torch.Tensor, TensorDict)): + raise TypeError("value must be a torch.Tensor or TensorDict.") + if not isinstance(self.metadata, Mapping): + raise TypeError("metadata must be a mapping.") + owned_value = self.value.clone() + owned_metadata = _json_safe_copy(self.metadata, field_name="metadata") + object.__setattr__(self, "value", owned_value) + object.__setattr__(self, "metadata", MappingProxyType(owned_metadata)) + + def snapshot(self) -> ProcessedEnvAction: + """Return an independently owned processed-action envelope.""" + return ProcessedEnvAction(value=self.value, metadata=self.metadata) + + @dataclass(frozen=True) class DemoSegment: """One semantic subtask inside a demonstration episode. @@ -69,6 +138,16 @@ class DemoSegment: parallel environment (or one scalar broadcast to every environment). Gym ``terminated`` and ``truncated`` remain episode-level signals; use this callback for subtask-level validation. + abort_actions: Optional callback invoked when the executor stops after + retrieving an action but before exhausting the iterable. It receives + a reason and ``last_action_consumed`` flag, and must return any + emergency controller actions that still need ordinary ``env.step`` + consumption. This is the explicit cancellation handshake for lazy + runtimes whose command acknowledgements only mean locally buffered. + failure_policy: ``"batch_abort"`` preserves legacy batch-atomic + behavior. ``"row_independent"`` permanently freezes only failed + environment rows while peers continue through the shared segment + and later lazy segments. """ actions: Iterable[Any] @@ -77,6 +156,20 @@ class DemoSegment: instruction: str | None = None metadata: Mapping[str, Any] = field(default_factory=dict) validator: Callable[[], Any] | None = field(default=None, repr=False, compare=False) + abort_actions: Callable[..., Iterable[Any]] | None = field( + default=None, + repr=False, + compare=False, + ) + failure_policy: Literal["batch_abort", "row_independent"] = "batch_abort" + + def __post_init__(self) -> None: + if self.abort_actions is not None and not callable(self.abort_actions): + raise TypeError("abort_actions must be callable or None.") + if self.failure_policy not in {"batch_abort", "row_independent"}: + raise ValueError( + "failure_policy must be 'batch_abort' or 'row_independent'." + ) @dataclass(frozen=True) @@ -118,6 +211,15 @@ class DemoSegmentResult: successes: tuple[bool, ...] = () failure_reasons: tuple[str | None, ...] = () + def __post_init__(self) -> None: + if not isinstance(self.metadata, Mapping): + raise TypeError("metadata must be a mapping.") + owned_metadata = _json_safe_copy( + self.metadata, + field_name="segment result metadata", + ) + object.__setattr__(self, "metadata", MappingProxyType(owned_metadata)) + def to_metadata(self, env_id: int | None = None) -> dict[str, Any]: """Return a JSON-compatible aggregate or per-environment representation. @@ -133,7 +235,10 @@ def to_metadata(self, env_id: int | None = None) -> dict[str, Any]: "name": self.name, "target_uid": self.target_uid, "instruction": self.instruction, - "metadata": dict(self.metadata), + "metadata": _json_safe_copy( + self.metadata, + field_name="segment result metadata", + ), } if env_id is not None and self.start_steps: metadata.update( @@ -269,6 +374,28 @@ def _as_bool_tuple(value: Any, num_envs: int) -> tuple[bool, ...]: return tuple(bool(item) for item in tensor.tolist()) +def _has_terminal_runtime_failure_trace(segment: DemoSegment) -> bool: + """Return whether a lazy segment recorded a canonical failed runtime. + + Expert-program action iterables may terminate before yielding a controller + command when planning fails. Their bridge finalizes the runtime trace while + exhausting the iterable and exposes a validator that commits row-local + failure. This marker distinguishes that outcome from an ordinary empty + ``DemoSegment``, whose existing ``empty_segment`` guard remains unchanged. + """ + runtime = segment.metadata.get("runtime") + if not isinstance(runtime, Mapping): + return False + return ( + runtime.get("kind") + in { + "skill_result", + "parallel_skill_result", + } + and runtime.get("status") == "failed" + ) + + def _dataset_instruction(env: Any) -> str: """Return the dataset-level instruction used for legacy demo segments.""" metadata = getattr(_env_target(env), "metadata", {}) @@ -448,7 +575,20 @@ def publish_active_mask() -> None: f"{segment.name}", ) - for action in actions: + action_iterator = iter(actions) + last_action_consumed: bool | None = None + action_error: Exception | None = None + while True: + try: + action = next(action_iterator) + except StopIteration: + break + except Exception as exc: + action_error = exc + actions_exhausted = False + segment_reason = "action_generation_failed" + break + last_action_consumed = False if should_stop is not None and should_stop(): actions_exhausted = False fatal_reason = "interrupted" @@ -461,18 +601,32 @@ def publish_active_mask() -> None: publish_active_mask() break - if normalize_action is not None: - action = normalize_action(action) - if not all(active): - if mask_action is None: - raise RuntimeError( - "A vector demo environment completed asynchronously but " - "does not implement _mask_demo_action(action, active_mask)." - ) - action = mask_action(action, tuple(active)) + try: + if normalize_action is not None: + action = normalize_action(action) + if not all(active): + if mask_action is None: + raise RuntimeError( + "A vector demo environment completed asynchronously " + "but does not implement " + "_mask_demo_action(action, active_mask)." + ) + action = mask_action(action, tuple(active)) + except Exception as exc: + action_error = exc + actions_exhausted = False + segment_reason = "action_processing_failed" + break active_before_step = tuple(active) - _, _, terminated_value, truncated_value, info = env.step(action) + try: + _, _, terminated_value, truncated_value, info = env.step(action) + except Exception as exc: + action_error = exc + actions_exhausted = False + segment_reason = "action_execution_failed" + break + last_action_consumed = True action_count += 1 last_info = info for env_id, was_active in enumerate(active_before_step): @@ -529,16 +683,21 @@ def publish_active_mask() -> None: step_failed = True if step_failed: - actions_exhausted = False - fatal_reason = "truncated" if active_step_truncated else "failure" - segment_reason = fatal_reason - for env_id, is_active in enumerate(active): - if is_active: - terminal_reasons[env_id] = "batch_aborted" - segment_failure_reasons[env_id] = "batch_aborted" - active[env_id] = False + if segment.failure_policy == "batch_abort": + actions_exhausted = False + fatal_reason = ( + "truncated" if active_step_truncated else "failure" + ) + segment_reason = fatal_reason + for env_id, is_active in enumerate(active): + if is_active: + terminal_reasons[env_id] = "batch_aborted" + segment_failure_reasons[env_id] = "batch_aborted" + active[env_id] = False publish_active_mask() - break + if segment.failure_policy == "batch_abort" or not any(active): + actions_exhausted = False + break publish_active_mask() if not any(active): @@ -558,7 +717,78 @@ def publish_active_mask() -> None: publish_active_mask() break - if action_count == 0 and segment_reason is None: + if not actions_exhausted: + if segment.abort_actions is not None: + reason = ( + segment_reason + or fatal_reason + or "demo segment execution stopped before exhaustion" + ) + try: + emergency_actions = segment.abort_actions( + reason, + last_action_consumed=bool(last_action_consumed), + ) + if isinstance(emergency_actions, (str, bytes)): + raise TypeError( + "abort_actions must return an iterable of actions." + ) + emergency_iterator = iter(emergency_actions) + try: + for emergency_action in emergency_iterator: + if normalize_action is not None: + emergency_action = normalize_action( + emergency_action + ) + try: + _, _, _, _, emergency_info = env.step( + emergency_action + ) + except Exception as exc: + raise RuntimeError( + "Emergency demo safe-stop action failed " + "during env.step()." + ) from exc + action_count += 1 + last_info = emergency_info + for env_id, is_participant in enumerate(participants): + if is_participant: + lengths[env_id] += 1 + finally: + close_emergency = getattr( + emergency_iterator, + "close", + None, + ) + if callable(close_emergency): + close_emergency() + finally: + close_actions = getattr(action_iterator, "close", None) + if callable(close_actions): + close_actions() + else: + close_actions = getattr(action_iterator, "close", None) + if callable(close_actions): + close_actions() + + if action_error is not None: + raise RuntimeError( + "Demo action generation, processing, or execution failed " + "after an emergency safe-stop attempt." + ) from action_error + + traced_terminal_runtime_failure = ( + action_count == 0 + and actions_exhausted + and segment_reason is None + and segment.validator is not None + and _has_terminal_runtime_failure_trace(segment) + ) + if ( + action_count == 0 + and segment_reason is None + and not traced_terminal_runtime_failure + ): fatal_reason = "empty_segment" segment_reason = fatal_reason for env_id, is_participant in enumerate(participants): @@ -588,15 +818,25 @@ def publish_active_mask() -> None: terminal_reasons[env_id] = "segment_validation_failed" if validation_failed: - fatal_reason = "segment_validation_failed" - segment_reason = fatal_reason - for env_id, is_active in enumerate(active): - if is_active: - if segment_failure_reasons[env_id] is None: - segment_failure_reasons[env_id] = "batch_aborted" - terminal_reasons[env_id] = "batch_aborted" - segment_successes[env_id] = False - active[env_id] = False + if segment.failure_policy == "batch_abort": + fatal_reason = "segment_validation_failed" + segment_reason = fatal_reason + for env_id, is_active in enumerate(active): + if is_active: + if segment_failure_reasons[env_id] is None: + segment_failure_reasons[env_id] = "batch_aborted" + terminal_reasons[env_id] = "batch_aborted" + segment_successes[env_id] = False + active[env_id] = False + else: + for env_id, is_active in enumerate(active): + if ( + is_active + and segment_failure_reasons[env_id] + == "segment_validation_failed" + ): + segment_successes[env_id] = False + active[env_id] = False publish_active_mask() participant_ids = [ diff --git a/embodichain/lab/gym/envs/embodied_env.py b/embodichain/lab/gym/envs/embodied_env.py index 51b2fa966..ba18a6a91 100644 --- a/embodichain/lab/gym/envs/embodied_env.py +++ b/embodichain/lab/gym/envs/embodied_env.py @@ -27,7 +27,17 @@ import gymnasium as gym from dataclasses import MISSING -from typing import Dict, Union, Sequence, Tuple, Any, Iterable, List, Optional +from typing import ( + TYPE_CHECKING, + Dict, + Union, + Sequence, + Tuple, + Any, + Iterable, + List, + Optional, +) from tensordict import TensorDict from embodichain.lab.sim.cfg import ( @@ -52,6 +62,7 @@ DemoEpisodeResult, DemoSegment, DemoSegmentResult, + ProcessedEnvAction, ) from embodichain.lab.gym.envs.managers import ( EventManager, @@ -70,6 +81,13 @@ from embodichain.data import get_data_path from embodichain.data.constants import EMBODICHAIN_DEFAULT_DATA_ROOT +if TYPE_CHECKING: + from embodichain.lab.gym.envs.expert_program import ( + CompiledProgram, + ExpertProgramCfg, + ) + from embodichain.lab.gym.envs.expert_program.bridge import AtomicDemoBridge + __all__ = ["EmbodiedEnvCfg", "EmbodiedEnv"] @@ -240,6 +258,14 @@ class EnvLightCfg: """If True (and record_trajectory is True), auto-save each env's trajectory to ``trajectory_save_dir`` at episode end and on close().""" + expert_program: ExpertProgramCfg | None = None + """Optional declarative Expert Program used to generate demo segments. + + The program remains inert until :meth:`EmbodiedEnv.create_demo_segments` + requests an explicit environment compiler and bridge through the dedicated + hooks. No live provider, planner, or callable is stored in this config. + """ + @register_env("EmbodiedEnv-v1") class EmbodiedEnv(BaseEnv): @@ -1248,14 +1274,30 @@ def _write_rl_rollout_step( : self.num_envs, self.current_rollout_step ].copy_(truncateds.to(buffer_device), non_blocking=True) - def _normalize_demo_action(self, action: EnvAction) -> EnvAction: - """Normalize one legacy or segment action to the environment action space.""" + def _normalize_demo_action( + self, action: EnvAction | ProcessedEnvAction + ) -> EnvAction | ProcessedEnvAction: + """Normalize one raw action or preserve a controller-ready envelope.""" + if isinstance(action, ProcessedEnvAction): + value = action.value + if value.ndim == 0: + raise ValueError( + "Processed demo actions must have a leading environment " + "dimension." + ) + if value.shape[0] != self.num_envs: + raise ValueError( + "Processed demo action batch size must match num_envs." + ) + return action.snapshot() expected_dim = int(np.prod(self.single_action_space.shape)) return self._normalize_demo_action_tensor(action, expected_dim) def _mask_demo_action( - self, action: EnvAction, active_mask: Sequence[bool] - ) -> EnvAction: + self, + action: EnvAction | ProcessedEnvAction, + active_mask: Sequence[bool], + ) -> EnvAction | ProcessedEnvAction: """Accept an asynchronously completed vector-demo action. Raw actions may still require :class:`ActionManager` preprocessing, so @@ -1548,15 +1590,18 @@ def evaluate(self, **kwargs) -> Dict[str, Any]: eval_dict[key] = value return eval_dict - def _preprocess_action(self, action: EnvAction) -> EnvAction: - """Delegate to ActionManager when configured; stash raw action for trajectory.""" + def _preprocess_action(self, action: EnvAction | ProcessedEnvAction) -> EnvAction: + """Apply raw preprocessing once and stash the executed controller action.""" + is_processed = isinstance(action, ProcessedEnvAction) + if is_processed: + action = action.value if self._traj_buffer is not None: self._traj_raw_action = ( action.clone() if hasattr(action, "clone") else action ) - if self.action_manager is not None: + if self.action_manager is not None and not is_processed: action = self.action_manager.process_action(action, mode="pre") - else: + elif not is_processed: action = super()._preprocess_action(action) if getattr(self, "_demo_no_auto_reset", False): action = self._mask_processed_demo_action(action) @@ -1763,13 +1808,65 @@ def create_demo_action_list(self, *args, **kwargs) -> Sequence[EnvAction] | None "The method 'create_demo_action_list' must be implemented in subclasses." ) + def compile_expert_program( + self, + program: ExpertProgramCfg, + ) -> CompiledProgram: + """Compile a configured Expert Program through explicit scene providers. + + Declarative environments override this hook to supply their authoritative + scene registry/resolver to :class:`ExpertProgramCompiler`. Keeping the + provider boundary explicit prevents the base environment from inferring + identities or scanning mutable simulator internals. + + Args: + program: Strict Expert Program configuration attached to ``cfg``. + + Returns: + Provider-free compiled program ready for runtime assembly. + + Raises: + NotImplementedError: If an environment enables ``expert_program`` + without supplying the compiler/provider integration. + """ + raise NotImplementedError( + "An environment with cfg.expert_program must implement " + "compile_expert_program() using an explicit scene resolver." + ) + + def create_expert_program_bridge( + self, + program: CompiledProgram, + ) -> AtomicDemoBridge: + """Create the Gym demo bridge through explicit runtime-port factories. + + Declarative environments override this hook to assemble ``SkillRuntime`` + and Gym-aware command/clock/post-policy/validator ports. The returned + bridge must emit commands through normal ``env.step()`` processing. + + Args: + program: Compiled provider-free Expert Program. + + Returns: + Atomic demo bridge whose segments are consumed lazily. + + Raises: + NotImplementedError: If no explicit runtime factory is available. + """ + raise NotImplementedError( + "An environment with cfg.expert_program must implement " + "create_expert_program_bridge() using explicit runtime ports." + ) + def create_demo_segments(self, *args, **kwargs) -> Iterable[DemoSegment] | None: """Create the semantic segments that make up one task episode. - The default adapter preserves existing tasks by wrapping their single - ``create_demo_action_list`` result in one segment. Multi-object tasks - should override this method and may return a lazy generator so each - segment can be planned from the scene state left by the previous one. + When ``cfg.expert_program`` is configured, the environment compiles it + through an explicit scene-provider hook and creates an atomic demo bridge + through an explicit runtime-port factory hook. Otherwise, the default + adapter wraps ``create_demo_action_list`` in one segment. Multi-object + tasks may return a lazy generator so each segment can be planned from + the scene state left by the previous one. Args: *args: Positional arguments forwarded to the legacy planner. @@ -1778,6 +1875,12 @@ def create_demo_segments(self, *args, **kwargs) -> Iterable[DemoSegment] | None: Returns: Segment sequence, or ``None`` when planning fails. """ + expert_program = getattr(getattr(self, "cfg", None), "expert_program", None) + if expert_program is not None: + compiled_program = self.compile_expert_program(expert_program) + bridge = self.create_expert_program_bridge(compiled_program) + return bridge.iter_segments() + actions = self.create_demo_action_list(*args, **kwargs) if actions is None: return None diff --git a/embodichain/lab/gym/envs/expert_program/__init__.py b/embodichain/lab/gym/envs/expert_program/__init__.py new file mode 100644 index 000000000..ca245f98b --- /dev/null +++ b/embodichain/lab/gym/envs/expert_program/__init__.py @@ -0,0 +1,251 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Versioned declarative Expert Program schema, compiler, and runtime types.""" + +from __future__ import annotations + +from .cfg import ( + EXPERT_PROGRAM_SCHEMA_VERSION, + EXPERT_PROGRAM_SCHEMA_VERSION_V2, + MAX_DECLARATIVE_DEPTH, + MAX_DECLARATIVE_NODES, + MAX_EXPANDED_CALLS, + MAX_PROGRAM_DEPTH, + MAX_PROGRAM_NODES, + MAX_REPEAT_COUNT, + SUPPORTED_EXPERT_PROGRAM_SCHEMA_VERSIONS, + BarrierCfg, + CyclicPoseTargetCfg, + DeclarativeCfgValue, + ExpertProgramCfg, + ExpertProgramIntegrationCfg, + HandOverCfg, + InvokeCfg, + ObjectNearTargetValidatorCfg, + OperateArticulationCfg, + ParallelCfg, + PickCfg, + PlaceCfg, + PoseCfg, + PostPolicyCfg, + ProgramNodeCfg, + RegisteredSemanticCallCfg, + RepeatCfg, + SegmentCfg, + SemanticCallCfg, + SequenceCfg, + TargetCfg, + TargetRefCfg, + ValidatorCfg, + WaitStablePostCfg, +) +from .decoder import ( + ConfigPath, + ConfigPathPart, + ExpertProgramConfigError, + ExpertProgramDecodeError, + ExpertProgramValidationContext, + ExpertProgramValidationError, + SceneReferenceRole, + decode_expert_program, + render_config_path, + validate_expert_program, +) +from .loader import ( + MAX_EXPERT_PROGRAM_BYTES, + load_expert_program, + loads_expert_program_json, + parse_expert_program_json, +) +from .bridge import ( + AcceptedRuntimeCommandObserver, + AtomicDemoBridge, + BufferedGymCommandSink, + DemoBridgeError, + EnvironmentStepClock, + EnvironmentStepTimingError, + GymPlanningObservationProvider, + RuntimeCommandFrameEncoder, + RuntimeTransportActionEncoder, + SegmentPostPolicyMetadataPort, + SegmentPostPolicyPort, + SegmentPostPolicyResultPort, + SegmentValidatorMetadataPort, + SegmentValidatorPort, + UnsupportedRuntimeTransportError, +) +from .compiler import ( + CompiledBarrier, + CompiledParallelBlock, + CompiledParallelBranch, + CompiledPostPolicy, + CompiledProgram, + CompiledProgramAnalysis, + CompiledProgramCall, + CompiledProgramSegment, + CompiledProgramValidator, + CompiledRepeatFrame, + CompiledTargetSelection, + ExpertProgramCompileError, + ExpertProgramCompiler, + ExpertProgramSceneResolver, + MaterializedCompiledProgram, + SceneRegistryProgramResolver, +) +from .environment import ( + AcceptedRuntimeCommandObserverFactory, + ExpertProgramEnvironmentAdapter, + ExpertProgramEnvironmentFactory, + ExpertProgramEnvironmentMixin, + ExpertProgramRuntimeAssembly, + PlanningObservationPort, +) +from .simulation import ( + AntipodalGraspAffordanceBinding, + ArticulationOperationAffordanceBinding, + ArticulationOperationTargetBinding, + ControlPartCommandPreset, + ControlPartEndpointBinding, + ControlPartResourceBinding, + RobotResourceBinding, + SimulationArticulationBinding, + SimulationArticulationLinkBinding, + SimulationRigidObjectBinding, + SimulationResourceEndpointBinding, + SimulationRobotResourceBinding, + SimulationRobotSkillProfileBinding, + SimulationSceneBinding, +) +from .simulation_environment import ( + ControlCommandStateEvidenceTracker, + MotionGeneratorFactory, + SharedTickSceneProvider, + SimulationExpertProgramEnvironment, + SimulationExpertProgramFactory, + SimulationPlanningObservationProvider, + create_simulation_expert_program_adapter, +) +from .simulation_policies import SimulationSegmentPolicyPort + +__all__ = [ + "AcceptedRuntimeCommandObserver", + "AcceptedRuntimeCommandObserverFactory", + "AntipodalGraspAffordanceBinding", + "ArticulationOperationAffordanceBinding", + "ArticulationOperationTargetBinding", + "AtomicDemoBridge", + "BarrierCfg", + "BufferedGymCommandSink", + "ConfigPath", + "ConfigPathPart", + "CompiledBarrier", + "CompiledParallelBlock", + "CompiledParallelBranch", + "CompiledPostPolicy", + "CompiledProgram", + "CompiledProgramAnalysis", + "CompiledProgramCall", + "CompiledProgramSegment", + "CompiledProgramValidator", + "CompiledRepeatFrame", + "CompiledTargetSelection", + "ControlCommandStateEvidenceTracker", + "ControlPartCommandPreset", + "ControlPartEndpointBinding", + "ControlPartResourceBinding", + "CyclicPoseTargetCfg", + "DeclarativeCfgValue", + "DemoBridgeError", + "EXPERT_PROGRAM_SCHEMA_VERSION", + "EXPERT_PROGRAM_SCHEMA_VERSION_V2", + "EnvironmentStepClock", + "EnvironmentStepTimingError", + "ExpertProgramCfg", + "ExpertProgramCompileError", + "ExpertProgramCompiler", + "ExpertProgramConfigError", + "ExpertProgramDecodeError", + "ExpertProgramEnvironmentAdapter", + "ExpertProgramEnvironmentFactory", + "ExpertProgramEnvironmentMixin", + "ExpertProgramIntegrationCfg", + "ExpertProgramRuntimeAssembly", + "ExpertProgramSceneResolver", + "ExpertProgramValidationContext", + "ExpertProgramValidationError", + "HandOverCfg", + "GymPlanningObservationProvider", + "InvokeCfg", + "MAX_DECLARATIVE_DEPTH", + "MAX_DECLARATIVE_NODES", + "MAX_EXPANDED_CALLS", + "MAX_EXPERT_PROGRAM_BYTES", + "MAX_PROGRAM_DEPTH", + "MAX_PROGRAM_NODES", + "MAX_REPEAT_COUNT", + "MaterializedCompiledProgram", + "MotionGeneratorFactory", + "ObjectNearTargetValidatorCfg", + "OperateArticulationCfg", + "ParallelCfg", + "PickCfg", + "PlaceCfg", + "PlanningObservationPort", + "PoseCfg", + "PostPolicyCfg", + "ProgramNodeCfg", + "RegisteredSemanticCallCfg", + "RepeatCfg", + "RobotResourceBinding", + "RuntimeCommandFrameEncoder", + "RuntimeTransportActionEncoder", + "SceneReferenceRole", + "SceneRegistryProgramResolver", + "SegmentPostPolicyMetadataPort", + "SegmentPostPolicyPort", + "SegmentPostPolicyResultPort", + "SegmentCfg", + "SegmentValidatorMetadataPort", + "SegmentValidatorPort", + "SemanticCallCfg", + "SequenceCfg", + "SharedTickSceneProvider", + "SimulationArticulationBinding", + "SimulationArticulationLinkBinding", + "SimulationExpertProgramEnvironment", + "SimulationExpertProgramFactory", + "SimulationPlanningObservationProvider", + "SimulationRigidObjectBinding", + "SimulationResourceEndpointBinding", + "SimulationRobotResourceBinding", + "SimulationRobotSkillProfileBinding", + "SimulationSceneBinding", + "SimulationSegmentPolicyPort", + "SUPPORTED_EXPERT_PROGRAM_SCHEMA_VERSIONS", + "TargetCfg", + "TargetRefCfg", + "UnsupportedRuntimeTransportError", + "ValidatorCfg", + "WaitStablePostCfg", + "create_simulation_expert_program_adapter", + "decode_expert_program", + "load_expert_program", + "loads_expert_program_json", + "parse_expert_program_json", + "render_config_path", + "validate_expert_program", +] diff --git a/embodichain/lab/gym/envs/expert_program/bridge.py b/embodichain/lab/gym/envs/expert_program/bridge.py new file mode 100644 index 000000000..f72e616d3 --- /dev/null +++ b/embodichain/lab/gym/envs/expert_program/bridge.py @@ -0,0 +1,1574 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Gym ports and a lazy demo adapter for compiled Expert Programs. + +This module deliberately stops at the Gym action boundary. It never calls +``env.step`` and never updates a simulator directly. The demo executor owns +the environment step; when it asks the action generator for the next value, +the bridge treats the previously yielded value as consumed and advances the +environment-backed execution clock by exactly one step. +""" + +from __future__ import annotations + +from collections import deque +from collections.abc import Callable, Iterable, Iterator, Mapping +from dataclasses import dataclass, field +import math +from typing import Any, Protocol, runtime_checkable + +import torch + +from embodichain.lab.gym.envs.demo import DemoSegment, ProcessedEnvAction +from embodichain.lab.sim.atomic_actions.bindings import ( + JointPositionTarget, + RuntimeEndpointTarget, +) +from embodichain.lab.sim.atomic_actions.runner import ( + CommandAcknowledgement, + ExecutionClock, +) +from embodichain.lab.sim.atomic_actions.runtime_commands import ( + EndpointCommand, + JointPositionPayload, + RuntimeCommandFrame, +) +from embodichain.lab.sim.atomic_actions.state import PlanningContext, TaskState +from embodichain.lab.sim.skills.parallel import ParallelTimingPolicy +from embodichain.lab.sim.skills.parallel_runtime import ( + ParallelCommandSafetyValidator, + ParallelSkillResult, + ParallelSkillRuntime, +) +from embodichain.lab.sim.skills.runtime import SkillResult, SkillRuntime, SkillStatus +from embodichain.lab.sim.types import EnvAction + +_SAFE_HOLD_ACTION_KINDS = frozenset( + {"runtime_safe_hold", "runtime_wait_hold", "runtime_abort_safe_hold"} +) + + +class DemoBridgeError(RuntimeError): + """Base error raised by the Expert Program Gym bridge.""" + + +class EnvironmentStepTimingError(DemoBridgeError, ValueError): + """Raised when runtime timing cannot be represented on the Gym step grid.""" + + +class UnsupportedRuntimeTransportError(DemoBridgeError, LookupError): + """Raised when a command frame names an unregistered transport.""" + + +def _validate_identifier(value: str, *, field_name: str) -> str: + """Validate and return one strict identifier.""" + if type(value) is not str or not value or value != value.strip(): + raise ValueError( + f"{field_name} must be a non-empty string without outer whitespace." + ) + return value + + +def _validate_timeout(timeout: float) -> None: + """Validate a runner-supplied acknowledgement timeout.""" + if not isinstance(timeout, (int, float)) or isinstance(timeout, bool): + raise TypeError("timeout must be a real number.") + if not math.isfinite(float(timeout)) or float(timeout) <= 0.0: + raise ValueError("timeout must be finite and positive.") + + +@runtime_checkable +class CurrentQposProvider(Protocol): + """Source of full robot positions aligned to explicit environment IDs.""" + + def current_qpos(self, env_ids: torch.Tensor) -> torch.Tensor: + """Return ``(batch_size, robot_dof)`` positions for ``env_ids``.""" + + +@runtime_checkable +class RuntimeTransportActionEncoder(Protocol): + """Extensible lowering boundary for one runtime transport kind. + + An encoder receives the action produced by earlier registered transports + and returns the next owned action value. This permits a future transport + to promote the built-in tensor action to a ``TensorDict`` when the Gym + action manager exposes a structured controller boundary. + """ + + @property + def transport_id(self) -> str: + """Return the exact runtime transport ID handled by this encoder.""" + + def encode( + self, + command: EndpointCommand, + *, + base_action: EnvAction, + active_mask: torch.Tensor, + ) -> EnvAction: + """Merge one addressed command into ``base_action``.""" + + def hold( + self, + targets: tuple[RuntimeEndpointTarget, ...], + *, + base_action: EnvAction, + context: PlanningContext, + ) -> EnvAction: + """Merge this transport's safe state into ``base_action``.""" + + +@runtime_checkable +class AcceptedRuntimeCommandObserver(Protocol): + """Transactional observer of commands accepted by the buffered Gym sink. + + Implementations may maintain runtime-local evidence state, but must not + control the robot or advance the environment. ``accepted`` is called only + after a complete frame was encoded and appended to the local buffer. + Cancellation and discard notifications are fail-closed reset boundaries. + """ + + def accepted(self, command: RuntimeCommandFrame) -> None: + """Record one independently owned accepted command frame.""" + + def cancelled(self, targets: tuple[RuntimeEndpointTarget, ...]) -> None: + """Clear state owned by the cancelled endpoint targets.""" + + def discarded(self) -> None: + """Clear every runtime-local state value after a buffer discard.""" + + +@runtime_checkable +class CompiledProgramPort(Protocol): + """Minimal provider-free compiled-program surface consumed by the bridge.""" + + schema_version: int + program_id: str + + def iter_segments(self) -> Iterator[Any]: + """Lazily yield compiled logical segments.""" + + def sequential_execution_analysis(self, segment_index: int) -> Any: + """Return current prefix plus downstream calls up to the next barrier.""" + + +@runtime_checkable +class SequentialSkillRuntimePort(Protocol): + """Nonblocking semantic runtime surface used by sequential segments.""" + + @property + def result(self) -> SkillResult: + """Return the current immutable runtime result.""" + + @property + def status(self) -> SkillStatus: + """Return the current runtime status.""" + + def start( + self, + *calls: Any, + workflow_id: str = "semantic_workflow", + eligible_mask: torch.Tensor | None = None, + execution_prefix_length: int | None = None, + ) -> SkillResult: + """Start one semantic workflow without blocking on motion.""" + + def step(self) -> SkillResult: + """Advance the workflow by at most one due runner cycle.""" + + def cancel(self, reason: str) -> SkillResult: + """Cancel one running workflow through the runner's safe-stop path.""" + + def adopt_verified_task_state(self, task_state: TaskState) -> SkillResult: + """Install state merged at an independent parallel barrier.""" + + +@runtime_checkable +class SegmentPostPolicyPort(Protocol): + """Environment-aware program post-policy boundary. + + Implementations may observe the environment after each resumed yield, but + must return every controller action to this iterable. The bridge then + routes those values through the ordinary demo executor and ``env.step``. + """ + + def validate_policy( + self, + policy: Any, + *, + segment: Any, + ) -> None: + """Validate one compiled policy without live observation or action.""" + + def actions( + self, + policy: Any, + *, + segment: Any, + active_mask: torch.Tensor, + ) -> Iterable[Any]: + """Yield holds until ``policy`` completes for the active rows only.""" + + +@runtime_checkable +class SegmentPostPolicyMetadataPort(Protocol): + """Optional result trace supplied by a segment post-policy port.""" + + def post_policy_metadata( + self, + policy: Any, + *, + segment: Any, + ) -> Mapping[str, Any]: + """Return JSON-safe metadata after one policy has run.""" + + +@runtime_checkable +class SegmentPostPolicyResultPort(Protocol): + """Optional row-local success result supplied by a post-policy port.""" + + def post_policy_result(self, policy: Any, *, segment: Any) -> Any: + """Return one boolean or one boolean per environment row.""" + + +@runtime_checkable +class SegmentValidatorPort(Protocol): + """Environment-aware boundary for compiled program validators.""" + + def validate_validator( + self, + validator: Any, + *, + segment: Any, + ) -> None: + """Validate one compiled validator without observing the environment.""" + + def validate(self, validator: Any, *, segment: Any) -> Any: + """Return one boolean or one boolean per environment row.""" + + +@runtime_checkable +class SegmentValidatorMetadataPort(Protocol): + """Optional result trace supplied by a segment validator port.""" + + def validator_metadata( + self, + validator: Any, + *, + segment: Any, + ) -> Mapping[str, Any]: + """Return JSON-safe metadata after one validator has run.""" + + +class GymPlanningObservationProvider: + """Callback-backed observation port that also exposes the latest qpos. + + Args: + capture: Callback accepting verified :class:`TaskState` and returning + one fresh :class:`PlanningContext` from the Gym environment. + + The callback is intentionally explicit: environment-specific scene, + simulator, and registry access remains in environment integration code. + """ + + def __init__(self, capture: Callable[[TaskState], PlanningContext]) -> None: + if not callable(capture): + raise TypeError("capture must be callable.") + self._capture = capture + self._latest: PlanningContext | None = None + + @property + def latest(self) -> PlanningContext | None: + """Return the latest immutable planning context, if one was captured.""" + return self._latest + + def observe(self, task_state: TaskState) -> PlanningContext: + """Capture and retain one fresh planning context.""" + if not isinstance(task_state, TaskState): + raise TypeError("task_state must be a TaskState.") + context = self._capture(task_state) + if not isinstance(context, PlanningContext): + raise TypeError("capture must return a PlanningContext.") + self._latest = context + return context + + def current_qpos(self, env_ids: torch.Tensor) -> torch.Tensor: + """Return latest full qpos rows in the requested stable-ID order.""" + context = self._latest + if context is None: + raise RuntimeError("No planning context has been observed yet.") + if not isinstance(env_ids, torch.Tensor): + raise TypeError("env_ids must be a torch.Tensor.") + if env_ids.dtype != torch.long or env_ids.dim() != 1 or env_ids.numel() == 0: + raise ValueError("env_ids must be a non-empty one-dimensional long tensor.") + if env_ids.device != context.env_ids.device: + raise ValueError("env_ids must share the latest context device.") + if torch.unique(env_ids).numel() != env_ids.numel(): + raise ValueError("env_ids must be unique.") + + row_by_id = { + int(env_id): row + for row, env_id in enumerate(context.env_ids.detach().cpu().tolist()) + } + try: + rows = [ + row_by_id[int(env_id)] for env_id in env_ids.detach().cpu().tolist() + ] + except KeyError as exc: + raise ValueError( + f"Environment ID {int(exc.args[0])} is absent from the latest context." + ) from exc + return context.robot.qpos[rows].clone() + + +class EnvironmentStepClock(ExecutionClock): + """Monotonic execution clock advanced only by explicit Gym steps. + + ``sleep`` intentionally raises. Calling synchronous ``SkillRuntime.run`` + with this clock would otherwise advance execution without an environment + transition. Demo integrations must use the nonblocking ``start``/``step`` + path and call :meth:`advance_after_env_step` only after a yielded action was + passed to ``env.step``. + """ + + def __init__(self, step_dt: float, *, initial_step: int = 0) -> None: + if not isinstance(step_dt, (int, float)) or isinstance(step_dt, bool): + raise TypeError("step_dt must be a real number.") + if not math.isfinite(float(step_dt)) or float(step_dt) <= 0.0: + raise ValueError("step_dt must be finite and positive.") + if type(initial_step) is not int or initial_step < 0: + raise ValueError("initial_step must be a non-negative integer.") + self._step_dt = float(step_dt) + self._step_index = initial_step + + @property + def step_dt(self) -> float: + """Return the authoritative Gym control cadence.""" + return self._step_dt + + @property + def step_index(self) -> int: + """Return the number of explicitly acknowledged environment steps.""" + return self._step_index + + def now(self) -> float: + """Return deterministic environment time in seconds.""" + return self._step_index * self._step_dt + + def sleep(self, duration: float) -> None: + """Reject implicit waiting that is not backed by ``env.step``.""" + self.steps_for_duration(duration, field_name="sleep duration") + raise RuntimeError( + "EnvironmentStepClock cannot sleep or advance implicitly; use the " + "nonblocking runtime and advance_after_env_step() after env.step()." + ) + + def steps_for_duration( + self, + duration: float, + *, + field_name: str = "duration", + ) -> int: + """Return an exact integer-grid representation of ``duration``. + + Float32 command tensors receive a small ratio-space tolerance, but an + incompatible cadence is never rounded or resampled. + """ + if not isinstance(duration, (int, float)) or isinstance(duration, bool): + raise TypeError(f"{field_name} must be a real number.") + duration = float(duration) + if not math.isfinite(duration) or duration < 0.0: + raise ValueError(f"{field_name} must be finite and non-negative.") + ratio = duration / self._step_dt + nearest = round(ratio) + tolerance = max(1.0e-6, abs(ratio) * 1.0e-6) + if not math.isclose(ratio, nearest, rel_tol=0.0, abs_tol=tolerance): + raise EnvironmentStepTimingError( + f"{field_name}={duration:.9g}s is not an integer multiple of " + f"step_dt={self._step_dt:.9g}s; explicit resampling is not supported." + ) + return int(nearest) + + def validate_frame(self, frame: RuntimeCommandFrame) -> None: + """Validate every row's command hold duration against the step grid.""" + if not isinstance(frame, RuntimeCommandFrame): + raise TypeError("frame must be a RuntimeCommandFrame.") + for row, duration in enumerate(frame.hold_duration.detach().cpu().tolist()): + self.steps_for_duration( + float(duration), + field_name=f"RuntimeCommandFrame.hold_duration[{row}]", + ) + + def advance_after_env_step(self, steps: int = 1) -> None: + """Advance time after ``steps`` completed Gym environment transitions.""" + if type(steps) is not int or steps <= 0: + raise ValueError("steps must be a positive integer.") + self._step_index += steps + + +class JointPositionGymTransportEncoder: + """Built-in ``robot.joint_position`` to full-qpos action encoder.""" + + @property + def transport_id(self) -> str: + """Return the built-in joint-position transport ID.""" + return JointPositionTarget.TRANSPORT_ID + + def encode( + self, + command: EndpointCommand, + *, + base_action: EnvAction, + active_mask: torch.Tensor, + ) -> EnvAction: + """Write addressed joints while holding every other qpos column.""" + if not isinstance(command.target, JointPositionTarget): + raise TypeError("Joint-position transport requires JointPositionTarget.") + if not isinstance(command.payload, JointPositionPayload): + raise TypeError("Joint-position transport requires JointPositionPayload.") + if not isinstance(base_action, torch.Tensor): + raise TypeError( + "The built-in joint-position encoder requires a tensor base action; " + "register structured transports after it or provide a compatible " + "custom composition encoder." + ) + if base_action.dim() != 2 or base_action.shape[0] != command.batch_size: + raise ValueError( + "The full-qpos base action must have shape (batch_size, robot_dof)." + ) + if active_mask.dtype != torch.bool or active_mask.shape != ( + command.batch_size, + ): + raise ValueError("active_mask must be bool with one value per command row.") + if active_mask.device != base_action.device: + raise ValueError("active_mask and base_action must share a device.") + joint_ids = command.target.joint_ids + if max(joint_ids) >= base_action.shape[1]: + raise ValueError( + f"Joint ID {max(joint_ids)} exceeds full qpos width " + f"{base_action.shape[1]}." + ) + positions = command.payload.positions + if positions.device != base_action.device: + raise ValueError("Joint payload and base action must share a device.") + if not base_action.is_floating_point(): + raise TypeError("The full-qpos base action must be floating point.") + + action = base_action.clone() + columns = torch.tensor(joint_ids, dtype=torch.long, device=action.device) + selected = action.index_select(1, columns) + selected[active_mask] = positions[active_mask].to(dtype=action.dtype) + action[:, columns] = selected + return action + + def hold( + self, + targets: tuple[RuntimeEndpointTarget, ...], + *, + base_action: EnvAction, + context: PlanningContext, + ) -> EnvAction: + """Keep observed full qpos unchanged for addressed joint targets.""" + del context + if not all(isinstance(target, JointPositionTarget) for target in targets): + raise TypeError("Joint-position hold received an incompatible target.") + return base_action.clone() + + +class RuntimeCommandFrameEncoder: + """Encode transport-neutral command frames to controller-ready Gym actions. + + Args: + qpos_provider: Full-qpos source aligned to a frame's explicit ``env_ids``. + transports: Optional additional transport encoders. The built-in + joint-position encoder is always installed first. + """ + + def __init__( + self, + qpos_provider: CurrentQposProvider, + *, + transports: Iterable[RuntimeTransportActionEncoder] = (), + ) -> None: + if not isinstance(qpos_provider, CurrentQposProvider): + raise TypeError("qpos_provider must implement CurrentQposProvider.") + self._qpos_provider = qpos_provider + self._transports: dict[str, RuntimeTransportActionEncoder] = {} + self.register_transport(JointPositionGymTransportEncoder()) + for transport in transports: + self.register_transport(transport) + + @property + def transport_ids(self) -> tuple[str, ...]: + """Return registered transport IDs in deterministic encoding order.""" + return tuple(self._transports) + + def register_transport( + self, + transport: RuntimeTransportActionEncoder, + *, + replace: bool = False, + ) -> None: + """Register one shared transport-to-Gym action encoder.""" + if not isinstance(transport, RuntimeTransportActionEncoder): + raise TypeError("transport must implement RuntimeTransportActionEncoder.") + transport_id = _validate_identifier( + transport.transport_id, + field_name="RuntimeTransportActionEncoder.transport_id", + ) + if type(replace) is not bool: + raise TypeError("replace must be a bool.") + if transport_id in self._transports and not replace: + raise ValueError(f"Transport {transport_id!r} is already registered.") + self._transports[transport_id] = transport + + def _base_qpos(self, env_ids: torch.Tensor) -> torch.Tensor: + """Capture and validate one owned full-qpos hold action.""" + qpos = self._qpos_provider.current_qpos(env_ids) + if not isinstance(qpos, torch.Tensor): + raise TypeError("CurrentQposProvider.current_qpos() must return a tensor.") + if qpos.dim() != 2 or qpos.shape[0] != env_ids.shape[0] or qpos.shape[1] == 0: + raise ValueError( + "Current qpos must have shape (batch_size, robot_dof) with non-zero DOF." + ) + if qpos.device != env_ids.device: + raise ValueError("Current qpos and env_ids must share a device.") + if not qpos.is_floating_point() or not torch.isfinite(qpos).all().item(): + raise ValueError("Current qpos must contain finite floating-point values.") + return qpos.clone() + + def encode(self, frame: RuntimeCommandFrame) -> EnvAction: + """Encode one frame on top of a fresh full-qpos hold action.""" + if not isinstance(frame, RuntimeCommandFrame): + raise TypeError("frame must be a RuntimeCommandFrame.") + action: EnvAction = self._base_qpos(frame.env_ids) + for command in frame.commands: + transport = self._transports.get(command.transport_id) + if transport is None: + raise UnsupportedRuntimeTransportError( + f"No Gym action encoder is registered for runtime transport " + f"{command.transport_id!r}." + ) + action = transport.encode( + command, + base_action=action, + active_mask=frame.active_mask, + ) + return action + + def encode_hold( + self, + targets: tuple[RuntimeEndpointTarget, ...], + context: PlanningContext, + ) -> EnvAction: + """Encode an observed-position safe hold for addressed transports.""" + if not isinstance(context, PlanningContext): + raise TypeError("context must be a PlanningContext.") + action: EnvAction = context.robot.qpos.clone() + by_transport: dict[str, list[RuntimeEndpointTarget]] = {} + for target in targets: + if not isinstance(target, RuntimeEndpointTarget): + raise TypeError("targets must contain RuntimeEndpointTarget values.") + by_transport.setdefault(target.transport_id, []).append(target) + for transport_id, grouped in by_transport.items(): + transport = self._transports.get(transport_id) + if transport is None: + raise UnsupportedRuntimeTransportError( + f"No Gym action encoder is registered for runtime transport " + f"{transport_id!r}." + ) + action = transport.hold( + tuple(grouped), + base_action=action, + context=context, + ) + return action + + def encode_idle_hold(self, env_ids: torch.Tensor) -> EnvAction: + """Return a fresh full-qpos hold when no transport was armed yet.""" + return self._base_qpos(env_ids) + + +@dataclass(frozen=True, slots=True) +class _BufferedAction: + """One owned action plus command-boundary provenance.""" + + action: ProcessedEnvAction + + def snapshot(self) -> _BufferedAction: + """Return one independently owned buffered action.""" + return _BufferedAction(self.action.snapshot()) + + +class BufferedGymCommandSink: + """Runner command sink that buffers actions for the Gym demo generator. + + Acceptance means the command was validated and copied into the local + buffer; it does not claim that an environment transition already occurred. + """ + + def __init__( + self, + encoder: RuntimeCommandFrameEncoder, + clock: EnvironmentStepClock, + *, + accepted_command_observer: AcceptedRuntimeCommandObserver | None = None, + ) -> None: + if not isinstance(encoder, RuntimeCommandFrameEncoder): + raise TypeError("encoder must be a RuntimeCommandFrameEncoder.") + if not isinstance(clock, EnvironmentStepClock): + raise TypeError("clock must be an EnvironmentStepClock.") + if accepted_command_observer is not None and not isinstance( + accepted_command_observer, + AcceptedRuntimeCommandObserver, + ): + raise TypeError( + "accepted_command_observer must implement " + "AcceptedRuntimeCommandObserver or be None." + ) + self._encoder = encoder + self._clock = clock + self._accepted_command_observer = accepted_command_observer + self._pending: deque[_BufferedAction] = deque() + self._last_emitted: ProcessedEnvAction | None = None + self._accepted_action_count = 0 + + @property + def clock(self) -> EnvironmentStepClock: + """Return the exact environment-step clock used for timing checks.""" + return self._clock + + @property + def pending_count(self) -> int: + """Return the number of accepted actions not yet yielded to Gym.""" + return len(self._pending) + + @property + def accepted_action_count(self) -> int: + """Return the monotonic count of actions accepted by this sink.""" + return self._accepted_action_count + + def send( + self, + command: RuntimeCommandFrame, + *, + timeout: float, + ) -> CommandAcknowledgement: + """Validate, encode, and buffer one runtime command frame.""" + _validate_timeout(timeout) + if not isinstance(command, RuntimeCommandFrame): + raise TypeError("command must be a RuntimeCommandFrame.") + self._clock.validate_frame(command) + action = self._encoder.encode(command) + metadata = { + "bridge_action_kind": "runtime_command", + "runtime_destinations": [ + [item.transport_id, item.target.target_id] for item in command.commands + ], + "active_mask": command.active_mask.detach().cpu().tolist(), + "hold_duration": command.hold_duration.detach().cpu().tolist(), + } + self._pending.append( + _BufferedAction(ProcessedEnvAction(value=action, metadata=metadata)) + ) + observer = self._accepted_command_observer + if observer is not None: + try: + observer.accepted(command.snapshot()) + except Exception: + self._pending.clear() + self._discard_observer_state() + raise + self._accepted_action_count += 1 + return CommandAcknowledgement.accepted_ack("Buffered for the Gym step loop.") + + def hold( + self, + targets: tuple[RuntimeEndpointTarget, ...], + context: PlanningContext, + *, + timeout: float, + ) -> CommandAcknowledgement: + """Buffer one observed-position safe hold action.""" + _validate_timeout(timeout) + action = self._encoder.encode_hold(tuple(targets), context) + metadata = { + "bridge_action_kind": "runtime_safe_hold", + "runtime_destinations": [ + [target.transport_id, target.target_id] for target in targets + ], + } + self._pending.append( + _BufferedAction(ProcessedEnvAction(value=action, metadata=metadata)) + ) + self._accepted_action_count += 1 + return CommandAcknowledgement.accepted_ack("Safe hold buffered for Gym.") + + def cancel( + self, + targets: tuple[RuntimeEndpointTarget, ...], + *, + timeout: float, + ) -> CommandAcknowledgement: + """Discard accepted-but-not-yielded frames before a safe-stop hold.""" + _validate_timeout(timeout) + if not all(isinstance(target, RuntimeEndpointTarget) for target in targets): + raise TypeError("targets must contain RuntimeEndpointTarget values.") + self._pending.clear() + observer = self._accepted_command_observer + if observer is not None: + try: + observer.cancelled(tuple(target.snapshot() for target in targets)) + except Exception: + self._discard_observer_state() + raise + return CommandAcknowledgement.accepted_ack("Buffered commands cancelled.") + + def discard_pending(self) -> None: + """Discard actions that were accepted locally but never yielded.""" + self._pending.clear() + self._discard_observer_state() + + def drain_safe_stop_action( + self, + *, + fallback: ProcessedEnvAction | None = None, + ) -> ProcessedEnvAction | None: + """Select one buffered safe hold and discard every other local action. + + This method is used only by the demo abort handshake. A runtime + acknowledgement proves local buffering, not ``env.step`` consumption; + therefore an interrupted generator must explicitly surface the final + safe hold to the executor while dropping stale motion commands. + """ + candidates: list[ProcessedEnvAction] = [] + for candidate in (self._last_emitted, fallback): + if ( + candidate is not None + and candidate.metadata.get("bridge_action_kind") + in _SAFE_HOLD_ACTION_KINDS + ): + candidates.append(candidate.snapshot()) + while self._pending: + candidate = self._pending.popleft().action + if candidate.metadata.get("bridge_action_kind") in _SAFE_HOLD_ACTION_KINDS: + candidates.append(candidate.snapshot()) + self._discard_observer_state() + return None if not candidates else candidates[-1].snapshot() + + def _discard_observer_state(self) -> None: + """Reset observer state after any fail-closed local discard.""" + observer = self._accepted_command_observer + if observer is not None: + observer.discarded() + + def pop(self) -> ProcessedEnvAction: + """Pop the next accepted action and remember it as the active hold.""" + if not self._pending: + raise RuntimeError("No buffered Gym command is available.") + action = self._pending.popleft().action.snapshot() + self._last_emitted = action.snapshot() + return action + + def wait_hold(self, env_ids: torch.Tensor) -> ProcessedEnvAction: + """Return an owned hold action for one runtime waiting step.""" + if self._last_emitted is None: + value = self._encoder.encode_idle_hold(env_ids) + else: + value = self._last_emitted.value + return ProcessedEnvAction( + value=value, + metadata={"bridge_action_kind": "runtime_wait_hold"}, + ) + + +@dataclass(slots=True) +class _SegmentLifecycle: + """Mutable state shared by one lazy action generator and validator.""" + + complete: bool = False + result: SkillResult | ParallelSkillResult | None = None + validation: torch.Tensor | None = None + runtime: SequentialSkillRuntimePort | ParallelSkillRuntime | None = None + pending_action: ProcessedEnvAction | None = None + actions_started: bool = False + sink_acceptance_baseline: int | None = None + yielded_action_count: int = 0 + abort_started: bool = False + abort_complete: bool = False + metadata: dict[str, Any] = field(default_factory=dict) + post_policy_success: torch.Tensor | None = None + + +def _validate_runtime_result( + result: SkillResult | ParallelSkillResult, +) -> SkillResult | ParallelSkillResult: + """Validate one exact sequential or parallel runtime boundary.""" + if not isinstance(result, (SkillResult, ParallelSkillResult)): + raise TypeError( + "Runtime methods must return SkillResult or ParallelSkillResult values." + ) + return result + + +def _normalize_validation( + value: Any, + *, + batch_size: int, + device: torch.device, +) -> torch.Tensor: + """Normalize one validator output to an owned row-local boolean tensor.""" + tensor = torch.as_tensor(value, dtype=torch.bool, device=device).reshape(-1) + if tensor.numel() == 1 and batch_size > 1: + tensor = tensor.repeat(batch_size) + if tensor.numel() != batch_size: + raise ValueError( + f"Segment validator returned {tensor.numel()} flags, expected " + f"{batch_size}." + ) + return tensor.clone() + + +def _json_safe_copy(value: Any, *, field_name: str) -> Any: + """Return an owned JSON value while rejecting lossy coercions.""" + if value is None or type(value) in {bool, int, str}: + return value + if type(value) is float: + if not math.isfinite(value): + raise ValueError(f"{field_name} contains a non-finite float.") + return value + if isinstance(value, Mapping): + result: dict[str, Any] = {} + for key, item in value.items(): + if type(key) is not str or not key or key != key.strip(): + raise ValueError( + f"{field_name} mapping keys must be non-empty strings " + "without outer whitespace." + ) + result[key] = _json_safe_copy( + item, + field_name=f"{field_name}.{key}", + ) + return result + if isinstance(value, (list, tuple)): + return [ + _json_safe_copy(item, field_name=f"{field_name}[{index}]") + for index, item in enumerate(value) + ] + raise TypeError(f"{field_name} contains non-JSON value {type(value).__name__}.") + + +def _runtime_result_metadata( + result: SkillResult | ParallelSkillResult, +) -> dict[str, Any]: + """Snapshot one core runtime result through its canonical serializer.""" + serializer = getattr(result, "to_metadata", None) + if not callable(serializer): + raise TypeError( + f"{type(result).__name__} must provide to_metadata() for demo tracing." + ) + metadata = _json_safe_copy(serializer(), field_name="runtime result metadata") + if not isinstance(metadata, dict): + raise TypeError("Runtime result to_metadata() must return a mapping.") + return metadata + + +class AtomicDemoBridge: + """Adapt sequential compiled program segments to lazy Gym demonstrations. + + Args: + program: Provider-free compiled Expert Program. + runtime: Nonblocking semantic :class:`SkillRuntime` surface. + command_sink: The same buffered sink installed in ``runtime``. + clock: The same environment-step clock installed in ``runtime``. + post_policy_port: Optional environment-aware post-policy executor. + validator_port: Optional environment-aware validator executor. + parallel_safety_validator: Optional authoritative physical-safety gate + required before any parallel branch can start. + + Schema-v2 parallel blocks retain their branch lanes and explicit barrier. + They are lowered through :class:`ParallelSkillRuntime`; they are never + flattened into a sequential semantic-call list. + """ + + def __init__( + self, + program: CompiledProgramPort, + runtime: SequentialSkillRuntimePort, + command_sink: BufferedGymCommandSink, + clock: EnvironmentStepClock, + *, + post_policy_port: SegmentPostPolicyPort | None = None, + validator_port: SegmentValidatorPort | None = None, + parallel_safety_validator: ParallelCommandSafetyValidator | None = None, + ) -> None: + if not isinstance(program, CompiledProgramPort): + raise TypeError("program must implement CompiledProgramPort.") + _validate_identifier(program.program_id, field_name="program.program_id") + if type(program.schema_version) is not int or program.schema_version < 1: + raise ValueError("program.schema_version must be a positive integer.") + if not isinstance(runtime, SequentialSkillRuntimePort): + raise TypeError("runtime must implement SequentialSkillRuntimePort.") + if not isinstance(command_sink, BufferedGymCommandSink): + raise TypeError("command_sink must be a BufferedGymCommandSink.") + if not isinstance(clock, EnvironmentStepClock): + raise TypeError("clock must be an EnvironmentStepClock.") + if command_sink.clock is not clock: + raise ValueError("command_sink and bridge must share the exact clock.") + if post_policy_port is not None and not isinstance( + post_policy_port, SegmentPostPolicyPort + ): + raise TypeError("post_policy_port must implement SegmentPostPolicyPort.") + if validator_port is not None and not isinstance( + validator_port, SegmentValidatorPort + ): + raise TypeError("validator_port must implement SegmentValidatorPort.") + if parallel_safety_validator is not None and not isinstance( + parallel_safety_validator, ParallelCommandSafetyValidator + ): + raise TypeError( + "parallel_safety_validator must implement " + "ParallelCommandSafetyValidator." + ) + self._program = program + self._runtime = runtime + self._sink = command_sink + self._clock = clock + self._post_policy_port = post_policy_port + self._validator_port = validator_port + self._parallel_safety_validator = parallel_safety_validator + self._active_segment_id: str | None = None + self._eligible_mask: torch.Tensor | None = None + + @property + def clock(self) -> EnvironmentStepClock: + """Return the environment-step clock used by this bridge.""" + return self._clock + + def iter_segments(self) -> Iterator[DemoSegment]: + """Lazily adapt compiled program segments to ``DemoSegment`` values. + + Consumers must exhaust each segment's actions and invoke its validator + before requesting the next segment. Skipping either lifecycle boundary + raises :class:`DemoBridgeError` instead of silently carrying stale row + eligibility into downstream execution. + """ + for segment in self._program.iter_segments(): + metadata = self._segment_metadata(segment) + lifecycle = _SegmentLifecycle(metadata=metadata) + validator = self._segment_validator(segment, lifecycle) + yield DemoSegment( + actions=self._segment_actions(segment, lifecycle), + name=segment.name, + metadata=metadata, + validator=validator, + abort_actions=self._segment_abort_actions(segment, lifecycle), + failure_policy="row_independent", + ) + self._require_consumed_segment_lifecycle(segment, lifecycle) + + def __iter__(self) -> Iterator[DemoSegment]: + """Delegate iteration to :meth:`iter_segments`.""" + return self.iter_segments() + + @staticmethod + def _require_consumed_segment_lifecycle( + segment: Any, + lifecycle: _SegmentLifecycle, + ) -> None: + """Reject advancing past a segment with an unconsumed lifecycle. + + The public demo executor exhausts ``actions`` and then invokes the + segment validator before requesting the next lazy segment. Direct + bridge consumers must preserve the same ordering because validation is + also the commit point for runtime and post-policy row eligibility. + """ + if not lifecycle.complete: + raise DemoBridgeError( + f"Segment {segment.segment_id!r} actions must be exhausted before " + "requesting the next compiled segment." + ) + if lifecycle.validation is None: + raise DemoBridgeError( + f"Segment {segment.segment_id!r} validator must be called after " + "its actions are exhausted and before requesting the next " + "compiled segment." + ) + + def _segment_metadata(self, segment: Any) -> dict[str, Any]: + """Build mutable JSON-safe metadata completed at lifecycle boundaries.""" + return { + "expert_program_schema_version": self._program.schema_version, + "expert_program_id": self._program.program_id, + "program_segment_id": segment.segment_id, + "program_segment_index": segment.segment_index, + "program_segment_source_path": list(segment.source_path), + "program_segment_implicit": bool(segment.implicit), + "semantic_call_indices": [call.call_index for call in segment.calls], + "post_policy_count": len(segment.post_policies), + "validator_count": len(segment.validators), + "parallel": getattr(segment, "parallel_block", None) is not None, + "runtime": None, + "post_policies": [], + "validation": None, + } + + @staticmethod + def _record_runtime_result( + lifecycle: _SegmentLifecycle, + result: SkillResult | ParallelSkillResult, + ) -> None: + """Snapshot one runtime boundary into its owning segment metadata.""" + lifecycle.result = result + lifecycle.metadata["runtime"] = _runtime_result_metadata(result) + + def _decorate_action( + self, + action: Any, + *, + segment: Any, + result: SkillResult | ParallelSkillResult, + action_kind: str | None = None, + ) -> ProcessedEnvAction: + """Own one action and attach stable program/runtime provenance.""" + if isinstance(action, ProcessedEnvAction): + value = action.value + metadata = dict(action.metadata) + else: + value = action + metadata = {} + if action_kind is not None: + metadata["bridge_action_kind"] = action_kind + metadata.update( + { + "expert_program_id": self._program.program_id, + "program_segment_id": segment.segment_id, + "program_segment_index": segment.segment_index, + "environment_step": self._clock.step_index, + "runtime_status": result.status.value, + "runtime_call_index": getattr(result, "current_call_index", None), + } + ) + return ProcessedEnvAction(value=value, metadata=metadata) + + def _yield_and_advance( + self, + action: ProcessedEnvAction, + lifecycle: _SegmentLifecycle, + ) -> Iterator[ProcessedEnvAction]: + """Yield once and advance only after explicit consumption acknowledgement.""" + if lifecycle.pending_action is not None: + raise RuntimeError("A prior demo action is still awaiting acknowledgement.") + lifecycle.pending_action = action.snapshot() + lifecycle.yielded_action_count += 1 + yield action + if lifecycle.pending_action is not None: + self._clock.advance_after_env_step() + lifecycle.pending_action = None + + def _segment_actions( + self, + segment: Any, + lifecycle: _SegmentLifecycle, + ) -> Iterator[ProcessedEnvAction]: + """Drive one semantic segment without bypassing the Gym step loop.""" + segment_id = segment.segment_id + lifecycle.actions_started = True + lifecycle.sink_acceptance_baseline = self._sink.accepted_action_count + if self._active_segment_id is not None: + raise RuntimeError( + f"Segment {self._active_segment_id!r} is still active; exhaust or " + "close it before starting another lazy segment." + ) + self._active_segment_id = segment_id + result: SkillResult | ParallelSkillResult | None = None + segment_runtime: SequentialSkillRuntimePort | ParallelSkillRuntime = ( + self._runtime + ) + is_parallel = getattr(segment, "parallel_block", None) is not None + try: + if is_parallel: + segment_runtime = self._parallel_runtime(segment) + lifecycle.runtime = segment_runtime + result = _validate_runtime_result( + segment_runtime.start( + workflow_id=f"{self._program.program_id}/{segment_id}", + eligible_mask=self._eligible_mask, + ) + ) + else: + lifecycle.runtime = segment_runtime + analysis = self._program.sequential_execution_analysis( + segment.segment_index + ) + calls = tuple(analysis.calls) + if not calls: + raise DemoBridgeError( + f"Compiled segment {segment_id!r} contains no semantic calls." + ) + execution_prefix_length = analysis.execution_prefix_length + if execution_prefix_length != len(segment.calls): + raise DemoBridgeError( + f"Compiled segment {segment_id!r} analysis prefix length " + "does not match its owned semantic calls." + ) + result = _validate_runtime_result( + segment_runtime.start( + calls, + workflow_id=f"{self._program.program_id}/{segment_id}", + eligible_mask=self._eligible_mask, + execution_prefix_length=execution_prefix_length, + ) + ) + + while True: + emitted = False + while self._sink.pending_count: + action = self._decorate_action( + self._sink.pop(), + segment=segment, + result=result, + ) + yield from self._yield_and_advance(action, lifecycle) + emitted = True + + if emitted and not result.terminal: + # The result's wait duration was measured before the action + # just consumed by Gym. Refresh it against the advanced + # environment clock before deciding whether another hold is due. + result = _validate_runtime_result(segment_runtime.step()) + continue + + if result.terminal: + break + + if result.wait_duration > 0.0: + self._clock.steps_for_duration( + result.wait_duration, + field_name="SkillResult.wait_duration", + ) + hold = self._decorate_action( + self._sink.wait_hold(result.env_ids), + segment=segment, + result=result, + action_kind="runtime_wait_hold", + ) + yield from self._yield_and_advance(hold, lifecycle) + + result = _validate_runtime_result(segment_runtime.step()) + + self._record_runtime_result(lifecycle, result) + self._retain_eligible_rows(result.success_mask) + if is_parallel: + self._runtime.adopt_verified_task_state(result.task_state) + if result.status is SkillStatus.COMPLETED: + yield from self._post_policy_actions(segment, result, lifecycle) + lifecycle.complete = True + finally: + if not lifecycle.abort_started and lifecycle.pending_action is not None: + if result is not None and not result.terminal: + segment_runtime.cancel( + f"Demo segment {segment_id!r} action iteration stopped early." + ) + raise DemoBridgeError( + f"Demo segment {segment_id!r} was closed with an unacknowledged " + "action. Consume DemoSegment.abort_actions through env.step() " + "before closing the action iterator." + ) + self._active_segment_id = None + + def _segment_abort_actions( + self, + segment: Any, + lifecycle: _SegmentLifecycle, + ) -> Callable[..., Iterator[ProcessedEnvAction]]: + """Create the explicit executor-to-runtime cancellation handshake.""" + + def abort( + reason: str, + *, + last_action_consumed: bool, + ) -> Iterator[ProcessedEnvAction]: + return self._abort_segment( + segment, + lifecycle, + reason=reason, + last_action_consumed=last_action_consumed, + ) + + return abort + + def _abort_segment( + self, + segment: Any, + lifecycle: _SegmentLifecycle, + *, + reason: str, + last_action_consumed: bool, + ) -> Iterator[ProcessedEnvAction]: + """Abort one segment, surfacing a safe hold only after controller activity.""" + if type(reason) is not str or not reason: + raise ValueError("abort reason must be a non-empty string.") + if type(last_action_consumed) is not bool: + raise TypeError("last_action_consumed must be a bool.") + if lifecycle.abort_started: + raise RuntimeError( + f"Segment {segment.segment_id!r} abort handshake already started." + ) + if not lifecycle.actions_started: + raise RuntimeError( + f"Segment {segment.segment_id!r} has no started action iteration " + "to abort." + ) + baseline = lifecycle.sink_acceptance_baseline + if baseline is None: + raise RuntimeError( + f"Segment {segment.segment_id!r} has no sink lifecycle baseline." + ) + controller_activity_started = ( + lifecycle.yielded_action_count > 0 + or lifecycle.pending_action is not None + or self._sink.accepted_action_count > baseline + ) + if not controller_activity_started: + # Runtime construction and preflight are deliberately observation- and + # command-free. If either fails before the first accepted or yielded + # action, there is no physical controller state to safe-stop. Mark the + # handshake complete without touching the partially constructed runtime + # so the original action-generation exception remains authoritative. + lifecycle.abort_started = True + lifecycle.abort_complete = True + return + runtime = lifecycle.runtime + pending = lifecycle.pending_action + if runtime is None: + raise DemoBridgeError( + f"Segment {segment.segment_id!r} accepted or yielded a controller " + "action without retaining a runtime capable of strict safe-stop." + ) + lifecycle.abort_started = True + if pending is not None: + pending = pending.snapshot() + if pending is not None and last_action_consumed: + self._clock.advance_after_env_step() + lifecycle.pending_action = None + + result = _validate_runtime_result(runtime.result) + if not result.terminal: + result = _validate_runtime_result(runtime.cancel(reason)) + self._record_runtime_result(lifecycle, result) + + pending_kind = ( + None if pending is None else pending.metadata.get("bridge_action_kind") + ) + if ( + pending is not None + and last_action_consumed + and pending_kind in _SAFE_HOLD_ACTION_KINDS + ): + self._sink.discard_pending() + lifecycle.abort_complete = True + return + + safe_action = self._sink.drain_safe_stop_action( + fallback=None if last_action_consumed else pending, + ) + if safe_action is None: + raise DemoBridgeError( + f"Segment {segment.segment_id!r} stopped before exhaustion, but " + "no controller safe-hold action was available for env.step()." + ) + processed = self._decorate_action( + safe_action, + segment=segment, + result=result, + action_kind="runtime_abort_safe_hold", + ) + yield processed + self._clock.advance_after_env_step() + lifecycle.abort_complete = True + + def _parallel_runtime(self, segment: Any) -> ParallelSkillRuntime: + """Build one one-shot coordinator from a compiled explicit barrier.""" + if self._parallel_safety_validator is None: + raise DemoBridgeError( + f"Parallel segment {segment.segment_id!r} requires an explicit " + "ParallelCommandSafetyValidator; resource claims alone do not " + "establish physical collision safety." + ) + if not isinstance(self._runtime, SkillRuntime): + # Production integration always supplies SkillRuntime. Keeping the + # sequential protocol permits lightweight tests and alternate + # frontends, but the canonical parallel factory requires forkable + # runtime internals by design. + raise TypeError( + "Parallel compiled segments require a concrete SkillRuntime " + "template." + ) + block = segment.parallel_block + branches = tuple(block.branches) + if len(branches) < 2: + raise DemoBridgeError( + f"Parallel segment {segment.segment_id!r} requires at least two " + "compiled branches." + ) + branch_calls = { + f"branch_{branch.branch_index}": tuple( + compiled.call for compiled in branch.calls + ) + for branch in branches + } + branch_paths = { + f"branch_{branch.branch_index}": tuple( + getattr(branch, "source_path", segment.source_path) + ) + for branch in branches + } + if any(not calls for calls in branch_calls.values()): + raise DemoBridgeError( + f"Parallel segment {segment.segment_id!r} contains an empty branch." + ) + barrier = block.barrier + return ParallelSkillRuntime.from_template( + self._runtime, + branch_calls, + self._sink, + ParallelTimingPolicy(self._clock.step_dt), + self._parallel_safety_validator, + timeout_steps=barrier.timeout_steps, + failure_policy=barrier.failure_policy, + workflow_id=( + f"{self._program.program_id}/{segment.segment_id}:parallel_analysis" + ), + branch_paths=branch_paths, + ) + + def _retain_eligible_rows(self, accepted: torch.Tensor) -> None: + """Permanently remove failed rows before a later lazy segment starts.""" + if not isinstance(accepted, torch.Tensor): + raise TypeError("accepted must be a torch.Tensor.") + if accepted.dtype != torch.bool or accepted.dim() != 1: + raise ValueError("accepted must be a one-dimensional bool tensor.") + if self._eligible_mask is None: + self._eligible_mask = torch.ones_like(accepted) + elif ( + self._eligible_mask.shape != accepted.shape + or self._eligible_mask.device != accepted.device + ): + raise ValueError("Environment rows changed across program segments.") + self._eligible_mask &= accepted + + def _post_policy_actions( + self, + segment: Any, + result: SkillResult | ParallelSkillResult, + lifecycle: _SegmentLifecycle, + ) -> Iterator[ProcessedEnvAction]: + """Route environment-aware post-policy actions through the same generator.""" + policies = tuple(segment.post_policies) + if policies and self._post_policy_port is None: + raise DemoBridgeError( + f"Segment {segment.segment_id!r} declares post-policies, but no " + "SegmentPostPolicyPort was installed." + ) + traces = lifecycle.metadata["post_policies"] + if not isinstance(traces, list): + raise TypeError("Segment post-policy metadata storage must be a list.") + for policy_index, policy in enumerate(policies): + assert self._post_policy_port is not None + active_mask = ( + result.success_mask.clone() + if self._eligible_mask is None + else self._eligible_mask.clone() + ) + if lifecycle.post_policy_success is not None: + active_mask &= lifecycle.post_policy_success + actions = self._post_policy_port.actions( + policy, + segment=segment, + active_mask=active_mask, + ) + if isinstance(actions, (str, bytes)): + raise TypeError("Post-policy actions must be an iterable of actions.") + action_iterator = iter(actions) + iteration_error: BaseException | None = None + try: + for action in action_iterator: + processed = self._decorate_action( + action, + segment=segment, + result=result, + action_kind="program_post_policy", + ) + yield from self._yield_and_advance(processed, lifecycle) + except BaseException as exc: + iteration_error = exc + raise + finally: + close = getattr(action_iterator, "close", None) + if callable(close): + close() + cfg = getattr(policy, "cfg", None) + trace: dict[str, Any] = { + "policy_index": policy_index, + "kind": getattr(cfg, "kind", type(policy).__name__), + "source_path": list(getattr(policy, "source_path", ())), + "result_mask": result.success_mask.detach().cpu().tolist(), + "result": None, + } + port = self._post_policy_port + policy_success = active_mask.clone() + if isinstance(port, SegmentPostPolicyResultPort): + try: + policy_success &= _normalize_validation( + port.post_policy_result(policy, segment=segment), + batch_size=result.env_ids.numel(), + device=result.env_ids.device, + ) + except Exception: + if iteration_error is None: + raise + if lifecycle.post_policy_success is None: + lifecycle.post_policy_success = policy_success.clone() + else: + lifecycle.post_policy_success &= policy_success + trace["result_mask"] = policy_success.detach().cpu().tolist() + if isinstance(port, SegmentPostPolicyMetadataPort): + try: + trace["result"] = port.post_policy_metadata( + policy, + segment=segment, + ) + except Exception: + if iteration_error is None: + raise + traces.append( + _json_safe_copy( + trace, + field_name=f"post-policy {policy_index} metadata", + ) + ) + + def _segment_validator( + self, + segment: Any, + lifecycle: _SegmentLifecycle, + ) -> Callable[[], torch.Tensor]: + """Create a demo-boundary validator including runtime row success.""" + + def validate() -> torch.Tensor: + if not lifecycle.complete or lifecycle.result is None: + raise RuntimeError( + f"Segment {segment.segment_id!r} cannot be validated before its " + "action iterable is exhausted." + ) + if lifecycle.validation is not None: + return lifecycle.validation.clone() + result = lifecycle.result + accepted = result.success_mask.clone() + runtime_success = result.success_mask.clone() + eligible_before = ( + torch.ones_like(accepted) + if self._eligible_mask is None + else self._eligible_mask.clone() + ) + if self._eligible_mask is not None: + accepted &= self._eligible_mask + if lifecycle.post_policy_success is not None: + accepted &= lifecycle.post_policy_success + validators = tuple(segment.validators) + if validators and self._validator_port is None: + raise DemoBridgeError( + f"Segment {segment.segment_id!r} declares validators, but no " + "SegmentValidatorPort was installed." + ) + validator_traces: list[dict[str, Any]] = [] + for validator_index, validator in enumerate(validators): + assert self._validator_port is not None + value = self._validator_port.validate(validator, segment=segment) + validator_result = _normalize_validation( + value, + batch_size=result.env_ids.numel(), + device=result.env_ids.device, + ) + accepted &= validator_result + cfg = getattr(validator, "cfg", None) + trace: dict[str, Any] = { + "validator_index": validator_index, + "kind": getattr(cfg, "kind", type(validator).__name__), + "source_path": list(getattr(validator, "source_path", ())), + "result_mask": validator_result.detach().cpu().tolist(), + "result": None, + } + port = self._validator_port + if isinstance(port, SegmentValidatorMetadataPort): + trace["result"] = port.validator_metadata( + validator, + segment=segment, + ) + validator_traces.append( + _json_safe_copy( + trace, + field_name=f"validator {validator_index} metadata", + ) + ) + lifecycle.metadata["validation"] = _json_safe_copy( + { + "env_ids": result.env_ids.detach().cpu().tolist(), + "runtime_success_mask": runtime_success.detach().cpu().tolist(), + "eligible_mask_before_validation": eligible_before.detach() + .cpu() + .tolist(), + "post_policy_success_mask": ( + None + if lifecycle.post_policy_success is None + else lifecycle.post_policy_success.detach().cpu().tolist() + ), + "validators": validator_traces, + "accepted_mask": accepted.detach().cpu().tolist(), + }, + field_name="segment validation metadata", + ) + self._retain_eligible_rows(accepted) + lifecycle.validation = accepted.clone() + return accepted.clone() + + return validate + + +__all__ = [ + "AcceptedRuntimeCommandObserver", + "AtomicDemoBridge", + "BufferedGymCommandSink", + "CompiledProgramPort", + "CurrentQposProvider", + "DemoBridgeError", + "EnvironmentStepClock", + "EnvironmentStepTimingError", + "GymPlanningObservationProvider", + "JointPositionGymTransportEncoder", + "ParallelCommandSafetyValidator", + "RuntimeCommandFrameEncoder", + "RuntimeTransportActionEncoder", + "SegmentPostPolicyMetadataPort", + "SegmentPostPolicyPort", + "SegmentPostPolicyResultPort", + "SegmentValidatorMetadataPort", + "SegmentValidatorPort", + "SequentialSkillRuntimePort", + "UnsupportedRuntimeTransportError", +] diff --git a/embodichain/lab/gym/envs/expert_program/cfg.py b/embodichain/lab/gym/envs/expert_program/cfg.py new file mode 100644 index 000000000..1e9b43278 --- /dev/null +++ b/embodichain/lab/gym/envs/expert_program/cfg.py @@ -0,0 +1,857 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Typed configuration values for declarative Expert Programs.""" + +from __future__ import annotations + +import math +import re +from dataclasses import MISSING, field +from typing import TypeAlias + +from embodichain.utils import configclass + +EXPERT_PROGRAM_SCHEMA_VERSION = 1 +"""Stable sequential Expert Program schema version.""" + +EXPERT_PROGRAM_SCHEMA_VERSION_V2 = 2 +"""Schema version adding deterministic ``Parallel`` and ``Barrier`` nodes.""" + +SUPPORTED_EXPERT_PROGRAM_SCHEMA_VERSIONS = ( + EXPERT_PROGRAM_SCHEMA_VERSION, + EXPERT_PROGRAM_SCHEMA_VERSION_V2, +) +"""Exact schema revisions accepted by the strict decoder.""" + +MAX_REPEAT_COUNT = 1_000 +"""Maximum repeat count accepted by one Expert Program repeat node.""" + +MAX_EXPANDED_CALLS = 10_000 +"""Maximum statically expanded semantic calls in one Expert Program.""" + +MAX_PROGRAM_DEPTH = 64 +"""Maximum nesting depth of a supported Expert Program AST.""" + +MAX_PROGRAM_NODES = 10_000 +"""Maximum number of stored nodes in a supported Expert Program AST.""" + +MAX_DECLARATIVE_DEPTH = 32 +"""Maximum nesting depth of a registered-call declarative payload.""" + +MAX_DECLARATIVE_NODES = 10_000 +"""Maximum number of values in a registered-call declarative payload.""" + +_REGISTERED_CALL_ID_PATTERN = re.compile(r"[a-z][a-z0-9_]*(?:\.[a-z][a-z0-9_]*)+") +_ENV_TRAVERSAL_PATTERN = re.compile( + r"(?:\$?(?:env|environment)(?:\.[A-Za-z_][A-Za-z0-9_]*)+|" + r"\$\{(?:env|environment)(?:\.[A-Za-z_][A-Za-z0-9_]*)+\})" +) +_FORBIDDEN_DECLARATIVE_KEYS = frozenset( + { + "__import__", + "attribute_path", + "callable", + "environment_path", + "env_path", + "eval", + "exec", + "expression", + "getattr", + "import", + "module", + "python", + } +) + +DeclarativeCfgValue: TypeAlias = ( + None + | bool + | int + | float + | str + | tuple["DeclarativeCfgValue", ...] + | dict[str, "DeclarativeCfgValue"] +) +"""Executable-free value accepted by a registered semantic call config.""" + + +def _validate_identifier(value: object, *, field_name: str) -> str: + """Return one exact non-empty identifier.""" + if type(value) is not str or not value or value != value.strip(): + raise ValueError( + f"{field_name} must be a non-empty string without outer whitespace." + ) + return value + + +def _validate_kind(value: object, *, expected: str, field_name: str) -> None: + """Require one exact discriminator value.""" + if type(value) is not str or value != expected: + raise ValueError(f"{field_name} must be exactly {expected!r}.") + + +def _validate_number(value: object, *, field_name: str) -> float: + """Return one finite number while rejecting bool values.""" + if type(value) not in (int, float): + raise TypeError(f"{field_name} must be an int or float.") + try: + normalized = float(value) + except OverflowError as error: + raise ValueError(f"{field_name} must be finite.") from error + if not math.isfinite(normalized): + raise ValueError(f"{field_name} must be finite.") + return normalized + + +def _validate_resources(value: object, *, field_name: str) -> dict[str, str]: + """Own one strict slot-to-resource mapping.""" + if type(value) is not dict: + raise TypeError(f"{field_name} must be an exact dict.") + resources: dict[str, str] = {} + for slot_id, resource_id in value.items(): + resources[ + _validate_identifier(slot_id, field_name=f"{field_name} slot IDs") + ] = _validate_identifier( + resource_id, + field_name=f"{field_name} resource IDs", + ) + return resources + + +def _validate_declarative_string(value: str, *, path: str) -> str: + """Reject strings that request executable or environment traversal behavior.""" + stripped = value.strip() + lowered = stripped.lower() + forbidden_prefixes = ( + "__import__(", + "eval(", + "exec(", + "import ", + "from ", + ) + if lowered.startswith(forbidden_prefixes): + raise ValueError(f"{path} contains an executable import/eval expression.") + if _ENV_TRAVERSAL_PATTERN.fullmatch(stripped) is not None: + raise ValueError(f"{path} contains dotted environment attribute traversal.") + return value + + +def _snapshot_declarative_value( + value: object, + *, + path: str, + _active: set[int] | None = None, + _budget: list[int] | None = None, + _depth: int = 0, +) -> DeclarativeCfgValue: + """Validate and own one bounded executable-free declarative value.""" + active = set() if _active is None else _active + budget = [MAX_DECLARATIVE_NODES] if _budget is None else _budget + if _depth > MAX_DECLARATIVE_DEPTH: + raise ValueError( + f"{path} exceeds declarative depth limit {MAX_DECLARATIVE_DEPTH}." + ) + budget[0] -= 1 + if budget[0] < 0: + raise ValueError( + f"{path} exceeds declarative node limit {MAX_DECLARATIVE_NODES}." + ) + if value is None or type(value) in (bool, int): + return value # type: ignore[return-value] + if type(value) is float: + if not math.isfinite(value): + raise ValueError(f"{path} contains a non-finite float.") + return value + if type(value) is str: + return _validate_declarative_string(value, path=path) + if type(value) in (list, tuple): + identity = id(value) + if identity in active: + raise ValueError(f"{path} contains a cyclic sequence.") + active.add(identity) + try: + return tuple( + _snapshot_declarative_value( + item, + path=f"{path}[{index}]", + _active=active, + _budget=budget, + _depth=_depth + 1, + ) + for index, item in enumerate(value) + ) + finally: + active.remove(identity) + if type(value) is dict: + identity = id(value) + if identity in active: + raise ValueError(f"{path} contains a cyclic mapping.") + active.add(identity) + try: + result: dict[str, DeclarativeCfgValue] = {} + for key, item in value.items(): + if type(key) is not str: + raise TypeError(f"{path} keys must be exact strings.") + if key.lower() in _FORBIDDEN_DECLARATIVE_KEYS: + raise ValueError( + f"{path}.{key} requests forbidden executable behavior." + ) + result[key] = _snapshot_declarative_value( + item, + path=f"{path}.{key}", + _active=active, + _budget=budget, + _depth=_depth + 1, + ) + return result + finally: + active.remove(identity) + raise TypeError( + f"{path} contains non-declarative {type(value).__name__}; callables, " + "classes, modules, tensors, and live objects are not allowed." + ) + + +@configclass +class ExpertProgramIntegrationCfg: + """Static integration references selected by one Expert Program.""" + + robot_profile: str = MISSING + scene_registry: str = MISSING + runtime_preset: str = MISSING + + def __post_init__(self) -> None: + """Validate stable integration identifiers.""" + _validate_identifier(self.robot_profile, field_name="robot_profile") + _validate_identifier(self.scene_registry, field_name="scene_registry") + _validate_identifier(self.runtime_preset, field_name="runtime_preset") + + +@configclass +class PoseCfg: + """One declarative Cartesian pose using a WXYZ quaternion.""" + + position: tuple[float, float, float] = MISSING + quaternion_wxyz: tuple[float, float, float, float] = MISSING + + def __post_init__(self) -> None: + """Validate pose shape, finiteness, and quaternion magnitude.""" + if type(self.position) not in (list, tuple) or len(self.position) != 3: + raise ValueError("position must contain exactly three numbers.") + if ( + type(self.quaternion_wxyz) not in (list, tuple) + or len(self.quaternion_wxyz) != 4 + ): + raise ValueError("quaternion_wxyz must contain exactly four numbers.") + position = tuple( + _validate_number(value, field_name=f"position[{index}]") + for index, value in enumerate(self.position) + ) + quaternion = tuple( + _validate_number(value, field_name=f"quaternion_wxyz[{index}]") + for index, value in enumerate(self.quaternion_wxyz) + ) + norm = math.sqrt(sum(value * value for value in quaternion)) + if norm <= 1.0e-12: + raise ValueError("quaternion_wxyz must have non-zero magnitude.") + self.position = position # type: ignore[assignment] + self.quaternion_wxyz = quaternion # type: ignore[assignment] + + +@configclass +class TargetRefCfg: + """Reference to one top-level typed target provider.""" + + target: str = MISSING + kind: str = "target_ref" + + def __post_init__(self) -> None: + """Validate the target identifier and discriminator.""" + _validate_identifier(self.target, field_name="target") + _validate_kind(self.kind, expected="target_ref", field_name="kind") + + +@configclass +class CyclicPoseTargetCfg: + """Finite pose values selected cyclically by the enclosing repeat index.""" + + values: tuple[PoseCfg, ...] = MISSING + kind: str = "cyclic_pose" + + def __post_init__(self) -> None: + """Validate a non-empty owned pose sequence.""" + _validate_kind(self.kind, expected="cyclic_pose", field_name="kind") + if type(self.values) not in (list, tuple) or not self.values: + raise ValueError("values must contain at least one PoseCfg.") + values = tuple(self.values) + if not all(type(value) is PoseCfg for value in values): + raise TypeError("values must contain exact PoseCfg values.") + self.values = values # type: ignore[assignment] + + +TargetCfg: TypeAlias = CyclicPoseTargetCfg + + +@configclass +class PickCfg: + """Declarative request to acquire one registered object.""" + + object: str = MISSING + grasp: str | None = None + resources: dict[str, str] = field(default_factory=dict) + kind: str = "pick" + + def __post_init__(self) -> None: + """Validate object, optional affordance, resources, and kind.""" + _validate_identifier(self.object, field_name="object") + if self.grasp is not None: + _validate_identifier(self.grasp, field_name="grasp") + self.resources = _validate_resources(self.resources, field_name="resources") + _validate_kind(self.kind, expected="pick", field_name="kind") + + +@configclass +class PlaceCfg: + """Declarative request to place one held object at one destination.""" + + object: str = MISSING + at: TargetRefCfg | None = None + on: str | None = None + inside: str | None = None + resources: dict[str, str] = field(default_factory=dict) + kind: str = "place" + + def __post_init__(self) -> None: + """Require exactly one typed destination.""" + _validate_identifier(self.object, field_name="object") + selected = sum(value is not None for value in (self.at, self.on, self.inside)) + if selected != 1: + raise ValueError("Place requires exactly one of at, on, or inside.") + if self.at is not None and type(self.at) is not TargetRefCfg: + raise TypeError("at must be exactly TargetRefCfg or None.") + if self.on is not None: + _validate_identifier(self.on, field_name="on") + if self.inside is not None: + _validate_identifier(self.inside, field_name="inside") + self.resources = _validate_resources(self.resources, field_name="resources") + _validate_kind(self.kind, expected="place", field_name="kind") + + +@configclass +class HandOverCfg: + """Declarative request to transfer one held object between resources.""" + + object: str = MISSING + receiver: str | None = None + final_target: TargetRefCfg | None = None + resources: dict[str, str] = field(default_factory=dict) + kind: str = "hand_over" + + def __post_init__(self) -> None: + """Validate object, destination resource, and optional target.""" + _validate_identifier(self.object, field_name="object") + if self.receiver is not None: + _validate_identifier(self.receiver, field_name="receiver") + if ( + self.final_target is not None + and type(self.final_target) is not TargetRefCfg + ): + raise TypeError("final_target must be exactly TargetRefCfg or None.") + resources = _validate_resources(self.resources, field_name="resources") + if self.receiver is not None: + selected = resources.get("destination") + if selected is not None and selected != self.receiver: + raise ValueError("receiver conflicts with resources['destination'].") + resources["destination"] = self.receiver + self.resources = resources + _validate_kind(self.kind, expected="hand_over", field_name="kind") + + +@configclass +class OperateArticulationCfg: + """Declarative request to operate one articulated joint through a handle.""" + + articulation: str = MISSING + handle: str | None = None + target: str | None = None + target_position: float | None = None + target_displacement: float | None = None + resources: dict[str, str] = field(default_factory=dict) + kind: str = "operate_articulation" + + def __post_init__(self) -> None: + """Require one named target or one complete explicit target pair.""" + _validate_identifier(self.articulation, field_name="articulation") + if self.handle is not None: + _validate_identifier(self.handle, field_name="handle") + named = self.target is not None + explicit_position = self.target_position is not None + explicit_displacement = self.target_displacement is not None + if named: + _validate_identifier(self.target, field_name="target") + if explicit_position or explicit_displacement: + raise ValueError( + "target is mutually exclusive with target_position and " + "target_displacement." + ) + elif not (explicit_position and explicit_displacement): + raise ValueError( + "OperateArticulation requires either target or the explicit " + "target_position and target_displacement pair." + ) + else: + self.target_position = _validate_number( + self.target_position, + field_name="target_position", + ) + self.target_displacement = _validate_number( + self.target_displacement, + field_name="target_displacement", + ) + self.resources = _validate_resources(self.resources, field_name="resources") + _validate_kind( + self.kind, + expected="operate_articulation", + field_name="kind", + ) + + +@configclass +class RegisteredSemanticCallCfg: + """Safe declarative payload for one catalog-registered semantic call.""" + + call_id: str = MISSING + schema_version: int = EXPERT_PROGRAM_SCHEMA_VERSION + arguments: dict[str, DeclarativeCfgValue] = field(default_factory=dict) + resources: dict[str, str] = field(default_factory=dict) + kind: str = "registered" + + def __post_init__(self) -> None: + """Validate versioned ID and recursively executable-free arguments.""" + _validate_identifier(self.call_id, field_name="call_id") + if _REGISTERED_CALL_ID_PATTERN.fullmatch(self.call_id) is None: + raise ValueError( + "call_id must contain two or more lowercase identifier segments " + "separated by single dots." + ) + if type(self.schema_version) is not int or self.schema_version != 1: + raise ValueError("Registered call schema_version must be exactly 1.") + if type(self.arguments) is not dict: + raise TypeError("arguments must be an exact dict.") + arguments = _snapshot_declarative_value( + self.arguments, + path="arguments", + ) + assert type(arguments) is dict + self.arguments = arguments + self.resources = _validate_resources(self.resources, field_name="resources") + _validate_kind(self.kind, expected="registered", field_name="kind") + + +SemanticCallCfg: TypeAlias = ( + PickCfg + | PlaceCfg + | HandOverCfg + | OperateArticulationCfg + | RegisteredSemanticCallCfg +) + + +@configclass +class WaitStablePostCfg: + """Wait for one registered entity to satisfy a named stability preset.""" + + entity: str = MISSING + preset: str = "rigid_object" + kind: str = "wait_stable" + + def __post_init__(self) -> None: + """Validate entity, preset, and discriminator.""" + _validate_identifier(self.entity, field_name="entity") + _validate_identifier(self.preset, field_name="preset") + _validate_kind(self.kind, expected="wait_stable", field_name="kind") + + +PostPolicyCfg: TypeAlias = WaitStablePostCfg + + +@configclass +class ObjectNearTargetValidatorCfg: + """Validate an object's position against one resolved target.""" + + object: str = MISSING + target: str = MISSING + position_tolerance: float = 0.03 + kind: str = "object_near_target" + + def __post_init__(self) -> None: + """Validate reference IDs and a positive finite tolerance.""" + _validate_identifier(self.object, field_name="object") + _validate_identifier(self.target, field_name="target") + tolerance = _validate_number( + self.position_tolerance, + field_name="position_tolerance", + ) + if tolerance <= 0.0: + raise ValueError("position_tolerance must be positive.") + self.position_tolerance = tolerance + _validate_kind( + self.kind, + expected="object_near_target", + field_name="kind", + ) + + +ValidatorCfg: TypeAlias = ObjectNearTargetValidatorCfg + + +@configclass +class InvokeCfg: + """Invoke exactly one semantic call at the current program boundary.""" + + call: SemanticCallCfg = MISSING + kind: str = "invoke" + + def __post_init__(self) -> None: + """Validate the semantic-call union and discriminator.""" + if type(self.call) not in _SEMANTIC_CALL_TYPES: + raise TypeError("call must be an exact SemanticCallCfg value.") + _validate_kind(self.kind, expected="invoke", field_name="kind") + + +@configclass +class BarrierCfg: + """Explicit synchronization boundary owned by one parallel node.""" + + name: str = "join" + timeout_steps: int = 1_000 + failure_policy: str = "fail_fast" + kind: str = "barrier" + + def __post_init__(self) -> None: + """Validate deterministic timeout and cancellation semantics.""" + _validate_kind(self.kind, expected="barrier", field_name="kind") + _validate_identifier(self.name, field_name="name") + if type(self.timeout_steps) is not int or self.timeout_steps <= 0: + raise ValueError("timeout_steps must be a positive integer.") + if self.failure_policy != "fail_fast": + raise ValueError("failure_policy must be exactly 'fail_fast'.") + + +@configclass +class SequenceCfg: + """Execute one non-empty ordered tuple of program nodes.""" + + items: tuple[ProgramNodeCfg, ...] = MISSING + kind: str = "sequence" + + def __post_init__(self) -> None: + """Validate ordered child nodes and discriminator.""" + _validate_kind(self.kind, expected="sequence", field_name="kind") + if type(self.items) not in (list, tuple) or not self.items: + raise ValueError("items must contain at least one program node.") + items = tuple(self.items) + if not all(type(item) in _PROGRAM_NODE_TYPES for item in items): + raise TypeError("items must contain exact ProgramNodeCfg values.") + self.items = items # type: ignore[assignment] + + +@configclass +class RepeatCfg: + """Repeat one child node a finite validated number of times.""" + + count: int = MISSING + body: ProgramNodeCfg = MISSING + kind: str = "repeat" + + def __post_init__(self) -> None: + """Validate a bounded positive repeat and its child node.""" + if type(self.count) is not int or not 1 <= self.count <= MAX_REPEAT_COUNT: + raise ValueError(f"count must be an integer in [1, {MAX_REPEAT_COUNT}].") + if type(self.body) not in _PROGRAM_NODE_TYPES: + raise TypeError("body must be an exact ProgramNodeCfg value.") + _validate_kind(self.kind, expected="repeat", field_name="kind") + + +@configclass +class SegmentCfg: + """Logical program transaction with post-policies and validators.""" + + name: str = MISSING + steps: ProgramNodeCfg = MISSING + post: tuple[PostPolicyCfg, ...] = field(default_factory=tuple) + validators: tuple[ValidatorCfg, ...] = field(default_factory=tuple) + kind: str = "segment" + + def __post_init__(self) -> None: + """Validate the segment boundary and its declarative hooks.""" + _validate_identifier(self.name, field_name="name") + if type(self.steps) not in _PROGRAM_NODE_TYPES: + raise TypeError("steps must be an exact ProgramNodeCfg value.") + if type(self.post) not in (list, tuple): + raise TypeError("post must be a list or tuple.") + if type(self.validators) not in (list, tuple): + raise TypeError("validators must be a list or tuple.") + post = tuple(self.post) + validators = tuple(self.validators) + if not all(type(value) in _POST_POLICY_TYPES for value in post): + raise TypeError("post must contain exact PostPolicyCfg values.") + if not all(type(value) in _VALIDATOR_TYPES for value in validators): + raise TypeError("validators must contain exact ValidatorCfg values.") + self.post = post # type: ignore[assignment] + self.validators = validators # type: ignore[assignment] + _validate_kind(self.kind, expected="segment", field_name="kind") + + +@configclass +class ParallelCfg: + """Execute two or more branches concurrently and join at one barrier.""" + + branches: tuple[ProgramNodeCfg, ...] = MISSING + barrier: BarrierCfg = MISSING + kind: str = "parallel" + + def __post_init__(self) -> None: + """Validate branch ownership and an explicit synchronization node.""" + _validate_kind(self.kind, expected="parallel", field_name="kind") + if type(self.branches) not in (list, tuple) or len(self.branches) < 2: + raise ValueError("branches must contain at least two program nodes.") + branches = tuple(self.branches) + if not all(type(branch) in _PROGRAM_NODE_TYPES for branch in branches): + raise TypeError("branches must contain exact ProgramNodeCfg values.") + if any(type(branch) in (ParallelCfg, BarrierCfg) for branch in branches): + raise ValueError( + "Nested Parallel and standalone Barrier branches are forbidden." + ) + if type(self.barrier) is not BarrierCfg: + raise TypeError("barrier must be exactly BarrierCfg.") + self.branches = branches # type: ignore[assignment] + + +ProgramNodeCfg: TypeAlias = ( + SequenceCfg | RepeatCfg | SegmentCfg | InvokeCfg | ParallelCfg | BarrierCfg +) + +_SEMANTIC_CALL_TYPES = ( + PickCfg, + PlaceCfg, + HandOverCfg, + OperateArticulationCfg, + RegisteredSemanticCallCfg, +) +_POST_POLICY_TYPES = (WaitStablePostCfg,) +_VALIDATOR_TYPES = (ObjectNearTargetValidatorCfg,) +_PROGRAM_NODE_TYPES = ( + SequenceCfg, + RepeatCfg, + SegmentCfg, + InvokeCfg, + ParallelCfg, + BarrierCfg, +) + + +def _validate_target_reference(target: str, targets: dict[str, TargetCfg]) -> None: + """Require one target reference to exist in the top-level registry.""" + if target not in targets: + raise ValueError(f"Unknown target reference {target!r}.") + + +def _validate_program( + node: ProgramNodeCfg, + *, + targets: dict[str, TargetCfg], + depth: int, + budget: list[int], + schema_version: int, + inside_parallel: bool = False, +) -> int: + """Validate references and return the statically expanded call count.""" + if depth > MAX_PROGRAM_DEPTH: + raise ValueError(f"Program exceeds depth limit {MAX_PROGRAM_DEPTH}.") + budget[0] -= 1 + if budget[0] < 0: + raise ValueError(f"Program exceeds node limit {MAX_PROGRAM_NODES}.") + if type(node) is InvokeCfg: + call = node.call + if type(call) is PlaceCfg and call.at is not None: + _validate_target_reference(call.at.target, targets) + if type(call) is HandOverCfg and call.final_target is not None: + _validate_target_reference(call.final_target.target, targets) + return 1 + if type(node) is BarrierCfg: + if schema_version < EXPERT_PROGRAM_SCHEMA_VERSION_V2: + raise ValueError("Barrier requires Expert Program schema version 2.") + if not inside_parallel: + raise ValueError("Barrier nodes may only be owned by Parallel.") + return 0 + if type(node) is SequenceCfg: + expanded = sum( + _validate_program( + child, + targets=targets, + depth=depth + 1, + budget=budget, + schema_version=schema_version, + inside_parallel=inside_parallel, + ) + for child in node.items + ) + elif type(node) is RepeatCfg: + expanded = node.count * _validate_program( + node.body, + targets=targets, + depth=depth + 1, + budget=budget, + schema_version=schema_version, + inside_parallel=inside_parallel, + ) + elif type(node) is SegmentCfg: + if inside_parallel: + raise ValueError( + "Parallel branches may contain only Invoke, Sequence, and Repeat " + "nodes; wrap the Parallel node in one Segment instead." + ) + for validator in node.validators: + _validate_target_reference(validator.target, targets) + expanded = _validate_program( + node.steps, + targets=targets, + depth=depth + 1, + budget=budget, + schema_version=schema_version, + inside_parallel=inside_parallel, + ) + elif type(node) is ParallelCfg: + if schema_version < EXPERT_PROGRAM_SCHEMA_VERSION_V2: + raise ValueError("Parallel requires Expert Program schema version 2.") + if inside_parallel: + raise ValueError("Nested Parallel nodes are forbidden in schema version 2.") + branch_counts = tuple( + _validate_program( + branch, + targets=targets, + depth=depth + 1, + budget=budget, + schema_version=schema_version, + inside_parallel=True, + ) + for branch in node.branches + ) + if any(count <= 0 for count in branch_counts): + raise ValueError("Every Parallel branch must contain a semantic call.") + _validate_program( + node.barrier, + targets=targets, + depth=depth + 1, + budget=budget, + schema_version=schema_version, + inside_parallel=True, + ) + expanded = sum(branch_counts) + else: # pragma: no cover - exact construction prevents this branch + raise TypeError("program must contain exact ProgramNodeCfg values.") + if expanded > MAX_EXPANDED_CALLS: + raise ValueError( + f"Program expands to more than {MAX_EXPANDED_CALLS} semantic calls." + ) + return expanded + + +@configclass +class ExpertProgramCfg: + """Strict, versioned, executable-free Expert Program configuration.""" + + schema_version: int = MISSING + program_id: str = MISSING + integration: ExpertProgramIntegrationCfg = MISSING + program: ProgramNodeCfg = MISSING + targets: dict[str, TargetCfg] = field(default_factory=dict) + + def __post_init__(self) -> None: + """Validate the complete static configuration and target graph.""" + if ( + type(self.schema_version) is not int + or self.schema_version not in SUPPORTED_EXPERT_PROGRAM_SCHEMA_VERSIONS + ): + raise ValueError( + "schema_version must be one of " + f"{SUPPORTED_EXPERT_PROGRAM_SCHEMA_VERSIONS}." + ) + _validate_identifier(self.program_id, field_name="program_id") + if type(self.integration) is not ExpertProgramIntegrationCfg: + raise TypeError("integration must be ExpertProgramIntegrationCfg.") + if type(self.targets) is not dict: + raise TypeError("targets must be an exact dict.") + targets: dict[str, TargetCfg] = {} + for target_id, target in self.targets.items(): + normalized_id = _validate_identifier( + target_id, + field_name="target IDs", + ) + if type(target) is not CyclicPoseTargetCfg: + raise TypeError("targets must contain exact TargetCfg values.") + targets[normalized_id] = target + if type(self.program) not in _PROGRAM_NODE_TYPES: + raise TypeError("program must be an exact ProgramNodeCfg value.") + expanded = _validate_program( + self.program, + targets=targets, + depth=0, + budget=[MAX_PROGRAM_NODES], + schema_version=self.schema_version, + ) + if expanded <= 0: + raise ValueError("program must contain at least one semantic call.") + self.targets = targets + + +__all__ = [ + "BarrierCfg", + "CyclicPoseTargetCfg", + "DeclarativeCfgValue", + "EXPERT_PROGRAM_SCHEMA_VERSION", + "EXPERT_PROGRAM_SCHEMA_VERSION_V2", + "ExpertProgramCfg", + "ExpertProgramIntegrationCfg", + "HandOverCfg", + "InvokeCfg", + "MAX_DECLARATIVE_DEPTH", + "MAX_DECLARATIVE_NODES", + "MAX_EXPANDED_CALLS", + "MAX_PROGRAM_DEPTH", + "MAX_PROGRAM_NODES", + "MAX_REPEAT_COUNT", + "ObjectNearTargetValidatorCfg", + "OperateArticulationCfg", + "ParallelCfg", + "PickCfg", + "PlaceCfg", + "PoseCfg", + "PostPolicyCfg", + "ProgramNodeCfg", + "RegisteredSemanticCallCfg", + "RepeatCfg", + "SegmentCfg", + "SemanticCallCfg", + "SequenceCfg", + "SUPPORTED_EXPERT_PROGRAM_SCHEMA_VERSIONS", + "TargetCfg", + "TargetRefCfg", + "ValidatorCfg", + "WaitStablePostCfg", +] diff --git a/embodichain/lab/gym/envs/expert_program/compiler.py b/embodichain/lab/gym/envs/expert_program/compiler.py new file mode 100644 index 000000000..cc12cea0d --- /dev/null +++ b/embodichain/lab/gym/envs/expert_program/compiler.py @@ -0,0 +1,1912 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Provider-free compilation and lazy expansion of Expert Program ASTs.""" + +from __future__ import annotations + +from collections.abc import Iterator, Mapping +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Protocol, runtime_checkable + +from embodichain.lab.sim.skills.calls import ( + DeclarativeValue, + HandOver, + OperateArticulation, + Pick, + Place, + RegisteredSemanticCall, + SemanticCallSpec, + SemanticPose, +) +from embodichain.lab.sim.skills.scene import ( + SceneAffordanceRef, + SceneArticulationRef, + SceneEntityRef, + SceneLinkRef, + SceneObjectRef, + SceneRegistry, +) + +from .cfg import ( + EXPERT_PROGRAM_SCHEMA_VERSION, + EXPERT_PROGRAM_SCHEMA_VERSION_V2, + MAX_EXPANDED_CALLS, + MAX_REPEAT_COUNT, + SUPPORTED_EXPERT_PROGRAM_SCHEMA_VERSIONS, + BarrierCfg, + CyclicPoseTargetCfg, + ExpertProgramCfg, + ExpertProgramIntegrationCfg, + HandOverCfg, + InvokeCfg, + ObjectNearTargetValidatorCfg, + OperateArticulationCfg, + ParallelCfg, + PickCfg, + PlaceCfg, + PoseCfg, + ProgramNodeCfg, + RegisteredSemanticCallCfg, + RepeatCfg, + SegmentCfg, + SequenceCfg, + TargetRefCfg, + WaitStablePostCfg, +) +from .decoder import ConfigPath, ExpertProgramConfigError, render_config_path + +_SEMANTIC_CALL_TYPES = ( + Pick, + Place, + HandOver, + OperateArticulation, + RegisteredSemanticCall, +) +_SCENE_REF_TYPES = ( + SceneEntityRef, + SceneObjectRef, + SceneArticulationRef, + SceneLinkRef, + SceneAffordanceRef, +) + + +class ExpertProgramCompileError(ExpertProgramConfigError): + """Raised when a validated AST cannot lower to canonical semantic calls.""" + + +@runtime_checkable +class ExpertProgramSceneResolver(Protocol): + """Provider-free typed resolver for canonical static scene references.""" + + def resolve( + self, + reference: str, + *, + expected_types: tuple[type[SceneEntityRef], ...], + path: ConfigPath, + ) -> SceneEntityRef: + """Resolve one canonical or aliased ID without observing scene state.""" + + +def _copy_scene_ref(reference: SceneEntityRef) -> SceneEntityRef: + """Return one independent exact typed scene reference.""" + if type(reference) not in _SCENE_REF_TYPES: + raise TypeError(f"Unsupported scene reference {type(reference).__name__}.") + return type(reference)(reference.entity_id) + + +class SceneRegistryProgramResolver: + """Provider-free static resolver snapshotted from one SceneRegistry. + + The resolver copies only canonical typed references and aliases. It does not + retain registrations, state providers, geometry providers, or the registry + itself, so compilation cannot observe dynamic scene state. + """ + + def __init__(self, registry: SceneRegistry) -> None: + """Snapshot the registry's static identity table. + + Args: + registry: Authoritative registry used only for static identity data. + """ + if type(registry) is not SceneRegistry: + raise TypeError("registry must be exactly SceneRegistry.") + references = { + reference.entity_id: _copy_scene_ref(reference) + for reference in registry.entity_refs + } + self._references = MappingProxyType(references) + self._aliases = MappingProxyType(dict(registry.aliases)) + + @property + def canonical_references(self) -> Mapping[str, SceneEntityRef]: + """Return an independent canonical typed-reference mapping.""" + return MappingProxyType( + { + entity_id: _copy_scene_ref(reference) + for entity_id, reference in self._references.items() + } + ) + + def resolve( + self, + reference: str, + *, + expected_types: tuple[type[SceneEntityRef], ...], + path: ConfigPath, + ) -> SceneEntityRef: + """Resolve one ID or alias through the snapshotted type table.""" + if ( + type(reference) is not str + or not reference + or reference != reference.strip() + ): + raise ExpertProgramCompileError( + "invalid_scene_reference", + path, + "Scene references must be non-empty strings without outer whitespace.", + ) + if ( + type(expected_types) is not tuple + or not expected_types + or not all( + isinstance(expected_type, type) + and issubclass(expected_type, SceneEntityRef) + for expected_type in expected_types + ) + ): + raise TypeError( + "expected_types must be a non-empty tuple of SceneEntityRef types." + ) + canonical_id = self._aliases.get(reference, reference) + resolved = self._references.get(canonical_id) + if resolved is None: + raise ExpertProgramCompileError( + "unknown_scene_reference", + path, + f"Unknown scene reference {reference!r}.", + ) + if type(resolved) not in expected_types: + expected_names = tuple(value.__name__ for value in expected_types) + raise ExpertProgramCompileError( + "scene_reference_type_mismatch", + path, + f"Scene reference {reference!r} resolves to " + f"{type(resolved).__name__}, expected one of {expected_names}.", + ) + return _copy_scene_ref(resolved) + + +@dataclass(frozen=True, slots=True) +class CompiledRepeatFrame: + """One lexical repeat occurrence in a compiled call or segment path.""" + + path: ConfigPath + iteration_index: int + count: int + + def __post_init__(self) -> None: + if type(self.path) is not tuple: + raise TypeError("path must be a ConfigPath tuple.") + if type(self.iteration_index) is not int or not 0 <= self.iteration_index: + raise ValueError("iteration_index must be a non-negative integer.") + if type(self.count) is not int or self.count <= 0: + raise ValueError("count must be a positive integer.") + if self.iteration_index >= self.count: + raise ValueError("iteration_index must be smaller than count.") + + +@dataclass(frozen=True, slots=True) +class CompiledTargetSelection: + """Deterministic cyclic-target selection metadata for one occurrence.""" + + target_id: str + value_index: int + repeat_path: ConfigPath | None + repeat_iteration_index: int | None + + def __post_init__(self) -> None: + if type(self.target_id) is not str or not self.target_id: + raise ValueError("target_id must be a non-empty string.") + if type(self.value_index) is not int or self.value_index < 0: + raise ValueError("value_index must be a non-negative integer.") + if (self.repeat_path is None) != (self.repeat_iteration_index is None): + raise ValueError( + "repeat_path and repeat_iteration_index must both be set or unset." + ) + if self.repeat_path is not None and type(self.repeat_path) is not tuple: + raise TypeError("repeat_path must be a ConfigPath tuple or None.") + if self.repeat_iteration_index is not None and ( + type(self.repeat_iteration_index) is not int + or self.repeat_iteration_index < 0 + ): + raise ValueError("repeat_iteration_index must be non-negative or None.") + + +def _snapshot_semantic_call(call: SemanticCallSpec) -> SemanticCallSpec: + """Return one independently owned exact semantic-call value.""" + if type(call) is Pick: + return Pick( + object=_copy_scene_ref(call.object), + grasp=(None if call.grasp is None else _copy_scene_ref(call.grasp)), + resources=dict(call.resources), + ) + if type(call) is Place: + return Place( + object=_copy_scene_ref(call.object), + at=None if call.at is None else call.at.snapshot(), + on=None if call.on is None else _copy_scene_ref(call.on), + inside=None if call.inside is None else _copy_scene_ref(call.inside), + resources=dict(call.resources), + ) + if type(call) is HandOver: + return HandOver( + object=_copy_scene_ref(call.object), + receiver=call.receiver, + final_target=( + None if call.final_target is None else call.final_target.snapshot() + ), + resources=dict(call.resources), + ) + if type(call) is OperateArticulation: + return OperateArticulation( + articulation=_copy_scene_ref(call.articulation), + handle=(None if call.handle is None else _copy_scene_ref(call.handle)), + target=call.target, + target_position=call.target_position, + target_displacement=call.target_displacement, + resources=dict(call.resources), + ) + if type(call) is RegisteredSemanticCall: + return RegisteredSemanticCall( + call_id=call.call_id, + arguments=call.arguments, + resources=dict(call.resources), + ) + raise TypeError("call must be an exact supported SemanticCallSpec value.") + + +@dataclass(frozen=True, slots=True) +class CompiledProgramCall: + """One owned semantic call occurrence emitted by lazy program expansion.""" + + call_index: int + segment_call_index: int + call: SemanticCallSpec + source_path: ConfigPath + repeat_frames: tuple[CompiledRepeatFrame, ...] = () + target_selections: tuple[CompiledTargetSelection, ...] = () + + def __post_init__(self) -> None: + if type(self.call_index) is not int or self.call_index < 0: + raise ValueError("call_index must be a non-negative integer.") + if type(self.segment_call_index) is not int or self.segment_call_index < 0: + raise ValueError("segment_call_index must be a non-negative integer.") + if type(self.call) not in _SEMANTIC_CALL_TYPES: + raise TypeError("call must be an exact supported SemanticCallSpec value.") + if type(self.source_path) is not tuple: + raise TypeError("source_path must be a ConfigPath tuple.") + frames = tuple(self.repeat_frames) + selections = tuple(self.target_selections) + if not all(type(frame) is CompiledRepeatFrame for frame in frames): + raise TypeError("repeat_frames must contain CompiledRepeatFrame values.") + if not all( + type(selection) is CompiledTargetSelection for selection in selections + ): + raise TypeError( + "target_selections must contain CompiledTargetSelection values." + ) + object.__setattr__(self, "call", _snapshot_semantic_call(self.call)) + object.__setattr__(self, "repeat_frames", frames) + object.__setattr__(self, "target_selections", selections) + + +@dataclass(frozen=True, slots=True) +class CompiledPostPolicy: + """Owned post-policy config plus its canonical scene entity and source path.""" + + cfg: WaitStablePostCfg + entity: SceneEntityRef + source_path: ConfigPath + + def __post_init__(self) -> None: + if type(self.cfg) is not WaitStablePostCfg: + raise TypeError("cfg must be exactly WaitStablePostCfg.") + if type(self.entity) not in _SCENE_REF_TYPES: + raise TypeError("entity must be an exact SceneEntityRef value.") + if type(self.source_path) is not tuple: + raise TypeError("source_path must be a ConfigPath tuple.") + object.__setattr__( + self, + "cfg", + WaitStablePostCfg( + entity=self.cfg.entity, + preset=self.cfg.preset, + kind=self.cfg.kind, + ), + ) + object.__setattr__(self, "entity", _copy_scene_ref(self.entity)) + + +@dataclass(frozen=True, slots=True) +class CompiledProgramValidator: + """Owned validator config with canonical object and resolved target pose.""" + + cfg: ObjectNearTargetValidatorCfg + object: SceneObjectRef + target_pose: SemanticPose + target_selection: CompiledTargetSelection + source_path: ConfigPath + + def __post_init__(self) -> None: + if type(self.cfg) is not ObjectNearTargetValidatorCfg: + raise TypeError("cfg must be exactly ObjectNearTargetValidatorCfg.") + if type(self.object) is not SceneObjectRef: + raise TypeError("object must be exactly SceneObjectRef.") + if type(self.target_pose) is not SemanticPose: + raise TypeError("target_pose must be exactly SemanticPose.") + if type(self.target_selection) is not CompiledTargetSelection: + raise TypeError("target_selection must be exactly CompiledTargetSelection.") + if type(self.source_path) is not tuple: + raise TypeError("source_path must be a ConfigPath tuple.") + object.__setattr__( + self, + "cfg", + ObjectNearTargetValidatorCfg( + object=self.cfg.object, + target=self.cfg.target, + position_tolerance=self.cfg.position_tolerance, + kind=self.cfg.kind, + ), + ) + object.__setattr__(self, "object", _copy_scene_ref(self.object)) + object.__setattr__(self, "target_pose", self.target_pose.snapshot()) + + +@dataclass(frozen=True, slots=True) +class CompiledBarrier: + """Explicit schema-v2 join semantics for one compiled parallel block.""" + + name: str + timeout_steps: int + failure_policy: str + source_path: ConfigPath + + def __post_init__(self) -> None: + if type(self.name) is not str or not self.name: + raise ValueError("barrier name must be non-empty.") + if type(self.timeout_steps) is not int or self.timeout_steps <= 0: + raise ValueError("barrier timeout_steps must be positive.") + if self.failure_policy != "fail_fast": + raise ValueError("barrier failure_policy must be 'fail_fast'.") + if type(self.source_path) is not tuple: + raise TypeError("barrier source_path must be a ConfigPath tuple.") + + +@dataclass(frozen=True, slots=True) +class CompiledParallelBranch: + """One ordered semantic-call lane inside a parallel block.""" + + branch_index: int + calls: tuple[CompiledProgramCall, ...] + source_path: ConfigPath + + def __post_init__(self) -> None: + if type(self.branch_index) is not int or self.branch_index < 0: + raise ValueError("branch_index must be non-negative.") + calls = tuple(self.calls) + if not calls or not all(type(call) is CompiledProgramCall for call in calls): + raise TypeError("parallel branch calls must be non-empty compiled calls.") + if type(self.source_path) is not tuple: + raise TypeError("parallel branch source_path must be a ConfigPath tuple.") + object.__setattr__(self, "calls", calls) + + +@dataclass(frozen=True, slots=True) +class CompiledParallelBlock: + """Two or more call lanes joined by an explicit deterministic barrier.""" + + branches: tuple[CompiledParallelBranch, ...] + barrier: CompiledBarrier + source_path: ConfigPath + + def __post_init__(self) -> None: + branches = tuple(self.branches) + if len(branches) < 2 or not all( + type(branch) is CompiledParallelBranch for branch in branches + ): + raise TypeError("parallel blocks require at least two compiled branches.") + if tuple(branch.branch_index for branch in branches) != tuple( + range(len(branches)) + ): + raise ValueError("parallel branch indices must be contiguous from zero.") + if type(self.barrier) is not CompiledBarrier: + raise TypeError("barrier must be exactly CompiledBarrier.") + if type(self.source_path) is not tuple: + raise TypeError("parallel source_path must be a ConfigPath tuple.") + object.__setattr__(self, "branches", branches) + + +@dataclass(frozen=True, slots=True) +class CompiledProgramSegment: + """One independent explicit or implicit logical program segment.""" + + segment_index: int + segment_id: str + name: str + calls: tuple[CompiledProgramCall, ...] + source_path: ConfigPath + repeat_frames: tuple[CompiledRepeatFrame, ...] = () + post_policies: tuple[CompiledPostPolicy, ...] = () + validators: tuple[CompiledProgramValidator, ...] = () + parallel_block: CompiledParallelBlock | None = None + implicit: bool = False + + def __post_init__(self) -> None: + if type(self.segment_index) is not int or self.segment_index < 0: + raise ValueError("segment_index must be a non-negative integer.") + for field_name in ("segment_id", "name"): + value = getattr(self, field_name) + if type(value) is not str or not value: + raise ValueError(f"{field_name} must be a non-empty string.") + calls = tuple(self.calls) + if not calls or not all(type(call) is CompiledProgramCall for call in calls): + raise TypeError( + "calls must contain at least one exact CompiledProgramCall." + ) + if tuple(call.segment_call_index for call in calls) != tuple(range(len(calls))): + raise ValueError("segment call indices must be contiguous from zero.") + if type(self.source_path) is not tuple: + raise TypeError("source_path must be a ConfigPath tuple.") + frames = tuple(self.repeat_frames) + post = tuple(self.post_policies) + validators = tuple(self.validators) + if not all(type(frame) is CompiledRepeatFrame for frame in frames): + raise TypeError("repeat_frames must contain CompiledRepeatFrame values.") + if not all(type(value) is CompiledPostPolicy for value in post): + raise TypeError("post_policies must contain CompiledPostPolicy values.") + if not all(type(value) is CompiledProgramValidator for value in validators): + raise TypeError("validators must contain CompiledProgramValidator values.") + if type(self.implicit) is not bool: + raise TypeError("implicit must be a bool.") + if self.implicit and (post or validators): + raise ValueError( + "Implicit segments cannot own post-policies or validators." + ) + if self.parallel_block is not None: + if type(self.parallel_block) is not CompiledParallelBlock: + raise TypeError("parallel_block must be CompiledParallelBlock or None.") + flattened = tuple( + call for branch in self.parallel_block.branches for call in branch.calls + ) + if flattened != calls: + raise ValueError( + "segment calls must equal parallel branch calls in branch order." + ) + object.__setattr__(self, "calls", calls) + object.__setattr__(self, "repeat_frames", frames) + object.__setattr__(self, "post_policies", post) + object.__setattr__(self, "validators", validators) + + +@dataclass(frozen=True, slots=True) +class CompiledProgramAnalysis: + """One owned canonical semantic-analysis window for a compiled program. + + ``execution_prefix_length`` separates calls that the current segment owns + from downstream calls included only for static state-flow and target + look-ahead. Preflight analyses set the prefix to the complete window. + """ + + analysis_id: str + kind: str + calls: tuple[SemanticCallSpec, ...] + source_path: ConfigPath + segment_indices: tuple[int, ...] + execution_prefix_length: int + + def __post_init__(self) -> None: + if type(self.analysis_id) is not str or not self.analysis_id: + raise ValueError("analysis_id must be a non-empty string.") + if self.kind not in { + "sequential_stretch", + "parallel_branch", + "sequential_suffix", + }: + raise ValueError("kind must identify a supported program analysis.") + calls = tuple(self.calls) + if not calls or not all(type(call) in _SEMANTIC_CALL_TYPES for call in calls): + raise TypeError("calls must contain supported semantic call values.") + if type(self.source_path) is not tuple: + raise TypeError("source_path must be a ConfigPath tuple.") + indices = tuple(self.segment_indices) + if not indices or any(type(index) is not int or index < 0 for index in indices): + raise ValueError("segment_indices must contain non-negative integers.") + if len(set(indices)) != len(indices) or tuple(sorted(indices)) != indices: + raise ValueError("segment_indices must be unique and ordered.") + if type( + self.execution_prefix_length + ) is not int or not 1 <= self.execution_prefix_length <= len(calls): + raise ValueError( + "execution_prefix_length must select a non-empty prefix of calls." + ) + object.__setattr__( + self, + "calls", + tuple(_snapshot_semantic_call(call) for call in calls), + ) + object.__setattr__(self, "segment_indices", indices) + + +@dataclass(frozen=True, slots=True) +class _CallTemplate: + kind: str + source_path: ConfigPath + object: SceneObjectRef | None = None + grasp: SceneAffordanceRef | None = None + at_target_id: str | None = None + on: SceneObjectRef | SceneAffordanceRef | None = None + inside: SceneObjectRef | SceneAffordanceRef | None = None + receiver: str | None = None + final_target_id: str | None = None + articulation: SceneArticulationRef | None = None + handle: SceneAffordanceRef | None = None + articulation_target: str | None = None + target_position: float | None = None + target_displacement: float | None = None + call_id: str | None = None + arguments: Mapping[str, DeclarativeValue] | None = None + resources: tuple[tuple[str, str], ...] = () + + +@dataclass(frozen=True, slots=True) +class _InvokeTemplate: + call: _CallTemplate + source_path: ConfigPath + + +@dataclass(frozen=True, slots=True) +class _SequenceTemplate: + items: tuple[_NodeTemplate, ...] + source_path: ConfigPath + + +@dataclass(frozen=True, slots=True) +class _RepeatTemplate: + count: int + body: _NodeTemplate + source_path: ConfigPath + + +@dataclass(frozen=True, slots=True) +class _BarrierTemplate: + name: str + timeout_steps: int + failure_policy: str + source_path: ConfigPath + + +@dataclass(frozen=True, slots=True) +class _ParallelTemplate: + branches: tuple[_NodeTemplate, ...] + barrier: _BarrierTemplate + source_path: ConfigPath + + +@dataclass(frozen=True, slots=True) +class _PostTemplate: + cfg: WaitStablePostCfg + entity: SceneEntityRef + source_path: ConfigPath + + +@dataclass(frozen=True, slots=True) +class _ValidatorTemplate: + cfg: ObjectNearTargetValidatorCfg + object: SceneObjectRef + target_id: str + source_path: ConfigPath + + +@dataclass(frozen=True, slots=True) +class _SegmentTemplate: + name: str + steps: _NodeTemplate + post: tuple[_PostTemplate, ...] + validators: tuple[_ValidatorTemplate, ...] + source_path: ConfigPath + + +_NodeTemplate = ( + _InvokeTemplate + | _SequenceTemplate + | _RepeatTemplate + | _SegmentTemplate + | _ParallelTemplate + | _BarrierTemplate +) + + +def _contains_parallel(template: _NodeTemplate) -> bool: + """Return whether a compiled subtree owns a parallel block.""" + if type(template) is _ParallelTemplate: + return True + if type(template) is _SequenceTemplate: + return any(_contains_parallel(child) for child in template.items) + if type(template) is _RepeatTemplate: + return _contains_parallel(template.body) + if type(template) is _SegmentTemplate: + return _contains_parallel(template.steps) + return False + + +@dataclass(slots=True) +class _ExpansionState: + segment_index: int = 0 + call_index: int = 0 + + +def _resolve_target( + target_id: str, + *, + targets: Mapping[str, tuple[SemanticPose, ...]], + repeat_frames: tuple[CompiledRepeatFrame, ...], +) -> tuple[SemanticPose, CompiledTargetSelection]: + """Select one cyclic target from the nearest lexical repeat frame.""" + values = targets[target_id] + repeat = repeat_frames[-1] if repeat_frames else None + value_index = 0 if repeat is None else repeat.iteration_index % len(values) + selection = CompiledTargetSelection( + target_id=target_id, + value_index=value_index, + repeat_path=None if repeat is None else repeat.path, + repeat_iteration_index=None if repeat is None else repeat.iteration_index, + ) + return values[value_index].snapshot(), selection + + +def _instantiate_call( + template: _CallTemplate, + *, + call_index: int, + segment_call_index: int, + targets: Mapping[str, tuple[SemanticPose, ...]], + repeat_frames: tuple[CompiledRepeatFrame, ...], +) -> CompiledProgramCall: + """Instantiate one semantic call occurrence from static templates.""" + resources = dict(template.resources) + selections: list[CompiledTargetSelection] = [] + if template.kind == "pick": + assert template.object is not None + call: SemanticCallSpec = Pick( + object=_copy_scene_ref(template.object), + grasp=(None if template.grasp is None else _copy_scene_ref(template.grasp)), + resources=resources, + ) + elif template.kind == "place": + assert template.object is not None + at: SemanticPose | None = None + if template.at_target_id is not None: + at, selection = _resolve_target( + template.at_target_id, + targets=targets, + repeat_frames=repeat_frames, + ) + selections.append(selection) + call = Place( + object=_copy_scene_ref(template.object), + at=at, + on=None if template.on is None else _copy_scene_ref(template.on), + inside=( + None if template.inside is None else _copy_scene_ref(template.inside) + ), + resources=resources, + ) + elif template.kind == "hand_over": + assert template.object is not None + final_target: SemanticPose | None = None + if template.final_target_id is not None: + final_target, selection = _resolve_target( + template.final_target_id, + targets=targets, + repeat_frames=repeat_frames, + ) + selections.append(selection) + call = HandOver( + object=_copy_scene_ref(template.object), + receiver=template.receiver, + final_target=final_target, + resources=resources, + ) + elif template.kind == "operate_articulation": + assert template.articulation is not None + call = OperateArticulation( + articulation=_copy_scene_ref(template.articulation), + handle=( + None if template.handle is None else _copy_scene_ref(template.handle) + ), + target=template.articulation_target, + target_position=template.target_position, + target_displacement=template.target_displacement, + resources=resources, + ) + elif template.kind == "registered": + assert template.call_id is not None and template.arguments is not None + call = RegisteredSemanticCall( + call_id=template.call_id, + arguments=template.arguments, + resources=resources, + ) + else: # pragma: no cover - compiler-owned templates prevent this + raise AssertionError(f"Unknown call template {template.kind!r}.") + return CompiledProgramCall( + call_index=call_index, + segment_call_index=segment_call_index, + call=call, + source_path=template.source_path, + repeat_frames=repeat_frames, + target_selections=tuple(selections), + ) + + +def _iter_call_templates( + template: _NodeTemplate, + *, + repeat_frames: tuple[CompiledRepeatFrame, ...], +) -> Iterator[tuple[_CallTemplate, tuple[CompiledRepeatFrame, ...]]]: + """Expand call templates inside one explicit segment without segment splits.""" + if type(template) is _InvokeTemplate: + yield template.call, repeat_frames + elif type(template) is _SequenceTemplate: + for child in template.items: + yield from _iter_call_templates(child, repeat_frames=repeat_frames) + elif type(template) is _RepeatTemplate: + for iteration_index in range(template.count): + frame = CompiledRepeatFrame( + path=template.source_path, + iteration_index=iteration_index, + count=template.count, + ) + yield from _iter_call_templates( + template.body, + repeat_frames=(*repeat_frames, frame), + ) + else: # pragma: no cover - nested segments are rejected during compilation + raise AssertionError("A nested segment reached call-only expansion.") + + +def _instantiate_parallel_block( + template: _ParallelTemplate, + *, + targets: Mapping[str, tuple[SemanticPose, ...]], + repeat_frames: tuple[CompiledRepeatFrame, ...], + state: _ExpansionState, +) -> tuple[CompiledParallelBlock, tuple[CompiledProgramCall, ...]]: + """Instantiate branch-local call order without serializing branch semantics.""" + branches: list[CompiledParallelBranch] = [] + flattened: list[CompiledProgramCall] = [] + segment_call_index = 0 + for branch_index, branch_template in enumerate(template.branches): + calls: list[CompiledProgramCall] = [] + for call_template, call_repeat_frames in _iter_call_templates( + branch_template, + repeat_frames=repeat_frames, + ): + call = _instantiate_call( + call_template, + call_index=state.call_index, + segment_call_index=segment_call_index, + targets=targets, + repeat_frames=call_repeat_frames, + ) + calls.append(call) + flattened.append(call) + state.call_index += 1 + segment_call_index += 1 + branches.append( + CompiledParallelBranch( + branch_index=branch_index, + calls=tuple(calls), + source_path=template.branches[branch_index].source_path, + ) + ) + barrier = CompiledBarrier( + name=template.barrier.name, + timeout_steps=template.barrier.timeout_steps, + failure_policy=template.barrier.failure_policy, + source_path=template.barrier.source_path, + ) + return ( + CompiledParallelBlock( + branches=tuple(branches), + barrier=barrier, + source_path=template.source_path, + ), + tuple(flattened), + ) + + +def _segment_identity( + program_id: str, + *, + source_path: ConfigPath, + repeat_frames: tuple[CompiledRepeatFrame, ...], + implicit: bool, +) -> str: + """Build one deterministic segment identity from lexical occurrence data.""" + repeat_suffix = "".join( + f"@{render_config_path(frame.path)}[{frame.iteration_index}]" + for frame in repeat_frames + ) + boundary = "implicit" if implicit else "segment" + return f"{program_id}:{boundary}:{render_config_path(source_path)}{repeat_suffix}" + + +def _iter_segments( + template: _NodeTemplate, + *, + program_id: str, + targets: Mapping[str, tuple[SemanticPose, ...]], + repeat_frames: tuple[CompiledRepeatFrame, ...], + state: _ExpansionState, +) -> Iterator[CompiledProgramSegment]: + """Lazily expand outer program structure into independent segments.""" + if type(template) is _SequenceTemplate: + for child in template.items: + yield from _iter_segments( + child, + program_id=program_id, + targets=targets, + repeat_frames=repeat_frames, + state=state, + ) + return + if type(template) is _RepeatTemplate: + for iteration_index in range(template.count): + frame = CompiledRepeatFrame( + path=template.source_path, + iteration_index=iteration_index, + count=template.count, + ) + yield from _iter_segments( + template.body, + program_id=program_id, + targets=targets, + repeat_frames=(*repeat_frames, frame), + state=state, + ) + return + if type(template) is _InvokeTemplate: + call = _instantiate_call( + template.call, + call_index=state.call_index, + segment_call_index=0, + targets=targets, + repeat_frames=repeat_frames, + ) + state.call_index += 1 + segment = CompiledProgramSegment( + segment_index=state.segment_index, + segment_id=_segment_identity( + program_id, + source_path=template.source_path, + repeat_frames=repeat_frames, + implicit=True, + ), + name=f"invoke:{call.call.semantic_id}", + calls=(call,), + source_path=template.source_path, + repeat_frames=repeat_frames, + implicit=True, + ) + state.segment_index += 1 + yield segment + return + + if type(template) is _ParallelTemplate: + parallel_block, calls = _instantiate_parallel_block( + template, + targets=targets, + repeat_frames=repeat_frames, + state=state, + ) + segment = CompiledProgramSegment( + segment_index=state.segment_index, + segment_id=_segment_identity( + program_id, + source_path=template.source_path, + repeat_frames=repeat_frames, + implicit=True, + ), + name=f"parallel:{parallel_block.barrier.name}", + calls=calls, + source_path=template.source_path, + repeat_frames=repeat_frames, + parallel_block=parallel_block, + implicit=True, + ) + state.segment_index += 1 + yield segment + return + + assert type(template) is _SegmentTemplate + parallel_block: CompiledParallelBlock | None = None + if type(template.steps) is _ParallelTemplate: + parallel_block, instantiated_calls = _instantiate_parallel_block( + template.steps, + targets=targets, + repeat_frames=repeat_frames, + state=state, + ) + calls = list(instantiated_calls) + else: + calls = [] + for segment_call_index, (call_template, call_repeat_frames) in enumerate( + _iter_call_templates(template.steps, repeat_frames=repeat_frames) + ): + calls.append( + _instantiate_call( + call_template, + call_index=state.call_index, + segment_call_index=segment_call_index, + targets=targets, + repeat_frames=call_repeat_frames, + ) + ) + state.call_index += 1 + post_policies = tuple( + CompiledPostPolicy( + cfg=post.cfg, + entity=post.entity, + source_path=post.source_path, + ) + for post in template.post + ) + validators: list[CompiledProgramValidator] = [] + for validator in template.validators: + target_pose, selection = _resolve_target( + validator.target_id, + targets=targets, + repeat_frames=repeat_frames, + ) + validators.append( + CompiledProgramValidator( + cfg=validator.cfg, + object=validator.object, + target_pose=target_pose, + target_selection=selection, + source_path=validator.source_path, + ) + ) + segment = CompiledProgramSegment( + segment_index=state.segment_index, + segment_id=_segment_identity( + program_id, + source_path=template.source_path, + repeat_frames=repeat_frames, + implicit=False, + ), + name=template.name, + calls=tuple(calls), + source_path=template.source_path, + repeat_frames=repeat_frames, + post_policies=post_policies, + validators=tuple(validators), + parallel_block=parallel_block, + implicit=False, + ) + state.segment_index += 1 + yield segment + + +@dataclass(frozen=True, slots=True, init=False) +class CompiledProgram: + """Owned provider-free program template with lazy deterministic expansion.""" + + schema_version: int + program_id: str + _integration: ExpertProgramIntegrationCfg = field(repr=False, compare=False) + _targets: Mapping[str, tuple[SemanticPose, ...]] = field( + repr=False, + compare=False, + ) + _root: _NodeTemplate = field(repr=False, compare=False) + + def __init__(self, *args: object, **kwargs: object) -> None: + """Reject construction outside :class:`ExpertProgramCompiler`.""" + del args, kwargs + raise TypeError("CompiledProgram values are created by ExpertProgramCompiler.") + + @classmethod + def _create( + cls, + *, + schema_version: int, + program_id: str, + integration: ExpertProgramIntegrationCfg, + targets: Mapping[str, tuple[SemanticPose, ...]], + root: _NodeTemplate, + ) -> CompiledProgram: + """Create one compiler-owned lazy program template.""" + instance = object.__new__(cls) + object.__setattr__(instance, "schema_version", schema_version) + object.__setattr__(instance, "program_id", program_id) + object.__setattr__(instance, "_integration", integration) + object.__setattr__( + instance, + "_targets", + MappingProxyType( + { + target_id: tuple(pose.snapshot() for pose in values) + for target_id, values in targets.items() + } + ), + ) + object.__setattr__(instance, "_root", root) + return instance + + @property + def integration(self) -> ExpertProgramIntegrationCfg: + """Return an independent integration-selection snapshot.""" + return ExpertProgramIntegrationCfg( + robot_profile=self._integration.robot_profile, + scene_registry=self._integration.scene_registry, + runtime_preset=self._integration.runtime_preset, + ) + + @property + def targets(self) -> Mapping[str, tuple[SemanticPose, ...]]: + """Return independent static target-pose snapshots.""" + return MappingProxyType( + { + target_id: tuple(pose.snapshot() for pose in values) + for target_id, values in self._targets.items() + } + ) + + def iter_segments(self) -> Iterator[CompiledProgramSegment]: + """Lazily expand a fresh deterministic segment stream.""" + return _iter_segments( + self._root, + program_id=self.program_id, + targets=self._targets, + repeat_frames=(), + state=_ExpansionState(), + ) + + def materialize(self) -> MaterializedCompiledProgram: + """Expand the bounded provider-free segment stream exactly once. + + Materialization never observes a scene provider. It also re-enforces + the public expanded-call bound so a configuration mutated after its + initial validation cannot create an unbounded bridge-preflight pass. + + Returns: + Immutable materialized program with deterministic analysis windows. + + Raises: + ExpertProgramCompileError: If expansion exceeds the configured + semantic-call bound. + """ + segments: list[CompiledProgramSegment] = [] + expanded_calls = 0 + for segment in self.iter_segments(): + expanded_calls += len(segment.calls) + if expanded_calls > MAX_EXPANDED_CALLS: + raise ExpertProgramCompileError( + "expanded_call_limit", + segment.source_path, + "Program materialization exceeds the static limit of " + f"{MAX_EXPANDED_CALLS} semantic calls.", + ) + segments.append(segment) + return MaterializedCompiledProgram._create( + schema_version=self.schema_version, + program_id=self.program_id, + integration=self._integration, + segments=tuple(segments), + ) + + def __iter__(self) -> Iterator[CompiledProgramSegment]: + return self.iter_segments() + + +@dataclass(frozen=True, slots=True, init=False) +class MaterializedCompiledProgram: + """Bounded provider-free segment snapshot used by preflight and execution.""" + + schema_version: int + program_id: str + _integration: ExpertProgramIntegrationCfg = field(repr=False, compare=False) + _segments: tuple[CompiledProgramSegment, ...] = field( + repr=False, + compare=False, + ) + + def __init__(self, *args: object, **kwargs: object) -> None: + """Reject construction outside :meth:`CompiledProgram.materialize`.""" + del args, kwargs + raise TypeError( + "MaterializedCompiledProgram values are created by " + "CompiledProgram.materialize()." + ) + + @classmethod + def _create( + cls, + *, + schema_version: int, + program_id: str, + integration: ExpertProgramIntegrationCfg, + segments: tuple[CompiledProgramSegment, ...], + ) -> MaterializedCompiledProgram: + """Create one compiler-owned materialized program.""" + if type(schema_version) is not int or schema_version < 1: + raise ValueError("schema_version must be a positive integer.") + if type(program_id) is not str or not program_id: + raise ValueError("program_id must be a non-empty string.") + if type(integration) is not ExpertProgramIntegrationCfg: + raise TypeError("integration must be ExpertProgramIntegrationCfg.") + values = tuple(segments) + if not values or not all( + type(segment) is CompiledProgramSegment for segment in values + ): + raise TypeError( + "segments must contain at least one CompiledProgramSegment." + ) + if tuple(segment.segment_index for segment in values) != tuple( + range(len(values)) + ): + raise ValueError("Materialized segment indices must be contiguous.") + flattened_calls = tuple(call for segment in values for call in segment.calls) + if len(flattened_calls) > MAX_EXPANDED_CALLS: + raise ValueError( + f"Materialized program exceeds {MAX_EXPANDED_CALLS} calls." + ) + if tuple(call.call_index for call in flattened_calls) != tuple( + range(len(flattened_calls)) + ): + raise ValueError("Materialized call indices must be contiguous.") + + instance = object.__new__(cls) + object.__setattr__(instance, "schema_version", schema_version) + object.__setattr__(instance, "program_id", program_id) + object.__setattr__( + instance, + "_integration", + ExpertProgramIntegrationCfg( + robot_profile=integration.robot_profile, + scene_registry=integration.scene_registry, + runtime_preset=integration.runtime_preset, + ), + ) + object.__setattr__(instance, "_segments", values) + return instance + + @property + def integration(self) -> ExpertProgramIntegrationCfg: + """Return an independent integration-selection snapshot.""" + return ExpertProgramIntegrationCfg( + robot_profile=self._integration.robot_profile, + scene_registry=self._integration.scene_registry, + runtime_preset=self._integration.runtime_preset, + ) + + @property + def segment_count(self) -> int: + """Return the number of materialized logical segments.""" + return len(self._segments) + + def iter_segments(self) -> Iterator[CompiledProgramSegment]: + """Iterate the already materialized provider-free segments.""" + return iter(self._segments) + + def preflight_analyses(self) -> tuple[CompiledProgramAnalysis, ...]: + """Return full-program analyses split only at parallel barriers. + + Consecutive sequential segments form one static workflow, preserving + their object-state flow and cross-segment target look-ahead. Each + parallel branch is analyzed independently; no state or target inference + crosses the barrier in either direction. + """ + analyses: list[CompiledProgramAnalysis] = [] + stretch: list[CompiledProgramSegment] = [] + + def flush_stretch() -> None: + if not stretch: + return + indices = tuple(segment.segment_index for segment in stretch) + calls = tuple(call.call for segment in stretch for call in segment.calls) + analyses.append( + CompiledProgramAnalysis( + analysis_id=( + f"{self.program_id}:preflight:sequential:" + f"{indices[0]}-{indices[-1]}" + ), + kind="sequential_stretch", + calls=calls, + source_path=stretch[0].source_path, + segment_indices=indices, + execution_prefix_length=len(calls), + ) + ) + stretch.clear() + + for segment in self._segments: + block = segment.parallel_block + if block is None: + stretch.append(segment) + continue + flush_stretch() + for branch in block.branches: + calls = tuple(call.call for call in branch.calls) + analyses.append( + CompiledProgramAnalysis( + analysis_id=( + f"{self.program_id}:preflight:parallel:" + f"{segment.segment_index}:{branch.branch_index}" + ), + kind="parallel_branch", + calls=calls, + source_path=branch.source_path, + segment_indices=(segment.segment_index,), + execution_prefix_length=len(calls), + ) + ) + flush_stretch() + return tuple(analyses) + + def sequential_execution_analysis( + self, + segment_index: int, + ) -> CompiledProgramAnalysis: + """Return current-segment prefix plus downstream sequential look-ahead. + + Args: + segment_index: Index of the sequential segment about to execute. + + Returns: + Analysis beginning at the selected segment and ending immediately + before the next parallel barrier or the end of the program. + + Raises: + IndexError: If ``segment_index`` is outside this program. + ValueError: If the selected segment is itself parallel. + """ + if type(segment_index) is not int: + raise TypeError("segment_index must be an integer.") + if not 0 <= segment_index < len(self._segments): + raise IndexError(f"segment_index {segment_index!r} is outside the program.") + current = self._segments[segment_index] + if current.parallel_block is not None: + raise ValueError("Parallel segments do not have sequential look-ahead.") + window: list[CompiledProgramSegment] = [] + for segment in self._segments[segment_index:]: + if segment.parallel_block is not None: + break + window.append(segment) + calls = tuple(call.call for segment in window for call in segment.calls) + indices = tuple(segment.segment_index for segment in window) + return CompiledProgramAnalysis( + analysis_id=( + f"{self.program_id}:execution:sequential:" f"{indices[0]}-{indices[-1]}" + ), + kind="sequential_suffix", + calls=calls, + source_path=current.source_path, + segment_indices=indices, + execution_prefix_length=len(current.calls), + ) + + def __iter__(self) -> Iterator[CompiledProgramSegment]: + return self.iter_segments() + + +class ExpertProgramCompiler: + """Compile validated Expert Program ASTs through one typed resolver.""" + + def __init__(self, scene_resolver: ExpertProgramSceneResolver) -> None: + """Create one provider-free compiler. + + Args: + scene_resolver: Static typed scene identity resolver. + """ + if not isinstance(scene_resolver, ExpertProgramSceneResolver): + raise TypeError("scene_resolver must implement ExpertProgramSceneResolver.") + self._scene_resolver = scene_resolver + + @classmethod + def from_scene_registry(cls, registry: SceneRegistry) -> ExpertProgramCompiler: + """Create a compiler from a provider-free SceneRegistry identity snapshot.""" + return cls(SceneRegistryProgramResolver(registry)) + + def _resolve_scene( + self, + reference: str, + *, + expected_types: tuple[type[SceneEntityRef], ...], + path: ConfigPath, + ) -> SceneEntityRef: + """Resolve and validate one exact typed canonical scene reference.""" + try: + resolved = self._scene_resolver.resolve( + reference, + expected_types=expected_types, + path=path, + ) + except ExpertProgramConfigError: + raise + except (KeyError, TypeError, ValueError) as exc: + raise ExpertProgramCompileError( + "scene_resolution_failed", + path, + str(exc), + ) from exc + if type(resolved) not in expected_types: + raise ExpertProgramCompileError( + "scene_resolver_contract_violation", + path, + "Scene resolver returned an incompatible typed reference.", + ) + return _copy_scene_ref(resolved) + + @staticmethod + def _target_id( + reference: TargetRefCfg, + *, + targets: Mapping[str, tuple[SemanticPose, ...]], + path: ConfigPath, + ) -> str: + """Resolve one statically registered target ID.""" + if type(reference) is not TargetRefCfg or reference.kind != "target_ref": + raise ExpertProgramCompileError( + "invalid_target_reference", + path, + "Expected an exact target_ref configuration.", + ) + if reference.target not in targets: + raise ExpertProgramCompileError( + "unknown_target", + (*path, "target"), + f"Unknown target {reference.target!r}.", + ) + return reference.target + + def _compile_call( + self, + cfg: object, + *, + targets: Mapping[str, tuple[SemanticPose, ...]], + path: ConfigPath, + ) -> _CallTemplate: + """Lower one config call into a provider-free canonical template.""" + if type(cfg) is PickCfg: + if cfg.kind != "pick": + raise ExpertProgramCompileError( + "invalid_discriminator", (*path, "kind"), "Expected 'pick'." + ) + object_ref = self._resolve_scene( + cfg.object, + expected_types=(SceneObjectRef,), + path=(*path, "object"), + ) + grasp_ref = ( + None + if cfg.grasp is None + else self._resolve_scene( + cfg.grasp, + expected_types=(SceneAffordanceRef,), + path=(*path, "grasp"), + ) + ) + return _CallTemplate( + kind="pick", + source_path=path, + object=object_ref, + grasp=grasp_ref, + resources=tuple(sorted(cfg.resources.items())), + ) + if type(cfg) is PlaceCfg: + if cfg.kind != "place": + raise ExpertProgramCompileError( + "invalid_discriminator", (*path, "kind"), "Expected 'place'." + ) + object_ref = self._resolve_scene( + cfg.object, + expected_types=(SceneObjectRef,), + path=(*path, "object"), + ) + at_target_id = ( + None + if cfg.at is None + else self._target_id( + cfg.at, + targets=targets, + path=(*path, "at"), + ) + ) + on = ( + None + if cfg.on is None + else self._resolve_scene( + cfg.on, + expected_types=(SceneObjectRef, SceneAffordanceRef), + path=(*path, "on"), + ) + ) + inside = ( + None + if cfg.inside is None + else self._resolve_scene( + cfg.inside, + expected_types=(SceneObjectRef, SceneAffordanceRef), + path=(*path, "inside"), + ) + ) + return _CallTemplate( + kind="place", + source_path=path, + object=object_ref, + at_target_id=at_target_id, + on=on, + inside=inside, + resources=tuple(sorted(cfg.resources.items())), + ) + if type(cfg) is HandOverCfg: + if cfg.kind != "hand_over": + raise ExpertProgramCompileError( + "invalid_discriminator", + (*path, "kind"), + "Expected 'hand_over'.", + ) + object_ref = self._resolve_scene( + cfg.object, + expected_types=(SceneObjectRef,), + path=(*path, "object"), + ) + final_target_id = ( + None + if cfg.final_target is None + else self._target_id( + cfg.final_target, + targets=targets, + path=(*path, "final_target"), + ) + ) + return _CallTemplate( + kind="hand_over", + source_path=path, + object=object_ref, + receiver=cfg.receiver, + final_target_id=final_target_id, + resources=tuple(sorted(cfg.resources.items())), + ) + if type(cfg) is OperateArticulationCfg: + if cfg.kind != "operate_articulation": + raise ExpertProgramCompileError( + "invalid_discriminator", + (*path, "kind"), + "Expected 'operate_articulation'.", + ) + articulation = self._resolve_scene( + cfg.articulation, + expected_types=(SceneArticulationRef,), + path=(*path, "articulation"), + ) + handle = ( + None + if cfg.handle is None + else self._resolve_scene( + cfg.handle, + expected_types=(SceneAffordanceRef,), + path=(*path, "handle"), + ) + ) + snapshot = OperateArticulation( + articulation=articulation, + handle=handle, + target=cfg.target, + target_position=cfg.target_position, + target_displacement=cfg.target_displacement, + resources=cfg.resources, + ) + return _CallTemplate( + kind="operate_articulation", + source_path=path, + articulation=snapshot.articulation, + handle=snapshot.handle, + articulation_target=snapshot.target, + target_position=snapshot.target_position, + target_displacement=snapshot.target_displacement, + resources=tuple(sorted(snapshot.resources.items())), + ) + if type(cfg) is RegisteredSemanticCallCfg: + if cfg.kind != "registered": + raise ExpertProgramCompileError( + "invalid_discriminator", + (*path, "kind"), + "Expected 'registered'.", + ) + if cfg.schema_version != EXPERT_PROGRAM_SCHEMA_VERSION: + raise ExpertProgramCompileError( + "unsupported_registered_schema", + (*path, "schema_version"), + "Registered call schema_version must be exactly 1.", + ) + snapshot = RegisteredSemanticCall( + call_id=cfg.call_id, + arguments=cfg.arguments, + resources=cfg.resources, + ) + return _CallTemplate( + kind="registered", + source_path=path, + call_id=snapshot.call_id, + arguments=snapshot.arguments, + resources=tuple(sorted(snapshot.resources.items())), + ) + raise ExpertProgramCompileError( + "unsupported_call", + path, + f"Unsupported semantic call config {type(cfg).__name__}.", + ) + + def _compile_node( + self, + node: ProgramNodeCfg, + *, + targets: Mapping[str, tuple[SemanticPose, ...]], + path: ConfigPath, + inside_segment: bool, + inside_parallel: bool, + ) -> _NodeTemplate: + """Compile static AST structure without expanding repeats.""" + if type(node) is InvokeCfg: + if node.kind != "invoke": + raise ExpertProgramCompileError( + "invalid_discriminator", (*path, "kind"), "Expected 'invoke'." + ) + return _InvokeTemplate( + call=self._compile_call( + node.call, + targets=targets, + path=(*path, "call"), + ), + source_path=path, + ) + if type(node) is SequenceCfg: + if node.kind != "sequence": + raise ExpertProgramCompileError( + "invalid_discriminator", + (*path, "kind"), + "Expected 'sequence'.", + ) + if not node.items: + raise ExpertProgramCompileError( + "empty_sequence", + (*path, "items"), + "Sequence items must contain at least one program node.", + ) + return _SequenceTemplate( + items=tuple( + self._compile_node( + child, + targets=targets, + path=(*path, "items", index), + inside_segment=inside_segment, + inside_parallel=inside_parallel, + ) + for index, child in enumerate(node.items) + ), + source_path=path, + ) + if type(node) is RepeatCfg: + if node.kind != "repeat": + raise ExpertProgramCompileError( + "invalid_discriminator", (*path, "kind"), "Expected 'repeat'." + ) + if type(node.count) is not int or not 1 <= node.count <= MAX_REPEAT_COUNT: + raise ExpertProgramCompileError( + "invalid_repeat_count", + (*path, "count"), + f"Repeat count must be an integer in [1, {MAX_REPEAT_COUNT}].", + ) + return _RepeatTemplate( + count=node.count, + body=self._compile_node( + node.body, + targets=targets, + path=(*path, "body"), + inside_segment=inside_segment, + inside_parallel=inside_parallel, + ), + source_path=path, + ) + if type(node) is SegmentCfg: + if inside_parallel: + raise ExpertProgramCompileError( + "segment_inside_parallel", + path, + "Parallel branches may contain only Invoke, Sequence, and " + "Repeat nodes; wrap the Parallel node in one Segment instead.", + ) + if inside_segment: + raise ExpertProgramCompileError( + "nested_segment", + path, + "Nested Segment nodes are ambiguous and forbidden.", + ) + if node.kind != "segment": + raise ExpertProgramCompileError( + "invalid_discriminator", (*path, "kind"), "Expected 'segment'." + ) + post: list[_PostTemplate] = [] + for index, cfg in enumerate(node.post): + post_path = (*path, "post", index) + if type(cfg) is not WaitStablePostCfg or cfg.kind != "wait_stable": + raise ExpertProgramCompileError( + "unsupported_post_policy", + post_path, + "Supported schemas accept only exact wait_stable post policies.", + ) + entity = self._resolve_scene( + cfg.entity, + expected_types=_SCENE_REF_TYPES, + path=(*post_path, "entity"), + ) + post.append( + _PostTemplate( + cfg=WaitStablePostCfg( + entity=cfg.entity, + preset=cfg.preset, + kind=cfg.kind, + ), + entity=entity, + source_path=post_path, + ) + ) + validators: list[_ValidatorTemplate] = [] + for index, cfg in enumerate(node.validators): + validator_path = (*path, "validators", index) + if ( + type(cfg) is not ObjectNearTargetValidatorCfg + or cfg.kind != "object_near_target" + ): + raise ExpertProgramCompileError( + "unsupported_validator", + validator_path, + "Supported schemas accept only exact object_near_target " + "validators.", + ) + if cfg.target not in targets: + raise ExpertProgramCompileError( + "unknown_target", + (*validator_path, "target"), + f"Unknown target {cfg.target!r}.", + ) + object_ref = self._resolve_scene( + cfg.object, + expected_types=(SceneObjectRef,), + path=(*validator_path, "object"), + ) + validators.append( + _ValidatorTemplate( + cfg=ObjectNearTargetValidatorCfg( + object=cfg.object, + target=cfg.target, + position_tolerance=cfg.position_tolerance, + kind=cfg.kind, + ), + object=object_ref, + target_id=cfg.target, + source_path=validator_path, + ) + ) + steps = self._compile_node( + node.steps, + targets=targets, + path=(*path, "steps"), + inside_segment=True, + inside_parallel=False, + ) + if type(steps) is not _ParallelTemplate and _contains_parallel(steps): + raise ExpertProgramCompileError( + "mixed_parallel_segment", + (*path, "steps"), + "A Segment may contain either a call-only program or one direct " + "Parallel node, not a mixed sequential/parallel tree.", + ) + return _SegmentTemplate( + name=node.name, + steps=steps, + post=tuple(post), + validators=tuple(validators), + source_path=path, + ) + if type(node) is ParallelCfg: + if inside_parallel: + raise ExpertProgramCompileError( + "nested_parallel", + path, + "Nested Parallel nodes are forbidden in schema version 2.", + ) + if node.kind != "parallel": + raise ExpertProgramCompileError( + "invalid_discriminator", + (*path, "kind"), + "Expected 'parallel'.", + ) + if len(node.branches) < 2: + raise ExpertProgramCompileError( + "parallel_branch_count", + (*path, "branches"), + "Parallel requires at least two branches.", + ) + if type(node.barrier) is not BarrierCfg: + raise ExpertProgramCompileError( + "parallel_barrier_required", + (*path, "barrier"), + "Parallel.barrier must be an exact BarrierCfg.", + ) + branches = tuple( + self._compile_node( + branch, + targets=targets, + path=(*path, "branches", index), + inside_segment=inside_segment, + inside_parallel=True, + ) + for index, branch in enumerate(node.branches) + ) + if any(_contains_parallel(branch) for branch in branches): + raise ExpertProgramCompileError( + "nested_parallel", + (*path, "branches"), + "Nested Parallel nodes are forbidden in schema version 2.", + ) + barrier = node.barrier + if barrier.kind != "barrier": + raise ExpertProgramCompileError( + "invalid_discriminator", + (*path, "barrier", "kind"), + "Expected 'barrier'.", + ) + if barrier.failure_policy != "fail_fast": + raise ExpertProgramCompileError( + "unsupported_failure_policy", + (*path, "barrier", "failure_policy"), + "Barrier failure_policy must be exactly 'fail_fast'.", + ) + return _ParallelTemplate( + branches=branches, + barrier=_BarrierTemplate( + name=barrier.name, + timeout_steps=barrier.timeout_steps, + failure_policy=barrier.failure_policy, + source_path=(*path, "barrier"), + ), + source_path=path, + ) + if type(node) is BarrierCfg: + raise ExpertProgramCompileError( + "standalone_barrier", + path, + "Barrier nodes may only be owned by Parallel.", + ) + raise ExpertProgramCompileError( + "unsupported_program_node", + path, + f"Unsupported program node {type(node).__name__}.", + ) + + @staticmethod + def _compile_targets( + targets: Mapping[str, CyclicPoseTargetCfg], + ) -> Mapping[str, tuple[SemanticPose, ...]]: + """Compile static pose providers without selecting repeat values.""" + compiled: dict[str, tuple[SemanticPose, ...]] = {} + for target_id, target in targets.items(): + path = ("targets", target_id) + if type(target) is not CyclicPoseTargetCfg or target.kind != "cyclic_pose": + raise ExpertProgramCompileError( + "unsupported_target", + path, + "Supported schemas accept only exact cyclic_pose targets.", + ) + poses: list[SemanticPose] = [] + if not target.values: + raise ExpertProgramCompileError( + "empty_target_values", + (*path, "values"), + "Cyclic target values must contain at least one pose.", + ) + for index, pose in enumerate(target.values): + if type(pose) is not PoseCfg: + raise ExpertProgramCompileError( + "invalid_pose", + (*path, "values", index), + "Target values must be exact PoseCfg values.", + ) + poses.append(SemanticPose(pose.position, pose.quaternion_wxyz)) + compiled[target_id] = tuple(poses) + return MappingProxyType(compiled) + + def compile(self, config: ExpertProgramCfg) -> CompiledProgram: + """Compile one validated AST into a provider-free lazy program. + + Args: + config: Strict, supported-version Expert Program configuration. + + Returns: + Owned static templates whose iteration resolves repeat-local targets + and emits independent logical segments. + + Raises: + ExpertProgramCompileError: If typed scene resolution or AST lowering + fails. + """ + if type(config) is not ExpertProgramCfg: + raise TypeError("config must be exactly ExpertProgramCfg.") + if config.schema_version not in SUPPORTED_EXPERT_PROGRAM_SCHEMA_VERSIONS: + raise ExpertProgramCompileError( + "unsupported_schema_version", + ("schema_version",), + "Supported Expert Program schema versions are " + f"{SUPPORTED_EXPERT_PROGRAM_SCHEMA_VERSIONS}.", + ) + targets = self._compile_targets(config.targets) + root = self._compile_node( + config.program, + targets=targets, + path=("program",), + inside_segment=False, + inside_parallel=False, + ) + integration = ExpertProgramIntegrationCfg( + robot_profile=config.integration.robot_profile, + scene_registry=config.integration.scene_registry, + runtime_preset=config.integration.runtime_preset, + ) + return CompiledProgram._create( + schema_version=config.schema_version, + program_id=config.program_id, + integration=integration, + targets=targets, + root=root, + ) + + +__all__ = [ + "CompiledBarrier", + "CompiledParallelBlock", + "CompiledParallelBranch", + "CompiledPostPolicy", + "CompiledProgram", + "CompiledProgramAnalysis", + "CompiledProgramCall", + "CompiledProgramSegment", + "CompiledProgramValidator", + "CompiledRepeatFrame", + "CompiledTargetSelection", + "ExpertProgramCompileError", + "ExpertProgramCompiler", + "ExpertProgramSceneResolver", + "MaterializedCompiledProgram", + "SceneRegistryProgramResolver", +] diff --git a/embodichain/lab/gym/envs/expert_program/decoder.py b/embodichain/lab/gym/envs/expert_program/decoder.py new file mode 100644 index 000000000..88ba825df --- /dev/null +++ b/embodichain/lab/gym/envs/expert_program/decoder.py @@ -0,0 +1,1240 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Strict JSON/YAML-value decoder for Expert Program schema versions 1 and 2.""" + +from __future__ import annotations + +import math +import re +from collections.abc import Callable +from typing import Literal, Protocol, TypeAlias, runtime_checkable + +from .cfg import ( + BarrierCfg, + EXPERT_PROGRAM_SCHEMA_VERSION, + EXPERT_PROGRAM_SCHEMA_VERSION_V2, + MAX_PROGRAM_DEPTH, + MAX_REPEAT_COUNT, + CyclicPoseTargetCfg, + ExpertProgramCfg, + ExpertProgramIntegrationCfg, + HandOverCfg, + InvokeCfg, + ObjectNearTargetValidatorCfg, + OperateArticulationCfg, + ParallelCfg, + PickCfg, + PlaceCfg, + PoseCfg, + PostPolicyCfg, + ProgramNodeCfg, + RegisteredSemanticCallCfg, + RepeatCfg, + SegmentCfg, + SemanticCallCfg, + SequenceCfg, + SUPPORTED_EXPERT_PROGRAM_SCHEMA_VERSIONS, + TargetCfg, + TargetRefCfg, + ValidatorCfg, + WaitStablePostCfg, +) + +ConfigPathPart: TypeAlias = str | int +ConfigPath: TypeAlias = tuple[ConfigPathPart, ...] +SceneReferenceRole: TypeAlias = Literal[ + "entity", + "object", + "articulation", + "affordance", + "object_or_affordance", +] + +_MAX_INPUT_DEPTH = 128 +_MAX_INPUT_NODES = 100_000 +_ENV_TRAVERSAL_PATTERN = re.compile( + r"(?:\$?(?:env|environment)(?:\.[A-Za-z_][A-Za-z0-9_]*)+|" + r"\$\{(?:env|environment)(?:\.[A-Za-z_][A-Za-z0-9_]*)+\})" +) +_FORBIDDEN_KEYS = frozenset( + { + "__import__", + "attribute_path", + "callable", + "environment_path", + "env_path", + "eval", + "exec", + "expression", + "getattr", + "import", + "module", + "python", + } +) + + +def render_config_path(path: ConfigPath) -> str: + """Render one configuration path using JSONPath-like notation. + + Args: + path: Tuple of mapping keys and sequence indices. + + Returns: + Stable human-readable path beginning at ``$``. + """ + rendered = "$" + for part in path: + if type(part) is int: + rendered += f"[{part}]" + elif re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", part) is not None: + rendered += f".{part}" + else: + rendered += f"[{part!r}]" + return rendered + + +class ExpertProgramConfigError(ValueError): + """Base pathful diagnostic for Expert Program configuration failures.""" + + def __init__(self, code: str, path: ConfigPath, message: str) -> None: + """Create one stable pathful diagnostic. + + Args: + code: Machine-readable failure code. + path: Exact configuration location. + message: Human-readable explanation. + """ + self.code = code + self.path = tuple(path) + self.message = message + super().__init__(f"{render_config_path(self.path)}: {message} [{code}]") + + +class ExpertProgramDecodeError(ExpertProgramConfigError): + """Raised when untrusted data does not match a supported strict schema.""" + + +class ExpertProgramValidationError(ExpertProgramConfigError): + """Raised when an explicit static integration context rejects a reference.""" + + +@runtime_checkable +class ExpertProgramValidationContext(Protocol): + """Provider-free static validation boundary for external references. + + Implementations may resolve profile, scene, preset, catalog, affordance, and + resource IDs, but must not observe simulation state, construct planners, or + execute calls. + """ + + def validate_integration( + self, + integration: ExpertProgramIntegrationCfg, + *, + path: ConfigPath, + ) -> None: + """Validate integration references at ``path``.""" + + def validate_semantic_call( + self, + call: SemanticCallCfg, + *, + path: ConfigPath, + ) -> None: + """Validate catalog identity, schema revision, and resource overrides.""" + + def validate_scene_reference( + self, + reference: str, + *, + role: SceneReferenceRole, + path: ConfigPath, + ) -> None: + """Validate one canonical scene reference with its semantic role.""" + + def validate_post_policy( + self, + policy: PostPolicyCfg, + *, + path: ConfigPath, + ) -> None: + """Validate a post-policy kind and its named preset.""" + + def validate_validator( + self, + validator: ValidatorCfg, + *, + path: ConfigPath, + ) -> None: + """Validate one registered segment-validator contract.""" + + +def _error(code: str, path: ConfigPath, message: str) -> ExpertProgramDecodeError: + """Build one decoder diagnostic.""" + return ExpertProgramDecodeError(code, path, message) + + +def _clone_untrusted_value( + value: object, + *, + path: ConfigPath, + active: set[int], + budget: list[int], + depth: int, +) -> object: + """Own and validate one bounded JSON-compatible value tree.""" + if depth > _MAX_INPUT_DEPTH: + raise _error( + "input_too_deep", + path, + f"Input exceeds nesting depth limit {_MAX_INPUT_DEPTH}.", + ) + budget[0] -= 1 + if budget[0] < 0: + raise _error( + "input_too_large", + path, + f"Input exceeds node limit {_MAX_INPUT_NODES}.", + ) + if value is None or type(value) in (bool, int): + return value + if type(value) is float: + if not math.isfinite(value): + raise _error("non_finite_number", path, "Floats must be finite.") + return value + if type(value) is str: + stripped = value.strip() + lowered = stripped.lower() + if lowered.startswith(("__import__(", "eval(", "exec(", "import ", "from ")): + raise _error( + "executable_expression", + path, + "Imports, eval, exec, and executable expressions are forbidden.", + ) + if _ENV_TRAVERSAL_PATTERN.fullmatch(stripped) is not None: + raise _error( + "environment_traversal", + path, + "Dotted environment attribute traversal is forbidden.", + ) + return value + if type(value) is list: + identity = id(value) + if identity in active: + raise _error("cyclic_input", path, "Input contains a cyclic list.") + active.add(identity) + try: + return [ + _clone_untrusted_value( + item, + path=(*path, index), + active=active, + budget=budget, + depth=depth + 1, + ) + for index, item in enumerate(value) + ] + finally: + active.remove(identity) + if type(value) is dict: + identity = id(value) + if identity in active: + raise _error("cyclic_input", path, "Input contains a cyclic mapping.") + active.add(identity) + try: + result: dict[str, object] = {} + for key, item in value.items(): + if type(key) is not str: + raise _error( + "invalid_mapping_key", + path, + "Mapping keys must be exact strings.", + ) + if key.lower() in _FORBIDDEN_KEYS: + raise _error( + "forbidden_construct", + (*path, key), + f"Field {key!r} requests executable or traversal behavior.", + ) + result[key] = _clone_untrusted_value( + item, + path=(*path, key), + active=active, + budget=budget, + depth=depth + 1, + ) + return result + finally: + active.remove(identity) + raise _error( + "non_declarative_value", + path, + f"{type(value).__name__} is not JSON-compatible declarative data; " + "callables, classes, modules, tensors, and live objects are forbidden.", + ) + + +def _expect_mapping(value: object, *, path: ConfigPath) -> dict[str, object]: + """Require one exact mapping.""" + if type(value) is not dict: + raise _error("expected_mapping", path, "Expected an object mapping.") + return value + + +def _expect_list(value: object, *, path: ConfigPath) -> list[object]: + """Require one exact JSON list.""" + if type(value) is not list: + raise _error("expected_list", path, "Expected a list.") + return value + + +def _validate_fields( + value: dict[str, object], + *, + allowed: frozenset[str], + required: frozenset[str], + path: ConfigPath, +) -> None: + """Reject unknown fields and report the first missing required field.""" + unknown = sorted(set(value).difference(allowed)) + if unknown: + field_name = unknown[0] + raise _error( + "unknown_field", + (*path, field_name), + f"Unknown field {field_name!r}; allowed fields are {sorted(allowed)}.", + ) + missing = sorted(required.difference(value)) + if missing: + field_name = missing[0] + raise _error( + "missing_field", + (*path, field_name), + f"Missing required field {field_name!r}.", + ) + + +def _expect_identifier(value: object, *, path: ConfigPath) -> str: + """Require one exact non-empty identifier.""" + if type(value) is not str or not value or value != value.strip(): + raise _error( + "invalid_identifier", + path, + "Expected a non-empty string without outer whitespace.", + ) + return value + + +def _expect_discriminator( + value: dict[str, object], + *, + path: ConfigPath, + supported: tuple[str, ...], +) -> str: + """Read one required exact string discriminator.""" + if "kind" not in value: + raise _error( + "missing_discriminator", + (*path, "kind"), + "Missing required discriminator 'kind'.", + ) + kind = value["kind"] + if type(kind) is not str or kind not in supported: + raise _error( + "unknown_discriminator", + (*path, "kind"), + f"Unsupported discriminator {kind!r}; expected one of {supported}.", + ) + return kind + + +def _decode_resources(value: object, *, path: ConfigPath) -> dict[str, str]: + """Decode one strict slot-to-resource mapping.""" + mapping = _expect_mapping(value, path=path) + return { + _expect_identifier(slot_id, path=(*path, slot_id)): _expect_identifier( + resource_id, + path=(*path, slot_id), + ) + for slot_id, resource_id in mapping.items() + } + + +def _construct( + constructor: Callable[..., object], + *, + path: ConfigPath, + **kwargs: object, +) -> object: + """Construct one config value and wrap invariant failures pathfully.""" + try: + return constructor(**kwargs) + except ExpertProgramConfigError: + raise + except (TypeError, ValueError) as exc: + raise _error("invalid_value", path, str(exc)) from exc + + +def _decode_pose(value: object, *, path: ConfigPath) -> PoseCfg: + """Decode one finite pose value.""" + mapping = _expect_mapping(value, path=path) + _validate_fields( + mapping, + allowed=frozenset({"position", "quaternion_wxyz"}), + required=frozenset({"position", "quaternion_wxyz"}), + path=path, + ) + position_values = _expect_list(mapping["position"], path=(*path, "position")) + quaternion_values = _expect_list( + mapping["quaternion_wxyz"], + path=(*path, "quaternion_wxyz"), + ) + if len(position_values) != 3: + raise _error( + "invalid_pose_shape", + (*path, "position"), + "position must contain exactly three numbers.", + ) + if len(quaternion_values) != 4: + raise _error( + "invalid_pose_shape", + (*path, "quaternion_wxyz"), + "quaternion_wxyz must contain exactly four numbers.", + ) + for name, values in ( + ("position", position_values), + ("quaternion_wxyz", quaternion_values), + ): + for index, number in enumerate(values): + if type(number) not in (int, float): + raise _error( + "invalid_number", + (*path, name, index), + "Pose components must be finite numbers, not bool values.", + ) + return _construct( + PoseCfg, + path=path, + position=tuple(position_values), + quaternion_wxyz=tuple(quaternion_values), + ) # type: ignore[return-value] + + +def _decode_target(value: object, *, path: ConfigPath) -> TargetCfg: + """Decode one target provider shared by the supported schema versions.""" + mapping = _expect_mapping(value, path=path) + kind = _expect_discriminator( + mapping, + path=path, + supported=("cyclic_pose",), + ) + assert kind == "cyclic_pose" + _validate_fields( + mapping, + allowed=frozenset({"kind", "values"}), + required=frozenset({"kind", "values"}), + path=path, + ) + values = tuple( + _decode_pose(item, path=(*path, "values", index)) + for index, item in enumerate( + _expect_list(mapping["values"], path=(*path, "values")) + ) + ) + return _construct( + CyclicPoseTargetCfg, + path=path, + kind=kind, + values=values, + ) # type: ignore[return-value] + + +def _decode_target_ref( + value: object, + *, + path: ConfigPath, + target_ids: frozenset[str], +) -> TargetRefCfg: + """Decode and statically resolve one target reference.""" + mapping = _expect_mapping(value, path=path) + kind = _expect_discriminator(mapping, path=path, supported=("target_ref",)) + _validate_fields( + mapping, + allowed=frozenset({"kind", "target"}), + required=frozenset({"kind", "target"}), + path=path, + ) + target = _expect_identifier(mapping["target"], path=(*path, "target")) + if target not in target_ids: + raise _error( + "unknown_target", + (*path, "target"), + f"Unknown target {target!r}; available targets are {sorted(target_ids)}.", + ) + return _construct( + TargetRefCfg, + path=path, + kind=kind, + target=target, + ) # type: ignore[return-value] + + +def _decode_call( + value: object, + *, + path: ConfigPath, + target_ids: frozenset[str], +) -> SemanticCallCfg: + """Decode one discriminated semantic call.""" + mapping = _expect_mapping(value, path=path) + kind = _expect_discriminator( + mapping, + path=path, + supported=( + "pick", + "place", + "hand_over", + "operate_articulation", + "registered", + ), + ) + resources = _decode_resources( + mapping.get("resources", {}), + path=(*path, "resources"), + ) + if kind == "pick": + _validate_fields( + mapping, + allowed=frozenset({"kind", "object", "grasp", "resources"}), + required=frozenset({"kind", "object"}), + path=path, + ) + grasp = mapping.get("grasp") + if grasp is not None: + grasp = _expect_identifier(grasp, path=(*path, "grasp")) + return _construct( + PickCfg, + path=path, + kind=kind, + object=_expect_identifier(mapping["object"], path=(*path, "object")), + grasp=grasp, + resources=resources, + ) # type: ignore[return-value] + if kind == "place": + _validate_fields( + mapping, + allowed=frozenset({"kind", "object", "at", "on", "inside", "resources"}), + required=frozenset({"kind", "object"}), + path=path, + ) + at = ( + None + if mapping.get("at") is None + else _decode_target_ref( + mapping["at"], + path=(*path, "at"), + target_ids=target_ids, + ) + ) + on = mapping.get("on") + inside = mapping.get("inside") + if on is not None: + on = _expect_identifier(on, path=(*path, "on")) + if inside is not None: + inside = _expect_identifier(inside, path=(*path, "inside")) + return _construct( + PlaceCfg, + path=path, + kind=kind, + object=_expect_identifier(mapping["object"], path=(*path, "object")), + at=at, + on=on, + inside=inside, + resources=resources, + ) # type: ignore[return-value] + if kind == "hand_over": + _validate_fields( + mapping, + allowed=frozenset( + {"kind", "object", "receiver", "final_target", "resources"} + ), + required=frozenset({"kind", "object"}), + path=path, + ) + receiver = mapping.get("receiver") + if receiver is not None: + receiver = _expect_identifier(receiver, path=(*path, "receiver")) + final_target = ( + None + if mapping.get("final_target") is None + else _decode_target_ref( + mapping["final_target"], + path=(*path, "final_target"), + target_ids=target_ids, + ) + ) + return _construct( + HandOverCfg, + path=path, + kind=kind, + object=_expect_identifier(mapping["object"], path=(*path, "object")), + receiver=receiver, + final_target=final_target, + resources=resources, + ) # type: ignore[return-value] + if kind == "operate_articulation": + _validate_fields( + mapping, + allowed=frozenset( + { + "kind", + "articulation", + "handle", + "target", + "target_position", + "target_displacement", + "resources", + } + ), + required=frozenset({"kind", "articulation"}), + path=path, + ) + handle = mapping.get("handle") + target = mapping.get("target") + if handle is not None: + handle = _expect_identifier(handle, path=(*path, "handle")) + if target is not None: + target = _expect_identifier(target, path=(*path, "target")) + target_position = mapping.get("target_position") + target_displacement = mapping.get("target_displacement") + for field_name, value in ( + ("target_position", target_position), + ("target_displacement", target_displacement), + ): + if value is not None and type(value) not in (int, float): + raise _error( + "invalid_number", + (*path, field_name), + f"{field_name} must be a finite number, not bool.", + ) + named = target is not None + explicit_position = target_position is not None + explicit_displacement = target_displacement is not None + if named and (explicit_position or explicit_displacement): + raise _error( + "conflicting_articulation_target", + path, + "target is mutually exclusive with target_position and " + "target_displacement.", + ) + if not named and not (explicit_position and explicit_displacement): + raise _error( + "incomplete_articulation_target", + path, + "Specify target or both target_position and target_displacement.", + ) + return _construct( + OperateArticulationCfg, + path=path, + kind=kind, + articulation=_expect_identifier( + mapping["articulation"], + path=(*path, "articulation"), + ), + handle=handle, + target=target, + target_position=target_position, + target_displacement=target_displacement, + resources=resources, + ) # type: ignore[return-value] + + _validate_fields( + mapping, + allowed=frozenset( + {"kind", "call_id", "schema_version", "arguments", "resources"} + ), + required=frozenset({"kind", "call_id", "schema_version"}), + path=path, + ) + arguments = _expect_mapping( + mapping.get("arguments", {}), + path=(*path, "arguments"), + ) + schema_version = mapping["schema_version"] + if type(schema_version) is not int or schema_version != 1: + raise _error( + "invalid_schema_version", + (*path, "schema_version"), + "Registered call schema_version must be exactly 1.", + ) + return _construct( + RegisteredSemanticCallCfg, + path=path, + kind=kind, + call_id=_expect_identifier(mapping["call_id"], path=(*path, "call_id")), + schema_version=schema_version, + arguments=arguments, + resources=resources, + ) # type: ignore[return-value] + + +def _decode_post_policy(value: object, *, path: ConfigPath) -> PostPolicyCfg: + """Decode one segment post-policy shared by the supported schemas.""" + mapping = _expect_mapping(value, path=path) + kind = _expect_discriminator(mapping, path=path, supported=("wait_stable",)) + _validate_fields( + mapping, + allowed=frozenset({"kind", "entity", "preset"}), + required=frozenset({"kind", "entity"}), + path=path, + ) + return _construct( + WaitStablePostCfg, + path=path, + kind=kind, + entity=_expect_identifier(mapping["entity"], path=(*path, "entity")), + preset=_expect_identifier( + mapping.get("preset", "rigid_object"), + path=(*path, "preset"), + ), + ) # type: ignore[return-value] + + +def _decode_validator( + value: object, + *, + path: ConfigPath, + target_ids: frozenset[str], +) -> ValidatorCfg: + """Decode one segment validator shared by the supported schemas.""" + mapping = _expect_mapping(value, path=path) + kind = _expect_discriminator( + mapping, + path=path, + supported=("object_near_target",), + ) + _validate_fields( + mapping, + allowed=frozenset({"kind", "object", "target", "position_tolerance"}), + required=frozenset({"kind", "object", "target"}), + path=path, + ) + target = _expect_identifier(mapping["target"], path=(*path, "target")) + if target not in target_ids: + raise _error( + "unknown_target", + (*path, "target"), + f"Unknown target {target!r}; available targets are {sorted(target_ids)}.", + ) + tolerance = mapping.get("position_tolerance", 0.03) + if type(tolerance) not in (int, float): + raise _error( + "invalid_number", + (*path, "position_tolerance"), + "position_tolerance must be a finite number, not bool.", + ) + return _construct( + ObjectNearTargetValidatorCfg, + path=path, + kind=kind, + object=_expect_identifier(mapping["object"], path=(*path, "object")), + target=target, + position_tolerance=tolerance, + ) # type: ignore[return-value] + + +def _decode_program_node( + value: object, + *, + path: ConfigPath, + target_ids: frozenset[str], + depth: int, + schema_version: int, +) -> ProgramNodeCfg: + """Recursively decode one bounded versioned program node.""" + if depth > MAX_PROGRAM_DEPTH: + raise _error( + "program_too_deep", + path, + "Program AST exceeds the configured nesting depth.", + ) + mapping = _expect_mapping(value, path=path) + supported_kinds = ["sequence", "repeat", "segment", "invoke"] + if schema_version >= EXPERT_PROGRAM_SCHEMA_VERSION_V2: + supported_kinds.extend(("parallel", "barrier")) + kind = _expect_discriminator( + mapping, + path=path, + supported=tuple(supported_kinds), + ) + if kind == "sequence": + _validate_fields( + mapping, + allowed=frozenset({"kind", "items"}), + required=frozenset({"kind", "items"}), + path=path, + ) + items = tuple( + _decode_program_node( + item, + path=(*path, "items", index), + target_ids=target_ids, + depth=depth + 1, + schema_version=schema_version, + ) + for index, item in enumerate( + _expect_list(mapping["items"], path=(*path, "items")) + ) + ) + return _construct( + SequenceCfg, + path=path, + kind=kind, + items=items, + ) # type: ignore[return-value] + if kind == "repeat": + _validate_fields( + mapping, + allowed=frozenset({"kind", "count", "body"}), + required=frozenset({"kind", "count", "body"}), + path=path, + ) + count = mapping["count"] + if type(count) is not int or not 1 <= count <= MAX_REPEAT_COUNT: + raise _error( + "invalid_repeat_count", + (*path, "count"), + f"Repeat count must be an integer in [1, {MAX_REPEAT_COUNT}].", + ) + return _construct( + RepeatCfg, + path=path, + kind=kind, + count=count, + body=_decode_program_node( + mapping["body"], + path=(*path, "body"), + target_ids=target_ids, + depth=depth + 1, + schema_version=schema_version, + ), + ) # type: ignore[return-value] + if kind == "segment": + _validate_fields( + mapping, + allowed=frozenset({"kind", "name", "steps", "post", "validators"}), + required=frozenset({"kind", "name", "steps"}), + path=path, + ) + post = tuple( + _decode_post_policy(item, path=(*path, "post", index)) + for index, item in enumerate( + _expect_list(mapping.get("post", []), path=(*path, "post")) + ) + ) + validators = tuple( + _decode_validator( + item, + path=(*path, "validators", index), + target_ids=target_ids, + ) + for index, item in enumerate( + _expect_list( + mapping.get("validators", []), + path=(*path, "validators"), + ) + ) + ) + return _construct( + SegmentCfg, + path=path, + kind=kind, + name=_expect_identifier(mapping["name"], path=(*path, "name")), + steps=_decode_program_node( + mapping["steps"], + path=(*path, "steps"), + target_ids=target_ids, + depth=depth + 1, + schema_version=schema_version, + ), + post=post, + validators=validators, + ) # type: ignore[return-value] + + if kind == "parallel": + _validate_fields( + mapping, + allowed=frozenset({"kind", "branches", "barrier"}), + required=frozenset({"kind", "branches", "barrier"}), + path=path, + ) + branches_values = _expect_list( + mapping["branches"], + path=(*path, "branches"), + ) + if len(branches_values) < 2: + raise _error( + "parallel_branch_count", + (*path, "branches"), + "Parallel requires at least two branches.", + ) + barrier = _decode_program_node( + mapping["barrier"], + path=(*path, "barrier"), + target_ids=target_ids, + depth=depth + 1, + schema_version=schema_version, + ) + if type(barrier) is not BarrierCfg: + raise _error( + "parallel_barrier_required", + (*path, "barrier"), + "Parallel.barrier must be an explicit barrier node.", + ) + return _construct( + ParallelCfg, + path=path, + kind=kind, + branches=tuple( + _decode_program_node( + branch, + path=(*path, "branches", index), + target_ids=target_ids, + depth=depth + 1, + schema_version=schema_version, + ) + for index, branch in enumerate(branches_values) + ), + barrier=barrier, + ) # type: ignore[return-value] + if kind == "barrier": + _validate_fields( + mapping, + allowed=frozenset({"kind", "name", "timeout_steps", "failure_policy"}), + required=frozenset({"kind", "name"}), + path=path, + ) + timeout_steps = mapping.get("timeout_steps", 1_000) + if type(timeout_steps) is not int or timeout_steps <= 0: + raise _error( + "invalid_barrier_timeout", + (*path, "timeout_steps"), + "Barrier timeout_steps must be a positive integer.", + ) + failure_policy = mapping.get("failure_policy", "fail_fast") + if failure_policy != "fail_fast": + raise _error( + "unsupported_failure_policy", + (*path, "failure_policy"), + "Barrier failure_policy must be exactly 'fail_fast'.", + ) + return _construct( + BarrierCfg, + path=path, + kind=kind, + name=_expect_identifier(mapping["name"], path=(*path, "name")), + timeout_steps=timeout_steps, + failure_policy=failure_policy, + ) # type: ignore[return-value] + + _validate_fields( + mapping, + allowed=frozenset({"kind", "call"}), + required=frozenset({"kind", "call"}), + path=path, + ) + return _construct( + InvokeCfg, + path=path, + kind=kind, + call=_decode_call( + mapping["call"], + path=(*path, "call"), + target_ids=target_ids, + ), + ) # type: ignore[return-value] + + +def _walk_program( + node: ProgramNodeCfg, + *, + path: ConfigPath, +) -> list[tuple[ProgramNodeCfg, ConfigPath]]: + """Return deterministic node/path pairs for static context validation.""" + values = [(node, path)] + if type(node) is SequenceCfg: + for index, child in enumerate(node.items): + values.extend(_walk_program(child, path=(*path, "items", index))) + elif type(node) is RepeatCfg: + values.extend(_walk_program(node.body, path=(*path, "body"))) + elif type(node) is SegmentCfg: + values.extend(_walk_program(node.steps, path=(*path, "steps"))) + elif type(node) is ParallelCfg: + for index, branch in enumerate(node.branches): + values.extend(_walk_program(branch, path=(*path, "branches", index))) + values.extend(_walk_program(node.barrier, path=(*path, "barrier"))) + return values + + +def _call_context( + callback: Callable[..., None], + *args: object, + path: ConfigPath, + **kwargs: object, +) -> None: + """Call one static validation hook and preserve pathful failures.""" + try: + callback(*args, path=path, **kwargs) + except ExpertProgramConfigError: + raise + except (KeyError, TypeError, ValueError) as exc: + raise ExpertProgramValidationError( + "reference_validation_failed", + path, + str(exc), + ) from exc + + +def validate_expert_program( + config: ExpertProgramCfg, + context: ExpertProgramValidationContext, +) -> None: + """Resolve external references without observing or executing an environment. + + Args: + config: Fully decoded and internally validated Expert Program. + context: Provider-free static integration/catalog/scene validator. + + Raises: + TypeError: If either argument has the wrong contract. + ExpertProgramValidationError: If an external reference is unavailable. + """ + if type(config) is not ExpertProgramCfg: + raise TypeError("config must be exactly ExpertProgramCfg.") + if not isinstance(context, ExpertProgramValidationContext): + raise TypeError( + "context must implement ExpertProgramValidationContext exactly." + ) + _call_context( + context.validate_integration, + config.integration, + path=("integration",), + ) + for node, path in _walk_program(config.program, path=("program",)): + if type(node) is InvokeCfg: + call = node.call + call_path = (*path, "call") + _call_context( + context.validate_semantic_call, + call, + path=call_path, + ) + if type(call) in (PickCfg, PlaceCfg, HandOverCfg): + _call_context( + context.validate_scene_reference, + call.object, + role="object", + path=(*call_path, "object"), + ) + if type(call) is PickCfg and call.grasp is not None: + _call_context( + context.validate_scene_reference, + call.grasp, + role="affordance", + path=(*call_path, "grasp"), + ) + if type(call) is PlaceCfg: + if call.on is not None: + _call_context( + context.validate_scene_reference, + call.on, + role="object_or_affordance", + path=(*call_path, "on"), + ) + if call.inside is not None: + _call_context( + context.validate_scene_reference, + call.inside, + role="object_or_affordance", + path=(*call_path, "inside"), + ) + if type(call) is OperateArticulationCfg: + _call_context( + context.validate_scene_reference, + call.articulation, + role="articulation", + path=(*call_path, "articulation"), + ) + if call.handle is not None: + _call_context( + context.validate_scene_reference, + call.handle, + role="affordance", + path=(*call_path, "handle"), + ) + elif type(node) is SegmentCfg: + for index, post in enumerate(node.post): + post_path = (*path, "post", index) + _call_context( + context.validate_post_policy, + post, + path=post_path, + ) + _call_context( + context.validate_scene_reference, + post.entity, + role="entity", + path=(*post_path, "entity"), + ) + for index, validator in enumerate(node.validators): + validator_path = (*path, "validators", index) + _call_context( + context.validate_validator, + validator, + path=validator_path, + ) + _call_context( + context.validate_scene_reference, + validator.object, + role="object", + path=(*validator_path, "object"), + ) + + +def decode_expert_program( + data: object, + *, + validation_context: ExpertProgramValidationContext | None = None, +) -> ExpertProgramCfg: + """Decode untrusted JSON/YAML-shaped values into strict versioned config. + + Schema versions 1 and 2 are supported. Version 2 adds deterministic + parallel blocks with explicit barriers while preserving the Version 1 + sequential nodes and semantic calls. + + Args: + data: Exact JSON-compatible mapping produced by a trusted parser. + validation_context: Optional provider-free static reference validator. + + Returns: + Fully owned and internally validated Expert Program configuration. + + Raises: + ExpertProgramDecodeError: If data is unsafe or violates the schema. + ExpertProgramValidationError: If an explicit context rejects a reference. + """ + owned = _clone_untrusted_value( + data, + path=(), + active=set(), + budget=[_MAX_INPUT_NODES], + depth=0, + ) + mapping = _expect_mapping(owned, path=()) + _validate_fields( + mapping, + allowed=frozenset( + {"schema_version", "program_id", "integration", "targets", "program"} + ), + required=frozenset( + {"schema_version", "program_id", "integration", "targets", "program"} + ), + path=(), + ) + schema_version = mapping["schema_version"] + if ( + type(schema_version) is not int + or schema_version not in SUPPORTED_EXPERT_PROGRAM_SCHEMA_VERSIONS + ): + raise _error( + "unsupported_schema_version", + ("schema_version",), + "Supported schema versions are " + f"{list(SUPPORTED_EXPERT_PROGRAM_SCHEMA_VERSIONS)}.", + ) + + integration_mapping = _expect_mapping( + mapping["integration"], + path=("integration",), + ) + _validate_fields( + integration_mapping, + allowed=frozenset({"robot_profile", "scene_registry", "runtime_preset"}), + required=frozenset({"robot_profile", "scene_registry", "runtime_preset"}), + path=("integration",), + ) + integration = _construct( + ExpertProgramIntegrationCfg, + path=("integration",), + robot_profile=_expect_identifier( + integration_mapping["robot_profile"], + path=("integration", "robot_profile"), + ), + scene_registry=_expect_identifier( + integration_mapping["scene_registry"], + path=("integration", "scene_registry"), + ), + runtime_preset=_expect_identifier( + integration_mapping["runtime_preset"], + path=("integration", "runtime_preset"), + ), + ) + + target_mapping = _expect_mapping(mapping["targets"], path=("targets",)) + targets: dict[str, TargetCfg] = {} + for target_id, target_value in target_mapping.items(): + normalized_id = _expect_identifier(target_id, path=("targets", target_id)) + targets[normalized_id] = _decode_target( + target_value, + path=("targets", normalized_id), + ) + target_ids = frozenset(targets) + program = _decode_program_node( + mapping["program"], + path=("program",), + target_ids=target_ids, + depth=0, + schema_version=schema_version, + ) + config = _construct( + ExpertProgramCfg, + path=(), + schema_version=schema_version, + program_id=_expect_identifier(mapping["program_id"], path=("program_id",)), + integration=integration, + targets=targets, + program=program, + ) + assert type(config) is ExpertProgramCfg + if validation_context is not None: + validate_expert_program(config, validation_context) + return config + + +__all__ = [ + "ConfigPath", + "ConfigPathPart", + "ExpertProgramConfigError", + "ExpertProgramDecodeError", + "ExpertProgramValidationContext", + "ExpertProgramValidationError", + "SceneReferenceRole", + "decode_expert_program", + "render_config_path", + "validate_expert_program", +] diff --git a/embodichain/lab/gym/envs/expert_program/environment.py b/embodichain/lab/gym/envs/expert_program/environment.py new file mode 100644 index 000000000..95b2294f3 --- /dev/null +++ b/embodichain/lab/gym/envs/expert_program/environment.py @@ -0,0 +1,831 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Explicit production assembly for environment-backed Expert Programs. + +The adapter in this module is deliberately strict. It does not scan a +simulation, infer robot resources, or manufacture task semantics from naming +conventions. An environment supplies one typed factory that owns all live +provider choices; the adapter validates those declarations and wires the +shared semantic compiler, runtime, and Gym bridge. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping +from dataclasses import dataclass +import math +from typing import Protocol, runtime_checkable + +from embodichain.lab.sim.atomic_actions.engine import AtomicActionEngine +from embodichain.lab.sim.atomic_actions.runner import ( + ExecutionRunnerCfg, + ObservationProvider, +) +from embodichain.lab.sim.skills.calls import ( + SemanticCallCatalog, + SemanticCallSpec, + builtin_semantic_call_catalog, +) +from embodichain.lab.sim.skills.compiler import ( + HandOverPoseProvider, + RegisteredSemanticLowerer, + RelationTargetGrounder, + SemanticSkillCompiler, +) +from embodichain.lab.sim.skills.effects import EffectMonitorRegistry +from embodichain.lab.sim.skills.evidence import ( + EffectEvidenceCollector, + EffectEvidenceProvider, + EffectEvidenceProviderRegistry, +) +from embodichain.lab.sim.skills.integration import ( + SceneManifest, + SemanticIntegrationManifest, +) +from embodichain.lab.sim.skills.parallel_runtime import ( + ParallelCommandSafetyValidator, + analyze_parallel_branches, +) +from embodichain.lab.sim.skills.profiles import ( + ResourceEndpoint, + ResourceEndpointAdapter, + RobotSkillProfile, +) +from embodichain.lab.sim.skills.runtime import SkillRuntime +from embodichain.lab.sim.skills.scene import SceneRegistry + +from .bridge import ( + AcceptedRuntimeCommandObserver, + AtomicDemoBridge, + BufferedGymCommandSink, + CurrentQposProvider, + DemoBridgeError, + EnvironmentStepClock, + RuntimeCommandFrameEncoder, + RuntimeTransportActionEncoder, + SegmentPostPolicyPort, + SegmentValidatorPort, +) +from .cfg import ExpertProgramCfg, ExpertProgramIntegrationCfg +from .compiler import ( + CompiledProgram, + ExpertProgramCompiler, + MaterializedCompiledProgram, +) + + +def _validate_identifier(value: str, *, field_name: str) -> str: + """Validate one stable integration identifier.""" + if type(value) is not str or not value or value != value.strip(): + raise ValueError( + f"{field_name} must be a non-empty string without outer whitespace." + ) + return value + + +@runtime_checkable +class PlanningObservationPort( + ObservationProvider, + CurrentQposProvider, + Protocol, +): + """Combined observation and full-qpos port required by the Gym runtime.""" + + +@runtime_checkable +class ExpertProgramEnvironmentFactory(Protocol): + """Environment-owned factories for one explicit semantic integration. + + Implementations normally live in reusable robot/task integration modules, + not in individual task motion planners. Every method is passed the exact + objects selected earlier in the assembly so a factory cannot silently bind + a different scene, robot profile, or engine. + """ + + @property + def scene_registry_id(self) -> str: + """Return the configuration ID selecting this scene declaration. + + Returns: + Stable scene-registry identifier. + """ + + @property + def robot_profile_id(self) -> str: + """Return the configuration ID selecting this robot profile. + + Returns: + Stable robot-profile identifier. + """ + + def create_scene_registry(self) -> SceneRegistry: + """Create the authoritative explicitly registered live scene. + + Returns: + Fresh registry containing only explicitly selected entities. + """ + + def create_robot_skill_profile(self) -> RobotSkillProfile: + """Create the authoritative declarative robot skill profile. + + Returns: + Profile whose ID matches :attr:`robot_profile_id`. + """ + + def create_atomic_action_engine( + self, + profile: RobotSkillProfile, + ) -> AtomicActionEngine: + """Create an engine for exactly ``profile`` and its motion backend. + + Args: + profile: Profile selected and validated by the adapter. + + Returns: + Atomic engine connected to the environment's robot and planner. + """ + + def create_planning_observation_provider( + self, + *, + scene_registry: SceneRegistry, + engine: AtomicActionEngine, + clock: EnvironmentStepClock, + ) -> PlanningObservationPort: + """Create fresh planning observations and aligned full-qpos reads. + + Args: + scene_registry: Exact registry selected for this runtime. + engine: Exact atomic engine selected for this runtime. + clock: Shared environment-step execution clock. + + Returns: + Combined planning-observation and qpos provider. + """ + + def create_effect_evidence_providers( + self, + *, + scene_registry: SceneRegistry, + engine: AtomicActionEngine, + observation_provider: PlanningObservationPort, + ) -> Iterable[EffectEvidenceProvider]: + """Create exact-version providers used by semantic effect monitors. + + Args: + scene_registry: Exact registry selected for this runtime. + engine: Exact atomic engine selected for this runtime. + observation_provider: Shared planning observation provider. + + Returns: + Explicit provider set; an empty iterable is permitted. + """ + + +@runtime_checkable +class AcceptedRuntimeCommandObserverFactory(Protocol): + """Optional factory capability for runtime-local accepted-command state. + + The observer is created from the exact observation provider used by the + runtime so command-derived evidence cannot leak across bridge instances or + bind to a different simulation batch. + """ + + def create_accepted_runtime_command_observer( + self, + *, + scene_registry: SceneRegistry, + engine: AtomicActionEngine, + observation_provider: PlanningObservationPort, + ) -> AcceptedRuntimeCommandObserver: + """Return the observer shared by the command sink and evidence ports.""" + + +@dataclass(frozen=True, slots=True) +class ExpertProgramRuntimeAssembly: + """Auditable result of one fresh environment runtime assembly. + + Attributes: + integration: Owned integration-selection snapshot. + scene_registry: Authoritative live scene registry. + robot_profile: Declarative robot resource profile. + manifest: Static scene/profile/call integration manifest. + engine: Bound atomic action engine. + compiler: Bound semantic skill compiler. + observation_provider: Shared planning and full-qpos provider. + evidence_collector: Exact-version semantic evidence collector. + clock: Shared environment-step clock. + command_encoder: Runtime-frame to Gym-action encoder. + command_sink: Buffered Gym command sink. + accepted_command_observer: Optional transactional command-state owner. + runtime: Nonblocking semantic skill runtime. + """ + + integration: ExpertProgramIntegrationCfg + scene_registry: SceneRegistry + robot_profile: RobotSkillProfile + manifest: SemanticIntegrationManifest + engine: AtomicActionEngine + compiler: SemanticSkillCompiler + observation_provider: PlanningObservationPort + evidence_collector: EffectEvidenceCollector + clock: EnvironmentStepClock + command_encoder: RuntimeCommandFrameEncoder + command_sink: BufferedGymCommandSink + accepted_command_observer: AcceptedRuntimeCommandObserver | None + runtime: SkillRuntime + + +@dataclass(frozen=True, slots=True) +class _ExpertProgramSemanticAssembly: + """Observation-free semantic components prepared for program preflight.""" + + integration: ExpertProgramIntegrationCfg + scene_registry: SceneRegistry + robot_profile: RobotSkillProfile + manifest: SemanticIntegrationManifest + engine: AtomicActionEngine + compiler: SemanticSkillCompiler + + +class ExpertProgramEnvironmentAdapter: + """Compile and run Expert Programs through explicit environment factories. + + Args: + factory: Environment-owned live-provider and engine factory. + step_dt: Authoritative Gym control cadence in seconds. + call_catalog: Optional immutable semantic call catalog. The built-in + catalog is used when omitted. + endpoint_adapters: Optional custom robot endpoint adapters. + registered_lowerers: Explicit lowerers for registered semantic calls. + relation_grounders: Explicit relation-target grounding providers. + handover_pose_providers: Explicit embodiment hand-over providers. + effect_monitor_registry: Optional exact-version monitor registry. + runtime_transports: Additional runtime-command-to-Gym encoders. + runner_cfg: Optional execution-runner policy. + post_policy_port: Optional environment post-policy executor. + validator_port: Optional environment segment validator. + parallel_safety_validator: Optional authoritative parallel safety gate. + + A call to :meth:`compile` snapshots only scene identities. A call to + :meth:`assemble_runtime` creates a fresh live runtime, which makes reset and + episode ownership explicit and avoids retaining providers in compiled data. + """ + + def __init__( + self, + factory: ExpertProgramEnvironmentFactory, + *, + step_dt: float, + call_catalog: SemanticCallCatalog | None = None, + endpoint_adapters: ( + Mapping[type[ResourceEndpoint], ResourceEndpointAdapter] | None + ) = None, + registered_lowerers: Iterable[RegisteredSemanticLowerer] = (), + relation_grounders: Iterable[RelationTargetGrounder] = (), + handover_pose_providers: Iterable[HandOverPoseProvider] = (), + effect_monitor_registry: EffectMonitorRegistry | None = None, + runtime_transports: Iterable[RuntimeTransportActionEncoder] = (), + runner_cfg: ExecutionRunnerCfg | None = None, + post_policy_port: SegmentPostPolicyPort | None = None, + validator_port: SegmentValidatorPort | None = None, + parallel_safety_validator: ParallelCommandSafetyValidator | None = None, + ) -> None: + if not isinstance(factory, ExpertProgramEnvironmentFactory): + raise TypeError("factory must implement ExpertProgramEnvironmentFactory.") + if not isinstance(step_dt, (int, float)) or isinstance(step_dt, bool): + raise TypeError("step_dt must be a real number.") + if not math.isfinite(float(step_dt)) or float(step_dt) <= 0.0: + raise ValueError("step_dt must be finite and positive.") + scene_registry_id = _validate_identifier( + factory.scene_registry_id, + field_name="factory.scene_registry_id", + ) + robot_profile_id = _validate_identifier( + factory.robot_profile_id, + field_name="factory.robot_profile_id", + ) + selected_catalog = call_catalog or builtin_semantic_call_catalog() + if type(selected_catalog) is not SemanticCallCatalog: + raise TypeError("call_catalog must be exactly SemanticCallCatalog or None.") + if endpoint_adapters is not None and not isinstance(endpoint_adapters, Mapping): + raise TypeError("endpoint_adapters must be a mapping or None.") + if runner_cfg is not None and not isinstance(runner_cfg, ExecutionRunnerCfg): + raise TypeError("runner_cfg must be an ExecutionRunnerCfg or None.") + if post_policy_port is not None and not isinstance( + post_policy_port, + SegmentPostPolicyPort, + ): + raise TypeError( + "post_policy_port must implement SegmentPostPolicyPort or be None." + ) + if validator_port is not None and not isinstance( + validator_port, + SegmentValidatorPort, + ): + raise TypeError( + "validator_port must implement SegmentValidatorPort or be None." + ) + if parallel_safety_validator is not None and not isinstance( + parallel_safety_validator, + ParallelCommandSafetyValidator, + ): + raise TypeError( + "parallel_safety_validator must implement " + "ParallelCommandSafetyValidator or be None." + ) + + self._factory = factory + self._scene_registry_id = scene_registry_id + self._robot_profile_id = robot_profile_id + self._step_dt = float(step_dt) + self._call_catalog = selected_catalog + self._endpoint_adapters = ( + None if endpoint_adapters is None else dict(endpoint_adapters) + ) + self._registered_lowerers = tuple(registered_lowerers) + self._relation_grounders = tuple(relation_grounders) + self._handover_pose_providers = tuple(handover_pose_providers) + self._effect_monitor_registry = effect_monitor_registry + self._runtime_transports = tuple(runtime_transports) + self._runner_cfg = runner_cfg + self._post_policy_port = post_policy_port + self._validator_port = validator_port + self._parallel_safety_validator = parallel_safety_validator + + @property + def scene_registry_id(self) -> str: + """Return the exact scene integration ID accepted by this adapter. + + Returns: + Stable scene-registry identifier. + """ + return self._scene_registry_id + + @property + def robot_profile_id(self) -> str: + """Return the exact robot profile ID accepted by this adapter. + + Returns: + Stable robot-profile identifier. + """ + return self._robot_profile_id + + @property + def step_dt(self) -> float: + """Return the authoritative environment-step cadence. + + Returns: + Positive control step duration in seconds. + """ + return self._step_dt + + def compile(self, program: ExpertProgramCfg) -> CompiledProgram: + """Compile one program after exact integration-selection validation. + + Args: + program: Strict declarative program configuration. + + Returns: + Provider-free lazily expanded compiled program. + """ + if type(program) is not ExpertProgramCfg: + raise TypeError("program must be exactly ExpertProgramCfg.") + self._validate_selection(program.integration) + registry = self._create_scene_registry() + return ExpertProgramCompiler.from_scene_registry(registry).compile(program) + + def assemble_runtime( + self, + integration: ExpertProgramIntegrationCfg, + ) -> ExpertProgramRuntimeAssembly: + """Create a fresh fully connected semantic runtime. + + Args: + integration: Exact scene, profile, and runtime-preset selection. + + Returns: + Owned assembly containing every validated runtime boundary. + """ + semantic = self._assemble_semantic_components(integration) + return self._assemble_execution_runtime(semantic) + + def _assemble_semantic_components( + self, + integration: ExpertProgramIntegrationCfg, + ) -> _ExpertProgramSemanticAssembly: + """Bind compiler dependencies without observation or evidence ports.""" + self._validate_selection(integration) + registry = self._create_scene_registry() + current_profile_id = _validate_identifier( + self._factory.robot_profile_id, + field_name="factory.robot_profile_id", + ) + if current_profile_id != self._robot_profile_id: + raise ValueError( + "Factory robot profile declaration drifted: expected " + f"{self._robot_profile_id!r}, got {current_profile_id!r}." + ) + profile = self._factory.create_robot_skill_profile() + if type(profile) is not RobotSkillProfile: + raise TypeError( + "create_robot_skill_profile() must return exactly RobotSkillProfile." + ) + if profile.profile_id != self._robot_profile_id: + raise ValueError( + "Factory robot profile declaration drifted: expected " + f"{self._robot_profile_id!r}, got {profile.profile_id!r}." + ) + + engine = self._factory.create_atomic_action_engine(profile) + if not isinstance(engine, AtomicActionEngine): + raise TypeError( + "create_atomic_action_engine() must return an AtomicActionEngine." + ) + + manifest = self._create_manifest( + registry, + profile, + runtime_preset=integration.runtime_preset, + ) + bound = manifest.bind( + registry, + engine, + endpoint_adapters=self._endpoint_adapters, + ) + compiler = SemanticSkillCompiler( + bound, + registered_lowerers=self._registered_lowerers, + relation_grounders=self._relation_grounders, + handover_pose_providers=self._handover_pose_providers, + effect_monitor_registry=self._effect_monitor_registry, + ) + + selection = ExpertProgramIntegrationCfg( + robot_profile=integration.robot_profile, + scene_registry=integration.scene_registry, + runtime_preset=integration.runtime_preset, + ) + return _ExpertProgramSemanticAssembly( + integration=selection, + scene_registry=registry, + robot_profile=profile, + manifest=manifest, + engine=engine, + compiler=compiler, + ) + + def _assemble_execution_runtime( + self, + semantic: _ExpertProgramSemanticAssembly, + ) -> ExpertProgramRuntimeAssembly: + """Attach live observation, evidence, command, and runtime boundaries.""" + if type(semantic) is not _ExpertProgramSemanticAssembly: + raise TypeError("semantic must be exactly _ExpertProgramSemanticAssembly.") + + clock = EnvironmentStepClock(self._step_dt) + observation_provider = self._factory.create_planning_observation_provider( + scene_registry=semantic.scene_registry, + engine=semantic.engine, + clock=clock, + ) + if not isinstance(observation_provider, PlanningObservationPort): + raise TypeError( + "create_planning_observation_provider() must return a port " + "implementing both ObservationProvider and CurrentQposProvider." + ) + providers = self._factory.create_effect_evidence_providers( + scene_registry=semantic.scene_registry, + engine=semantic.engine, + observation_provider=observation_provider, + ) + if isinstance(providers, (str, bytes)): + raise TypeError( + "create_effect_evidence_providers() must return an iterable of " + "EffectEvidenceProvider values." + ) + try: + provider_values = tuple(providers) + except TypeError as exc: + raise TypeError( + "create_effect_evidence_providers() must return an iterable of " + "EffectEvidenceProvider values." + ) from exc + evidence_collector = EffectEvidenceCollector( + EffectEvidenceProviderRegistry(provider_values) + ) + command_encoder = RuntimeCommandFrameEncoder( + observation_provider, + transports=self._runtime_transports, + ) + accepted_command_observer: AcceptedRuntimeCommandObserver | None = None + if isinstance(self._factory, AcceptedRuntimeCommandObserverFactory): + accepted_command_observer = ( + self._factory.create_accepted_runtime_command_observer( + scene_registry=semantic.scene_registry, + engine=semantic.engine, + observation_provider=observation_provider, + ) + ) + if not isinstance( + accepted_command_observer, + AcceptedRuntimeCommandObserver, + ): + raise TypeError( + "create_accepted_runtime_command_observer() must return an " + "AcceptedRuntimeCommandObserver." + ) + command_sink = BufferedGymCommandSink( + command_encoder, + clock, + accepted_command_observer=accepted_command_observer, + ) + runtime = SkillRuntime.from_components( + semantic.compiler, + observation_provider, + command_sink, + evidence_collector, + clock=clock, + runner_cfg=self._runner_cfg, + ) + return ExpertProgramRuntimeAssembly( + integration=semantic.integration, + scene_registry=semantic.scene_registry, + robot_profile=semantic.robot_profile, + manifest=semantic.manifest, + engine=semantic.engine, + compiler=semantic.compiler, + observation_provider=observation_provider, + evidence_collector=evidence_collector, + clock=clock, + command_encoder=command_encoder, + command_sink=command_sink, + accepted_command_observer=accepted_command_observer, + runtime=runtime, + ) + + def create_bridge(self, program: CompiledProgram) -> AtomicDemoBridge: + """Create a fresh Gym bridge for one provider-free compiled program. + + Args: + program: Program compiled for this adapter's exact integration IDs. + + Returns: + Lazy bridge sharing one newly assembled runtime, clock, and sink. + """ + if type(program) is not CompiledProgram: + raise TypeError("program must be exactly CompiledProgram.") + materialized = program.materialize() + self._validate_selection(materialized.integration) + self._preflight_program_surfaces(materialized) + semantic = self._assemble_semantic_components(materialized.integration) + self._preflight_program(materialized, semantic.compiler) + assembly = self._assemble_execution_runtime(semantic) + return AtomicDemoBridge( + materialized, + assembly.runtime, + assembly.command_sink, + assembly.clock, + post_policy_port=self._post_policy_port, + validator_port=self._validator_port, + parallel_safety_validator=self._parallel_safety_validator, + ) + + def _preflight_program_surfaces( + self, + program: MaterializedCompiledProgram, + ) -> None: + """Validate every segment hook without live observation or action.""" + if type(program) is not MaterializedCompiledProgram: + raise TypeError("program must be exactly MaterializedCompiledProgram.") + for segment in program.iter_segments(): + if segment.post_policies and self._post_policy_port is None: + raise DemoBridgeError( + f"Segment {segment.segment_id!r} declares post-policies, but no " + "SegmentPostPolicyPort was installed." + ) + for policy in segment.post_policies: + assert self._post_policy_port is not None + self._post_policy_port.validate_policy(policy, segment=segment) + if segment.validators and self._validator_port is None: + raise DemoBridgeError( + f"Segment {segment.segment_id!r} declares validators, but no " + "SegmentValidatorPort was installed." + ) + for validator in segment.validators: + assert self._validator_port is not None + self._validator_port.validate_validator( + validator, + segment=segment, + ) + + def _preflight_program( + self, + program: MaterializedCompiledProgram, + compiler: SemanticSkillCompiler, + ) -> None: + """Analyze every program workflow before any physical action can run. + + Sequential stretches retain cross-segment state flow and target + look-ahead. A parallel barrier cuts that flow; each branch is checked + independently through the same canonical semantic compiler used by the + runtime. This boundary materializes no observations and starts no + execution session. + """ + if type(program) is not MaterializedCompiledProgram: + raise TypeError("program must be exactly MaterializedCompiledProgram.") + if not isinstance(compiler, SemanticSkillCompiler): + raise TypeError("compiler must be a SemanticSkillCompiler.") + analyses = program.preflight_analyses() + if any(analysis.kind == "parallel_branch" for analysis in analyses) and ( + self._parallel_safety_validator is None + ): + raise ValueError( + "Expert Programs containing parallel blocks require an explicit " + "ParallelCommandSafetyValidator before bridge creation." + ) + index = 0 + while index < len(analyses): + analysis = analyses[index] + if analysis.kind != "parallel_branch": + compiler.analyze( + analysis.calls, + workflow_id=analysis.analysis_id, + path=analysis.source_path, + ) + index += 1 + continue + segment_index = analysis.segment_indices[0] + branches: dict[str, tuple[SemanticCallSpec, ...]] = {} + branch_paths: dict[str, tuple[str | int, ...]] = {} + while index < len(analyses): + branch = analyses[index] + if branch.kind != "parallel_branch" or branch.segment_indices != ( + segment_index, + ): + break + branch_id = f"branch_{len(branches)}" + branches[branch_id] = branch.calls + branch_paths[branch_id] = branch.source_path + index += 1 + analyze_parallel_branches( + compiler, + branches, + workflow_id=( + f"{program.program_id}:preflight:parallel:{segment_index}" + ), + branch_paths=branch_paths, + ) + + def _validate_selection( + self, + integration: ExpertProgramIntegrationCfg, + ) -> None: + """Reject an integration selection owned by another adapter.""" + if type(integration) is not ExpertProgramIntegrationCfg: + raise TypeError("integration must be exactly ExpertProgramIntegrationCfg.") + current_scene_id = _validate_identifier( + self._factory.scene_registry_id, + field_name="factory.scene_registry_id", + ) + current_profile_id = _validate_identifier( + self._factory.robot_profile_id, + field_name="factory.robot_profile_id", + ) + if current_scene_id != self._scene_registry_id: + raise ValueError( + "Factory scene registry declaration drifted: expected " + f"{self._scene_registry_id!r}, got {current_scene_id!r}." + ) + if current_profile_id != self._robot_profile_id: + raise ValueError( + "Factory robot profile declaration drifted: expected " + f"{self._robot_profile_id!r}, got {current_profile_id!r}." + ) + if integration.scene_registry != self._scene_registry_id: + raise ValueError( + f"Expert Program selects scene_registry " + f"{integration.scene_registry!r}, but this environment exposes " + f"only {self._scene_registry_id!r}." + ) + if integration.robot_profile != self._robot_profile_id: + raise ValueError( + f"Expert Program selects robot_profile " + f"{integration.robot_profile!r}, but this environment exposes " + f"only {self._robot_profile_id!r}." + ) + + def _create_scene_registry(self) -> SceneRegistry: + """Create and validate one exact live scene registry.""" + current_id = _validate_identifier( + self._factory.scene_registry_id, + field_name="factory.scene_registry_id", + ) + if current_id != self._scene_registry_id: + raise ValueError( + "Factory scene registry declaration drifted: expected " + f"{self._scene_registry_id!r}, got {current_id!r}." + ) + registry = self._factory.create_scene_registry() + if type(registry) is not SceneRegistry: + raise TypeError( + "create_scene_registry() must return exactly SceneRegistry." + ) + return registry + + def _create_manifest( + self, + registry: SceneRegistry, + profile: RobotSkillProfile, + *, + runtime_preset: str, + ) -> SemanticIntegrationManifest: + """Create one static manifest from exact selected declarations.""" + return SemanticIntegrationManifest( + scene=SceneManifest.from_registry(registry), + robot_profile=profile, + call_catalog=self._call_catalog, + runtime_preset=runtime_preset, + ) + + +class ExpertProgramEnvironmentMixin: + """Delegate environment hooks to one reusable explicit adapter. + + Environment classes place this mixin before their normal environment base + and implement only :attr:`expert_program_adapter`. Motion generation and + runtime stepping remain in shared components. + """ + + @property + def expert_program_adapter(self) -> ExpertProgramEnvironmentAdapter: + """Return the environment-owned reusable adapter. + + Returns: + Exact shared Expert Program environment adapter. + """ + raise NotImplementedError( + "Expert Program environments must expose expert_program_adapter." + ) + + def compile_expert_program( + self, + program: ExpertProgramCfg, + ) -> CompiledProgram: + """Delegate provider-free compilation to the explicit adapter. + + Args: + program: Strict declarative program configuration. + + Returns: + Provider-free compiled program. + """ + return self._checked_expert_program_adapter().compile(program) + + def create_expert_program_bridge( + self, + program: CompiledProgram, + ) -> AtomicDemoBridge: + """Delegate live runtime and Gym bridge assembly to the adapter. + + Args: + program: Provider-free compiled program. + + Returns: + Fresh lazy Gym bridge. + """ + return self._checked_expert_program_adapter().create_bridge(program) + + def _checked_expert_program_adapter(self) -> ExpertProgramEnvironmentAdapter: + """Return the exact adapter or fail before any provider is touched.""" + adapter = self.expert_program_adapter + if type(adapter) is not ExpertProgramEnvironmentAdapter: + raise TypeError( + "expert_program_adapter must be exactly " + "ExpertProgramEnvironmentAdapter." + ) + return adapter + + +__all__ = [ + "AcceptedRuntimeCommandObserverFactory", + "ExpertProgramEnvironmentAdapter", + "ExpertProgramEnvironmentFactory", + "ExpertProgramEnvironmentMixin", + "ExpertProgramRuntimeAssembly", + "PlanningObservationPort", +] diff --git a/embodichain/lab/gym/envs/expert_program/loader.py b/embodichain/lab/gym/envs/expert_program/loader.py new file mode 100644 index 000000000..f47e60940 --- /dev/null +++ b/embodichain/lab/gym/envs/expert_program/loader.py @@ -0,0 +1,337 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Safe file and strict JSON loading for declarative Expert Programs.""" + +from __future__ import annotations + +import json +import math +import os +from pathlib import Path + +import yaml + +from .cfg import ExpertProgramCfg +from .decoder import ( + ExpertProgramDecodeError, + ExpertProgramValidationContext, + decode_expert_program, +) + +__all__ = [ + "MAX_EXPERT_PROGRAM_BYTES", + "load_expert_program", + "loads_expert_program_json", + "parse_expert_program_json", +] + +MAX_EXPERT_PROGRAM_BYTES = 4 * 1024 * 1024 +"""Maximum serialized Expert Program size accepted by the file loader.""" + + +class _StrictJsonValueError(ValueError): + """Carry one stable strict-JSON failure into the public decode boundary.""" + + def __init__(self, code: str, message: str) -> None: + self.code = code + self.message = message + super().__init__(message) + + +def _reject_duplicate_json_keys( + pairs: list[tuple[str, object]], +) -> dict[str, object]: + """Build a JSON mapping while rejecting ambiguous duplicate keys.""" + mapping: dict[str, object] = {} + for key, value in pairs: + if key in mapping: + raise _StrictJsonValueError( + "duplicate_json_key", + f"Duplicate JSON key {key!r}.", + ) + mapping[key] = value + return mapping + + +def _reject_non_finite_json_constant(token: str) -> object: + """Reject the non-standard NaN and Infinity JSON constants.""" + raise _StrictJsonValueError( + "non_finite_number", + f"Non-finite JSON number {token!r} is forbidden.", + ) + + +def _parse_finite_json_float(token: str) -> float: + """Parse one JSON float while rejecting overflow to infinity.""" + value = float(token) + if not math.isfinite(value): + raise _StrictJsonValueError( + "non_finite_number", + f"JSON number {token!r} is not finite.", + ) + return value + + +def _validate_decoded_json_unicode(value: object) -> None: + """Reject decoded JSON strings that cannot be represented as UTF-8.""" + if type(value) is str: + try: + value.encode("utf-8") + except UnicodeEncodeError as error: + raise _StrictJsonValueError( + "invalid_utf8", + "Expert Program JSON contains an unpaired Unicode surrogate.", + ) from error + return + if type(value) is list: + for item in value: + _validate_decoded_json_unicode(item) + return + if type(value) is dict: + for key, item in value.items(): + _validate_decoded_json_unicode(key) + _validate_decoded_json_unicode(item) + + +def _loads_strict_json_value( + text: str, + *, + max_bytes: int = MAX_EXPERT_PROGRAM_BYTES, +) -> object: + """Parse one bounded JSON document into exact JSON-compatible values.""" + + if type(text) is not str: + raise TypeError("text must be exactly str.") + if type(max_bytes) is not int: + raise TypeError("max_bytes must be exactly int.") + if max_bytes <= 0: + raise ValueError("max_bytes must be positive.") + try: + payload = text.encode("utf-8") + except UnicodeEncodeError as error: + raise ExpertProgramDecodeError( + "invalid_utf8", + (), + "Expert Program JSON must be valid UTF-8 text.", + ) from error + if len(payload) > max_bytes: + raise ExpertProgramDecodeError( + "input_too_large", + (), + f"Expert Program JSON exceeds the {max_bytes}-byte input limit.", + ) + try: + value = json.loads( + text, + object_pairs_hook=_reject_duplicate_json_keys, + parse_constant=_reject_non_finite_json_constant, + parse_float=_parse_finite_json_float, + ) + _validate_decoded_json_unicode(value) + return value + except _StrictJsonValueError as error: + raise ExpertProgramDecodeError(error.code, (), error.message) from error + except json.JSONDecodeError as error: + raise ExpertProgramDecodeError( + "invalid_json", + (), + "Invalid Expert Program JSON at " + f"line {error.lineno}, column {error.colno}.", + ) from error + except RecursionError as error: + raise ExpertProgramDecodeError( + "input_too_deep", + (), + "Expert Program JSON exceeds the parser nesting limit.", + ) from error + except ValueError as error: + raise ExpertProgramDecodeError( + "invalid_json", + (), + "Expert Program JSON contains an invalid numeric value.", + ) from error + + +def parse_expert_program_json( + text: str, + *, + max_bytes: int = MAX_EXPERT_PROGRAM_BYTES, +) -> dict[str, object]: + """Parse one bounded Expert Program JSON object without decoding its schema. + + This parse-only boundary lets a host-controlled frontend inspect or inject + fields before calling :func:`decode_expert_program`. It rejects duplicate + keys, non-finite numbers, trailing content, invalid Unicode, excessive + nesting, oversized UTF-8 input, and non-object top-level values. It does + not validate the Expert Program schema. + + Args: + text: Untrusted JSON document text. + max_bytes: Maximum accepted UTF-8 encoded input size. + + Returns: + Exact JSON object mapping ready for explicit schema decoding. + + Raises: + TypeError: If ``text`` or ``max_bytes`` has the wrong exact type. + ValueError: If ``max_bytes`` is not positive. + ExpertProgramDecodeError: If strict JSON parsing fails. + """ + value = _loads_strict_json_value(text, max_bytes=max_bytes) + if type(value) is not dict: + raise ExpertProgramDecodeError( + "expected_mapping", + (), + "Expected an object mapping.", + ) + return value + + +def loads_expert_program_json( + text: str, + *, + validation_context: ExpertProgramValidationContext | None = None, + max_bytes: int = MAX_EXPERT_PROGRAM_BYTES, +) -> ExpertProgramCfg: + """Strictly parse and decode one untrusted Expert Program JSON document. + + The input must be one plain JSON document. Markdown fences, trailing text, + multiple documents, duplicate keys, non-finite numbers, and oversized input + are rejected before the existing Expert Program decoder is called. + + Args: + text: Untrusted JSON response text. + validation_context: Optional provider-free static reference validator. + max_bytes: Maximum UTF-8 encoded response size. + + Returns: + Fully owned and internally validated Expert Program configuration. + + Raises: + TypeError: If ``text`` or ``max_bytes`` has the wrong exact type. + ValueError: If ``max_bytes`` is not positive. + ExpertProgramDecodeError: If parsing or strict decoding fails. + """ + data = parse_expert_program_json(text, max_bytes=max_bytes) + return decode_expert_program(data, validation_context=validation_context) + + +class _UniqueKeySafeLoader(yaml.SafeLoader): + """YAML safe loader that also rejects ambiguous duplicate keys.""" + + +def _construct_unique_yaml_mapping( + loader: _UniqueKeySafeLoader, + node: yaml.MappingNode, + deep: bool = False, +) -> dict[object, object]: + """Construct one YAML mapping with unique, hashable keys.""" + loader.flatten_mapping(node) + mapping: dict[object, object] = {} + for key_node, value_node in node.value: + key = loader.construct_object(key_node, deep=deep) + try: + duplicate = key in mapping + except TypeError as error: + raise yaml.constructor.ConstructorError( + "while constructing a mapping", + node.start_mark, + "found an unhashable key", + key_node.start_mark, + ) from error + if duplicate: + raise yaml.constructor.ConstructorError( + "while constructing a mapping", + node.start_mark, + f"found duplicate key {key!r}", + key_node.start_mark, + ) + mapping[key] = loader.construct_object(value_node, deep=deep) + return mapping + + +_UniqueKeySafeLoader.add_constructor( + yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, + _construct_unique_yaml_mapping, +) + + +def load_expert_program( + path: str | os.PathLike[str], + *, + base_dir: str | os.PathLike[str] | None = None, + validation_context: ExpertProgramValidationContext | None = None, +) -> ExpertProgramCfg: + """Safely load and strictly decode one JSON or YAML Expert Program file. + + Relative paths are resolved from ``base_dir`` when provided. Otherwise, + they retain normal :class:`pathlib.Path` semantics and therefore resolve + from the process working directory when opened. + + Args: + path: JSON, YAML, or YML file to load. + base_dir: Optional directory used to resolve a relative ``path``. + validation_context: Optional provider-free static reference validator + applied after decoding either serialized format. + + Returns: + An owned, validated Expert Program configuration. + + Raises: + FileNotFoundError: If the resolved path is not a regular file. + ValueError: If the file is too large, has an unsupported extension, or + contains ambiguous or invalid serialized data. + ExpertProgramValidationError: If ``validation_context`` rejects an + external reference. + UnicodeDecodeError: If the file is not valid UTF-8. + """ + program_path = Path(path).expanduser() + if base_dir is not None and not program_path.is_absolute(): + program_path = Path(base_dir).expanduser() / program_path + if not program_path.is_file(): + raise FileNotFoundError(f"Expert Program path is not a file: {program_path}.") + suffix = program_path.suffix.lower() + if suffix not in {".json", ".yaml", ".yml"}: + raise ValueError( + "Expert Program must use a .json, .yaml, or .yml extension; " + f"got {program_path.name!r}." + ) + + payload = program_path.read_bytes() + if len(payload) > MAX_EXPERT_PROGRAM_BYTES: + raise ExpertProgramDecodeError( + "input_too_large", + (), + "Expert Program exceeds the " + f"{MAX_EXPERT_PROGRAM_BYTES}-byte input limit.", + ) + text = payload.decode("utf-8") + if suffix == ".json": + return loads_expert_program_json( + text, + validation_context=validation_context, + ) + try: + data = yaml.load(text, Loader=_UniqueKeySafeLoader) + except yaml.YAMLError as error: + raise ValueError( + f"Invalid Expert Program YAML in {program_path}: {error}" + ) from error + return decode_expert_program( + data, + validation_context=validation_context, + ) diff --git a/embodichain/lab/gym/envs/expert_program/simulation.py b/embodichain/lab/gym/envs/expert_program/simulation.py new file mode 100644 index 000000000..5317dc091 --- /dev/null +++ b/embodichain/lab/gym/envs/expert_program/simulation.py @@ -0,0 +1,1240 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Explicit simulation bindings for declarative Expert Programs. + +The values in this module bridge task-owned, executable-free declarations to +the existing :class:`SceneRegistry` and :class:`RobotSkillProfile` contracts. +They deliberately do not scan the simulation or infer semantic capabilities +from names. Every simulation entity, articulation member, control part, and +semantic command is selected explicitly and validated while the binding is +built. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from copy import deepcopy +from dataclasses import dataclass, field, replace +import math +from types import MappingProxyType +from typing import Any, Protocol, runtime_checkable, TYPE_CHECKING + +import torch + +from embodichain.lab.sim.atomic_actions import ( + AntipodalAffordance, + ArticulationOperationAffordance, + ArticulationOperationTarget, + ControlPartCommandProfile, + EntityState, +) +from embodichain.lab.sim.skills.profiles import ( + ControlPartEndpoint, + ResourceEndpoint, + ResourceBinding, + RobotResource, + RobotSkillProfile, + SkillPolicyPreset, +) +from embodichain.lab.sim.skills.scene import ( + ARTICULATION_OPERATION_AFFORDANCE_CAPABILITY, + GRASP_AFFORDANCE_CAPABILITY, + SceneAffordanceRef, + SceneArticulationRef, + SceneCollisionRole, + SceneCollisionWorldMode, + SceneDynamics, + SceneEntityRegistration, + SceneGeometryProvider, + SceneLinkRef, + SceneObjectRef, + SceneRegistry, +) +from embodichain.toolkits.graspkit.pg_grasp import GraspGeneratorCfg +from embodichain.toolkits.graspkit.pg_grasp.gripper_collision_checker import ( + GripperCollisionCfg, +) + +if TYPE_CHECKING: + from embodichain.lab.sim.objects import Robot + from embodichain.lab.sim.sim_manager import SimulationManager + + +_IDENTITY_POSE = ( + 1.0, + 0.0, + 0.0, + 0.0, + 0.0, + 1.0, + 0.0, + 0.0, + 0.0, + 0.0, + 1.0, + 0.0, + 0.0, + 0.0, + 0.0, + 1.0, +) + + +def _identifier(value: str, *, field_name: str) -> str: + """Return one exact non-empty identifier.""" + if type(value) is not str or not value or value != value.strip(): + raise ValueError( + f"{field_name} must be a non-empty string without outer whitespace." + ) + return value + + +def _optional_identifier(value: str | None, *, field_name: str) -> str | None: + """Validate one optional identifier.""" + if value is not None: + _identifier(value, field_name=field_name) + return value + + +def _identifier_tuple( + values: tuple[str, ...], + *, + field_name: str, +) -> tuple[str, ...]: + """Own a duplicate-free tuple of exact identifiers.""" + if isinstance(values, (str, bytes)): + raise TypeError(f"{field_name} must be an iterable of identifiers.") + normalized = tuple(values) + for value in normalized: + _identifier(value, field_name=field_name) + if len(set(normalized)) != len(normalized): + raise ValueError(f"{field_name} must contain unique identifiers.") + return normalized + + +def _finite(value: float, *, field_name: str) -> float: + """Return one finite non-boolean float.""" + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError(f"{field_name} must be a finite number.") + normalized = float(value) + if not math.isfinite(normalized): + raise ValueError(f"{field_name} must be finite.") + return normalized + + +def _pose_tuple( + values: tuple[float, ...], + *, + field_name: str, +) -> tuple[float, ...]: + """Own and validate one flattened SE(3) matrix.""" + if isinstance(values, (str, bytes)): + raise TypeError(f"{field_name} must contain 16 finite numbers.") + normalized = tuple( + _finite(value, field_name=f"{field_name}[{index}]") + for index, value in enumerate(values) + ) + if len(normalized) != 16: + raise ValueError(f"{field_name} must contain exactly 16 numbers.") + pose = torch.tensor(normalized, dtype=torch.float64).reshape(4, 4) + bottom = torch.tensor((0.0, 0.0, 0.0, 1.0), dtype=torch.float64) + if not torch.allclose(pose[3], bottom, atol=1.0e-6, rtol=0.0): + raise ValueError(f"{field_name} must have bottom row [0, 0, 0, 1].") + rotation = pose[:3, :3] + if not torch.allclose( + rotation.T @ rotation, + torch.eye(3, dtype=torch.float64), + atol=1.0e-6, + rtol=0.0, + ) or not torch.isclose( + torch.linalg.det(rotation), + torch.tensor(1.0, dtype=torch.float64), + atol=1.0e-6, + rtol=0.0, + ): + raise ValueError(f"{field_name} must contain a proper SE(3) rotation.") + return normalized + + +def _pose_tensor(values: tuple[float, ...]) -> torch.Tensor: + """Materialize an owned float32 pose matrix.""" + return torch.tensor(values, dtype=torch.float32).reshape(4, 4) + + +def _validate_scene_classification( + dynamics: SceneDynamics, + collision_role: SceneCollisionRole, +) -> None: + """Validate exact scene-enum values.""" + if not isinstance(dynamics, SceneDynamics): + raise TypeError("dynamics must be a SceneDynamics value.") + if not isinstance(collision_role, SceneCollisionRole): + raise TypeError("collision_role must be a SceneCollisionRole value.") + + +@dataclass(frozen=True, slots=True) +class SimulationRigidObjectBinding: + """Explicit binding for one simulation rigid object.""" + + entity_id: str + simulation_uid: str + aliases: tuple[str, ...] = () + dynamics: SceneDynamics = SceneDynamics.UNKNOWN + collision_role: SceneCollisionRole = SceneCollisionRole.NONE + semantic_type: str | None = None + default_grasp_affordance: str | None = None + geometry_provider: SceneGeometryProvider | None = None + + def __post_init__(self) -> None: + _identifier(self.entity_id, field_name="entity_id") + _identifier(self.simulation_uid, field_name="simulation_uid") + object.__setattr__( + self, + "aliases", + _identifier_tuple(self.aliases, field_name="aliases"), + ) + _validate_scene_classification(self.dynamics, self.collision_role) + _optional_identifier(self.semantic_type, field_name="semantic_type") + _optional_identifier( + self.default_grasp_affordance, + field_name="default_grasp_affordance", + ) + + +@dataclass(frozen=True, slots=True) +class SimulationArticulationBinding: + """Explicit binding for one simulation articulation.""" + + entity_id: str + simulation_uid: str + aliases: tuple[str, ...] = () + dynamics: SceneDynamics = SceneDynamics.UNKNOWN + collision_role: SceneCollisionRole = SceneCollisionRole.NONE + semantic_type: str | None = None + default_operation_affordance: str | None = None + geometry_provider: SceneGeometryProvider | None = None + + def __post_init__(self) -> None: + _identifier(self.entity_id, field_name="entity_id") + _identifier(self.simulation_uid, field_name="simulation_uid") + object.__setattr__( + self, + "aliases", + _identifier_tuple(self.aliases, field_name="aliases"), + ) + _validate_scene_classification(self.dynamics, self.collision_role) + _optional_identifier(self.semantic_type, field_name="semantic_type") + _optional_identifier( + self.default_operation_affordance, + field_name="default_operation_affordance", + ) + + +@dataclass(frozen=True, slots=True) +class SimulationArticulationLinkBinding: + """Explicit canonical link backed by one native articulation link.""" + + entity_id: str + articulation_id: str + native_link_name: str + aliases: tuple[str, ...] = () + dynamics: SceneDynamics = SceneDynamics.UNKNOWN + semantic_type: str | None = None + + def __post_init__(self) -> None: + _identifier(self.entity_id, field_name="entity_id") + _identifier(self.articulation_id, field_name="articulation_id") + _identifier(self.native_link_name, field_name="native_link_name") + object.__setattr__( + self, + "aliases", + _identifier_tuple(self.aliases, field_name="aliases"), + ) + if not isinstance(self.dynamics, SceneDynamics): + raise TypeError("dynamics must be a SceneDynamics value.") + _optional_identifier(self.semantic_type, field_name="semantic_type") + + +@dataclass(frozen=True, slots=True) +class AntipodalGraspAffordanceBinding: + """Build one antipodal grasp affordance from a selected rigid-object mesh.""" + + entity_id: str + object_id: str + native_name: str + revision: str + aliases: tuple[str, ...] = () + relative_pose: tuple[float, ...] = _IDENTITY_POSE + mesh_env_id: int = 0 + generator_cfg: GraspGeneratorCfg | None = None + gripper_collision_cfg: GripperCollisionCfg | None = None + force_reannotate: bool = False + + def __post_init__(self) -> None: + for field_name in ("entity_id", "object_id", "native_name", "revision"): + _identifier(getattr(self, field_name), field_name=field_name) + object.__setattr__( + self, + "aliases", + _identifier_tuple(self.aliases, field_name="aliases"), + ) + object.__setattr__( + self, + "relative_pose", + _pose_tuple(self.relative_pose, field_name="relative_pose"), + ) + if ( + isinstance(self.mesh_env_id, bool) + or not isinstance(self.mesh_env_id, int) + or self.mesh_env_id < 0 + ): + raise ValueError("mesh_env_id must be a non-negative integer.") + if self.generator_cfg is not None and not isinstance( + self.generator_cfg, + GraspGeneratorCfg, + ): + raise TypeError("generator_cfg must be GraspGeneratorCfg or None.") + if self.gripper_collision_cfg is not None and not isinstance( + self.gripper_collision_cfg, + GripperCollisionCfg, + ): + raise TypeError( + "gripper_collision_cfg must be GripperCollisionCfg or None." + ) + if not isinstance(self.force_reannotate, bool): + raise TypeError("force_reannotate must be a bool.") + object.__setattr__(self, "generator_cfg", deepcopy(self.generator_cfg)) + object.__setattr__( + self, + "gripper_collision_cfg", + deepcopy(self.gripper_collision_cfg), + ) + + +@dataclass(frozen=True, slots=True) +class ArticulationOperationTargetBinding: + """Declarative named target for one articulation operation.""" + + target_position: float + displacement: float + + def __post_init__(self) -> None: + object.__setattr__( + self, + "target_position", + _finite(self.target_position, field_name="target_position"), + ) + object.__setattr__( + self, + "displacement", + _finite(self.displacement, field_name="displacement"), + ) + + def build(self) -> ArticulationOperationTarget: + """Build the existing atomic-action target value.""" + return ArticulationOperationTarget( + target_position=self.target_position, + displacement=self.displacement, + ) + + +@dataclass(frozen=True, slots=True) +class ArticulationOperationAffordanceBinding: + """Bind one handle operation to an explicit native link and joint.""" + + entity_id: str + articulation_id: str + link_id: str + joint_id: str + revision: str + semantic_targets: Mapping[str, ArticulationOperationTargetBinding] + aliases: tuple[str, ...] = () + handle_pose_offset: tuple[float, ...] = _IDENTITY_POSE + approach_offset: tuple[float, ...] = _IDENTITY_POSE + contact_offset: tuple[float, ...] = _IDENTITY_POSE + operation_offset: tuple[float, ...] = _IDENTITY_POSE + retract_offset: tuple[float, ...] = _IDENTITY_POSE + operation_axis: tuple[float, float, float] = (1.0, 0.0, 0.0) + position_scale: float = 1.0 + + def __post_init__(self) -> None: + for field_name in ( + "entity_id", + "articulation_id", + "link_id", + "joint_id", + "revision", + ): + _identifier(getattr(self, field_name), field_name=field_name) + object.__setattr__( + self, + "aliases", + _identifier_tuple(self.aliases, field_name="aliases"), + ) + for field_name in ( + "handle_pose_offset", + "approach_offset", + "contact_offset", + "operation_offset", + "retract_offset", + ): + object.__setattr__( + self, + field_name, + _pose_tuple(getattr(self, field_name), field_name=field_name), + ) + axis = tuple( + _finite(value, field_name=f"operation_axis[{index}]") + for index, value in enumerate(self.operation_axis) + ) + if len(axis) != 3 or math.sqrt(sum(value * value for value in axis)) <= 0.0: + raise ValueError("operation_axis must contain three non-zero values.") + object.__setattr__(self, "operation_axis", axis) + position_scale = _finite(self.position_scale, field_name="position_scale") + if position_scale <= 0.0: + raise ValueError("position_scale must be positive.") + object.__setattr__(self, "position_scale", position_scale) + if not isinstance(self.semantic_targets, Mapping): + raise TypeError("semantic_targets must be a mapping.") + targets: dict[str, ArticulationOperationTargetBinding] = {} + for target_id, target in self.semantic_targets.items(): + _identifier(target_id, field_name="semantic target IDs") + if type(target) is not ArticulationOperationTargetBinding: + raise TypeError( + "semantic_targets values must be exact " + "ArticulationOperationTargetBinding values." + ) + targets[target_id] = target + object.__setattr__(self, "semantic_targets", MappingProxyType(targets)) + + +@dataclass(frozen=True, slots=True) +class _SimulationArticulationLinkStateProvider: + """Read one selected native link pose with an optional local offset.""" + + articulation: Any + native_link_name: str + local_offset: torch.Tensor = field(repr=False) + + def observe( + self, + *, + timestamp: float, + env_ids: torch.Tensor, + ) -> EntityState: + del timestamp + getter = getattr(self.articulation, "get_link_pose", None) + if not callable(getter): + raise TypeError("Simulation articulation must provide get_link_pose().") + pose = getter( + self.native_link_name, + env_ids=env_ids.detach().to("cpu").tolist(), + to_matrix=True, + ) + if not isinstance(pose, torch.Tensor): + raise TypeError( + "Simulation articulation get_link_pose() must return a tensor." + ) + offset = self.local_offset.to(device=pose.device, dtype=pose.dtype) + return EntityState(torch.matmul(pose, offset)) + + +def _require_native_entity( + simulation: SimulationManager, + *, + getter_name: str, + registry_id: str, + simulation_uid: str, +) -> Any: + """Resolve one explicitly selected native entity or fail closed.""" + getter = getattr(simulation, getter_name, None) + if not callable(getter): + raise TypeError(f"simulation must provide {getter_name}().") + entity = getter(simulation_uid) + if entity is None: + raise KeyError( + f"Simulation UID {simulation_uid!r} selected for registry entity " + f"{registry_id!r} was not found." + ) + return entity + + +def _native_names(entity: Any, *, attribute: str, owner: str) -> tuple[str, ...]: + """Read and validate one existing native-name collection.""" + values = getattr(entity, attribute, None) + if values is None: + raise TypeError(f"{owner} must expose {attribute}.") + if isinstance(values, (str, bytes)): + raise TypeError(f"{owner}.{attribute} must be an iterable of names.") + try: + names = tuple(values) + except TypeError as exc: + raise TypeError(f"{owner}.{attribute} must be an iterable of names.") from exc + for name in names: + _identifier(name, field_name=f"{owner}.{attribute}") + if len(set(names)) != len(names): + raise ValueError(f"{owner}.{attribute} must contain unique names.") + return names + + +def _mesh_tensor( + entity: Any, + *, + getter_name: str, + mesh_env_id: int, + vertices: bool, +) -> torch.Tensor: + """Read one explicitly selected mesh row with strict shape validation.""" + getter = getattr(entity, getter_name, None) + if not callable(getter): + raise TypeError(f"Simulation rigid object must provide {getter_name}().") + if vertices: + value = getter(env_ids=[mesh_env_id], scale=True) + else: + value = getter(env_ids=[mesh_env_id]) + if not isinstance(value, torch.Tensor): + raise TypeError( + f"Simulation rigid object {getter_name}() must return a tensor." + ) + if value.dim() != 3 or value.shape[0] != 1 or value.shape[2] != 3: + raise ValueError( + f"Simulation rigid object {getter_name}() must return shape (1, N, 3)." + ) + selected = value[0].detach().clone() + if selected.shape[0] == 0: + raise ValueError(f"Simulation rigid object {getter_name}() returned no data.") + if vertices: + if not selected.is_floating_point() or not torch.isfinite(selected).all(): + raise ValueError("Antipodal mesh vertices must be finite floating values.") + elif selected.dtype == torch.bool or selected.is_floating_point(): + raise TypeError("Antipodal mesh triangles must use an integer dtype.") + return selected + + +def _antipodal_affordance( + binding: AntipodalGraspAffordanceBinding, + entity: Any, +) -> AntipodalAffordance: + """Build and validate one owned antipodal affordance payload.""" + vertices = _mesh_tensor( + entity, + getter_name="get_vertices", + mesh_env_id=binding.mesh_env_id, + vertices=True, + ) + triangles = _mesh_tensor( + entity, + getter_name="get_triangles", + mesh_env_id=binding.mesh_env_id, + vertices=False, + ) + if bool((triangles < 0).any()) or int(triangles.max().item()) >= vertices.shape[0]: + raise ValueError("Antipodal mesh triangles reference invalid vertex indices.") + return AntipodalAffordance( + mesh_vertices=vertices, + mesh_triangles=triangles, + generator_cfg=deepcopy(binding.generator_cfg), + gripper_collision_cfg=deepcopy(binding.gripper_collision_cfg), + force_reannotate=binding.force_reannotate, + ) + + +@dataclass(frozen=True, slots=True) +class SimulationSceneBinding: + """Build one authoritative registry from explicit simulation bindings.""" + + registry_id: str + rigid_objects: tuple[SimulationRigidObjectBinding, ...] = () + articulations: tuple[SimulationArticulationBinding, ...] = () + links: tuple[SimulationArticulationLinkBinding, ...] = () + antipodal_grasps: tuple[AntipodalGraspAffordanceBinding, ...] = () + articulation_operations: tuple[ArticulationOperationAffordanceBinding, ...] = () + collision_world_mode: SceneCollisionWorldMode | None = None + + def __post_init__(self) -> None: + _identifier(self.registry_id, field_name="registry_id") + expected_types = { + "rigid_objects": SimulationRigidObjectBinding, + "articulations": SimulationArticulationBinding, + "links": SimulationArticulationLinkBinding, + "antipodal_grasps": AntipodalGraspAffordanceBinding, + "articulation_operations": ArticulationOperationAffordanceBinding, + } + all_ids: list[str] = [] + for field_name, expected_type in expected_types.items(): + values = tuple(getattr(self, field_name)) + if not all(type(value) is expected_type for value in values): + raise TypeError( + f"{field_name} must contain exact {expected_type.__name__} values." + ) + object.__setattr__(self, field_name, values) + all_ids.extend(value.entity_id for value in values) + duplicates = sorted( + entity_id for entity_id in set(all_ids) if all_ids.count(entity_id) > 1 + ) + if duplicates: + raise ValueError(f"Scene binding entity IDs must be unique: {duplicates}.") + if self.collision_world_mode is not None and not isinstance( + self.collision_world_mode, + SceneCollisionWorldMode, + ): + raise TypeError( + "collision_world_mode must be SceneCollisionWorldMode or None." + ) + + def build(self, simulation: SimulationManager) -> SceneRegistry: + """Build the existing authoritative scene registry. + + Args: + simulation: Live simulation used only for explicitly named lookups. + + Returns: + Immutable registry with typed roots, links, and affordances. + """ + objects = {item.entity_id: item for item in self.rigid_objects} + articulations = {item.entity_id: item for item in self.articulations} + geometry = { + item.entity_id: item.geometry_provider + for item in (*self.rigid_objects, *self.articulations) + if item.geometry_provider is not None + } + roles = { + item.entity_id: item.collision_role + for item in (*self.rigid_objects, *self.articulations) + } + base = SceneRegistry.from_simulation( + simulation, + rigid_objects={ + item.entity_id: item.simulation_uid for item in self.rigid_objects + }, + articulations={ + item.entity_id: item.simulation_uid for item in self.articulations + }, + collision_roles=roles, + geometry_providers=geometry, + collision_world_mode=self.collision_world_mode, + ) + + registrations: list[SceneEntityRegistration] = [] + for registration in base.registrations: + entity_id = registration.ref.entity_id + if isinstance(registration.ref, SceneObjectRef): + binding = objects[entity_id] + defaults = ( + {} + if binding.default_grasp_affordance is None + else { + GRASP_AFFORDANCE_CAPABILITY: SceneAffordanceRef( + binding.default_grasp_affordance + ) + } + ) + else: + binding = articulations[entity_id] + defaults = ( + {} + if binding.default_operation_affordance is None + else { + ARTICULATION_OPERATION_AFFORDANCE_CAPABILITY: ( + SceneAffordanceRef(binding.default_operation_affordance) + ) + } + ) + registrations.append( + replace( + registration, + aliases=(*registration.aliases, *binding.aliases), + dynamics=binding.dynamics, + semantic_type=binding.semantic_type, + default_affordances=defaults, + ) + ) + + native_articulations: dict[str, Any] = {} + links: dict[str, SimulationArticulationLinkBinding] = {} + for binding in self.links: + articulation_binding = articulations.get(binding.articulation_id) + if articulation_binding is None: + raise KeyError( + f"Link {binding.entity_id!r} references unbound articulation " + f"{binding.articulation_id!r}." + ) + articulation = native_articulations.setdefault( + binding.articulation_id, + _require_native_entity( + simulation, + getter_name="get_articulation", + registry_id=binding.articulation_id, + simulation_uid=articulation_binding.simulation_uid, + ), + ) + native_links = _native_names( + articulation, + attribute="link_names", + owner=f"articulation {binding.articulation_id!r}", + ) + if binding.native_link_name not in native_links: + raise KeyError( + f"Native link {binding.native_link_name!r} selected for " + f"{binding.entity_id!r} was not found; available links are " + f"{sorted(native_links)}." + ) + links[binding.entity_id] = binding + registrations.append( + SceneEntityRegistration( + ref=SceneLinkRef(binding.entity_id), + state_provider=_SimulationArticulationLinkStateProvider( + articulation, + binding.native_link_name, + _pose_tensor(_IDENTITY_POSE), + ), + aliases=binding.aliases, + parent=SceneArticulationRef(binding.articulation_id), + native_name=binding.native_link_name, + dynamics=binding.dynamics, + semantic_type=binding.semantic_type, + ) + ) + + native_objects: dict[str, Any] = {} + for binding in self.antipodal_grasps: + object_binding = objects.get(binding.object_id) + if object_binding is None: + raise KeyError( + f"Grasp affordance {binding.entity_id!r} references unbound " + f"object {binding.object_id!r}." + ) + entity = native_objects.setdefault( + binding.object_id, + _require_native_entity( + simulation, + getter_name="get_rigid_object", + registry_id=binding.object_id, + simulation_uid=object_binding.simulation_uid, + ), + ) + registrations.append( + SceneEntityRegistration( + ref=SceneAffordanceRef(binding.entity_id), + aliases=binding.aliases, + parent=SceneObjectRef(binding.object_id), + native_name=binding.native_name, + affordance=_antipodal_affordance(binding, entity), + affordance_capabilities=frozenset({GRASP_AFFORDANCE_CAPABILITY}), + affordance_revision=binding.revision, + relative_pose=_pose_tensor(binding.relative_pose), + ) + ) + + for binding in self.articulation_operations: + articulation_binding = articulations.get(binding.articulation_id) + if articulation_binding is None: + raise KeyError( + f"Operation affordance {binding.entity_id!r} references " + f"unbound articulation {binding.articulation_id!r}." + ) + link_binding = links.get(binding.link_id) + if link_binding is None: + raise KeyError( + f"Operation affordance {binding.entity_id!r} references " + f"unbound link {binding.link_id!r}." + ) + if link_binding.articulation_id != binding.articulation_id: + raise ValueError( + f"Operation affordance {binding.entity_id!r} and link " + f"{binding.link_id!r} select different articulations." + ) + articulation = native_articulations[binding.articulation_id] + native_joints = _native_names( + articulation, + attribute="joint_names", + owner=f"articulation {binding.articulation_id!r}", + ) + if binding.joint_id not in native_joints: + raise KeyError( + f"Native joint {binding.joint_id!r} selected for " + f"{binding.entity_id!r} was not found; available joints are " + f"{sorted(native_joints)}." + ) + payload = ArticulationOperationAffordance( + joint_id=binding.joint_id, + approach_offset=_pose_tensor(binding.approach_offset), + contact_offset=_pose_tensor(binding.contact_offset), + operation_offset=_pose_tensor(binding.operation_offset), + retract_offset=_pose_tensor(binding.retract_offset), + operation_axis=torch.tensor( + binding.operation_axis, + dtype=torch.float32, + ), + position_scale=binding.position_scale, + semantic_targets={ + target_id: target.build() + for target_id, target in binding.semantic_targets.items() + }, + ) + registrations.append( + SceneEntityRegistration( + ref=SceneAffordanceRef(binding.entity_id), + state_provider=_SimulationArticulationLinkStateProvider( + articulation, + link_binding.native_link_name, + _pose_tensor(binding.handle_pose_offset), + ), + aliases=binding.aliases, + parent=SceneArticulationRef(binding.articulation_id), + native_name=link_binding.native_link_name, + affordance=payload, + affordance_capabilities=frozenset( + {ARTICULATION_OPERATION_AFFORDANCE_CAPABILITY} + ), + affordance_revision=binding.revision, + ) + ) + + return SceneRegistry( + registrations, + collision_world_mode=self.collision_world_mode, + ) + + +@dataclass(frozen=True, slots=True) +class ControlPartCommandPreset: + """Named one-dimensional joint commands for one exact control part.""" + + preset_id: str + control_part: str + commands: Mapping[str, tuple[float, ...]] + + def __post_init__(self) -> None: + _identifier(self.preset_id, field_name="preset_id") + _identifier(self.control_part, field_name="control_part") + if not isinstance(self.commands, Mapping): + raise TypeError("commands must be a mapping.") + commands: dict[str, tuple[float, ...]] = {} + for command_id, positions in self.commands.items(): + _identifier(command_id, field_name="command IDs") + if isinstance(positions, (str, bytes)): + raise TypeError("command positions must be an iterable of numbers.") + normalized = tuple( + _finite(value, field_name=f"commands[{command_id!r}][{index}]") + for index, value in enumerate(positions) + ) + if not normalized: + raise ValueError("command positions must not be empty.") + commands[command_id] = normalized + object.__setattr__(self, "commands", MappingProxyType(commands)) + + def build(self, *, control_dof: int) -> ControlPartCommandProfile: + """Build a command profile after validating the native control width.""" + for command_id, positions in self.commands.items(): + if len(positions) != control_dof: + raise ValueError( + f"Command {command_id!r} in preset {self.preset_id!r} has " + f"{len(positions)} positions, but control part " + f"{self.control_part!r} has {control_dof} joints." + ) + return ControlPartCommandProfile.joint_positions( + **{ + command_id: torch.tensor(positions, dtype=torch.float32) + for command_id, positions in self.commands.items() + } + ) + + +def _require_control_part_dof(robot: Robot, control_part: str) -> int: + """Validate one native joint-backed control part and return its width.""" + control_parts = getattr(robot, "control_parts", None) + if not isinstance(control_parts, Mapping): + raise TypeError("robot must expose a control_parts mapping.") + get_joint_ids = getattr(robot, "get_joint_ids", None) + if not callable(get_joint_ids): + raise TypeError("robot must provide get_joint_ids().") + if control_part not in control_parts: + raise KeyError( + f"Robot control part {control_part!r} was not found; available " + f"control parts are {sorted(str(key) for key in control_parts)}." + ) + joint_ids = tuple(get_joint_ids(name=control_part)) + if not joint_ids: + raise ValueError(f"Robot control part {control_part!r} contains no joints.") + if not all( + isinstance(joint_id, int) and not isinstance(joint_id, bool) and joint_id >= 0 + for joint_id in joint_ids + ): + raise ValueError( + f"Robot control part {control_part!r} returned invalid joint IDs." + ) + if len(set(joint_ids)) != len(joint_ids): + raise ValueError( + f"Robot control part {control_part!r} contains duplicate joint IDs." + ) + return len(joint_ids) + + +@runtime_checkable +class SimulationResourceEndpointBinding(Protocol): + """Build one typed resource endpoint from an explicitly selected robot. + + Implementations are reusable robot-integration declarations. They may + validate embodiment-specific controller surfaces, but must only return an + owned :class:`ResourceEndpoint`; live controller handles remain in the + endpoint adapter and runtime transport. + """ + + @property + def endpoint_id(self) -> str: + """Return the stable endpoint ID within its containing resource.""" + + def build(self, robot: Robot) -> ResourceEndpoint: + """Build and validate one endpoint declaration for ``robot``.""" + + +@runtime_checkable +class SimulationRobotResourceBinding(Protocol): + """Build one leaf or composite resource in the robot resource DAG.""" + + @property + def resource_id(self) -> str: + """Return the stable resource ID.""" + + @property + def members(self) -> tuple[str, ...]: + """Return explicitly declared child resource IDs.""" + + def build(self, robot: Robot) -> RobotResource: + """Build and validate one owned robot resource declaration.""" + + +@dataclass(frozen=True, slots=True) +class RobotResourceBinding: + """Generic simulation binding for arbitrary typed resource endpoints. + + This is the direct configuration path for mobile bases, whole-body + controllers, tools, and other non-joint transports. Endpoint-specific + validation remains in the registered :class:`ResourceEndpointAdapter`; + this value owns the declaration and preserves the resource DAG exactly. + """ + + resource_id: str + endpoints: Mapping[str, ResourceEndpoint] = field(default_factory=dict) + members: tuple[str, ...] = () + + def __post_init__(self) -> None: + resource = RobotResource( + resource_id=self.resource_id, + endpoints=self.endpoints, + members=self.members, + ) + object.__setattr__(self, "endpoints", resource.endpoints) + object.__setattr__(self, "members", resource.members) + + def build(self, robot: Robot) -> RobotResource: + """Build an independently owned resource without assuming robot joints.""" + del robot + return RobotResource( + resource_id=self.resource_id, + endpoints=self.endpoints, + members=self.members, + ) + + +@dataclass(frozen=True, slots=True) +class ControlPartEndpointBinding: + """Profile endpoint backed by one explicit robot control part.""" + + endpoint_id: str + control_part: str + capabilities: frozenset[str] + command_preset: str | None = None + + def __post_init__(self) -> None: + _identifier(self.endpoint_id, field_name="endpoint_id") + _identifier(self.control_part, field_name="control_part") + if isinstance(self.capabilities, (str, bytes)): + raise TypeError("capabilities must be an iterable of identifiers.") + capabilities = frozenset(self.capabilities) + for capability in capabilities: + _identifier(capability, field_name="capabilities") + object.__setattr__(self, "capabilities", capabilities) + _optional_identifier(self.command_preset, field_name="command_preset") + + def build(self, robot: Robot) -> ResourceEndpoint: + """Build a joint-backed endpoint after native control-part validation.""" + _require_control_part_dof(robot, self.control_part) + return ControlPartEndpoint( + control_part=self.control_part, + command_profile=self.command_preset, + capabilities=self.capabilities, + ) + + +@dataclass(frozen=True, slots=True) +class ControlPartResourceBinding: + """Joint-backed robot resource containing control-part endpoints.""" + + resource_id: str + endpoints: tuple[ControlPartEndpointBinding, ...] = () + members: tuple[str, ...] = () + + def __post_init__(self) -> None: + _identifier(self.resource_id, field_name="resource_id") + endpoints = tuple(self.endpoints) + if not all( + type(endpoint) is ControlPartEndpointBinding for endpoint in endpoints + ): + raise TypeError( + "endpoints must contain exact ControlPartEndpointBinding values." + ) + endpoint_ids = [endpoint.endpoint_id for endpoint in endpoints] + if len(set(endpoint_ids)) != len(endpoint_ids): + raise ValueError("endpoint_id values must be unique within a resource.") + object.__setattr__(self, "endpoints", endpoints) + object.__setattr__( + self, + "members", + _identifier_tuple(self.members, field_name="members"), + ) + + def build(self, robot: Robot) -> RobotResource: + """Build a resource containing strictly validated control-part endpoints.""" + endpoints: dict[str, ResourceEndpoint] = {} + for binding in self.endpoints: + endpoint = binding.build(robot) + if type(endpoint) is not ControlPartEndpoint: + raise TypeError( + "ControlPartEndpointBinding.build() must return exactly " + "ControlPartEndpoint." + ) + endpoints[binding.endpoint_id] = endpoint + return RobotResource( + resource_id=self.resource_id, + endpoints=endpoints, + members=self.members, + ) + + +def _owned_nested_identifier_mapping( + values: Mapping[str, Mapping[str, str]], + *, + field_name: str, +) -> Mapping[str, Mapping[str, str]]: + """Own a strict two-level identifier mapping.""" + if not isinstance(values, Mapping): + raise TypeError(f"{field_name} must be a mapping.") + outer: dict[str, Mapping[str, str]] = {} + for key, nested in values.items(): + _identifier(key, field_name=f"{field_name} keys") + if not isinstance(nested, Mapping): + raise TypeError(f"{field_name}[{key!r}] must be a mapping.") + normalized: dict[str, str] = {} + for nested_key, nested_value in nested.items(): + _identifier(nested_key, field_name=f"{field_name} slot IDs") + _identifier(nested_value, field_name=f"{field_name} resource IDs") + normalized[nested_key] = nested_value + outer[key] = MappingProxyType(normalized) + return MappingProxyType(outer) + + +def _owned_identifier_mapping( + values: Mapping[str, str], + *, + field_name: str, +) -> Mapping[str, str]: + """Own one strict identifier mapping.""" + if not isinstance(values, Mapping): + raise TypeError(f"{field_name} must be a mapping.") + normalized: dict[str, str] = {} + for key, value in values.items(): + _identifier(key, field_name=f"{field_name} keys") + _identifier(value, field_name=f"{field_name} values") + normalized[key] = value + return MappingProxyType(normalized) + + +@dataclass(frozen=True, slots=True) +class SimulationRobotSkillProfileBinding: + """Build a profile from typed resources with strict native validation.""" + + profile_id: str + resources: tuple[SimulationRobotResourceBinding, ...] + command_presets: tuple[ControlPartCommandPreset, ...] = () + defaults: Mapping[str, Mapping[str, str]] = field(default_factory=dict) + presets: tuple[SkillPolicyPreset, ...] = () + default_preset: str | None = None + skill_presets: Mapping[str, str] = field(default_factory=dict) + grounding_providers: Mapping[str, str] = field(default_factory=dict) + + def __post_init__(self) -> None: + _identifier(self.profile_id, field_name="profile_id") + resources = tuple(self.resources) + if not all( + isinstance(resource, SimulationRobotResourceBinding) + for resource in resources + ): + raise TypeError("resources must implement SimulationRobotResourceBinding.") + for resource in resources: + _identifier(resource.resource_id, field_name="resource_id") + _identifier_tuple(resource.members, field_name="resource members") + resource_ids = [resource.resource_id for resource in resources] + if len(set(resource_ids)) != len(resource_ids): + raise ValueError("resource_id values must be unique.") + object.__setattr__(self, "resources", resources) + command_presets = tuple(self.command_presets) + if not all( + type(preset) is ControlPartCommandPreset for preset in command_presets + ): + raise TypeError( + "command_presets must contain exact ControlPartCommandPreset values." + ) + command_preset_ids = [preset.preset_id for preset in command_presets] + if len(set(command_preset_ids)) != len(command_preset_ids): + raise ValueError("command preset IDs must be unique.") + object.__setattr__(self, "command_presets", command_presets) + object.__setattr__( + self, + "defaults", + _owned_nested_identifier_mapping(self.defaults, field_name="defaults"), + ) + presets = tuple(self.presets) + if not all(type(preset) is SkillPolicyPreset for preset in presets): + raise TypeError("presets must contain exact SkillPolicyPreset values.") + preset_ids = [preset.preset_id for preset in presets] + if len(set(preset_ids)) != len(preset_ids): + raise ValueError("policy preset IDs must be unique.") + object.__setattr__(self, "presets", presets) + _optional_identifier(self.default_preset, field_name="default_preset") + object.__setattr__( + self, + "skill_presets", + _owned_identifier_mapping( + self.skill_presets, + field_name="skill_presets", + ), + ) + object.__setattr__( + self, + "grounding_providers", + _owned_identifier_mapping( + self.grounding_providers, + field_name="grounding_providers", + ), + ) + + def build(self, robot: Robot) -> RobotSkillProfile: + """Build the existing profile after validating every typed resource. + + Args: + robot: Live robot selected by the simulation factory. + + Returns: + Reusable, engine-independent robot skill profile. + """ + control_dofs: dict[str, int] = {} + + def require_control_part(control_part: str) -> int: + if control_part not in control_dofs: + control_dofs[control_part] = _require_control_part_dof( + robot, + control_part, + ) + return control_dofs[control_part] + + command_presets = {preset.preset_id: preset for preset in self.command_presets} + command_profiles: dict[str, ControlPartCommandProfile] = {} + for preset in self.command_presets: + command_profiles[preset.preset_id] = preset.build( + control_dof=require_control_part(preset.control_part) + ) + + resources: dict[str, RobotResource] = {} + for resource_binding in self.resources: + resource = resource_binding.build(robot) + if type(resource) is not RobotResource: + raise TypeError( + f"Resource binding {resource_binding.resource_id!r} must build " + "exactly RobotResource." + ) + if resource.resource_id != resource_binding.resource_id: + raise ValueError( + f"Resource binding {resource_binding.resource_id!r} built " + f"resource ID {resource.resource_id!r}." + ) + if resource.members != tuple(resource_binding.members): + raise ValueError( + f"Resource binding {resource_binding.resource_id!r} changed its " + "declared resource members while building." + ) + for endpoint_id, endpoint in resource.endpoints.items(): + if not isinstance(endpoint, ControlPartEndpoint): + continue + require_control_part(endpoint.control_part) + profile_id = ( + endpoint.control_part + if endpoint.command_profile is None + else endpoint.command_profile + ) + command_preset = command_presets.get(profile_id) + if endpoint.command_profile is not None and command_preset is None: + raise KeyError( + f"Endpoint {resource.resource_id!r}.{endpoint_id!r} " + "references unknown command " + f"preset {profile_id!r}." + ) + if ( + command_preset is not None + and command_preset.control_part != endpoint.control_part + ): + raise ValueError( + f"Endpoint {resource.resource_id!r}.{endpoint_id!r} uses " + "control part " + f"{endpoint.control_part!r}, but command preset " + f"{profile_id!r} targets " + f"{command_preset.control_part!r}." + ) + resources[resource.resource_id] = resource + + return RobotSkillProfile( + profile_id=self.profile_id, + resources=resources, + command_profiles=command_profiles, + defaults={ + skill_id: ResourceBinding(resources=bindings) + for skill_id, bindings in self.defaults.items() + }, + presets={preset.preset_id: preset for preset in self.presets}, + default_preset=self.default_preset, + skill_presets=self.skill_presets, + grounding_providers=self.grounding_providers, + ) + + +__all__ = [ + "AntipodalGraspAffordanceBinding", + "ArticulationOperationAffordanceBinding", + "ArticulationOperationTargetBinding", + "ControlPartCommandPreset", + "ControlPartEndpointBinding", + "ControlPartResourceBinding", + "RobotResourceBinding", + "SimulationArticulationBinding", + "SimulationArticulationLinkBinding", + "SimulationRigidObjectBinding", + "SimulationResourceEndpointBinding", + "SimulationRobotResourceBinding", + "SimulationRobotSkillProfileBinding", + "SimulationSceneBinding", +] diff --git a/embodichain/lab/gym/envs/expert_program/simulation_environment.py b/embodichain/lab/gym/envs/expert_program/simulation_environment.py new file mode 100644 index 000000000..557b8f707 --- /dev/null +++ b/embodichain/lab/gym/envs/expert_program/simulation_environment.py @@ -0,0 +1,1223 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Production simulation assembly for Gym-backed Expert Programs. + +This module owns the reusable live wiring between declarative simulation +bindings and :class:`ExpertProgramEnvironmentAdapter`. A task declares a +scene binding and a robot profile binding; this factory constructs the motion +generator, atomic-action engine, planning observation port, effect-evidence +providers, and segment-policy port without task-local motion code. + +The resulting runtime is intentionally Gym-only. Its buffered command sink +must remain attached to :class:`AtomicDemoBridge`, which advances the shared +clock only after an ordinary ``env.step()`` consumes a yielded command. It is +therefore not a ``SkillRuntimeProvider`` for synchronous ``AtomicSkills`` use. +""" + +from __future__ import annotations + +from collections.abc import Callable, Iterable, Mapping +from copy import deepcopy +from dataclasses import replace +import math +from typing import Any, Protocol, TYPE_CHECKING + +import torch + +from embodichain.lab.gym.envs.settling import DynamicSettleMonitorCfg +from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + EntityState, + ObservedArticulationJointState, + PlanningContext, + RobotObservation, + SceneProvider, + SceneSnapshot, + TaskState, +) +from embodichain.lab.sim.atomic_actions.bindings import ( + JointPositionTarget, + RuntimeEndpointTarget, +) +from embodichain.lab.sim.atomic_actions.control import ( + GRASP_COMMAND, + OPEN_COMMAND, + ControlPartCommandProfile, + JointPositionCommand, +) +from embodichain.lab.sim.atomic_actions.runner import ExecutionRunnerCfg +from embodichain.lab.sim.atomic_actions.runtime_commands import ( + JointPositionPayload, + RuntimeCommandFrame, +) +from embodichain.lab.sim.planners import ( + BasePlannerCfg, + MotionGenCfg, + MotionGenerator, + ToppraPlannerCfg, +) +from embodichain.lab.sim.skills.calls import SemanticCallCatalog +from embodichain.lab.sim.skills.compiler import ( + HandOverPoseProvider, + RegisteredSemanticLowerer, + RelationTargetGrounder, +) +from embodichain.lab.sim.skills.effects import ( + ControlPartEvidenceAddress, + EffectMonitorRegistry, +) +from embodichain.lab.sim.skills.evidence import ( + BinaryEffectEvidenceQuery, + BinaryObservationCallback, + BinaryEffectObservation, + ControlPartRobotEvidenceSource, + ControlPartSimulationEvidenceProvider, + EffectEvidenceCollectionContext, + EffectEvidenceProvider, + ScalarObservationCallback, + SceneArticulationEvidenceProvider, +) +from embodichain.lab.sim.skills.parallel_runtime import ( + ParallelCommandSafetyValidator, +) +from embodichain.lab.sim.skills.profiles import ( + ResourceEndpoint, + ResourceEndpointAdapter, + RobotSkillProfile, + SkillPolicyPreset, +) +from embodichain.lab.sim.skills.scene import RegistrySceneProvider, SceneRegistry + +from .bridge import ( + AcceptedRuntimeCommandObserver, + EnvironmentStepClock, + GymPlanningObservationProvider, + RuntimeTransportActionEncoder, +) +from .environment import ( + ExpertProgramEnvironmentAdapter, + ExpertProgramEnvironmentFactory, + PlanningObservationPort, +) +from .simulation import ( + SimulationRobotSkillProfileBinding, + SimulationSceneBinding, +) +from .simulation_policies import SimulationSegmentPolicyPort + +if TYPE_CHECKING: + from embodichain.lab.sim.objects import Robot + from embodichain.lab.sim.sim_manager import SimulationManager + + +MotionGeneratorFactory = Callable[[], MotionGenerator] +"""Zero-argument factory that must return one fresh motion generator.""" + + +class ControlCommandStateEvidenceTracker(AcceptedRuntimeCommandObserver): + """Track row-local open/grasp state from accepted semantic commands. + + This is the explicit lightweight evidence option selected for simulations + that do not expose a typed contact sensor. It does not claim physical + contact by itself: the built-in effect contract still conjuncts this + binary command state with live object-to-endpoint pose evidence. State is + updated only after the complete command frame has been encoded and accepted + by :class:`BufferedGymCommandSink`. + + Args: + control_profiles: Exact semantic command profiles installed in the + atomic engine, keyed by concrete control-part name. + env_ids: Stable full simulation batch correlation IDs. + """ + + def __init__( + self, + control_profiles: Mapping[str, ControlPartCommandProfile], + env_ids: torch.Tensor, + ) -> None: + if not isinstance(control_profiles, Mapping): + raise TypeError("control_profiles must be a mapping.") + if ( + not isinstance(env_ids, torch.Tensor) + or env_ids.dtype != torch.long + or env_ids.dim() != 1 + or env_ids.numel() == 0 + ): + raise ValueError( + "env_ids must be a non-empty one-dimensional int64 tensor." + ) + if torch.unique(env_ids).numel() != env_ids.numel(): + raise ValueError("env_ids must be unique.") + + commands: dict[ + str, + tuple[JointPositionCommand, JointPositionCommand], + ] = {} + for control_part, profile in control_profiles.items(): + if type(control_part) is not str or not control_part: + raise ValueError("control_profiles keys must be non-empty strings.") + if not isinstance(profile, ControlPartCommandProfile): + raise TypeError( + "control_profiles values must be ControlPartCommandProfile values." + ) + open_command = profile.commands.get(OPEN_COMMAND) + grasp_command = profile.commands.get(GRASP_COMMAND) + if open_command is None and grasp_command is None: + continue + if not isinstance(open_command, JointPositionCommand) or not isinstance( + grasp_command, + JointPositionCommand, + ): + raise TypeError( + f"Control part {control_part!r} must define both open and grasp " + "as JointPositionCommand values for command-state evidence." + ) + if open_command.equivalent_to(grasp_command): + raise ValueError( + f"Control part {control_part!r} has indistinguishable open and " + "grasp commands." + ) + commands[control_part] = ( + open_command.snapshot(), + grasp_command.snapshot(), + ) + + self._commands = commands + self._env_ids = env_ids.clone() + self._row_by_env_id = { + int(env_id): row + for row, env_id in enumerate(env_ids.detach().cpu().tolist()) + } + batch_size = int(env_ids.numel()) + self._values = { + control_part: torch.zeros( + batch_size, + dtype=torch.bool, + device=env_ids.device, + ) + for control_part in commands + } + self._valid = { + control_part: torch.zeros_like(values) + for control_part, values in self._values.items() + } + + @property + def tracked_control_parts(self) -> tuple[str, ...]: + """Return control parts with exact open/grasp semantic commands.""" + return tuple(self._commands) + + def accepted(self, command: RuntimeCommandFrame) -> None: + """Commit exact open/grasp states for active rows in an accepted frame.""" + if not isinstance(command, RuntimeCommandFrame): + raise TypeError("command must be a RuntimeCommandFrame.") + rows = self._rows(command.env_ids) + for endpoint_command in command.commands: + target = endpoint_command.target + payload = endpoint_command.payload + if not isinstance(target, JointPositionTarget) or not isinstance( + payload, + JointPositionPayload, + ): + continue + semantic_commands = self._commands.get(target.control_part) + if semantic_commands is None: + continue + open_command, grasp_command = semantic_commands + open_positions = open_command.resolve( + n_envs=command.batch_size, + control_dof=payload.dof, + device=payload.device, + dtype=payload.positions.dtype, + ) + grasp_positions = grasp_command.resolve( + n_envs=command.batch_size, + control_dof=payload.dof, + device=payload.device, + dtype=payload.positions.dtype, + ) + is_open = torch.isclose( + payload.positions, + open_positions, + rtol=0.0, + atol=1.0e-7, + ).all(dim=1) + is_grasp = torch.isclose( + payload.positions, + grasp_positions, + rtol=0.0, + atol=1.0e-7, + ).all(dim=1) + if bool((is_open & is_grasp).any().item()): + raise ValueError( + "An accepted row matched both open and grasp commands." + ) + recognized = command.active_mask & (is_open | is_grasp) + if not bool(recognized.any().item()): + continue + destination_rows = torch.tensor( + rows, + dtype=torch.long, + device=self._env_ids.device, + ) + selected_rows = destination_rows[recognized.to(destination_rows.device)] + values = self._values[target.control_part] + valid = self._valid[target.control_part] + values[selected_rows] = is_grasp[recognized].to(values.device) + valid[selected_rows] = True + + def cancelled(self, targets: tuple[RuntimeEndpointTarget, ...]) -> None: + """Invalidate every row owned by cancelled control-part targets.""" + if not isinstance(targets, tuple) or not all( + isinstance(target, RuntimeEndpointTarget) for target in targets + ): + raise TypeError("targets must contain RuntimeEndpointTarget values.") + for target in targets: + if isinstance(target, JointPositionTarget): + self._clear_control_part(target.control_part) + + def discarded(self) -> None: + """Invalidate all command-derived state after a fail-closed discard.""" + for control_part in self._commands: + self._clear_control_part(control_part) + + def observe( + self, + query: BinaryEffectEvidenceQuery, + context: EffectEvidenceCollectionContext, + ) -> BinaryEffectObservation: + """Return selected command-state rows for one typed binary query.""" + if type(query) is not BinaryEffectEvidenceQuery: + raise TypeError("query must be exactly BinaryEffectEvidenceQuery.") + if type(context) is not EffectEvidenceCollectionContext: + raise TypeError("context must be exactly EffectEvidenceCollectionContext.") + address = query.source.address + if type(address) is not ControlPartEvidenceAddress: + raise TypeError( + "Command-state evidence requires ControlPartEvidenceAddress." + ) + rows = self._rows(context.env_ids) + values = self._values.get(address.control_part) + valid = self._valid.get(address.control_part) + if values is None or valid is None: + missing = torch.zeros( + context.env_ids.numel(), + dtype=torch.bool, + device=context.env_ids.device, + ) + return BinaryEffectObservation( + values=missing, + valid=missing, + acquisition_errors=( + f"Control part {address.control_part!r} has no exact open/grasp " + "command-state profile.", + ) + * int(context.env_ids.numel()), + ) + indices = torch.tensor(rows, dtype=torch.long, device=values.device) + selected_values = values.index_select(0, indices).to(context.env_ids.device) + selected_valid = valid.index_select(0, indices).to(context.env_ids.device) + errors = tuple( + ( + None + if bool(row_valid) + else "No accepted open/grasp command has established this row's state." + ) + for row_valid in selected_valid.detach().cpu().tolist() + ) + return BinaryEffectObservation( + values=selected_values, + valid=selected_valid, + acquisition_errors=errors, + ) + + def __call__( + self, + query: BinaryEffectEvidenceQuery, + context: EffectEvidenceCollectionContext, + ) -> BinaryEffectObservation: + """Delegate callback use to :meth:`observe`.""" + return self.observe(query, context) + + def _rows(self, env_ids: torch.Tensor) -> tuple[int, ...]: + """Resolve stable correlation IDs to full simulation row indices.""" + if ( + not isinstance(env_ids, torch.Tensor) + or env_ids.dtype != torch.long + or env_ids.dim() != 1 + or env_ids.numel() == 0 + ): + raise ValueError( + "env_ids must be a non-empty one-dimensional int64 tensor." + ) + if env_ids.device != self._env_ids.device: + raise ValueError("env_ids must share the tracker device.") + if torch.unique(env_ids).numel() != env_ids.numel(): + raise ValueError("env_ids must be unique.") + try: + return tuple( + self._row_by_env_id[int(env_id)] + for env_id in env_ids.detach().cpu().tolist() + ) + except KeyError as exc: + raise ValueError( + f"Environment ID {int(exc.args[0])} is absent from tracker env_ids." + ) from exc + + def _clear_control_part(self, control_part: str) -> None: + """Fail-closed reset one tracked control part when present.""" + values = self._values.get(control_part) + valid = self._valid.get(control_part) + if values is not None and valid is not None: + values.zero_() + valid.zero_() + + +class SimulationExpertProgramEnvironment(Protocol): + """Minimal Gym environment surface used by the simulation factory.""" + + sim: SimulationManager + robot: Robot + + @property + def step_dt(self) -> float: + """Return the authoritative Gym control cadence in seconds.""" + + +def _positive_finite(value: float, *, field_name: str) -> float: + """Validate one positive finite real number.""" + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError(f"{field_name} must be a real number.") + normalized = float(value) + if not math.isfinite(normalized) or normalized <= 0.0: + raise ValueError(f"{field_name} must be finite and positive.") + return normalized + + +def _non_negative_finite(value: float, *, field_name: str) -> float: + """Validate one non-negative finite real number.""" + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError(f"{field_name} must be a real number.") + normalized = float(value) + if not math.isfinite(normalized) or normalized < 0.0: + raise ValueError(f"{field_name} must be finite and non-negative.") + return normalized + + +def _robot_uid(robot: Robot) -> str: + """Return one strict live robot UID.""" + uid = getattr(robot, "uid", None) + if type(uid) is not str or not uid or uid != uid.strip(): + raise ValueError( + "robot.uid must be a non-empty string without outer whitespace." + ) + return uid + + +def _full_robot_tensor( + robot: Robot, + getter_name: str, + *, + required: bool, + reference: torch.Tensor | None = None, +) -> torch.Tensor | None: + """Read and validate one full-robot floating state tensor.""" + getter = getattr(robot, getter_name, None) + if not callable(getter): + if required: + raise TypeError(f"robot must provide {getter_name}().") + return None + value = getter() + if not isinstance(value, torch.Tensor): + raise TypeError(f"robot.{getter_name}() must return a torch.Tensor.") + if not value.is_floating_point() or value.dim() != 2: + raise ValueError( + f"robot.{getter_name}() must return floating shape (B, robot_dof)." + ) + if value.shape[0] == 0 or value.shape[1] == 0: + raise ValueError(f"robot.{getter_name}() dimensions must be non-zero.") + if reference is not None and ( + value.shape != reference.shape or value.device != reference.device + ): + raise ValueError( + f"robot.{getter_name}() must match robot.get_qpos() shape and device." + ) + if not bool(torch.isfinite(value).all().item()): + raise ValueError(f"robot.{getter_name}() must contain only finite values.") + return value.clone() + + +class SharedTickSceneProvider(SceneProvider): + """Share one immutable scene snapshot across consumers in the same tick. + + ``RegistrySceneProvider`` is stateful: every call observes native entities + and updates material-change baselines. Planning observations and multiple + evidence providers can legitimately request the same timestamp. This + wrapper always delegates one full-batch request per tick, then returns the + exact snapshot or an owned ordered-row projection to later consumers. + """ + + def __init__( + self, + delegate: RegistrySceneProvider, + full_env_ids: torch.Tensor, + ) -> None: + if type(delegate) is not RegistrySceneProvider: + raise TypeError("delegate must be exactly RegistrySceneProvider.") + if ( + not isinstance(full_env_ids, torch.Tensor) + or full_env_ids.dtype != torch.long + or full_env_ids.dim() != 1 + or full_env_ids.numel() == 0 + ): + raise ValueError("full_env_ids must be a non-empty 1D int64 tensor.") + if torch.unique(full_env_ids).numel() != full_env_ids.numel(): + raise ValueError("full_env_ids must be unique.") + self._delegate = delegate + self._full_env_ids = full_env_ids.clone() + self._row_by_env_id = { + int(env_id): row + for row, env_id in enumerate(full_env_ids.detach().cpu().tolist()) + } + self._timestamp: float | None = None + self._snapshot: SceneSnapshot | None = None + + @property + def delegate(self) -> RegistrySceneProvider: + """Return the authoritative stateful registry provider.""" + return self._delegate + + @property + def collision_entity_ids(self) -> tuple[str, ...]: + """Return canonical dynamic collision IDs from the delegate.""" + return self._delegate.collision_entity_ids + + @property + def full_env_ids(self) -> torch.Tensor: + """Return the authoritative full simulation batch order.""" + return self._full_env_ids.clone() + + def snapshot( + self, + *, + timestamp: float, + env_ids: torch.Tensor, + ) -> SceneSnapshot: + """Return the single shared snapshot for ``timestamp`` and ``env_ids``.""" + if isinstance(timestamp, bool) or not isinstance(timestamp, (int, float)): + raise TypeError("timestamp must be a real number.") + normalized_timestamp = float(timestamp) + if not math.isfinite(normalized_timestamp) or normalized_timestamp < 0.0: + raise ValueError("timestamp must be finite and non-negative.") + if ( + not isinstance(env_ids, torch.Tensor) + or env_ids.dtype != torch.long + or env_ids.dim() != 1 + or env_ids.numel() == 0 + ): + raise ValueError("env_ids must be a non-empty 1D int64 tensor.") + if torch.unique(env_ids).numel() != env_ids.numel(): + raise ValueError("env_ids must be unique.") + if env_ids.device != self._full_env_ids.device: + raise ValueError("env_ids must share the full simulation batch device.") + try: + rows = tuple( + self._row_by_env_id[int(env_id)] + for env_id in env_ids.detach().cpu().tolist() + ) + except KeyError as exc: + raise ValueError( + f"Environment ID {int(exc.args[0])} is absent from full_env_ids." + ) from exc + + if self._timestamp is not None: + if normalized_timestamp < self._timestamp: + raise ValueError("Shared scene snapshot timestamps must be monotonic.") + if normalized_timestamp == self._timestamp: + assert self._snapshot is not None + return self._select_rows(self._snapshot, rows) + + snapshot = self._delegate.snapshot( + timestamp=normalized_timestamp, + env_ids=self._full_env_ids.clone(), + ) + if not isinstance(snapshot, SceneSnapshot): + raise TypeError( + "RegistrySceneProvider.snapshot() must return SceneSnapshot." + ) + if snapshot.timestamp != normalized_timestamp: + raise ValueError("Scene snapshot timestamp must match the requested tick.") + self._timestamp = normalized_timestamp + self._snapshot = snapshot + return self._select_rows(snapshot, rows) + + def _select_rows( + self, + snapshot: SceneSnapshot, + rows: tuple[int, ...], + ) -> SceneSnapshot: + """Project one cached full-batch snapshot to an ordered row subset.""" + full_size = int(self._full_env_ids.numel()) + if rows == tuple(range(full_size)): + return snapshot + entities: dict[str, EntityState] = {} + for entity_id, state in snapshot.entities.items(): + pose = state.pose + if pose.dim() == 3: + if pose.shape[0] != full_size: + raise ValueError( + f"Scene entity {entity_id!r} batch does not match " + "full_env_ids." + ) + index = torch.tensor(rows, dtype=torch.long, device=pose.device) + pose = pose.index_select(0, index) + entities[entity_id] = EntityState(pose, confidence=state.confidence) + + articulation_joints: dict[tuple[str, str], ObservedArticulationJointState] = {} + for address, state in snapshot.articulation_joints.items(): + position = state.position + valid = state.valid_mask + if position.dim() == 2: + if position.shape[0] != full_size: + raise ValueError( + f"Scene articulation joint {address!r} batch does not " + "match full_env_ids." + ) + index = torch.tensor(rows, dtype=torch.long, device=position.device) + position = position.index_select(0, index) + if valid is not None: + valid = valid.index_select(0, index.to(valid.device)) + articulation_joints[address] = ObservedArticulationJointState( + position, + valid, + ) + + revisions = snapshot.collision_world_revisions(full_size) + return SceneSnapshot( + timestamp=snapshot.timestamp, + version=snapshot.version, + entities=entities, + collision_world_revision=tuple(revisions[row] for row in rows), + collision_entity_ids=snapshot.collision_entity_ids, + articulation_joints=articulation_joints, + ) + + +class SimulationPlanningObservationProvider(GymPlanningObservationProvider): + """Gym planning observations backed by live robot and shared scene state.""" + + def __init__( + self, + robot: Robot, + scene_provider: SharedTickSceneProvider, + clock: EnvironmentStepClock, + env_ids: torch.Tensor, + command_state_tracker: ControlCommandStateEvidenceTracker, + *, + owner_token: object, + ) -> None: + if type(scene_provider) is not SharedTickSceneProvider: + raise TypeError("scene_provider must be exactly SharedTickSceneProvider.") + if type(clock) is not EnvironmentStepClock: + raise TypeError("clock must be exactly EnvironmentStepClock.") + qpos = _full_robot_tensor(robot, "get_qpos", required=True) + assert qpos is not None + 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 qpos must share a device.") + if torch.unique(env_ids).numel() != env_ids.numel(): + raise ValueError("env_ids must be unique.") + if type(command_state_tracker) is not ControlCommandStateEvidenceTracker: + raise TypeError( + "command_state_tracker must be exactly " + "ControlCommandStateEvidenceTracker." + ) + self._robot = robot + self._scene_provider = scene_provider + self._clock = clock + self._env_ids = env_ids.clone() + self._command_state_tracker = command_state_tracker + self._owner_token = owner_token + super().__init__(self._capture) + + @property + def scene_provider(self) -> SharedTickSceneProvider: + """Return the snapshot-sharing scene provider used by evidence ports.""" + return self._scene_provider + + @property + def env_ids(self) -> torch.Tensor: + """Return stable ordered simulation row IDs.""" + return self._env_ids.clone() + + @property + def command_state_tracker(self) -> ControlCommandStateEvidenceTracker: + """Return the runtime-local accepted-command evidence owner.""" + return self._command_state_tracker + + def is_owned_by(self, owner_token: object) -> bool: + """Return whether this provider belongs to one factory instance.""" + return self._owner_token is owner_token + + def _capture(self, task_state: TaskState) -> PlanningContext: + """Capture one synchronized robot and scene observation.""" + qpos = _full_robot_tensor(self._robot, "get_qpos", required=True) + assert qpos is not None + if ( + qpos.shape[0] != self._env_ids.numel() + or qpos.device != self._env_ids.device + ): + raise ValueError("Robot batch shape or device changed after assembly.") + qvel = _full_robot_tensor( + self._robot, + "get_qvel", + required=False, + reference=qpos, + ) + if qvel is None: + qvel = torch.zeros_like(qpos) + qeffort = _full_robot_tensor( + self._robot, + "get_qf", + required=False, + reference=qpos, + ) + timestamp = self._clock.now() + scene = self._scene_provider.snapshot( + timestamp=timestamp, + env_ids=self._env_ids.clone(), + ) + return PlanningContext( + robot=RobotObservation( + timestamp=timestamp, + qpos=qpos, + qvel=qvel, + qeffort=qeffort, + ), + task=task_state, + scene=scene, + env_ids=self._env_ids, + ) + + +class SimulationExpertProgramFactory(ExpertProgramEnvironmentFactory): + """Build every live Expert Program component from explicit declarations. + + Args: + simulation: Exact live simulation that owns ``robot`` and scene UIDs. + robot: Exact robot selected for planning and evidence acquisition. + scene_binding: Canonical-to-native scene declaration. + robot_profile_binding: Typed robot resource and policy declaration. + step_dt: Authoritative Gym control cadence. + planner_cfg: Explicit planner configuration. ``None`` selects TOPPRA + for ``robot.uid``. + motion_generator_factory: Optional fresh-generator factory. It is + mutually exclusive with ``planner_cfg`` and intended for custom + planners and isolated tests. + endpoint_adapters: Explicit adapters for non-built-in resource endpoint + types. + settle_presets: Optional named segment settling policies. + translation_threshold: Material scene translation threshold. + rotation_threshold: Material scene rotation threshold. + contact_observer: Optional raw contact evidence callback. + constraint_observer: Optional raw constraint evidence callback. + force_observer: Optional raw force evidence callback. + wrench_observer: Optional raw wrench evidence callback. + + Every profile policy is rebuilt with ``control_dt == step_dt``. The Gym + cadence is authoritative because commands cannot be emitted between + environment steps; silently retaining a preset's unrelated fallback + cadence would make trajectory timing unrepresentable at the bridge. + """ + + def __init__( + self, + simulation: SimulationManager, + robot: Robot, + scene_binding: SimulationSceneBinding, + robot_profile_binding: SimulationRobotSkillProfileBinding, + *, + step_dt: float, + planner_cfg: BasePlannerCfg | None = None, + motion_generator_factory: MotionGeneratorFactory | None = None, + endpoint_adapters: ( + Mapping[type[ResourceEndpoint], ResourceEndpointAdapter] | None + ) = None, + settle_presets: Mapping[str, DynamicSettleMonitorCfg] | None = None, + translation_threshold: float = 1.0e-4, + rotation_threshold: float = 1.0e-3, + contact_observer: BinaryObservationCallback | None = None, + constraint_observer: BinaryObservationCallback | None = None, + force_observer: ScalarObservationCallback | None = None, + wrench_observer: ScalarObservationCallback | None = None, + ) -> None: + if type(scene_binding) is not SimulationSceneBinding: + raise TypeError("scene_binding must be exactly SimulationSceneBinding.") + if type(robot_profile_binding) is not SimulationRobotSkillProfileBinding: + raise TypeError( + "robot_profile_binding must be exactly " + "SimulationRobotSkillProfileBinding." + ) + if planner_cfg is not None and motion_generator_factory is not None: + raise ValueError( + "planner_cfg and motion_generator_factory are mutually exclusive." + ) + if planner_cfg is not None and not isinstance(planner_cfg, BasePlannerCfg): + raise TypeError("planner_cfg must be a BasePlannerCfg or None.") + if motion_generator_factory is not None and not callable( + motion_generator_factory + ): + raise TypeError("motion_generator_factory must be callable or None.") + if endpoint_adapters is not None and not isinstance(endpoint_adapters, Mapping): + raise TypeError("endpoint_adapters must be a mapping or None.") + for name, callback in ( + ("contact_observer", contact_observer), + ("constraint_observer", constraint_observer), + ("force_observer", force_observer), + ("wrench_observer", wrench_observer), + ): + if callback is not None and not callable(callback): + raise TypeError(f"{name} must be callable or None.") + + robot_uid = _robot_uid(robot) + get_robot = getattr(simulation, "get_robot", None) + if not callable(get_robot): + raise TypeError("simulation must provide get_robot().") + if get_robot(robot_uid) is not robot: + raise ValueError( + f"simulation.get_robot({robot_uid!r}) must return the exact " + "selected robot." + ) + selected_planner_cfg = deepcopy(planner_cfg) + if ( + selected_planner_cfg is not None + and selected_planner_cfg.robot_uid != robot_uid + ): + raise ValueError( + f"planner_cfg.robot_uid must equal selected robot UID {robot_uid!r}." + ) + + self._simulation = simulation + self._robot = robot + self._scene_binding = scene_binding + self._robot_profile_binding = robot_profile_binding + self._step_dt = _positive_finite(step_dt, field_name="step_dt") + self._planner_cfg = selected_planner_cfg + self._motion_generator_factory = motion_generator_factory + self._endpoint_adapters = ( + None if endpoint_adapters is None else dict(endpoint_adapters) + ) + self._translation_threshold = _non_negative_finite( + translation_threshold, + field_name="translation_threshold", + ) + self._rotation_threshold = _non_negative_finite( + rotation_threshold, + field_name="rotation_threshold", + ) + self._contact_observer = contact_observer + self._constraint_observer = constraint_observer + self._force_observer = force_observer + self._wrench_observer = wrench_observer + self._owner_token = object() + + qpos = _full_robot_tensor(robot, "get_qpos", required=True) + assert qpos is not None + self._env_ids = torch.arange( + qpos.shape[0], + dtype=torch.long, + device=qpos.device, + ) + self._segment_policy_port = SimulationSegmentPolicyPort( + simulation, + robot, + scene_binding, + settle_presets=settle_presets, + env_ids=self._env_ids, + ) + + @classmethod + def from_environment( + cls, + environment: SimulationExpertProgramEnvironment, + *, + scene_binding: SimulationSceneBinding, + robot_profile_binding: SimulationRobotSkillProfileBinding, + planner_cfg: BasePlannerCfg | None = None, + motion_generator_factory: MotionGeneratorFactory | None = None, + endpoint_adapters: ( + Mapping[type[ResourceEndpoint], ResourceEndpointAdapter] | None + ) = None, + settle_presets: Mapping[str, DynamicSettleMonitorCfg] | None = None, + translation_threshold: float = 1.0e-4, + rotation_threshold: float = 1.0e-3, + contact_observer: BinaryObservationCallback | None = None, + constraint_observer: BinaryObservationCallback | None = None, + force_observer: ScalarObservationCallback | None = None, + wrench_observer: ScalarObservationCallback | None = None, + ) -> SimulationExpertProgramFactory: + """Create a factory from the explicit standard Gym environment surface.""" + simulation = getattr(environment, "sim", None) + robot = getattr(environment, "robot", None) + try: + step_dt = environment.step_dt + except AttributeError as exc: + raise TypeError("environment must expose step_dt.") from exc + if simulation is None or robot is None: + raise TypeError("environment must expose non-None sim and robot values.") + return cls( + simulation, + robot, + scene_binding, + robot_profile_binding, + step_dt=step_dt, + planner_cfg=planner_cfg, + motion_generator_factory=motion_generator_factory, + endpoint_adapters=endpoint_adapters, + settle_presets=settle_presets, + translation_threshold=translation_threshold, + rotation_threshold=rotation_threshold, + contact_observer=contact_observer, + constraint_observer=constraint_observer, + force_observer=force_observer, + wrench_observer=wrench_observer, + ) + + @property + def scene_registry_id(self) -> str: + """Return the exact configured scene-registry ID.""" + return self._scene_binding.registry_id + + @property + def robot_profile_id(self) -> str: + """Return the exact configured robot-profile ID.""" + return self._robot_profile_binding.profile_id + + @property + def step_dt(self) -> float: + """Return the authoritative Gym control cadence.""" + return self._step_dt + + @property + def segment_policy_port(self) -> SimulationSegmentPolicyPort: + """Return the shared simulation post-policy and validator port.""" + return self._segment_policy_port + + @property + def endpoint_adapters( + self, + ) -> Mapping[type[ResourceEndpoint], ResourceEndpointAdapter] | None: + """Return an owned copy of installed custom endpoint adapters.""" + return ( + None if self._endpoint_adapters is None else dict(self._endpoint_adapters) + ) + + def create_scene_registry(self) -> SceneRegistry: + """Build one fresh authoritative registry from explicit bindings.""" + return self._scene_binding.build(self._simulation) + + def create_robot_skill_profile(self) -> RobotSkillProfile: + """Build a profile whose every motion policy uses the Gym cadence.""" + profile = self._robot_profile_binding.build(self._robot) + aligned_presets = { + preset_id: SkillPolicyPreset( + preset_id=preset.preset_id, + schema_version=preset.schema_version, + motion_policy=replace( + preset.motion_policy, + control_dt=self._step_dt, + ), + recovery_policy=preset.recovery_policy, + runner_cfg=preset.runner_cfg, + effect_monitors=preset.effect_monitors, + ) + for preset_id, preset in profile.presets.items() + } + aligned = replace(profile, presets=aligned_presets) + if any( + preset.motion_policy.control_dt != self._step_dt + for preset in aligned.presets.values() + ): + raise AssertionError("Profile motion policies were not cadence-aligned.") + return aligned + + def create_atomic_action_engine( + self, + profile: RobotSkillProfile, + ) -> AtomicActionEngine: + """Create a fresh engine around the selected planner and exact profile.""" + if not isinstance(profile, RobotSkillProfile): + raise TypeError("profile must be a RobotSkillProfile.") + if profile.profile_id != self.robot_profile_id: + raise ValueError( + f"profile ID must be {self.robot_profile_id!r}, got " + f"{profile.profile_id!r}." + ) + motion_generator = self._create_motion_generator() + if motion_generator.robot is not self._robot: + raise ValueError( + "Motion generator must own the exact robot selected by the factory." + ) + return AtomicActionEngine( + motion_generator, + skill_profile=profile, + endpoint_adapters=self._endpoint_adapters, + ) + + def create_planning_observation_provider( + self, + *, + scene_registry: SceneRegistry, + engine: AtomicActionEngine, + clock: EnvironmentStepClock, + ) -> PlanningObservationPort: + """Create one planning port and planner-validated shared scene provider.""" + if type(scene_registry) is not SceneRegistry: + raise TypeError("scene_registry must be exactly SceneRegistry.") + if not isinstance(engine, AtomicActionEngine): + raise TypeError("engine must be an AtomicActionEngine.") + if engine.robot is not self._robot: + raise ValueError("engine must own the exact factory robot.") + if type(clock) is not EnvironmentStepClock: + raise TypeError("clock must be exactly EnvironmentStepClock.") + if clock.step_dt != self._step_dt: + raise ValueError("clock.step_dt must equal the factory Gym cadence.") + provider = scene_registry.make_planning_scene_provider( + engine.motion_generator, + batch_size=int(self._env_ids.numel()), + translation_threshold=self._translation_threshold, + rotation_threshold=self._rotation_threshold, + ) + shared = SharedTickSceneProvider(provider, self._env_ids) + command_state_tracker = ControlCommandStateEvidenceTracker( + engine.control_profiles, + self._env_ids, + ) + return SimulationPlanningObservationProvider( + self._robot, + shared, + clock, + self._env_ids, + command_state_tracker, + owner_token=self._owner_token, + ) + + def create_effect_evidence_providers( + self, + *, + scene_registry: SceneRegistry, + engine: AtomicActionEngine, + observation_provider: PlanningObservationPort, + ) -> Iterable[EffectEvidenceProvider]: + """Create built-in control-part and articulation evidence providers.""" + if type(scene_registry) is not SceneRegistry: + raise TypeError("scene_registry must be exactly SceneRegistry.") + if not isinstance(engine, AtomicActionEngine): + raise TypeError("engine must be an AtomicActionEngine.") + if engine.robot is not self._robot: + raise ValueError("engine must own the exact factory robot.") + if type(observation_provider) is not SimulationPlanningObservationProvider: + raise TypeError( + "observation_provider must be exactly " + "SimulationPlanningObservationProvider." + ) + if not observation_provider.is_owned_by(self._owner_token): + raise ValueError("observation_provider belongs to another factory.") + scene_provider = observation_provider.scene_provider + command_state_tracker = observation_provider.command_state_tracker + contact_observer = self._contact_observer or command_state_tracker + constraint_observer = self._constraint_observer or command_state_tracker + providers: list[EffectEvidenceProvider] = [] + if isinstance(self._robot, ControlPartRobotEvidenceSource): + providers.append( + ControlPartSimulationEvidenceProvider( + self._robot, + scene_provider=scene_provider, + contact_observer=contact_observer, + constraint_observer=constraint_observer, + force_observer=self._force_observer, + wrench_observer=self._wrench_observer, + ) + ) + providers.append( + SceneArticulationEvidenceProvider(scene_provider=scene_provider) + ) + return tuple(providers) + + def create_accepted_runtime_command_observer( + self, + *, + scene_registry: SceneRegistry, + engine: AtomicActionEngine, + observation_provider: PlanningObservationPort, + ) -> AcceptedRuntimeCommandObserver: + """Return the tracker already shared with this runtime's evidence ports.""" + if type(scene_registry) is not SceneRegistry: + raise TypeError("scene_registry must be exactly SceneRegistry.") + if not isinstance(engine, AtomicActionEngine): + raise TypeError("engine must be an AtomicActionEngine.") + if engine.robot is not self._robot: + raise ValueError("engine must own the exact factory robot.") + if type(observation_provider) is not SimulationPlanningObservationProvider: + raise TypeError( + "observation_provider must be exactly " + "SimulationPlanningObservationProvider." + ) + if not observation_provider.is_owned_by(self._owner_token): + raise ValueError("observation_provider belongs to another factory.") + return observation_provider.command_state_tracker + + def create_adapter( + self, + *, + call_catalog: SemanticCallCatalog | None = None, + registered_lowerers: Iterable[RegisteredSemanticLowerer] = (), + relation_grounders: Iterable[RelationTargetGrounder] = (), + handover_pose_providers: Iterable[HandOverPoseProvider] = (), + effect_monitor_registry: EffectMonitorRegistry | None = None, + runtime_transports: Iterable[RuntimeTransportActionEncoder] = (), + runner_cfg: ExecutionRunnerCfg | None = None, + parallel_safety_validator: ParallelCommandSafetyValidator | None = None, + ) -> ExpertProgramEnvironmentAdapter: + """Create the exact Gym adapter with shared simulation policy ports.""" + return ExpertProgramEnvironmentAdapter( + self, + step_dt=self._step_dt, + call_catalog=call_catalog, + endpoint_adapters=self._endpoint_adapters, + registered_lowerers=registered_lowerers, + relation_grounders=relation_grounders, + handover_pose_providers=handover_pose_providers, + effect_monitor_registry=effect_monitor_registry, + runtime_transports=runtime_transports, + runner_cfg=runner_cfg, + post_policy_port=self._segment_policy_port, + validator_port=self._segment_policy_port, + parallel_safety_validator=parallel_safety_validator, + ) + + def _create_motion_generator(self) -> MotionGenerator: + """Create and validate one exact motion generator.""" + if self._motion_generator_factory is not None: + generator = self._motion_generator_factory() + else: + planner_cfg = ( + ToppraPlannerCfg(robot_uid=_robot_uid(self._robot)) + if self._planner_cfg is None + else deepcopy(self._planner_cfg) + ) + generator = MotionGenerator(MotionGenCfg(planner_cfg=planner_cfg)) + if not isinstance(generator, MotionGenerator): + raise TypeError( + "motion_generator_factory must return a MotionGenerator instance." + ) + return generator + + +def create_simulation_expert_program_adapter( + environment: SimulationExpertProgramEnvironment, + *, + scene_binding: SimulationSceneBinding, + robot_profile_binding: SimulationRobotSkillProfileBinding, + planner_cfg: BasePlannerCfg | None = None, + motion_generator_factory: MotionGeneratorFactory | None = None, + endpoint_adapters: ( + Mapping[type[ResourceEndpoint], ResourceEndpointAdapter] | None + ) = None, + relation_grounders: Iterable[RelationTargetGrounder] = (), + handover_pose_providers: Iterable[HandOverPoseProvider] = (), + runtime_transports: Iterable[RuntimeTransportActionEncoder] = (), + settle_presets: Mapping[str, DynamicSettleMonitorCfg] | None = None, + translation_threshold: float = 1.0e-4, + rotation_threshold: float = 1.0e-3, + contact_observer: BinaryObservationCallback | None = None, + constraint_observer: BinaryObservationCallback | None = None, + force_observer: ScalarObservationCallback | None = None, + wrench_observer: ScalarObservationCallback | None = None, + parallel_safety_validator: ParallelCommandSafetyValidator | None = None, +) -> ExpertProgramEnvironmentAdapter: + """Create a complete production adapter from one standard Gym environment. + + This is the intended task-side one-line integration. Relation-target + grounders and embodiment-owned handover pose providers are explicit and + default to empty collections, so calls that require an uninstalled provider + remain fail-closed during program preflight. Advanced callers can retain + :class:`SimulationExpertProgramFactory` and call ``create_adapter`` directly + to install registered semantic lowerers or custom monitors. Custom endpoint + adapters and their matching Gym runtime transports are accepted here so a + non-joint endpoint remains executable through the one-line path. + + Args: + environment: Standard Gym simulation environment exposing ``sim``, + ``robot``, and ``step_dt``. + scene_binding: Authoritative typed scene declaration. + robot_profile_binding: Typed robot resource and policy declaration. + planner_cfg: Optional planner configuration owned by the factory. + motion_generator_factory: Optional factory for one fresh motion generator. + endpoint_adapters: Optional exact-type custom endpoint adapters. + relation_grounders: Explicit typed relation-target grounders. + handover_pose_providers: Explicit embodiment-owned handover pose providers. + runtime_transports: Additional runtime-command-to-Gym encoders. + settle_presets: Optional named dynamic-settling policies. + translation_threshold: Scene translation revision threshold. + rotation_threshold: Scene rotation revision threshold. + contact_observer: Optional raw contact evidence callback. + constraint_observer: Optional raw constraint evidence callback. + force_observer: Optional raw force evidence callback. + wrench_observer: Optional raw wrench evidence callback. + parallel_safety_validator: Optional authoritative parallel-command gate. + + Returns: + Complete production Expert Program environment adapter. + """ + factory = SimulationExpertProgramFactory.from_environment( + environment, + scene_binding=scene_binding, + robot_profile_binding=robot_profile_binding, + planner_cfg=planner_cfg, + motion_generator_factory=motion_generator_factory, + endpoint_adapters=endpoint_adapters, + settle_presets=settle_presets, + translation_threshold=translation_threshold, + rotation_threshold=rotation_threshold, + contact_observer=contact_observer, + constraint_observer=constraint_observer, + force_observer=force_observer, + wrench_observer=wrench_observer, + ) + return factory.create_adapter( + relation_grounders=relation_grounders, + handover_pose_providers=handover_pose_providers, + runtime_transports=runtime_transports, + parallel_safety_validator=parallel_safety_validator, + ) + + +__all__ = [ + "ControlCommandStateEvidenceTracker", + "MotionGeneratorFactory", + "SharedTickSceneProvider", + "SimulationExpertProgramEnvironment", + "SimulationExpertProgramFactory", + "SimulationPlanningObservationProvider", + "create_simulation_expert_program_adapter", +] diff --git a/embodichain/lab/gym/envs/expert_program/simulation_policies.py b/embodichain/lab/gym/envs/expert_program/simulation_policies.py new file mode 100644 index 000000000..c18dace2c --- /dev/null +++ b/embodichain/lab/gym/envs/expert_program/simulation_policies.py @@ -0,0 +1,715 @@ +# ---------------------------------------------------------------------------- +# 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-backed Expert Program post-policies and validators. + +The port in this module deliberately consumes the same explicit +:class:`SimulationSceneBinding` used to construct the semantic scene registry. +It never scans a simulation or guesses a native entity from a canonical name. +Post-policy actions are full-qpos holds and therefore remain inside the normal +Gym ``env.step()`` path owned by :class:`AtomicDemoBridge`. +""" + +from __future__ import annotations + +from collections.abc import Iterator, Mapping +from copy import deepcopy +from dataclasses import dataclass +import math +from types import MappingProxyType +from typing import Any, TYPE_CHECKING + +import torch + +from embodichain.lab.gym.envs.settling import ( + DynamicSettleMonitor, + DynamicSettleMonitorCfg, + DynamicSettleSample, + DynamicSettleState, +) + +from .compiler import ( + CompiledPostPolicy, + CompiledProgramSegment, + CompiledProgramValidator, +) +from .simulation import SimulationSceneBinding + +if TYPE_CHECKING: + from embodichain.lab.sim.objects import Robot + from embodichain.lab.sim.sim_manager import SimulationManager + + +@dataclass(frozen=True, slots=True) +class _SimulationSettleTarget: + """One canonical entity resolved through an explicit native binding.""" + + canonical_id: str + kind: str + native_entity: Any + + +def _default_settle_presets() -> Mapping[str, DynamicSettleMonitorCfg]: + """Return independently owned built-in post-policy presets.""" + return MappingProxyType( + { + "rigid_object": DynamicSettleMonitorCfg( + linear_velocity_threshold=0.03, + angular_velocity_threshold=0.20, + min_steps=10, + max_steps=240, + check_interval_steps=2, + required_stable_checks=3, + ) + } + ) + + +def _json_speed_values(value: torch.Tensor) -> list[float | None]: + """Convert speed evidence to finite JSON numbers or explicit unknowns.""" + return [ + float(item) if math.isfinite(float(item)) else None + for item in value.detach().cpu().tolist() + ] + + +class SimulationSegmentPolicyPort: + """Execute built-in segment policies against explicitly bound simulation data. + + Args: + simulation: Live simulation used only for UIDs declared in + ``scene_binding``. + robot: Live robot used to produce controller-safe full-qpos holds. + scene_binding: Exact canonical-to-native scene declaration. + settle_presets: Named settling policies. ``None`` installs the shared + ``rigid_object`` preset. + env_ids: Optional stable logical row IDs. They describe correlation, + not simulator row indices; simulator rows remain ordered exactly as + returned by the robot and bound entities. + + The same instance implements both ``SegmentPostPolicyPort`` and + ``SegmentValidatorPort``. Unknown policy types, presets, canonical IDs, or + native entities fail before an action is emitted. + """ + + def __init__( + self, + simulation: SimulationManager, + robot: Robot, + scene_binding: SimulationSceneBinding, + *, + settle_presets: Mapping[str, DynamicSettleMonitorCfg] | None = None, + env_ids: torch.Tensor | None = None, + ) -> None: + if type(scene_binding) is not SimulationSceneBinding: + raise TypeError("scene_binding must be exactly SimulationSceneBinding.") + qpos = self._read_robot_qpos(robot) + 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): + raise TypeError("env_ids must be a torch.Tensor or None.") + if env_ids.dtype != torch.long or env_ids.shape != (qpos.shape[0],): + raise ValueError("env_ids must be int64 with one ID per simulator row.") + if env_ids.device != qpos.device: + raise ValueError("env_ids and robot qpos must share a device.") + if torch.unique(env_ids).numel() != env_ids.numel(): + raise ValueError("env_ids must contain unique values.") + + selected_presets = ( + _default_settle_presets() if settle_presets is None else settle_presets + ) + if not isinstance(selected_presets, Mapping) or not selected_presets: + raise ValueError("settle_presets must be a non-empty mapping.") + normalized_presets: dict[str, DynamicSettleMonitorCfg] = {} + for preset_id, cfg in selected_presets.items(): + if ( + type(preset_id) is not str + or not preset_id + or preset_id != preset_id.strip() + ): + raise ValueError( + "Settle preset IDs must be non-empty strings without outer " + "whitespace." + ) + if not isinstance(cfg, DynamicSettleMonitorCfg): + raise TypeError( + "settle_presets values must be DynamicSettleMonitorCfg values." + ) + normalized_presets[preset_id] = cfg.snapshot() + + self._simulation = simulation + self._robot = robot + self._scene_binding = scene_binding + self._env_ids = env_ids.clone() + self._row_indices = torch.arange( + qpos.shape[0], + dtype=torch.long, + device=qpos.device, + ) + self._settle_presets = MappingProxyType(normalized_presets) + self._settle_targets, self._rigid_objects = self._resolve_native_entities() + self._post_policy_results: dict[int, dict[str, object]] = {} + self._post_policy_success: dict[int, torch.Tensor] = {} + self._validator_results: dict[int, dict[str, object]] = {} + + @property + def settle_preset_ids(self) -> tuple[str, ...]: + """Return installed post-policy preset IDs in declaration order.""" + return tuple(self._settle_presets) + + def validate_policy( + self, + policy: Any, + *, + segment: Any, + ) -> None: + """Validate one post-policy against static bindings without observation. + + This method reads only the compiled declaration, installed preset + table, and entities resolved when the port was constructed. It never + samples velocity or qpos and never emits a controller action. + """ + if type(policy) is not CompiledPostPolicy: + raise TypeError("policy must be exactly CompiledPostPolicy.") + self._validate_segment_membership(segment, policy, kind="post policy") + if policy.cfg.kind != "wait_stable": + raise ValueError( + f"Unsupported compiled post-policy kind {policy.cfg.kind!r}." + ) + if policy.cfg.preset not in self._settle_presets: + raise KeyError( + f"Unknown settle preset {policy.cfg.preset!r}; available presets " + f"are {sorted(self._settle_presets)}." + ) + entity_id = policy.entity.entity_id + target = self._settle_targets.get(entity_id) + if target is None: + raise KeyError( + f"Canonical settle entity {entity_id!r} has no explicit native " + "dynamic binding." + ) + if target.kind == "rigid_object" and bool( + getattr(target.native_entity, "is_non_dynamic", False) + ): + raise ValueError( + f"Canonical settle entity {entity_id!r} is static or kinematic." + ) + + def actions( + self, + policy: Any, + *, + segment: Any, + active_mask: torch.Tensor, + ) -> Iterator[torch.Tensor]: + """Yield full-qpos hold actions until active rows settle or time out. + + Args: + policy: Exact compiled ``wait_stable`` policy. + segment: Exact segment that owns ``policy``. + active_mask: Rows that remain eligible after runtime execution and + preceding post-policies. Inactive rows are held safely but do + not participate in settling, timeout, or success results. + + Yields: + Fresh full-qpos hold commands consumed by ordinary ``env.step()``. + + Timeout is a normal row-local result boundary. Timed-out rows are + exposed through :meth:`post_policy_result` and + :meth:`post_policy_metadata`; no batch-level exception is raised. + """ + self.validate_policy(policy, segment=segment) + active_mask = self._validate_active_mask(active_mask) + preset = self._settle_presets[policy.cfg.preset] + entity_id = policy.entity.entity_id + target = self._settle_targets[entity_id] + + result_key = id(policy) + self._post_policy_results.pop(result_key, None) + self._post_policy_success.pop(result_key, None) + if not bool(active_mask.any().item()): + self._post_policy_success[result_key] = active_mask.clone() + self._post_policy_results[result_key] = { + "kind": policy.cfg.kind, + "entity_id": entity_id, + "preset": policy.cfg.preset, + "source_path": list(policy.source_path), + "status": "skipped", + "active_mask": active_mask.detach().cpu().tolist(), + "thresholds": self._settle_threshold_metadata(preset), + "state": self._empty_settle_state_metadata(active_mask), + } + return + + active_rows = self._row_indices[active_mask] + monitor = DynamicSettleMonitor(preset, self._env_ids[active_mask]) + elapsed_steps = 0 + while True: + state = monitor.observe( + (self._measure_settle_target(target, row_indices=active_rows),), + elapsed_steps=elapsed_steps, + ) + settled_mask = torch.zeros_like(active_mask) + settled_mask[active_mask] = state.settled_mask + self._post_policy_results[result_key] = { + "kind": policy.cfg.kind, + "entity_id": entity_id, + "preset": policy.cfg.preset, + "source_path": list(policy.source_path), + "active_mask": active_mask.detach().cpu().tolist(), + "status": ( + "settled" + if bool(state.settled_mask.all().item()) + else ( + "timed_out" + if bool(state.timeout_mask.any().item()) + else "running" + ) + ), + "thresholds": self._settle_threshold_metadata(preset), + "state": self._expand_settle_state_metadata(state, active_mask), + } + self._post_policy_success[result_key] = settled_mask + if bool(state.settled_mask.all().item()): + return + if bool(state.timeout_mask.any().item()): + return + yield self._read_robot_qpos(self._robot) + elapsed_steps += 1 + + def post_policy_result( + self, + policy: Any, + *, + segment: Any, + ) -> torch.Tensor: + """Return the latest independently owned per-row settling result.""" + if type(policy) is not CompiledPostPolicy: + raise TypeError("policy must be exactly CompiledPostPolicy.") + self._validate_segment_membership(segment, policy, kind="post policy") + result = self._post_policy_success.get(id(policy)) + if result is None: + raise RuntimeError("Post-policy result is unavailable before execution.") + return result.clone() + + def post_policy_metadata( + self, + policy: Any, + *, + segment: Any, + ) -> Mapping[str, object]: + """Return the latest JSON-safe settling trace for one policy. + + The trace is available after the policy generator has started. A + terminal trace has status ``"settled"``, ``"timed_out"``, or + ``"skipped"`` when no rows remain active; an early demo interruption + intentionally retains the latest ``"running"`` snapshot for diagnosis. + """ + if type(policy) is not CompiledPostPolicy: + raise TypeError("policy must be exactly CompiledPostPolicy.") + self._validate_segment_membership(segment, policy, kind="post policy") + metadata = self._post_policy_results.get(id(policy)) + if metadata is None: + raise RuntimeError("Post-policy metadata is unavailable before execution.") + return deepcopy(metadata) + + def validate_validator( + self, + validator: Any, + *, + segment: Any, + ) -> None: + """Validate one validator against static bindings without observation.""" + if type(validator) is not CompiledProgramValidator: + raise TypeError("validator must be exactly CompiledProgramValidator.") + self._validate_segment_membership(segment, validator, kind="validator") + if validator.cfg.kind != "object_near_target": + raise ValueError( + f"Unsupported compiled validator kind {validator.cfg.kind!r}." + ) + entity_id = validator.object.entity_id + if entity_id not in self._rigid_objects: + raise KeyError( + f"Canonical validator object {entity_id!r} has no explicit rigid-" + "object binding." + ) + + def validate(self, validator: Any, *, segment: Any) -> torch.Tensor: + """Observe an explicitly bound rigid object against a world target. + + Args: + validator: Exact compiled ``object_near_target`` validator. + segment: Exact segment that owns ``validator``. + + Returns: + Boolean tensor with one result per simulation row. + """ + self.validate_validator(validator, segment=segment) + entity_id = validator.object.entity_id + entity = self._rigid_objects[entity_id] + pose = self._read_pose(entity, entity_id=entity_id) + current_position = pose[:, :3, 3] + target_position = validator.target_pose.position.to( + device=current_position.device, + dtype=current_position.dtype, + ) + if target_position.dim() == 1: + target_position = target_position.unsqueeze(0).expand_as(current_position) + elif target_position.shape != current_position.shape: + raise ValueError( + "Validator target batch must be unbatched or match simulator rows." + ) + error = torch.linalg.vector_norm(current_position - target_position, dim=1) + accepted = torch.isfinite(error) & ( + error <= float(validator.cfg.position_tolerance) + ) + self._validator_results[id(validator)] = { + "kind": validator.cfg.kind, + "object_id": entity_id, + "target_id": validator.target_selection.target_id, + "target_value_index": validator.target_selection.value_index, + "source_path": list(validator.source_path), + "position_tolerance": float(validator.cfg.position_tolerance), + "env_ids": self._env_ids.detach().cpu().tolist(), + "object_position": current_position.detach().cpu().tolist(), + "target_position": target_position.detach().cpu().tolist(), + "position_error": error.detach().cpu().tolist(), + "accepted_mask": accepted.detach().cpu().tolist(), + } + return accepted + + def validator_metadata( + self, + validator: Any, + *, + segment: Any, + ) -> Mapping[str, object]: + """Return an owned JSON-safe trace for one completed validator.""" + if type(validator) is not CompiledProgramValidator: + raise TypeError("validator must be exactly CompiledProgramValidator.") + self._validate_segment_membership(segment, validator, kind="validator") + metadata = self._validator_results.get(id(validator)) + if metadata is None: + raise RuntimeError("Validator metadata is unavailable before validation.") + return deepcopy(metadata) + + @staticmethod + def _read_robot_qpos(robot: Robot) -> torch.Tensor: + """Capture one finite full-robot position batch.""" + get_qpos = getattr(robot, "get_qpos", None) + if not callable(get_qpos): + raise TypeError("robot must provide get_qpos().") + qpos = get_qpos() + if ( + not isinstance(qpos, torch.Tensor) + or not qpos.is_floating_point() + or qpos.dim() != 2 + or qpos.shape[0] == 0 + or qpos.shape[1] == 0 + ): + raise ValueError("robot.get_qpos() must return floating shape (B, J).") + if not bool(torch.isfinite(qpos).all().item()): + raise ValueError("robot.get_qpos() must contain finite values.") + return qpos.clone() + + def _validate_active_mask(self, active_mask: torch.Tensor) -> torch.Tensor: + """Return one owned row mask aligned with the simulator batch.""" + if not isinstance(active_mask, torch.Tensor): + raise TypeError("active_mask must be a torch.Tensor.") + if active_mask.dtype != torch.bool or active_mask.shape != self._env_ids.shape: + raise ValueError( + "active_mask must be bool with one value per simulator row." + ) + if active_mask.device != self._env_ids.device: + raise ValueError("active_mask and env_ids must share a device.") + return active_mask.clone() + + @staticmethod + def _settle_threshold_metadata( + preset: DynamicSettleMonitorCfg, + ) -> dict[str, float | int]: + """Serialize one settling preset without exposing mutable state.""" + return { + "linear_velocity": float(preset.linear_velocity_threshold), + "angular_velocity": float(preset.angular_velocity_threshold), + "min_steps": preset.min_steps, + "max_steps": preset.max_steps, + "check_interval_steps": preset.check_interval_steps, + "required_stable_checks": preset.required_stable_checks, + } + + def _empty_settle_state_metadata( + self, + active_mask: torch.Tensor, + ) -> dict[str, object]: + """Return a full-batch trace for a policy with no eligible rows.""" + batch_size = self._env_ids.numel() + return { + "elapsed_steps": 0, + "observation_count": 0, + "env_ids": self._env_ids.detach().cpu().tolist(), + "active_mask": active_mask.detach().cpu().tolist(), + "stable_counts": [0] * batch_size, + "settled_mask": [False] * batch_size, + "timeout_mask": [False] * batch_size, + "checked": False, + "max_linear_speed": [None] * batch_size, + "max_angular_speed": [None] * batch_size, + } + + def _expand_settle_state_metadata( + self, + state: DynamicSettleState, + active_mask: torch.Tensor, + ) -> dict[str, object]: + """Expand active-row monitor state to the stable full-batch ordering.""" + stable_counts = torch.zeros_like(self._env_ids) + settled_mask = torch.zeros_like(active_mask) + timeout_mask = torch.zeros_like(active_mask) + max_linear_speed = torch.full( + active_mask.shape, + float("inf"), + dtype=state.max_linear_speed.dtype, + device=active_mask.device, + ) + max_angular_speed = torch.full_like(max_linear_speed, float("inf")) + stable_counts[active_mask] = state.stable_counts + settled_mask[active_mask] = state.settled_mask + timeout_mask[active_mask] = state.timeout_mask + max_linear_speed[active_mask] = state.max_linear_speed + max_angular_speed[active_mask] = state.max_angular_speed + return { + "elapsed_steps": state.elapsed_steps, + "observation_count": state.observation_count, + "env_ids": self._env_ids.detach().cpu().tolist(), + "active_mask": active_mask.detach().cpu().tolist(), + "stable_counts": stable_counts.detach().cpu().tolist(), + "settled_mask": settled_mask.detach().cpu().tolist(), + "timeout_mask": timeout_mask.detach().cpu().tolist(), + "checked": state.checked, + "max_linear_speed": _json_speed_values(max_linear_speed), + "max_angular_speed": _json_speed_values(max_angular_speed), + } + + @staticmethod + def _validate_segment_membership( + segment: Any, + member: CompiledPostPolicy | CompiledProgramValidator, + *, + kind: str, + ) -> None: + """Require the supplied compiled value to belong to the exact segment.""" + if type(segment) is not CompiledProgramSegment: + raise TypeError("segment must be exactly CompiledProgramSegment.") + values = ( + segment.post_policies + if type(member) is CompiledPostPolicy + else segment.validators + ) + if not any(value is member for value in values): + raise ValueError( + f"Compiled {kind} does not belong to the supplied segment." + ) + + def _resolve_native_entities( + self, + ) -> tuple[ + Mapping[str, _SimulationSettleTarget], + Mapping[str, Any], + ]: + """Resolve only explicitly declared canonical/native pairs.""" + settle_targets: dict[str, _SimulationSettleTarget] = {} + rigid_objects: dict[str, Any] = {} + articulation_targets: dict[str, _SimulationSettleTarget] = {} + + for binding in self._scene_binding.rigid_objects: + entity = self._require_native( + "get_rigid_object", + canonical_id=binding.entity_id, + simulation_uid=binding.simulation_uid, + ) + target = _SimulationSettleTarget( + binding.entity_id, + "rigid_object", + entity, + ) + settle_targets[binding.entity_id] = target + rigid_objects[binding.entity_id] = entity + + for binding in self._scene_binding.articulations: + entity = self._require_native( + "get_articulation", + canonical_id=binding.entity_id, + simulation_uid=binding.simulation_uid, + ) + target = _SimulationSettleTarget( + binding.entity_id, + "articulation", + entity, + ) + settle_targets[binding.entity_id] = target + articulation_targets[binding.entity_id] = target + + for binding in self._scene_binding.links: + settle_targets[binding.entity_id] = self._require_parent_target( + articulation_targets, + child_id=binding.entity_id, + parent_id=binding.articulation_id, + ) + for binding in self._scene_binding.antipodal_grasps: + parent = settle_targets.get(binding.object_id) + if parent is None or parent.kind != "rigid_object": + raise KeyError( + f"Affordance {binding.entity_id!r} references unavailable rigid " + f"object {binding.object_id!r}." + ) + settle_targets[binding.entity_id] = parent + for binding in self._scene_binding.articulation_operations: + settle_targets[binding.entity_id] = self._require_parent_target( + articulation_targets, + child_id=binding.entity_id, + parent_id=binding.articulation_id, + ) + return MappingProxyType(settle_targets), MappingProxyType(rigid_objects) + + @staticmethod + def _require_parent_target( + targets: Mapping[str, _SimulationSettleTarget], + *, + child_id: str, + parent_id: str, + ) -> _SimulationSettleTarget: + """Resolve a child to one explicitly declared articulation root.""" + target = targets.get(parent_id) + if target is None: + raise KeyError( + f"Canonical entity {child_id!r} references unavailable parent " + f"{parent_id!r}." + ) + return target + + def _require_native( + self, + getter_name: str, + *, + canonical_id: str, + simulation_uid: str, + ) -> Any: + """Resolve one explicitly selected native simulation entity.""" + getter = getattr(self._simulation, getter_name, None) + if not callable(getter): + raise TypeError(f"simulation must provide {getter_name}().") + entity = getter(simulation_uid) + if entity is None: + raise KeyError( + f"Native entity {simulation_uid!r} selected for canonical entity " + f"{canonical_id!r} was not found." + ) + return entity + + def _measure_settle_target( + self, + target: _SimulationSettleTarget, + *, + row_indices: torch.Tensor, + ) -> DynamicSettleSample: + """Measure physical bodies for explicitly selected simulator rows.""" + if target.kind == "articulation": + body_data = getattr(target.native_entity, "body_data", None) + velocity = getattr(body_data, "body_link_vel", None) + if not isinstance(velocity, torch.Tensor): + raise RuntimeError( + f"Articulation settle target {target.canonical_id!r} has no " + "body_link_vel tensor." + ) + selected = velocity.index_select(0, row_indices.to(velocity.device)) + if selected.dim() != 3 or selected.shape[-1] != 6: + raise ValueError( + "Articulation body_link_vel must have shape (B, N, 6)." + ) + linear_velocity = selected[..., :3] + angular_velocity = selected[..., 3:] + else: + body_data = getattr(target.native_entity, "body_data", None) + linear_velocity = getattr(body_data, "lin_vel", None) + angular_velocity = getattr(body_data, "ang_vel", None) + if not isinstance(linear_velocity, torch.Tensor) or not isinstance( + angular_velocity, + torch.Tensor, + ): + raise RuntimeError( + f"Rigid settle target {target.canonical_id!r} has no linear/" + "angular velocity tensors." + ) + rows = row_indices.to(linear_velocity.device) + linear_velocity = linear_velocity.index_select(0, rows) + angular_velocity = angular_velocity.index_select( + 0, + row_indices.to(angular_velocity.device), + ) + if ( + linear_velocity.shape != angular_velocity.shape + or linear_velocity.dim() < 2 + or linear_velocity.shape[-1] != 3 + ): + raise ValueError( + "Rigid body velocities must have equal shape (B, ..., 3)." + ) + + linear_speed = torch.linalg.vector_norm(linear_velocity, dim=-1).reshape( + row_indices.numel(), + -1, + ) + angular_speed = torch.linalg.vector_norm(angular_velocity, dim=-1).reshape( + row_indices.numel(), + -1, + ) + device = self._env_ids.device + return DynamicSettleSample( + entity_id=target.canonical_id, + linear_speed=linear_speed.to(device=device), + angular_speed=angular_speed.to(device=device), + ) + + def _read_pose(self, entity: Any, *, entity_id: str) -> torch.Tensor: + """Read one rigid-object pose batch in simulator row order.""" + getter = getattr(entity, "get_local_pose", None) + if not callable(getter): + raise TypeError( + f"Native rigid object for {entity_id!r} must provide " + "get_local_pose()." + ) + pose = getter(to_matrix=True) + if not isinstance(pose, torch.Tensor) or not pose.is_floating_point(): + raise TypeError("get_local_pose(to_matrix=True) must return a tensor.") + batch_size = int(self._env_ids.numel()) + if pose.shape == (4, 4): + pose = pose.unsqueeze(0).expand(batch_size, -1, -1) + elif pose.shape != (batch_size, 4, 4): + raise ValueError( + f"Rigid object {entity_id!r} pose must have shape " + f"({batch_size}, 4, 4)." + ) + if not bool(torch.isfinite(pose).all().item()): + raise ValueError(f"Rigid object {entity_id!r} pose must be finite.") + return pose.clone() + + +__all__ = ["SimulationSegmentPolicyPort"] diff --git a/embodichain/lab/gym/envs/managers/_event_functors/dynamic_settling.py b/embodichain/lab/gym/envs/managers/_event_functors/dynamic_settling.py index b857f1078..9418d6182 100644 --- a/embodichain/lab/gym/envs/managers/_event_functors/dynamic_settling.py +++ b/embodichain/lab/gym/envs/managers/_event_functors/dynamic_settling.py @@ -18,14 +18,17 @@ from __future__ import annotations -import math from collections.abc import Sequence -from numbers import Real from typing import TYPE_CHECKING, Literal import torch from embodichain.lab.gym.envs.managers.cfg import SceneEntityCfg +from embodichain.lab.gym.envs.settling import ( + DynamicSettleMonitor, + DynamicSettleMonitorCfg, + DynamicSettleSample, +) from embodichain.lab.sim.objects import Articulation, RigidObject, RigidObjectGroup from embodichain.utils import logger @@ -37,7 +40,6 @@ _DynamicEntity = RigidObject | RigidObjectGroup | Articulation _SettleEntity = tuple[str, SceneEntityCfg, _DynamicEntity] -_SpeedSample = tuple[str, torch.Tensor, torch.Tensor] def _validate_settle_parameters( @@ -49,48 +51,21 @@ def _validate_settle_parameters( required_stable_checks: int, timeout_behavior: str, allow_partial_envs: bool, -) -> None: - """Validate dynamic-object settle parameters.""" - for name, value in ( - ("min_steps", min_steps), - ("max_steps", max_steps), - ("check_interval_steps", check_interval_steps), - ("required_stable_checks", required_stable_checks), - ): - if isinstance(value, bool) or not isinstance(value, int): - raise TypeError(f"{name} must be an integer, got {type(value).__name__}.") - - if min_steps < 0: - raise ValueError("min_steps must be non-negative.") - if max_steps < min_steps: - raise ValueError("max_steps must be greater than or equal to min_steps.") - if check_interval_steps < 1: - raise ValueError("check_interval_steps must be at least 1.") - if required_stable_checks < 1: - raise ValueError("required_stable_checks must be at least 1.") - - for name, value in ( - ("linear_velocity_threshold", linear_velocity_threshold), - ("angular_velocity_threshold", angular_velocity_threshold), - ): - if isinstance(value, bool) or not isinstance(value, Real): - raise TypeError(f"{name} must be a real number.") - if not math.isfinite(float(value)) or value < 0: - raise ValueError(f"{name} must be finite and non-negative.") - - available_checks = ( - 1 + (max_steps - min_steps + check_interval_steps - 1) // check_interval_steps +) -> DynamicSettleMonitorCfg: + """Validate parameters and return the reusable monitor policy.""" + cfg = DynamicSettleMonitorCfg( + linear_velocity_threshold=linear_velocity_threshold, + angular_velocity_threshold=angular_velocity_threshold, + min_steps=min_steps, + max_steps=max_steps, + check_interval_steps=check_interval_steps, + required_stable_checks=required_stable_checks, ) - if required_stable_checks > available_checks: - raise ValueError( - "required_stable_checks cannot be reached within the configured " - f"step budget; at most {available_checks} checks are possible." - ) - if timeout_behavior not in ("warn", "raise"): raise ValueError("timeout_behavior must be either 'warn' or 'raise'.") if not isinstance(allow_partial_envs, bool): raise TypeError("allow_partial_envs must be a boolean.") + return cfg def _normalize_settle_env_ids( @@ -210,9 +185,9 @@ def _resolve_settle_entities( def _measure_settle_speeds( entities: Sequence[_SettleEntity], env_ids: torch.Tensor, -) -> list[_SpeedSample]: +) -> list[DynamicSettleSample]: """Measure per-body linear and angular speeds for selected environments.""" - samples: list[_SpeedSample] = [] + samples: list[DynamicSettleSample] = [] for kind, entity_cfg, entity in entities: if kind == "articulation": velocity = entity.body_data.body_link_vel[env_ids] @@ -238,29 +213,35 @@ def _measure_settle_speeds( angular_speed = torch.linalg.vector_norm(angular_velocity, dim=-1).reshape( env_ids.numel(), -1 ) - samples.append((entity_cfg.uid, linear_speed, angular_speed)) + samples.append( + DynamicSettleSample( + entity_id=entity_cfg.uid, + linear_speed=linear_speed, + angular_speed=angular_speed, + ) + ) return samples def _settle_samples_are_stable( - samples: Sequence[_SpeedSample], + samples: Sequence[DynamicSettleSample], linear_velocity_threshold: float, angular_velocity_threshold: float, ) -> bool: """Return whether every measured body is finite and below both thresholds.""" stable = [] - for _, linear_speed, angular_speed in samples: + for sample in samples: stable.append( - torch.isfinite(linear_speed) - & torch.isfinite(angular_speed) - & (linear_speed <= linear_velocity_threshold) - & (angular_speed <= angular_velocity_threshold) + torch.isfinite(sample.linear_speed) + & torch.isfinite(sample.angular_speed) + & (sample.linear_speed <= linear_velocity_threshold) + & (sample.angular_speed <= angular_velocity_threshold) ) return bool(torch.cat([value.reshape(-1) for value in stable]).all().item()) def _format_settle_timeout( - samples: Sequence[_SpeedSample], + samples: Sequence[DynamicSettleSample], env_ids: torch.Tensor, linear_velocity_threshold: float, angular_velocity_threshold: float, @@ -272,7 +253,9 @@ def _format_settle_timeout( unsettled: list[str] = [] all_linear_speeds: list[torch.Tensor] = [] all_angular_speeds: list[torch.Tensor] = [] - for uid, linear_speed, angular_speed in samples: + for sample in samples: + linear_speed = sample.linear_speed + angular_speed = sample.angular_speed stable = ( torch.isfinite(linear_speed) & torch.isfinite(angular_speed) @@ -282,7 +265,7 @@ def _format_settle_timeout( unsettled_mask = ~stable.all(dim=1) if bool(unsettled_mask.any().item()): unsettled_env_ids = env_ids[unsettled_mask].detach().cpu().tolist() - unsettled.append(f"{uid}(env_ids={unsettled_env_ids})") + unsettled.append(f"{sample.entity_id}(env_ids={unsettled_env_ids})") all_linear_speeds.append(linear_speed.reshape(-1)) all_angular_speeds.append(angular_speed.reshape(-1)) @@ -364,7 +347,7 @@ def wait_for_dynamic_objects_to_settle( TypeError: If a parameter or entity configuration has the wrong type. ValueError: If parameters, targets, or environment selection are invalid. """ - _validate_settle_parameters( + monitor_cfg = _validate_settle_parameters( linear_velocity_threshold=linear_velocity_threshold, angular_velocity_threshold=angular_velocity_threshold, min_steps=min_steps, @@ -395,20 +378,14 @@ def wait_for_dynamic_objects_to_settle( env.sim.update(step=min_steps) step_count = min_steps - stable_checks = 0 - samples: list[_SpeedSample] + monitor = DynamicSettleMonitor(monitor_cfg, target_env_ids) + samples: list[DynamicSettleSample] + settle_state = None while True: samples = _measure_settle_speeds(entities, target_env_ids) - if _settle_samples_are_stable( - samples, - linear_velocity_threshold=linear_velocity_threshold, - angular_velocity_threshold=angular_velocity_threshold, - ): - stable_checks += 1 - if stable_checks >= required_stable_checks: - return - else: - stable_checks = 0 + settle_state = monitor.observe(samples, elapsed_steps=step_count) + if bool(settle_state.settled_mask.all().item()): + return if step_count >= max_steps: break @@ -422,7 +399,7 @@ def wait_for_dynamic_objects_to_settle( linear_velocity_threshold=linear_velocity_threshold, angular_velocity_threshold=angular_velocity_threshold, max_steps=max_steps, - stable_checks=stable_checks, + stable_checks=int(settle_state.stable_counts.min().item()), required_stable_checks=required_stable_checks, ) if timeout_behavior == "raise": diff --git a/embodichain/lab/gym/envs/settling.py b/embodichain/lab/gym/envs/settling.py new file mode 100644 index 000000000..39d5be171 --- /dev/null +++ b/embodichain/lab/gym/envs/settling.py @@ -0,0 +1,374 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Reusable per-environment dynamic-settling state machine.""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from dataclasses import dataclass +from numbers import Real + +import torch + +from embodichain.utils import configclass + + +@configclass +class DynamicSettleMonitorCfg: + """Threshold and cadence policy for :class:`DynamicSettleMonitor`. + + The monitor never advances an environment. Callers own the stepping path + and provide raw velocity samples after the configured minimum/cadence. + This lets reset events and demonstration post-policies share exactly the + same state transition rules while using different stepping ports. + """ + + linear_velocity_threshold: float = 0.03 + """Maximum stable linear speed in metres per second.""" + + angular_velocity_threshold: float = 0.20 + """Maximum stable angular speed in radians per second.""" + + min_steps: int = 10 + """Minimum number of environment steps before the first check.""" + + max_steps: int = 240 + """Maximum elapsed environment steps before unresolved rows time out.""" + + check_interval_steps: int = 2 + """Minimum number of steps between independent evidence checks.""" + + required_stable_checks: int = 3 + """Consecutive stable checks required independently for each row.""" + + def __post_init__(self) -> None: + for name in ( + "min_steps", + "max_steps", + "check_interval_steps", + "required_stable_checks", + ): + value = getattr(self, name) + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be an integer.") + if self.min_steps < 0: + raise ValueError("min_steps must be non-negative.") + if self.max_steps < self.min_steps: + raise ValueError("max_steps must be greater than or equal to min_steps.") + if self.check_interval_steps < 1: + raise ValueError("check_interval_steps must be at least 1.") + if self.required_stable_checks < 1: + raise ValueError("required_stable_checks must be at least 1.") + for name in ( + "linear_velocity_threshold", + "angular_velocity_threshold", + ): + value = getattr(self, name) + if isinstance(value, bool) or not isinstance(value, Real): + raise TypeError(f"{name} must be a real number.") + if not math.isfinite(float(value)) or float(value) < 0.0: + raise ValueError(f"{name} must be finite and non-negative.") + available_checks = ( + 1 + + (self.max_steps - self.min_steps + self.check_interval_steps - 1) + // self.check_interval_steps + ) + if self.required_stable_checks > available_checks: + raise ValueError( + "required_stable_checks cannot be reached within the configured " + f"step budget; at most {available_checks} checks are possible." + ) + + def snapshot(self) -> DynamicSettleMonitorCfg: + """Return an independently owned configuration value.""" + return DynamicSettleMonitorCfg( + linear_velocity_threshold=self.linear_velocity_threshold, + angular_velocity_threshold=self.angular_velocity_threshold, + min_steps=self.min_steps, + max_steps=self.max_steps, + check_interval_steps=self.check_interval_steps, + required_stable_checks=self.required_stable_checks, + ) + + +@dataclass(frozen=True, slots=True, eq=False) +class DynamicSettleSample: + """Raw per-body speed evidence for one registered scene entity. + + Args: + entity_id: Stable entity identifier used in metadata and diagnostics. + linear_speed: Per-row body speeds with shape ``(B, N)``. + angular_speed: Per-row body speeds with shape ``(B, N)``. + """ + + entity_id: str + linear_speed: torch.Tensor + angular_speed: torch.Tensor + + def __post_init__(self) -> None: + if ( + type(self.entity_id) is not str + or not self.entity_id + or self.entity_id != self.entity_id.strip() + ): + raise ValueError( + "entity_id must be a non-empty string without outer whitespace." + ) + for name in ("linear_speed", "angular_speed"): + value = getattr(self, name) + if not isinstance(value, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor.") + if not value.is_floating_point() or value.dim() != 2: + raise ValueError(f"{name} must be a floating tensor with shape (B, N).") + if value.shape[0] == 0 or value.shape[1] == 0: + raise ValueError(f"{name} must contain at least one row and body.") + if self.linear_speed.shape != self.angular_speed.shape: + raise ValueError("linear_speed and angular_speed must have equal shapes.") + if self.linear_speed.device != self.angular_speed.device: + raise ValueError("linear_speed and angular_speed must share a device.") + object.__setattr__(self, "linear_speed", self.linear_speed.clone()) + object.__setattr__(self, "angular_speed", self.angular_speed.clone()) + + def snapshot(self) -> DynamicSettleSample: + """Return an independently owned raw evidence sample.""" + return DynamicSettleSample( + entity_id=self.entity_id, + linear_speed=self.linear_speed, + angular_speed=self.angular_speed, + ) + + +@dataclass(frozen=True, slots=True, eq=False) +class DynamicSettleState: + """Owned state emitted after one monitor observation.""" + + env_ids: torch.Tensor + elapsed_steps: int + observation_count: int + checked: bool + stable_counts: torch.Tensor + settled_mask: torch.Tensor + timeout_mask: torch.Tensor + max_linear_speed: torch.Tensor + max_angular_speed: torch.Tensor + + def __post_init__(self) -> None: + if not isinstance(self.env_ids, torch.Tensor): + raise TypeError("env_ids must be a torch.Tensor.") + if self.env_ids.dtype != torch.long or self.env_ids.dim() != 1: + raise ValueError("env_ids must be a one-dimensional torch.long tensor.") + if self.env_ids.numel() == 0: + raise ValueError("env_ids must contain at least one row.") + if type(self.elapsed_steps) is not int or self.elapsed_steps < 0: + raise ValueError("elapsed_steps must be a non-negative integer.") + if type(self.observation_count) is not int or self.observation_count < 0: + raise ValueError("observation_count must be a non-negative integer.") + if type(self.checked) is not bool: + raise TypeError("checked must be a bool.") + row_count = self.env_ids.numel() + for name, dtype in ( + ("stable_counts", torch.long), + ("settled_mask", torch.bool), + ("timeout_mask", torch.bool), + ): + value = getattr(self, name) + if not isinstance(value, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor.") + if value.dtype != dtype or value.shape != (row_count,): + raise ValueError(f"{name} must have shape (B,) and dtype {dtype}.") + if value.device != self.env_ids.device: + raise ValueError(f"{name} and env_ids must share a device.") + if (self.settled_mask & self.timeout_mask).any(): + raise ValueError("settled_mask and timeout_mask must not overlap.") + for name in ("max_linear_speed", "max_angular_speed"): + value = getattr(self, name) + if not isinstance(value, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor.") + if not value.is_floating_point() or value.shape != (row_count,): + raise ValueError(f"{name} must be a floating tensor with shape (B,).") + if value.device != self.env_ids.device: + raise ValueError(f"{name} and env_ids must share a device.") + for name in ( + "env_ids", + "stable_counts", + "settled_mask", + "timeout_mask", + "max_linear_speed", + "max_angular_speed", + ): + object.__setattr__(self, name, getattr(self, name).clone()) + + @property + def complete(self) -> bool: + """Whether every row has either settled or timed out.""" + return bool((self.settled_mask | self.timeout_mask).all().item()) + + def to_metadata(self) -> dict[str, object]: + """Return deterministic, JSON-compatible post-policy metadata.""" + return { + "elapsed_steps": self.elapsed_steps, + "observation_count": self.observation_count, + "env_ids": self.env_ids.detach().to("cpu").tolist(), + "stable_counts": self.stable_counts.detach().to("cpu").tolist(), + "settled_mask": self.settled_mask.detach().to("cpu").tolist(), + "timeout_mask": self.timeout_mask.detach().to("cpu").tolist(), + "max_linear_speed": self.max_linear_speed.detach().to("cpu").tolist(), + "max_angular_speed": self.max_angular_speed.detach().to("cpu").tolist(), + } + + +class DynamicSettleMonitor: + """Track settling independently for stable environment IDs. + + Duplicate observations at the same ``elapsed_steps`` value are idempotent. + Regressing step counters are rejected, and a jump across multiple cadence + boundaries counts as one fresh observation rather than replaying one sample. + """ + + def __init__( + self, + cfg: DynamicSettleMonitorCfg, + env_ids: torch.Tensor, + ) -> None: + if not isinstance(cfg, DynamicSettleMonitorCfg): + raise TypeError("cfg must be a DynamicSettleMonitorCfg.") + if not isinstance(env_ids, torch.Tensor): + raise TypeError("env_ids must be a torch.Tensor.") + if env_ids.dtype != torch.long or env_ids.dim() != 1: + raise ValueError("env_ids must be a one-dimensional torch.long tensor.") + if env_ids.numel() == 0 or torch.unique(env_ids).numel() != env_ids.numel(): + raise ValueError("env_ids must contain unique environment IDs.") + self.cfg = cfg.snapshot() + self._env_ids = env_ids.clone() + self._stable_counts = torch.zeros_like(env_ids) + self._settled = torch.zeros_like(env_ids, dtype=torch.bool) + self._timeout = torch.zeros_like(env_ids, dtype=torch.bool) + self._max_linear = torch.full( + env_ids.shape, + float("inf"), + dtype=torch.float32, + device=env_ids.device, + ) + self._max_angular = self._max_linear.clone() + self._last_elapsed_steps = -1 + self._last_checked_steps = -1 + self._observation_count = 0 + + @property + def env_ids(self) -> torch.Tensor: + """Return the stable row IDs owned by this monitor.""" + return self._env_ids.clone() + + def observe( + self, + samples: Sequence[DynamicSettleSample], + *, + elapsed_steps: int, + ) -> DynamicSettleState: + """Consume one raw speed observation when the configured cadence is due. + + Args: + samples: One speed sample per monitored entity. + elapsed_steps: Steps advanced by the caller since post-policy start. + + Returns: + Per-row stable, settled, timeout, and velocity metadata. + """ + if type(elapsed_steps) is not int or elapsed_steps < 0: + raise ValueError("elapsed_steps must be a non-negative integer.") + if elapsed_steps < self._last_elapsed_steps: + raise ValueError("elapsed_steps must be monotonic.") + normalized = tuple(samples) + if not normalized or not all( + isinstance(sample, DynamicSettleSample) for sample in normalized + ): + raise ValueError("samples must contain DynamicSettleSample values.") + if len({sample.entity_id for sample in normalized}) != len(normalized): + raise ValueError("samples must use unique entity IDs.") + for sample in normalized: + if sample.linear_speed.shape[0] != self._env_ids.numel(): + raise ValueError("Every sample batch must match env_ids length.") + if sample.linear_speed.device != self._env_ids.device: + raise ValueError("Samples and env_ids must share a device.") + + duplicate = elapsed_steps == self._last_elapsed_steps + due = elapsed_steps >= self.cfg.min_steps and ( + self._last_checked_steps < 0 + or elapsed_steps - self._last_checked_steps >= self.cfg.check_interval_steps + or elapsed_steps >= self.cfg.max_steps + ) + checked = due and not duplicate and not self._timeout.all() + if checked: + linear = torch.cat([sample.linear_speed for sample in normalized], dim=1) + angular = torch.cat([sample.angular_speed for sample in normalized], dim=1) + finite = torch.isfinite(linear).all(dim=1) & torch.isfinite(angular).all( + dim=1 + ) + self._max_linear = torch.where( + torch.isfinite(linear), linear, torch.full_like(linear, float("inf")) + ).amax(dim=1) + self._max_angular = torch.where( + torch.isfinite(angular), + angular, + torch.full_like(angular, float("inf")), + ).amax(dim=1) + stable = ( + finite + & (self._max_linear <= self.cfg.linear_velocity_threshold) + & (self._max_angular <= self.cfg.angular_velocity_threshold) + ) + active = ~self._settled & ~self._timeout + self._stable_counts = torch.where( + active & stable, + self._stable_counts + 1, + torch.where( + active, torch.zeros_like(self._stable_counts), self._stable_counts + ), + ) + self._settled |= active & ( + self._stable_counts >= self.cfg.required_stable_checks + ) + self._observation_count += 1 + self._last_checked_steps = elapsed_steps + + if elapsed_steps >= self.cfg.max_steps: + self._timeout |= ~self._settled + self._last_elapsed_steps = elapsed_steps + return self._state(elapsed_steps=elapsed_steps, checked=checked) + + def _state(self, *, elapsed_steps: int, checked: bool) -> DynamicSettleState: + """Build an owned state snapshot.""" + return DynamicSettleState( + env_ids=self._env_ids, + elapsed_steps=elapsed_steps, + observation_count=self._observation_count, + checked=checked, + stable_counts=self._stable_counts, + settled_mask=self._settled, + timeout_mask=self._timeout, + max_linear_speed=self._max_linear, + max_angular_speed=self._max_angular, + ) + + +__all__ = [ + "DynamicSettleMonitor", + "DynamicSettleMonitorCfg", + "DynamicSettleSample", + "DynamicSettleState", +] diff --git a/embodichain/lab/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index cf3c1086b..e524b6765 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -17,6 +17,7 @@ from __future__ import annotations import os +from pathlib import Path import numpy as np import torch import dexsim @@ -393,13 +394,22 @@ def cat_tensor_with_ids( return out -def config_to_cfg(config: dict, manager_modules: list = None) -> "EmbodiedEnvCfg": +def config_to_cfg( + config: dict, + manager_modules: list | None = None, + *, + source_path: str | os.PathLike[str] | None = None, +) -> "EmbodiedEnvCfg": """Parser configuration file into cfgs for env initialization. Args: config (dict): The configuration dictionary containing robot, sensor, light, background, and interactive objects. manager_modules (list): List of module paths for dataset, event, observation, and reward managers. If not provided, uses default module paths. + source_path: Optional path of the Gym configuration source file. A + relative top-level ``expert_program_path`` is resolved from this + file's directory. Without it, relative paths use the current + working directory. Returns: EmbodiedEnvCfg: A configuration object for initializing the environment. @@ -446,6 +456,30 @@ class ComponentCfg: if key not in config: log_error(f"Missing required config key: {key}") + if "expert_program_path" in config: + expert_program_path = config["expert_program_path"] + if type(expert_program_path) is not str: + raise TypeError("expert_program_path must be an exact string.") + if ( + not expert_program_path + or expert_program_path != expert_program_path.strip() + ): + raise ValueError( + "expert_program_path must be a non-empty string without outer " + "whitespace." + ) + from embodichain.lab.gym.envs.expert_program.loader import ( + load_expert_program, + ) + + expert_program_base_dir = ( + None if source_path is None else Path(source_path).expanduser().parent + ) + env_cfg.expert_program = load_expert_program( + expert_program_path, + base_dir=expert_program_base_dir, + ) + env_cfg.max_episode_steps = config.get("max_episode_steps", 300) env_cfg.num_envs = config.get("num_envs", 1) @@ -1021,16 +1055,20 @@ def build_env_cfg_from_args( tuple[EmbodiedEnvCfg, dict, dict]: A tuple containing the environment configuration object, the merged gym configuration dictionary, and the action configuration dictionary. """ + from embodichain.utils.config_paths import resolve_config_path from embodichain.utils.utility import load_config from embodichain.lab.gym.envs import EmbodiedEnvCfg - gym_config = load_config(args.gym_config) + gym_config_source_path = resolve_config_path(args.gym_config) + gym_config = load_config(gym_config_source_path) gym_config = merge_args_with_gym_config(args, gym_config) if gym_config_modifier is not None: gym_config_modifier(gym_config) cfg: EmbodiedEnvCfg = config_to_cfg( - gym_config, manager_modules=get_manager_modules() + gym_config, + manager_modules=get_manager_modules(), + source_path=gym_config_source_path, ) cfg.filter_visual_rand = args.filter_visual_rand cfg.filter_dataset_saving = args.filter_dataset_saving diff --git a/embodichain/lab/scripts/run_env.py b/embodichain/lab/scripts/run_env.py index 880d040fc..8c99f098f 100644 --- a/embodichain/lab/scripts/run_env.py +++ b/embodichain/lab/scripts/run_env.py @@ -17,6 +17,7 @@ from __future__ import annotations import argparse +import json import os import select import sys @@ -31,6 +32,9 @@ import tqdm from embodichain.lab.gym.envs.demo import DemoEpisodeResult, execute_demo_episode +from embodichain.lab.gym.envs.expert_program.loader import ( + load_expert_program as _load_expert_program, +) from embodichain.lab.gym.envs.wrapper import ReplayWrapper from embodichain.lab.gym.utils.gym_utils import ( add_env_launcher_args_to_parser, @@ -289,6 +293,16 @@ def generate_function( f"Episode {time_id} attempt {attempt}/{max_attempts} failed: " f"{result.terminal_reason}. Discarding {result.length} frames." ) + if debug_mode: + log_warning( + "Failed demo trace: " + + json.dumps( + result.to_metadata(), + allow_nan=False, + sort_keys=True, + separators=(",", ":"), + ) + ) return False @@ -741,6 +755,18 @@ def _create_parser() -> argparse.ArgumentParser: add_env_launcher_args_to_parser(parser, require_gym_config=True) parser.set_defaults(viser_image_fps=None) + parser.add_argument( + "--expert-program", + type=str, + default=None, + help="Path to a declarative Expert Program (.json, .yaml, or .yml).", + ) + parser.add_argument( + "--debug-mode", + action="store_true", + help="Log the structured trace for each failed demo attempt.", + ) + parser.add_argument( "--replay", action="store_true", @@ -832,6 +858,9 @@ def cli(argv: Sequence[str] | None = None) -> None: execute_init_hooks() env_cfg, gym_config, action_config = build_env_cfg_from_args(args) + expert_program_path = getattr(args, "expert_program", None) + if expert_program_path is not None: + env_cfg.expert_program = _load_expert_program(expert_program_path) if args.replay and args.replay_mode == "control": log_info("Dataset saving disabled for control replay mode.", color="green") diff --git a/embodichain/utils/__init__.py b/embodichain/utils/__init__.py index f3dd6ba62..fd680446c 100644 --- a/embodichain/utils/__init__.py +++ b/embodichain/utils/__init__.py @@ -20,6 +20,15 @@ """ from .configclass import configclass, is_configclass +from .config_paths import resolve_config_path + +__all__ = [ + "GLOBAL_SEED", + "configclass", + "is_configclass", + "resolve_config_path", + "set_seed", +] GLOBAL_SEED = 1024 diff --git a/embodichain/utils/config_paths.py b/embodichain/utils/config_paths.py new file mode 100644 index 000000000..e97346c98 --- /dev/null +++ b/embodichain/utils/config_paths.py @@ -0,0 +1,56 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Stable path resolution for user and packaged configuration files.""" + +from __future__ import annotations + +from pathlib import Path + +__all__ = ["resolve_config_path"] + + +def resolve_config_path(path: str | Path) -> Path: + """Resolve one configuration path without opening the target file. + + Existing, absolute, and ordinary relative paths preserve their normal + filesystem meaning. Repository-style paths below + ``embodichain_tasks/configs`` are redirected to the packaged task-config + resource so the same configuration reference works from an installed + wheel. + + Args: + path: User path or repository-style official-task configuration path. + + Returns: + Expanded filesystem path, resolved through the packaged task resource + only when the input uses the official-task configuration prefix. + + Raises: + TypeError: If ``path`` is not path-like. + """ + resolved_path = Path(path).expanduser() + if resolved_path.exists() or resolved_path.is_absolute(): + return resolved_path + + task_prefix = ("embodichain_tasks", "configs") + if resolved_path.parts[: len(task_prefix)] != task_prefix: + return resolved_path + + from embodichain_tasks.configs import get_config_path + + relative_path = Path(*resolved_path.parts[len(task_prefix) :]) + return get_config_path(relative_path) diff --git a/embodichain/utils/utility.py b/embodichain/utils/utility.py index 2c6b3cd1d..e44c5633f 100644 --- a/embodichain/utils/utility.py +++ b/embodichain/utils/utility.py @@ -31,6 +31,7 @@ from pathlib import Path from typing import Any, Dict, List, Tuple, Callable +from embodichain.utils.config_paths import resolve_config_path as _resolve_config_path from embodichain.utils.string import callable_to_string @@ -375,22 +376,6 @@ def _config_format_from_path(path: str | Path) -> str: ) -def _resolve_config_path(path: str | Path) -> Path: - """Resolve repository-style official-task paths from an installed wheel.""" - resolved_path = Path(path).expanduser() - if resolved_path.exists() or resolved_path.is_absolute(): - return resolved_path - - task_prefix = ("embodichain_tasks", "configs") - if resolved_path.parts[: len(task_prefix)] != task_prefix: - return resolved_path - - from embodichain_tasks.configs import get_config_path - - relative_path = Path(*resolved_path.parts[len(task_prefix) :]) - return get_config_path(relative_path) - - def load_config(path: str | Path) -> Dict[str, Any]: """Load a gym or agent config file into a dictionary. diff --git a/tests/gym/envs/expert_program/test_articulation_program.py b/tests/gym/envs/expert_program/test_articulation_program.py new file mode 100644 index 000000000..6629724cc --- /dev/null +++ b/tests/gym/envs/expert_program/test_articulation_program.py @@ -0,0 +1,185 @@ +# ---------------------------------------------------------------------------- +# 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 declarative articulation calls in Expert Programs.""" + +from __future__ import annotations + +import pytest +import torch + +from embodichain.lab.gym.envs.expert_program import ( + ExpertProgramCfg, + ExpertProgramCompiler, + ExpertProgramDecodeError, + ExpertProgramIntegrationCfg, + InvokeCfg, + OperateArticulationCfg, + decode_expert_program, +) +from embodichain.lab.sim.atomic_actions import Affordance, EntityState +from embodichain.lab.sim.skills.calls import OperateArticulation +from embodichain.lab.sim.skills.scene import ( + SceneAffordanceRef, + SceneArticulationRef, + SceneEntityRegistration, + SceneRegistry, +) + + +class _NeverObserveProvider: + """Reject state observation during static program compilation.""" + + def observe( + self, + *, + timestamp: float, + env_ids: torch.Tensor, + ) -> EntityState: + del timestamp, env_ids + raise AssertionError("Compilation must not observe providers.") + + +def _payload(call: dict[str, object]) -> dict[str, object]: + return { + "schema_version": 1, + "program_id": "open_drawer", + "integration": { + "robot_profile": "manipulator", + "scene_registry": "scene", + "runtime_preset": "safe", + }, + "targets": {}, + "program": {"kind": "invoke", "call": call}, + } + + +def _compiler() -> ExpertProgramCompiler: + provider = _NeverObserveProvider() + drawer = SceneArticulationRef("drawer") + registry = SceneRegistry( + ( + SceneEntityRegistration(ref=drawer, state_provider=provider), + SceneEntityRegistration( + ref=SceneAffordanceRef("drawer_handle"), + parent=drawer, + native_name="handle", + relative_pose=torch.eye(4), + affordance=Affordance(), + ), + ) + ) + return ExpertProgramCompiler.from_scene_registry(registry) + + +def _integration() -> ExpertProgramIntegrationCfg: + return ExpertProgramIntegrationCfg( + robot_profile="manipulator", + scene_registry="scene", + runtime_preset="safe", + ) + + +def test_decoder_accepts_named_and_explicit_articulation_targets() -> None: + named = decode_expert_program( + _payload( + { + "kind": "operate_articulation", + "articulation": "drawer", + "handle": "drawer_handle", + "target": "open", + "resources": {"primary": "right_arm"}, + } + ) + ) + explicit = decode_expert_program( + _payload( + { + "kind": "operate_articulation", + "articulation": "drawer", + "target_position": 0.42, + "target_displacement": 0.40, + } + ) + ) + + assert type(named.program) is InvokeCfg + assert named.program.call == OperateArticulationCfg( + articulation="drawer", + handle="drawer_handle", + target="open", + resources={"primary": "right_arm"}, + ) + assert type(explicit.program) is InvokeCfg + assert explicit.program.call == OperateArticulationCfg( + articulation="drawer", + target_position=0.42, + target_displacement=0.40, + ) + + +@pytest.mark.parametrize( + ("fields", "code"), + ( + ({"target": "open", "target_position": 0.4}, "conflicting_articulation_target"), + ({"target_position": 0.4}, "incomplete_articulation_target"), + ({"target_displacement": 0.2}, "incomplete_articulation_target"), + ({"target_position": True, "target_displacement": 0.2}, "invalid_number"), + ), +) +def test_decoder_rejects_ambiguous_or_incomplete_articulation_targets( + fields: dict[str, object], + code: str, +) -> None: + call = { + "kind": "operate_articulation", + "articulation": "drawer", + **fields, + } + + with pytest.raises(ExpertProgramDecodeError) as error: + decode_expert_program(_payload(call)) + + assert error.value.code == code + + +def test_compiler_preserves_typed_articulation_call_without_observation() -> None: + config = ExpertProgramCfg( + schema_version=1, + program_id="open_drawer", + integration=_integration(), + targets={}, + program=InvokeCfg( + call=OperateArticulationCfg( + articulation="drawer", + handle="drawer_handle", + target="open", + ) + ), + ) + + segment = tuple(_compiler().compile(config))[0] + call = segment.calls[0].call + + assert type(call) is OperateArticulation + assert call.articulation == SceneArticulationRef("drawer") + assert call.handle == SceneAffordanceRef("drawer_handle") + assert call.target == "open" + assert call.target_position is None + assert call.target_displacement is None + + +__all__: list[str] = [] diff --git a/tests/gym/envs/expert_program/test_bridge.py b/tests/gym/envs/expert_program/test_bridge.py new file mode 100644 index 000000000..9389fd5a9 --- /dev/null +++ b/tests/gym/envs/expert_program/test_bridge.py @@ -0,0 +1,2049 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from dataclasses import dataclass +import json +from typing import Any + +import pytest +import torch + +from embodichain.lab.gym.envs.demo import ProcessedEnvAction, execute_demo_episode +from embodichain.lab.gym.envs.expert_program.bridge import ( + AtomicDemoBridge, + BufferedGymCommandSink, + DemoBridgeError, + EnvironmentStepClock, + EnvironmentStepTimingError, + GymPlanningObservationProvider, + RuntimeCommandFrameEncoder, + UnsupportedRuntimeTransportError, +) +import embodichain.lab.gym.envs.expert_program.bridge as bridge_module +from embodichain.lab.sim.atomic_actions.bindings import ( + JointPositionTarget, + RuntimeEndpointTarget, +) +from embodichain.lab.sim.atomic_actions.execution import ( + ExecutionEvent, + ExecutionEventKind, +) +from embodichain.lab.sim.atomic_actions.runtime_commands import ( + EndpointCommand, + JointPositionPayload, + RuntimeCommandFrame, + RuntimeCommandPayload, +) +from embodichain.lab.sim.atomic_actions.state import ( + PlanningContext, + RobotObservation, + SceneSnapshot, + TaskState, +) +from embodichain.lab.sim.skills.calls import RegisteredSemanticCall +from embodichain.lab.sim.skills.parallel import ParallelTimingPolicy +from embodichain.lab.sim.skills.runtime import SkillResult, SkillStatus +from embodichain.lab.sim.skills.parallel_runtime import ( + ParallelLaneCommandSink, + ParallelRuntimeBranch, + ParallelSkillResult, + ParallelSkillRuntime, +) +from embodichain.lab.sim.skills.profiles import ResourceClaim + +STEP_DT = 0.02 +BATCH_SIZE = 2 +ROBOT_DOF = 5 + + +class _QposProvider: + """Return an owned fixed full-qpos snapshot.""" + + def __init__(self, qpos: torch.Tensor) -> None: + self.qpos = qpos.clone() + + def current_qpos(self, env_ids: torch.Tensor) -> torch.Tensor: + assert env_ids.numel() == self.qpos.shape[0] + return self.qpos.clone() + + +def _context( + *, + qpos: torch.Tensor | None = None, + env_ids: torch.Tensor | None = None, +) -> PlanningContext: + qpos = torch.zeros(BATCH_SIZE, ROBOT_DOF) if qpos is None else qpos + env_ids = torch.tensor([7, 3], dtype=torch.long) if env_ids is None else env_ids + return PlanningContext( + robot=RobotObservation( + timestamp=0.0, + qpos=qpos, + qvel=torch.zeros_like(qpos), + ), + task=TaskState.empty(qpos.shape[0], qpos.device), + scene=SceneSnapshot.empty(), + env_ids=env_ids, + ) + + +def _joint_frame( + *, + duration: float, + active_mask: torch.Tensor | None = None, + positions: torch.Tensor | None = None, +) -> RuntimeCommandFrame: + active_mask = torch.tensor([True, True]) if active_mask is None else active_mask + positions = ( + torch.tensor([[10.0, 30.0], [11.0, 31.0]]) if positions is None else positions + ) + return RuntimeCommandFrame( + commands=( + EndpointCommand( + target=JointPositionTarget( + control_part="arm", + joint_ids=(1, 3), + ), + payload=JointPositionPayload(positions=positions), + ), + ), + active_mask=active_mask, + env_ids=torch.tensor([7, 3], dtype=torch.long), + hold_duration=torch.full((BATCH_SIZE,), duration), + ) + + +@dataclass(frozen=True, slots=True) +class _DummyTarget(RuntimeEndpointTarget): + """Test-only non-joint runtime target.""" + + name: str + + @property + def transport_id(self) -> str: + return "test.transport" + + @property + def target_id(self) -> str: + return self.name + + def snapshot(self) -> _DummyTarget: + return _DummyTarget(self.name) + + +@dataclass(frozen=True, slots=True, eq=False) +class _DummyPayload(RuntimeCommandPayload): + """Test-only scalar payload.""" + + values: torch.Tensor + + def __post_init__(self) -> None: + object.__setattr__(self, "values", self.values.clone()) + + @property + def batch_size(self) -> int: + return int(self.values.shape[0]) + + @property + def device(self) -> torch.device: + return self.values.device + + @property + def transport_id(self) -> str: + return "test.transport" + + def snapshot(self) -> _DummyPayload: + return _DummyPayload(self.values) + + +class _DummyTransportEncoder: + """Test registration proving the frame encoder is transport-extensible.""" + + @property + def transport_id(self) -> str: + return "test.transport" + + def encode( + self, + command: EndpointCommand, + *, + base_action: Any, + active_mask: torch.Tensor, + ) -> Any: + assert isinstance(command.payload, _DummyPayload) + action = base_action.clone() + action[active_mask, 0] = command.payload.values[active_mask] + return action + + def hold( + self, + targets: tuple[RuntimeEndpointTarget, ...], + *, + base_action: Any, + context: PlanningContext, + ) -> Any: + del targets, context + return base_action.clone() + + +class _RecordingAcceptedCommandObserver: + """Record transactional sink notifications and optional callback failures.""" + + def __init__( + self, + *, + fail_accept: bool = False, + fail_cancel: bool = False, + ) -> None: + self.fail_accept = fail_accept + self.fail_cancel = fail_cancel + self.sink: BufferedGymCommandSink | None = None + self.accepted_frames: list[RuntimeCommandFrame] = [] + self.accepted_pending_counts: list[int] = [] + self.cancelled_targets: list[tuple[RuntimeEndpointTarget, ...]] = [] + self.cancelled_pending_counts: list[int] = [] + self.discard_count = 0 + + def accepted(self, command: RuntimeCommandFrame) -> None: + """Record acceptance after observing the sink's committed buffer.""" + if self.sink is None: + raise AssertionError("Observer sink must be assigned before use.") + self.accepted_pending_counts.append(self.sink.pending_count) + self.accepted_frames.append(command) + if self.fail_accept: + raise RuntimeError("observer rejected accepted command") + + def cancelled(self, targets: tuple[RuntimeEndpointTarget, ...]) -> None: + """Record owned cancellation targets.""" + if self.sink is None: + raise AssertionError("Observer sink must be assigned before use.") + self.cancelled_pending_counts.append(self.sink.pending_count) + self.cancelled_targets.append(targets) + if self.fail_cancel: + raise RuntimeError("observer rejected cancellation") + + def discarded(self) -> None: + """Record one fail-closed observer reset.""" + self.discard_count += 1 + + +def _dummy_frame() -> RuntimeCommandFrame: + return RuntimeCommandFrame( + commands=( + EndpointCommand( + target=_DummyTarget("base"), + payload=_DummyPayload(torch.tensor([4.0, 5.0])), + ), + ), + active_mask=torch.tensor([True, False]), + env_ids=torch.tensor([7, 3], dtype=torch.long), + hold_duration=torch.full((BATCH_SIZE,), STEP_DT), + ) + + +@dataclass(frozen=True, slots=True) +class _FakeCompiledCall: + call_index: int + call: object + + +@dataclass(frozen=True, slots=True) +class _FakeSegment: + segment_index: int = 0 + segment_id: str = "segment-0" + name: str = "pick-and-place" + calls: tuple[_FakeCompiledCall, ...] = ( + _FakeCompiledCall(0, "pick"), + _FakeCompiledCall(1, "place"), + ) + source_path: tuple[object, ...] = ("program", "steps", 0) + post_policies: tuple[object, ...] = () + validators: tuple[object, ...] = () + parallel_block: object | None = None + implicit: bool = False + + +@dataclass(frozen=True, slots=True) +class _FakeParallelBranch: + branch_index: int + calls: tuple[_FakeCompiledCall, ...] + + +@dataclass(frozen=True, slots=True) +class _FakeBarrier: + timeout_steps: int = 17 + failure_policy: str = "fail_fast" + + +@dataclass(frozen=True, slots=True) +class _FakeParallelBlock: + branches: tuple[_FakeParallelBranch, ...] + barrier: _FakeBarrier = _FakeBarrier() + + +@dataclass(frozen=True, slots=True) +class _FakeProgramAnalysis: + calls: tuple[object, ...] + execution_prefix_length: int + + +class _FakeProgram: + schema_version = 1 + program_id = "demo-program" + + def __init__(self, *segments: _FakeSegment) -> None: + self.segments = segments + + def iter_segments(self): + yield from self.segments + + def sequential_execution_analysis( + self, + segment_index: int, + ) -> _FakeProgramAnalysis: + current = self.segments[segment_index] + if current.parallel_block is not None: + raise ValueError("Parallel segments have no sequential analysis.") + calls: list[object] = [] + for segment in self.segments[segment_index:]: + if segment.parallel_block is not None: + break + calls.extend(compiled.call for compiled in segment.calls) + return _FakeProgramAnalysis(tuple(calls), len(current.calls)) + + +def _skill_result( + status: SkillStatus, + *, + wait_duration: float = 0.0, + workflow_id: str = "demo-program/segment-0", +) -> SkillResult: + env_ids = torch.tensor([7, 3], dtype=torch.long) + eligible = torch.ones(BATCH_SIZE, dtype=torch.bool) + success = ( + torch.ones(BATCH_SIZE, dtype=torch.bool) + if status is SkillStatus.COMPLETED + else torch.zeros(BATCH_SIZE, dtype=torch.bool) + ) + failure = ( + torch.ones(BATCH_SIZE, dtype=torch.bool) + if status is SkillStatus.FAILED + else torch.zeros(BATCH_SIZE, dtype=torch.bool) + ) + if status is SkillStatus.FAILED: + eligible = torch.zeros_like(eligible) + return SkillResult( + status=status, + workflow_id=workflow_id, + current_call_index=0 if status is SkillStatus.RUNNING else None, + env_ids=env_ids, + success_mask=success, + failure_mask=failure, + cancelled_mask=torch.zeros(BATCH_SIZE, dtype=torch.bool), + eligible_mask=eligible, + task_state=TaskState.empty(BATCH_SIZE, "cpu"), + wait_duration=wait_duration, + ) + + +class _FakeRuntime: + """Clock-aware nonblocking runtime used to test the Gym boundary only.""" + + def __init__( + self, + sink: BufferedGymCommandSink, + clock: EnvironmentStepClock, + frame: RuntimeCommandFrame, + ) -> None: + self.sink = sink + self.clock = clock + self.frame = frame + self._status = SkillStatus.IDLE + self._result = _skill_result(SkillStatus.IDLE) + self._due_at = 0.0 + self._sent = False + self.start_count = 0 + self.step_count = 0 + self.cancel_count = 0 + self.calls: tuple[object, ...] = () + self.execution_prefix_lengths: list[int | None] = [] + self.eligible_masks: list[torch.Tensor | None] = [] + self.adopted_states: list[TaskState] = [] + + @property + def result(self) -> SkillResult: + return self._result + + @property + def status(self) -> SkillStatus: + return self._status + + def start( + self, + *calls: object, + workflow_id: str = "semantic_workflow", + eligible_mask: torch.Tensor | None = None, + execution_prefix_length: int | None = None, + ) -> SkillResult: + self.start_count += 1 + self.calls = tuple(calls[0]) if len(calls) == 1 else calls + self.execution_prefix_lengths.append(execution_prefix_length) + self.eligible_masks.append( + None if eligible_mask is None else eligible_mask.clone() + ) + self._sent = False + self._due_at = 0.0 + self._status = SkillStatus.RUNNING + self._result = _skill_result( + SkillStatus.RUNNING, + workflow_id=workflow_id, + ) + return self._result + + def adopt_verified_task_state(self, task_state: TaskState) -> SkillResult: + self.adopted_states.append(task_state) + return self._result + + def step(self) -> SkillResult: + self.step_count += 1 + if not self._sent: + self.sink.send(self.frame, timeout=1.0) + self._sent = True + self._due_at = self.clock.now() + float( + self.frame.hold_duration.max().item() + ) + remaining = max(self._due_at - self.clock.now(), 0.0) + if remaining > 1.0e-9: + self._result = _skill_result( + SkillStatus.RUNNING, + wait_duration=remaining, + workflow_id=self._result.workflow_id or "semantic_workflow", + ) + return self._result + self._status = SkillStatus.COMPLETED + self._result = _skill_result( + SkillStatus.COMPLETED, + workflow_id=self._result.workflow_id or "semantic_workflow", + ) + return self._result + + def cancel(self, reason: str) -> SkillResult: + del reason + self.cancel_count += 1 + self.sink.cancel(self.frame.targets, timeout=1.0) + self.sink.hold(self.frame.targets, _context(), timeout=1.0) + self._status = SkillStatus.CANCELLED + self._result = SkillResult( + status=SkillStatus.CANCELLED, + workflow_id=self._result.workflow_id, + current_call_index=None, + env_ids=torch.tensor([7, 3], dtype=torch.long), + success_mask=torch.zeros(BATCH_SIZE, dtype=torch.bool), + failure_mask=torch.zeros(BATCH_SIZE, dtype=torch.bool), + cancelled_mask=torch.ones(BATCH_SIZE, dtype=torch.bool), + eligible_mask=torch.zeros(BATCH_SIZE, dtype=torch.bool), + task_state=TaskState.empty(BATCH_SIZE, "cpu"), + ) + return self._result + + +class _StartFailingRuntime(_FakeRuntime): + """Fail semantic preflight before accepting any controller command.""" + + def start( + self, + *calls: object, + workflow_id: str = "semantic_workflow", + eligible_mask: torch.Tensor | None = None, + execution_prefix_length: int | None = None, + ) -> SkillResult: + del calls, workflow_id, eligible_mask, execution_prefix_length + self.start_count += 1 + raise RuntimeError("semantic runtime preflight failed") + + +class _TerminalHoldRuntime(_FakeRuntime): + """Emit a terminal safe hold after one consumed command.""" + + def step(self) -> SkillResult: + self.step_count += 1 + if not self._sent: + self.sink.send(self.frame, timeout=1.0) + self._sent = True + self._status = SkillStatus.RUNNING + self._result = _skill_result( + SkillStatus.RUNNING, + workflow_id=self._result.workflow_id or "semantic_workflow", + ) + return self._result + self.sink.hold(self.frame.targets, _context(), timeout=1.0) + self._status = SkillStatus.COMPLETED + self._result = _skill_result( + SkillStatus.COMPLETED, + workflow_id=self._result.workflow_id or "semantic_workflow", + ) + return self._result + + +class _TerminalFailedRuntime(_FakeRuntime): + """Fail terminally during planning without accepting a command.""" + + def start( + self, + *calls: object, + workflow_id: str = "semantic_workflow", + eligible_mask: torch.Tensor | None = None, + execution_prefix_length: int | None = None, + ) -> SkillResult: + running = super().start( + *calls, + workflow_id=workflow_id, + eligible_mask=eligible_mask, + execution_prefix_length=execution_prefix_length, + ) + failed_mask = torch.ones(BATCH_SIZE, dtype=torch.bool) + self._status = SkillStatus.FAILED + self._result = SkillResult( + status=SkillStatus.FAILED, + workflow_id=running.workflow_id, + current_call_index=None, + env_ids=running.env_ids, + success_mask=torch.zeros_like(failed_mask), + failure_mask=failed_mask, + cancelled_mask=torch.zeros_like(failed_mask), + eligible_mask=torch.zeros_like(failed_mask), + task_state=running.task_state, + events=( + ExecutionEvent( + kind=ExecutionEventKind.ACTION_PLANNING_FAILED, + timestamp=self.clock.now(), + skill_id="operate_articulation", + invocation_id="open-drawer-call", + invocation_revision=0, + invocation_index=0, + env_mask=failed_mask, + message="Articulation motion phase 'operate' failed.", + ), + ), + message="Motion planning failed before the first command.", + ) + return self._result + + +class _PartialSuccessRuntime(_FakeRuntime): + """Complete the workflow while retaining one failed environment row.""" + + def step(self) -> SkillResult: + result = super().step() + if result.status is SkillStatus.COMPLETED: + active_mask = torch.tensor([True, False]) + self._result = SkillResult( + status=SkillStatus.COMPLETED, + workflow_id=result.workflow_id, + current_call_index=None, + env_ids=result.env_ids, + success_mask=active_mask, + failure_mask=~active_mask, + cancelled_mask=torch.zeros_like(active_mask), + eligible_mask=active_mask, + task_state=result.task_state, + ) + return self._result + + +def _parallel_result( + status: SkillStatus, + *, + wait_duration: float = 0.0, +) -> ParallelSkillResult: + env_ids = torch.tensor([7, 3], dtype=torch.long) + terminal = status in { + SkillStatus.COMPLETED, + SkillStatus.FAILED, + SkillStatus.CANCELLED, + } + return ParallelSkillResult( + status=status, + env_ids=env_ids, + success_mask=( + torch.ones(BATCH_SIZE, dtype=torch.bool) + if status is SkillStatus.COMPLETED + else torch.zeros(BATCH_SIZE, dtype=torch.bool) + ), + failure_mask=( + torch.ones(BATCH_SIZE, dtype=torch.bool) + if status is SkillStatus.FAILED + else torch.zeros(BATCH_SIZE, dtype=torch.bool) + ), + cancelled_mask=( + torch.ones(BATCH_SIZE, dtype=torch.bool) + if status is SkillStatus.CANCELLED + else torch.zeros(BATCH_SIZE, dtype=torch.bool) + ), + pending_mask=( + torch.zeros(BATCH_SIZE, dtype=torch.bool) + if terminal + else torch.ones(BATCH_SIZE, dtype=torch.bool) + ), + task_state=TaskState.empty(BATCH_SIZE, "cpu"), + branch_results={}, + elapsed_steps=0, + command_count=0, + wait_duration=wait_duration, + ) + + +class _FakeParallelRuntime: + """One-grid-frame parallel coordinator used at the bridge boundary.""" + + def __init__( + self, + sink: BufferedGymCommandSink, + clock: EnvironmentStepClock, + ) -> None: + self.sink = sink + self.clock = clock + self._result = _parallel_result(SkillStatus.IDLE) + self._sent = False + self._due_at = 0.0 + self.eligible_mask: torch.Tensor | None = None + + @property + def result(self) -> ParallelSkillResult: + return self._result + + def start( + self, + *, + workflow_id: str = "parallel_workflow", + eligible_mask: torch.Tensor | None = None, + ) -> ParallelSkillResult: + del workflow_id + self.eligible_mask = None if eligible_mask is None else eligible_mask.clone() + self._result = _parallel_result(SkillStatus.RUNNING) + return self._result + + def step(self) -> ParallelSkillResult: + if not self._sent: + self.sink.send(_joint_frame(duration=STEP_DT), timeout=1.0) + self._sent = True + self._due_at = self.clock.now() + STEP_DT + remaining = max(self._due_at - self.clock.now(), 0.0) + self._result = ( + _parallel_result(SkillStatus.RUNNING, wait_duration=remaining) + if remaining > 1.0e-9 + else _parallel_result(SkillStatus.COMPLETED) + ) + return self._result + + def cancel(self, reason: str) -> ParallelSkillResult: + del reason + self._result = _parallel_result(SkillStatus.CANCELLED) + return self._result + + +class _GridLaneRuntime: + """Small branch runtime used with the real parallel coordinator and sink.""" + + def __init__( + self, + sink: ParallelLaneCommandSink, + script: tuple[tuple[SkillStatus, RuntimeCommandFrame | None], ...], + ) -> None: + self.sink = sink + self.script = script + self.step_count = 0 + self._result = _skill_result(SkillStatus.IDLE) + + @property + def result(self) -> SkillResult: + return self._result + + def start( + self, + *calls: object, + workflow_id: str, + eligible_mask: torch.Tensor | None = None, + ) -> SkillResult: + del calls, eligible_mask + self._result = _skill_result(SkillStatus.RUNNING, workflow_id=workflow_id) + return self._result + + def step(self) -> SkillResult: + status, frame = self.script[min(self.step_count, len(self.script) - 1)] + self.step_count += 1 + if frame is not None: + self.sink.send(frame, timeout=1.0) + if status is not SkillStatus.RUNNING: + last_frame = frame or self.sink.last_frame + assert last_frame is not None + self.sink.hold(last_frame.targets, _context(), timeout=1.0) + self._result = _skill_result( + status, workflow_id=self._result.workflow_id or "lane" + ) + return self._result + + def deactivate_rows( + self, + env_mask: torch.Tensor, + *, + reason: str, + ) -> SkillResult: + del env_mask, reason + return self._result + + def cancel(self, reason: str) -> SkillResult: + del reason + last_frame = self.sink.last_frame + if last_frame is not None: + self.sink.cancel(last_frame.targets, timeout=1.0) + self.sink.hold(last_frame.targets, _context(), timeout=1.0) + self._result = SkillResult( + status=SkillStatus.CANCELLED, + workflow_id=self._result.workflow_id, + current_call_index=None, + env_ids=torch.tensor([7, 3], dtype=torch.long), + success_mask=torch.zeros(BATCH_SIZE, dtype=torch.bool), + failure_mask=torch.zeros(BATCH_SIZE, dtype=torch.bool), + cancelled_mask=torch.ones(BATCH_SIZE, dtype=torch.bool), + eligible_mask=torch.zeros(BATCH_SIZE, dtype=torch.bool), + task_state=TaskState.empty(BATCH_SIZE, "cpu"), + ) + return self._result + + +def _grid_frame( + control_part: str, + joint_id: int, + value: float, +) -> RuntimeCommandFrame: + return RuntimeCommandFrame( + commands=( + EndpointCommand( + JointPositionTarget(control_part, (joint_id,)), + JointPositionPayload( + torch.full((BATCH_SIZE, 1), value, dtype=torch.float32) + ), + ), + ), + active_mask=torch.ones(BATCH_SIZE, dtype=torch.bool), + env_ids=torch.tensor([7, 3], dtype=torch.long), + hold_duration=torch.full((BATCH_SIZE,), STEP_DT), + ) + + +class _PostPolicyPort: + def __init__(self, action: torch.Tensor) -> None: + self.action = action + self.seen: list[object] = [] + self.active_masks: list[torch.Tensor] = [] + + def validate_policy(self, policy: object, *, segment: object) -> None: + del policy, segment + + def actions( + self, + policy: object, + *, + segment: object, + active_mask: torch.Tensor, + ): + del segment + self.seen.append(policy) + self.active_masks.append(active_mask.clone()) + yield self.action + + +class _ValidatorPort: + def __init__(self, result: torch.Tensor) -> None: + self.result = result + self.seen: list[object] = [] + + def validate_validator(self, validator: object, *, segment: object) -> None: + del validator, segment + + def validate(self, validator: object, *, segment: object) -> torch.Tensor: + del segment + self.seen.append(validator) + return self.result + + +class _MetadataPostPolicyPort(_PostPolicyPort): + """Post-policy test port exposing a deterministic result trace.""" + + def post_policy_metadata( + self, + policy: object, + *, + segment: object, + ) -> dict[str, object]: + del policy, segment + return { + "status": "timed_out", + "state": { + "elapsed_steps": 1, + "settled_mask": [True, False], + "timeout_mask": [False, True], + }, + } + + def post_policy_result( + self, + policy: object, + *, + segment: object, + ) -> torch.Tensor: + del policy, segment + return torch.tensor([True, False]) + + +class _MetadataValidatorPort(_ValidatorPort): + """Validator test port exposing observed error metadata.""" + + def validator_metadata( + self, + validator: object, + *, + segment: object, + ) -> dict[str, object]: + del validator, segment + return {"position_error": [0.01, 0.10]} + + +class _FailingPostPolicyPort: + """Raise from lazy policy iteration after the runtime reached a safe hold.""" + + def validate_policy(self, policy: object, *, segment: object) -> None: + del policy, segment + + def actions( + self, + policy: object, + *, + segment: object, + active_mask: torch.Tensor, + ): + del policy, segment, active_mask + if False: + yield torch.empty(0) + raise RuntimeError("post-policy observation failed") + + +class _AcceptParallelSafety: + """Test-only authoritative gate that accepts the supplied merged frame.""" + + def validate( + self, + *, + branch_frames: dict[str, RuntimeCommandFrame], + merged_frame: RuntimeCommandFrame, + ) -> None: + assert branch_frames + assert isinstance(merged_frame, RuntimeCommandFrame) + + +def _bridge( + *, + duration: float, + segment: _FakeSegment | None = None, + post_policy_port: object | None = None, + validator_port: object | None = None, + parallel_safety_validator: object | None = None, +) -> tuple[AtomicDemoBridge, _FakeRuntime, EnvironmentStepClock]: + clock = EnvironmentStepClock(STEP_DT) + encoder = RuntimeCommandFrameEncoder( + _QposProvider(torch.zeros(BATCH_SIZE, ROBOT_DOF)) + ) + sink = BufferedGymCommandSink(encoder, clock) + runtime = _FakeRuntime(sink, clock, _joint_frame(duration=duration)) + bridge = AtomicDemoBridge( + _FakeProgram(_FakeSegment() if segment is None else segment), + runtime, + sink, + clock, + post_policy_port=post_policy_port, + validator_port=validator_port, + parallel_safety_validator=parallel_safety_validator, + ) + return bridge, runtime, clock + + +def test_environment_step_clock_advances_only_explicitly() -> None: + clock = EnvironmentStepClock(STEP_DT) + + assert clock.now() == 0.0 + assert clock.steps_for_duration(3 * STEP_DT) == 3 + with pytest.raises(EnvironmentStepTimingError, match="not an integer multiple"): + clock.steps_for_duration(0.03) + with pytest.raises(RuntimeError, match="cannot sleep"): + clock.sleep(STEP_DT) + assert clock.now() == 0.0 + + clock.advance_after_env_step() + assert clock.step_index == 1 + assert clock.now() == pytest.approx(STEP_DT) + + +def test_observation_provider_reorders_qpos_by_stable_env_id() -> None: + context = _context( + qpos=torch.tensor([[7.0, 7.1, 7.2, 7.3, 7.4], [3.0, 3.1, 3.2, 3.3, 3.4]]) + ) + provider = GymPlanningObservationProvider(lambda task_state: context) + + observed = provider.observe(context.task) + reordered = provider.current_qpos(torch.tensor([3, 7], dtype=torch.long)) + + assert observed is context + assert torch.equal(reordered[0], context.robot.qpos[1]) + assert torch.equal(reordered[1], context.robot.qpos[0]) + + +def test_joint_encoder_emits_full_qpos_and_holds_inactive_rows() -> None: + qpos = torch.arange(BATCH_SIZE * ROBOT_DOF, dtype=torch.float32).reshape( + BATCH_SIZE, ROBOT_DOF + ) + encoder = RuntimeCommandFrameEncoder(_QposProvider(qpos)) + frame = _joint_frame( + duration=STEP_DT, + active_mask=torch.tensor([True, False]), + ) + + action = encoder.encode(frame) + + assert isinstance(action, torch.Tensor) + assert action.shape == qpos.shape + assert torch.equal(action[0, torch.tensor([1, 3])], torch.tensor([10.0, 30.0])) + assert torch.equal(action[0, torch.tensor([0, 2, 4])], qpos[0, [0, 2, 4]]) + assert torch.equal(action[1], qpos[1]) + + +def test_frame_encoder_supports_registered_future_transport() -> None: + encoder = RuntimeCommandFrameEncoder( + _QposProvider(torch.zeros(BATCH_SIZE, ROBOT_DOF)) + ) + frame = _dummy_frame() + with pytest.raises(UnsupportedRuntimeTransportError, match="test.transport"): + encoder.encode(frame) + + encoder.register_transport(_DummyTransportEncoder()) + action = encoder.encode(frame) + + assert isinstance(action, torch.Tensor) + assert action[0, 0].item() == 4.0 + assert action[1, 0].item() == 0.0 + + +def test_buffered_sink_rejects_off_grid_frame_before_buffering() -> None: + clock = EnvironmentStepClock(STEP_DT) + sink = BufferedGymCommandSink( + RuntimeCommandFrameEncoder(_QposProvider(torch.zeros(BATCH_SIZE, ROBOT_DOF))), + clock, + ) + + with pytest.raises(EnvironmentStepTimingError, match="hold_duration"): + sink.send(_joint_frame(duration=0.03), timeout=1.0) + + assert sink.pending_count == 0 + assert clock.step_index == 0 + + +def test_buffered_sink_buffers_command_hold_and_cancel_without_stepping() -> None: + clock = EnvironmentStepClock(STEP_DT) + sink = BufferedGymCommandSink( + RuntimeCommandFrameEncoder(_QposProvider(torch.zeros(BATCH_SIZE, ROBOT_DOF))), + clock, + ) + frame = _joint_frame(duration=STEP_DT) + + acknowledgement = sink.send(frame, timeout=1.0) + assert acknowledgement.accepted + assert sink.pending_count == 1 + action = sink.pop() + assert isinstance(action, ProcessedEnvAction) + assert action.metadata["bridge_action_kind"] == "runtime_command" + assert clock.step_index == 0 + + sink.hold(frame.targets, _context(), timeout=1.0) + assert sink.pending_count == 1 + sink.cancel(frame.targets, timeout=1.0) + assert sink.pending_count == 0 + assert clock.step_index == 0 + + +def test_buffered_sink_notifies_observer_only_after_successful_buffering() -> None: + """Observer acceptance follows encoding and owns a command snapshot.""" + clock = EnvironmentStepClock(STEP_DT) + observer = _RecordingAcceptedCommandObserver() + sink = BufferedGymCommandSink( + RuntimeCommandFrameEncoder(_QposProvider(torch.zeros(BATCH_SIZE, ROBOT_DOF))), + clock, + accepted_command_observer=observer, + ) + observer.sink = sink + frame = _joint_frame(duration=STEP_DT) + + sink.send(frame, timeout=1.0) + + assert observer.accepted_pending_counts == [1] + assert len(observer.accepted_frames) == 1 + observed = observer.accepted_frames[0] + assert observed is not frame + assert torch.equal(observed.env_ids, frame.env_ids) + assert observed.env_ids.data_ptr() != frame.env_ids.data_ptr() + + +def test_buffered_sink_does_not_notify_observer_when_encoding_fails() -> None: + """A frame that never reaches the buffer cannot establish evidence.""" + clock = EnvironmentStepClock(STEP_DT) + observer = _RecordingAcceptedCommandObserver() + sink = BufferedGymCommandSink( + RuntimeCommandFrameEncoder(_QposProvider(torch.zeros(BATCH_SIZE, ROBOT_DOF))), + clock, + accepted_command_observer=observer, + ) + observer.sink = sink + + with pytest.raises(UnsupportedRuntimeTransportError, match="test.transport"): + sink.send(_dummy_frame(), timeout=1.0) + + assert sink.pending_count == 0 + assert observer.accepted_frames == [] + assert observer.discard_count == 0 + + +def test_buffered_sink_rolls_back_buffer_when_observer_rejects_acceptance() -> None: + """Observer failure atomically clears the pending action and evidence state.""" + clock = EnvironmentStepClock(STEP_DT) + observer = _RecordingAcceptedCommandObserver(fail_accept=True) + sink = BufferedGymCommandSink( + RuntimeCommandFrameEncoder(_QposProvider(torch.zeros(BATCH_SIZE, ROBOT_DOF))), + clock, + accepted_command_observer=observer, + ) + observer.sink = sink + + with pytest.raises(RuntimeError, match="observer rejected accepted command"): + sink.send(_joint_frame(duration=STEP_DT), timeout=1.0) + + assert observer.accepted_pending_counts == [1] + assert sink.pending_count == 0 + assert observer.discard_count == 1 + + +def test_buffered_sink_notifies_observer_on_cancel_and_explicit_discard() -> None: + """Cancel is target-scoped while a local discard resets all evidence.""" + clock = EnvironmentStepClock(STEP_DT) + observer = _RecordingAcceptedCommandObserver() + sink = BufferedGymCommandSink( + RuntimeCommandFrameEncoder(_QposProvider(torch.zeros(BATCH_SIZE, ROBOT_DOF))), + clock, + accepted_command_observer=observer, + ) + observer.sink = sink + frame = _joint_frame(duration=STEP_DT) + sink.send(frame, timeout=1.0) + + sink.cancel(frame.targets, timeout=1.0) + + assert sink.pending_count == 0 + assert observer.cancelled_pending_counts == [0] + assert len(observer.cancelled_targets) == 1 + assert observer.cancelled_targets[0][0].address_fingerprint == ( + frame.targets[0].address_fingerprint + ) + assert observer.discard_count == 0 + + sink.send(frame, timeout=1.0) + sink.discard_pending() + assert sink.pending_count == 0 + assert observer.discard_count == 1 + + +def test_atomic_demo_bridge_is_lazy_and_waits_with_hold_actions() -> None: + bridge, runtime, clock = _bridge(duration=2 * STEP_DT) + + demo_segment = next(bridge.iter_segments()) + assert runtime.start_count == 0 + with pytest.raises(RuntimeError, match="before its action iterable"): + demo_segment.validator() + + actions = iter(demo_segment.actions) + command = next(actions) + assert runtime.start_count == 1 + assert runtime.calls == ("pick", "place") + assert clock.step_index == 0 + assert command.metadata["bridge_action_kind"] == "runtime_command" + assert command.metadata["environment_step"] == 0 + + wait_hold = next(actions) + assert clock.step_index == 1 + assert wait_hold.metadata["bridge_action_kind"] == "runtime_wait_hold" + assert torch.equal(wait_hold.value, command.value) + + with pytest.raises(StopIteration): + next(actions) + assert clock.step_index == 2 + assert runtime.status is SkillStatus.COMPLETED + assert runtime.cancel_count == 0 + assert demo_segment.validator().tolist() == [True, True] + + +def test_closing_without_abort_handshake_fails_loudly_and_does_not_ack() -> None: + bridge, runtime, clock = _bridge(duration=2 * STEP_DT) + actions = iter(next(bridge.iter_segments()).actions) + + next(actions) + with pytest.raises(DemoBridgeError, match="abort_actions"): + actions.close() + + assert clock.step_index == 0 + assert runtime.cancel_count == 1 + + +def test_abort_handshake_discards_unconsumed_command_and_yields_safe_hold() -> None: + bridge, runtime, clock = _bridge(duration=2 * STEP_DT) + segment = next(bridge.iter_segments()) + actions = iter(segment.actions) + + command = next(actions) + assert command.metadata["bridge_action_kind"] == "runtime_command" + assert segment.abort_actions is not None + emergency = iter(segment.abort_actions("operator stop", last_action_consumed=False)) + hold = next(emergency) + + assert hold.metadata["bridge_action_kind"] == "runtime_abort_safe_hold" + assert clock.step_index == 0 + with pytest.raises(StopIteration): + next(emergency) + assert clock.step_index == 1 + assert runtime.cancel_count == 1 + assert runtime.sink.pending_count == 0 + assert segment.metadata["runtime"]["status"] == "cancelled" + assert segment.metadata["runtime"]["masks"]["cancelled"] == [True, True] + with pytest.raises(RuntimeError, match="already started"): + next( + iter( + segment.abort_actions( + "duplicate stop", + last_action_consumed=False, + ) + ) + ) + actions.close() + + +def test_abort_handshake_acknowledges_consumed_command_exactly_once() -> None: + bridge, runtime, clock = _bridge(duration=2 * STEP_DT) + segment = next(bridge.iter_segments()) + actions = iter(segment.actions) + + next(actions) + assert segment.abort_actions is not None + emergency = iter( + segment.abort_actions("environment failure", last_action_consumed=True) + ) + hold = next(emergency) + + assert clock.step_index == 1 + assert hold.metadata["bridge_action_kind"] == "runtime_abort_safe_hold" + with pytest.raises(StopIteration): + next(emergency) + assert clock.step_index == 2 + assert runtime.cancel_count == 1 + actions.close() + + +def test_abort_replays_unconsumed_terminal_safe_hold_without_recancelling() -> None: + clock = EnvironmentStepClock(STEP_DT) + sink = BufferedGymCommandSink( + RuntimeCommandFrameEncoder(_QposProvider(torch.zeros(BATCH_SIZE, ROBOT_DOF))), + clock, + ) + runtime = _TerminalHoldRuntime(sink, clock, _joint_frame(duration=STEP_DT)) + bridge = AtomicDemoBridge(_FakeProgram(_FakeSegment()), runtime, sink, clock) + segment = next(bridge.iter_segments()) + actions = iter(segment.actions) + + command = next(actions) + assert command.metadata["bridge_action_kind"] == "runtime_command" + terminal_hold = next(actions) + assert terminal_hold.metadata["bridge_action_kind"] == "runtime_safe_hold" + assert runtime.status is SkillStatus.COMPLETED + assert clock.step_index == 1 + + assert segment.abort_actions is not None + emergency = iter( + segment.abort_actions("stop before hold", last_action_consumed=False) + ) + replay = next(emergency) + assert replay.metadata["bridge_action_kind"] == "runtime_abort_safe_hold" + assert torch.equal(replay.value, terminal_hold.value) + with pytest.raises(StopIteration): + next(emergency) + + assert clock.step_index == 2 + assert runtime.cancel_count == 0 + assert sink.pending_count == 0 + actions.close() + + +def test_post_policy_interruption_replays_last_runtime_safe_hold() -> None: + clock = EnvironmentStepClock(STEP_DT) + sink = BufferedGymCommandSink( + RuntimeCommandFrameEncoder(_QposProvider(torch.zeros(BATCH_SIZE, ROBOT_DOF))), + clock, + ) + runtime = _TerminalHoldRuntime(sink, clock, _joint_frame(duration=STEP_DT)) + post_policy = object() + segment_spec = _FakeSegment(post_policies=(post_policy,)) + bridge = AtomicDemoBridge( + _FakeProgram(segment_spec), + runtime, + sink, + clock, + post_policy_port=_PostPolicyPort(torch.ones(BATCH_SIZE, ROBOT_DOF)), + ) + segment = next(bridge.iter_segments()) + actions = iter(segment.actions) + + assert next(actions).metadata["bridge_action_kind"] == "runtime_command" + assert next(actions).metadata["bridge_action_kind"] == "runtime_safe_hold" + post_action = next(actions) + assert post_action.metadata["bridge_action_kind"] == "program_post_policy" + assert clock.step_index == 2 + + assert segment.abort_actions is not None + emergency = iter( + segment.abort_actions("post policy stop", last_action_consumed=False) + ) + replay = next(emergency) + assert replay.metadata["bridge_action_kind"] == "runtime_abort_safe_hold" + assert torch.equal(replay.value, torch.zeros(BATCH_SIZE, ROBOT_DOF)) + with pytest.raises(StopIteration): + next(emergency) + + assert clock.step_index == 3 + assert runtime.cancel_count == 0 + actions.close() + + +class _BridgeExecutorEnv: + """Minimal demo executor proving abort actions cross the Gym boundary.""" + + def __init__( + self, + bridge: AtomicDemoBridge, + *, + fail_first_mask: torch.Tensor | None = None, + raise_first_step: bool = False, + ) -> None: + self.bridge = bridge + self.fail_first_mask = ( + torch.zeros(BATCH_SIZE, dtype=torch.bool) + if fail_first_mask is None + else fail_first_mask.clone() + ) + self.raise_first_step = raise_first_step + self.num_envs = BATCH_SIZE + self.steps: list[ProcessedEnvAction] = [] + self._demo_no_auto_reset = False + + @property + def unwrapped(self) -> _BridgeExecutorEnv: + return self + + def create_demo_segments(self): + return self.bridge.iter_segments() + + def step(self, action: ProcessedEnvAction): + assert isinstance(action, ProcessedEnvAction) + self.steps.append(action.snapshot()) + if self.raise_first_step and len(self.steps) == 1: + raise RuntimeError("simulated environment failure") + failed = ( + self.fail_first_mask + if len(self.steps) == 1 + else torch.zeros(BATCH_SIZE, dtype=torch.bool) + ) + return ( + None, + torch.zeros(BATCH_SIZE), + torch.zeros(BATCH_SIZE, dtype=torch.bool), + torch.zeros(BATCH_SIZE, dtype=torch.bool), + {"fail": failed}, + ) + + def _mask_demo_action( + self, + action: ProcessedEnvAction, + active_mask: tuple[bool, ...], + ) -> ProcessedEnvAction: + del active_mask + return action.snapshot() + + +def test_zero_command_terminal_runtime_failure_preserves_trace_and_validates_once() -> ( + None +): + validator = object() + first = _FakeSegment( + calls=(_FakeCompiledCall(0, "operate_articulation"),), + validators=(validator,), + ) + second = _FakeSegment( + segment_index=1, + segment_id="segment-1", + name="must-not-start", + calls=(_FakeCompiledCall(1, "place"),), + source_path=("program", "steps", 1), + ) + clock = EnvironmentStepClock(STEP_DT) + sink = BufferedGymCommandSink( + RuntimeCommandFrameEncoder(_QposProvider(torch.zeros(BATCH_SIZE, ROBOT_DOF))), + clock, + ) + runtime = _TerminalFailedRuntime( + sink, + clock, + _joint_frame(duration=STEP_DT), + ) + validator_port = _ValidatorPort(torch.ones(BATCH_SIZE, dtype=torch.bool)) + bridge = AtomicDemoBridge( + _FakeProgram(first, second), + runtime, + sink, + clock, + validator_port=validator_port, + ) + env = _BridgeExecutorEnv(bridge) + + result = execute_demo_episode(env) + + assert env.steps == [] + assert clock.step_index == 0 + assert runtime.start_count == 1 + assert runtime.step_count == 0 + assert validator_port.seen == [validator] + assert not result.completed + assert result.terminal_reason == "segment_validation_failed" + assert len(result.segments) == 1 + segment_result = result.segments[0] + assert segment_result.failure_reason == "segment_validation_failed" + runtime_trace = segment_result.metadata["runtime"] + assert runtime_trace["status"] == "failed" + assert ( + runtime_trace["message"] == "Motion planning failed before the first command." + ) + assert runtime_trace["events"] == [ + { + "kind": "action_planning_failed", + "timestamp": 0.0, + "skill_id": "operate_articulation", + "invocation_id": "open-drawer-call", + "invocation_revision": 0, + "invocation_index": 0, + "env_mask": [True, True], + "message": "Articulation motion phase 'operate' failed.", + } + ] + assert segment_result.metadata["validation"] == { + "env_ids": [7, 3], + "runtime_success_mask": [False, False], + "eligible_mask_before_validation": [False, False], + "post_policy_success_mask": None, + "validators": [ + { + "validator_index": 0, + "kind": "object", + "source_path": [], + "result_mask": [True, True], + "result": None, + } + ], + "accepted_mask": [False, False], + } + + +def test_sequential_start_failure_before_first_command_preserves_cause() -> None: + clock = EnvironmentStepClock(STEP_DT) + sink = BufferedGymCommandSink( + RuntimeCommandFrameEncoder(_QposProvider(torch.zeros(BATCH_SIZE, ROBOT_DOF))), + clock, + ) + runtime = _StartFailingRuntime(sink, clock, _joint_frame(duration=STEP_DT)) + bridge = AtomicDemoBridge(_FakeProgram(_FakeSegment()), runtime, sink, clock) + env = _BridgeExecutorEnv(bridge) + + with pytest.raises(RuntimeError, match="action generation") as error: + execute_demo_episode(env) + + assert isinstance(error.value.__cause__, RuntimeError) + assert str(error.value.__cause__) == "semantic runtime preflight failed" + assert env.steps == [] + assert clock.step_index == 0 + assert runtime.start_count == 1 + assert runtime.cancel_count == 0 + assert sink.pending_count == 0 + + +def test_parallel_construction_failure_before_first_command_preserves_cause( + monkeypatch: pytest.MonkeyPatch, +) -> None: + block = _FakeParallelBlock( + branches=( + _FakeParallelBranch(0, (_FakeCompiledCall(0, "left"),)), + _FakeParallelBranch(1, (_FakeCompiledCall(1, "right"),)), + ) + ) + segment = _FakeSegment(parallel_block=block) + bridge, runtime, clock = _bridge( + duration=STEP_DT, + segment=segment, + parallel_safety_validator=_AcceptParallelSafety(), + ) + env = _BridgeExecutorEnv(bridge) + + def fail_construction(*args: object, **kwargs: object) -> _FakeParallelRuntime: + del args, kwargs + raise RuntimeError("parallel runtime construction failed") + + monkeypatch.setattr(bridge_module, "SkillRuntime", _FakeRuntime) + monkeypatch.setattr( + ParallelSkillRuntime, + "from_template", + classmethod(fail_construction), + ) + + with pytest.raises(RuntimeError, match="action generation") as error: + execute_demo_episode(env) + + assert isinstance(error.value.__cause__, RuntimeError) + assert str(error.value.__cause__) == "parallel runtime construction failed" + assert env.steps == [] + assert clock.step_index == 0 + assert runtime.start_count == 0 + assert runtime.cancel_count == 0 + assert runtime.sink.pending_count == 0 + + +def test_post_policy_timeout_is_row_local_and_preserved_in_segment_result() -> None: + segment = _FakeSegment(post_policies=(object(),)) + bridge, runtime, clock = _bridge( + duration=STEP_DT, + segment=segment, + post_policy_port=_MetadataPostPolicyPort(torch.ones(BATCH_SIZE, ROBOT_DOF)), + ) + env = _BridgeExecutorEnv(bridge) + + result = execute_demo_episode(env) + + assert result.segments[0].successes == (True, False) + assert result.segments[0].failure_reasons == ( + None, + "segment_validation_failed", + ) + assert result.segments[0].metadata["post_policies"][0]["result_mask"] == [ + True, + False, + ] + assert result.segments[0].metadata["validation"]["accepted_mask"] == [ + True, + False, + ] + assert [action.metadata["bridge_action_kind"] for action in env.steps] == [ + "runtime_command", + "program_post_policy", + ] + assert runtime.cancel_count == 0 + assert clock.step_index == 2 + + +def test_post_policy_generator_error_replays_safe_hold_before_propagating() -> None: + clock = EnvironmentStepClock(STEP_DT) + sink = BufferedGymCommandSink( + RuntimeCommandFrameEncoder(_QposProvider(torch.zeros(BATCH_SIZE, ROBOT_DOF))), + clock, + ) + runtime = _TerminalHoldRuntime(sink, clock, _joint_frame(duration=STEP_DT)) + segment_spec = _FakeSegment(post_policies=(object(),)) + bridge = AtomicDemoBridge( + _FakeProgram(segment_spec), + runtime, + sink, + clock, + post_policy_port=_FailingPostPolicyPort(), + ) + env = _BridgeExecutorEnv(bridge) + + with pytest.raises(RuntimeError, match="action generation") as error: + execute_demo_episode(env) + + assert isinstance(error.value.__cause__, RuntimeError) + assert [action.metadata["bridge_action_kind"] for action in env.steps] == [ + "runtime_command", + "runtime_safe_hold", + "runtime_abort_safe_hold", + ] + assert runtime.cancel_count == 0 + assert sink.pending_count == 0 + assert clock.step_index == 3 + + +def test_demo_executor_pre_step_stop_consumes_only_abort_hold() -> None: + bridge, runtime, clock = _bridge(duration=2 * STEP_DT) + env = _BridgeExecutorEnv(bridge) + checks = iter((False, True)) + + result = execute_demo_episode(env, should_stop=lambda: next(checks, True)) + + assert result.terminal_reason == "interrupted" + assert result.length == 1 + assert [action.metadata["bridge_action_kind"] for action in env.steps] == [ + "runtime_abort_safe_hold" + ] + assert clock.step_index == 1 + assert runtime.cancel_count == 1 + assert runtime.sink.pending_count == 0 + + +def test_demo_executor_post_step_failure_acknowledges_then_safe_stops() -> None: + bridge, runtime, clock = _bridge(duration=2 * STEP_DT) + env = _BridgeExecutorEnv( + bridge, + fail_first_mask=torch.ones(BATCH_SIZE, dtype=torch.bool), + ) + + result = execute_demo_episode(env) + + assert result.terminal_reason == "failure" + assert result.length == 2 + assert [action.metadata["bridge_action_kind"] for action in env.steps] == [ + "runtime_command", + "runtime_abort_safe_hold", + ] + assert clock.step_index == 2 + assert runtime.cancel_count == 1 + assert runtime.sink.pending_count == 0 + + +def test_demo_executor_safe_stops_when_regular_env_step_raises() -> None: + bridge, runtime, clock = _bridge(duration=2 * STEP_DT) + env = _BridgeExecutorEnv(bridge, raise_first_step=True) + + with pytest.raises(RuntimeError, match="emergency safe-stop"): + execute_demo_episode(env) + + assert [action.metadata["bridge_action_kind"] for action in env.steps] == [ + "runtime_command", + "runtime_abort_safe_hold", + ] + assert clock.step_index == 1 + assert runtime.cancel_count == 1 + assert runtime.sink.pending_count == 0 + + +def test_row_independent_partial_failure_does_not_abort_healthy_peer() -> None: + bridge, runtime, clock = _bridge(duration=2 * STEP_DT) + env = _BridgeExecutorEnv( + bridge, + fail_first_mask=torch.tensor([True, False]), + ) + + result = execute_demo_episode(env) + + assert result.lengths == (1, 2) + assert [action.metadata["bridge_action_kind"] for action in env.steps] == [ + "runtime_command", + "runtime_wait_hold", + ] + assert clock.step_index == 2 + assert runtime.cancel_count == 0 + assert runtime.sink.pending_count == 0 + + +def test_post_policy_and_validator_ports_stay_at_demo_boundary() -> None: + post_policy = object() + validator = object() + segment = _FakeSegment( + post_policies=(post_policy,), + validators=(validator,), + ) + post_port = _PostPolicyPort(torch.ones(BATCH_SIZE, ROBOT_DOF)) + validator_port = _ValidatorPort(torch.tensor([True, False])) + bridge, _, clock = _bridge( + duration=STEP_DT, + segment=segment, + post_policy_port=post_port, + validator_port=validator_port, + ) + demo_segment = next(bridge.iter_segments()) + actions = iter(demo_segment.actions) + + runtime_action = next(actions) + assert runtime_action.metadata["bridge_action_kind"] == "runtime_command" + post_action = next(actions) + assert clock.step_index == 1 + assert post_action.metadata["bridge_action_kind"] == "program_post_policy" + with pytest.raises(StopIteration): + next(actions) + + assert clock.step_index == 2 + assert post_port.seen == [post_policy] + assert len(post_port.active_masks) == 1 + assert post_port.active_masks[0].tolist() == [True, True] + assert demo_segment.validator().tolist() == [True, False] + assert validator_port.seen == [validator] + + +def test_post_policy_receives_only_rows_surviving_partial_runtime_failure() -> None: + """Post-policy completion cannot be blocked by a runtime-failed row.""" + post_policy = object() + segment = _FakeSegment(post_policies=(post_policy,)) + clock = EnvironmentStepClock(STEP_DT) + sink = BufferedGymCommandSink( + RuntimeCommandFrameEncoder(_QposProvider(torch.zeros(BATCH_SIZE, ROBOT_DOF))), + clock, + ) + runtime = _PartialSuccessRuntime( + sink, + clock, + _joint_frame(duration=STEP_DT), + ) + post_port = _PostPolicyPort(torch.ones(BATCH_SIZE, ROBOT_DOF)) + bridge = AtomicDemoBridge( + _FakeProgram(segment), + runtime, + sink, + clock, + post_policy_port=post_port, + ) + demo_segment = next(bridge.iter_segments()) + + tuple(demo_segment.actions) + + assert len(post_port.active_masks) == 1 + assert post_port.active_masks[0].tolist() == [True, False] + assert demo_segment.metadata["post_policies"][0]["result_mask"] == [True, False] + assert demo_segment.validator().tolist() == [True, False] + + +def test_later_post_policy_receives_only_rows_passing_prior_policy() -> None: + """Sequential post-policies monotonically narrow their active cohort.""" + segment = _FakeSegment(post_policies=(object(), object())) + post_port = _MetadataPostPolicyPort( + torch.ones(BATCH_SIZE, ROBOT_DOF), + ) + bridge, _, _ = _bridge( + duration=STEP_DT, + segment=segment, + post_policy_port=post_port, + ) + + tuple(next(bridge.iter_segments()).actions) + + assert [mask.tolist() for mask in post_port.active_masks] == [ + [True, True], + [True, False], + ] + + +def test_segment_lifecycle_metadata_records_runtime_post_and_validation() -> None: + post_policy = object() + validator = object() + segment = _FakeSegment( + post_policies=(post_policy,), + validators=(validator,), + ) + bridge, _, _ = _bridge( + duration=STEP_DT, + segment=segment, + post_policy_port=_MetadataPostPolicyPort(torch.ones(BATCH_SIZE, ROBOT_DOF)), + validator_port=_MetadataValidatorPort(torch.tensor([True, False])), + ) + demo_segment = next(bridge.iter_segments()) + + tuple(demo_segment.actions) + assert demo_segment.metadata["runtime"]["status"] == "completed" + assert demo_segment.metadata["post_policies"] == [ + { + "policy_index": 0, + "kind": "object", + "source_path": [], + "result_mask": [True, False], + "result": { + "status": "timed_out", + "state": { + "elapsed_steps": 1, + "settled_mask": [True, False], + "timeout_mask": [False, True], + }, + }, + } + ] + + assert demo_segment.validator().tolist() == [True, False] + assert demo_segment.metadata["validation"] == { + "env_ids": [7, 3], + "runtime_success_mask": [True, True], + "eligible_mask_before_validation": [True, True], + "post_policy_success_mask": [True, False], + "validators": [ + { + "validator_index": 0, + "kind": "object", + "source_path": [], + "result_mask": [True, False], + "result": {"position_error": [0.01, 0.1]}, + } + ], + "accepted_mask": [True, False], + } + json.dumps(demo_segment.metadata, allow_nan=False, sort_keys=True) + + +def test_declared_post_policy_requires_explicit_port() -> None: + segment = _FakeSegment(post_policies=(object(),)) + bridge, _, _ = _bridge(duration=STEP_DT, segment=segment) + + with pytest.raises(DemoBridgeError, match="no SegmentPostPolicyPort"): + tuple(next(bridge.iter_segments()).actions) + + +def test_bridge_marks_segments_row_independent_and_retains_failed_rows() -> None: + first_validator = object() + first = _FakeSegment(validators=(first_validator,)) + second = _FakeSegment( + segment_index=1, + segment_id="segment-1", + name="place-next", + calls=(_FakeCompiledCall(2, "place-next"),), + source_path=("program", "steps", 1), + ) + clock = EnvironmentStepClock(STEP_DT) + sink = BufferedGymCommandSink( + RuntimeCommandFrameEncoder(_QposProvider(torch.zeros(BATCH_SIZE, ROBOT_DOF))), + clock, + ) + runtime = _FakeRuntime(sink, clock, _joint_frame(duration=STEP_DT)) + bridge = AtomicDemoBridge( + _FakeProgram(first, second), + runtime, + sink, + clock, + validator_port=_ValidatorPort(torch.tensor([True, False])), + ) + segments = iter(bridge.iter_segments()) + + first_demo = next(segments) + assert first_demo.failure_policy == "row_independent" + tuple(first_demo.actions) + assert first_demo.validator().tolist() == [True, False] + + second_demo = next(segments) + tuple(second_demo.actions) + assert runtime.eligible_masks[0] is None + assert runtime.eligible_masks[1].tolist() == [True, False] + assert second_demo.validator().tolist() == [True, False] + + +def test_bridge_refuses_next_segment_when_validator_was_skipped() -> None: + first = _FakeSegment(calls=(_FakeCompiledCall(0, "pick"),)) + second = _FakeSegment( + segment_index=1, + segment_id="segment-1", + name="place", + calls=(_FakeCompiledCall(1, "place"),), + source_path=("program", "steps", 1), + ) + clock = EnvironmentStepClock(STEP_DT) + sink = BufferedGymCommandSink( + RuntimeCommandFrameEncoder(_QposProvider(torch.zeros(BATCH_SIZE, ROBOT_DOF))), + clock, + ) + runtime = _FakeRuntime(sink, clock, _joint_frame(duration=STEP_DT)) + segments = iter(AtomicDemoBridge(_FakeProgram(first, second), runtime, sink, clock)) + + first_demo = next(segments) + tuple(first_demo.actions) + + with pytest.raises(DemoBridgeError, match="validator must be called"): + next(segments) + assert runtime.start_count == 1 + + +def test_demo_executor_consumes_validation_before_requesting_next_segment() -> None: + first_validator = object() + first = _FakeSegment( + calls=(_FakeCompiledCall(0, "pick"),), + validators=(first_validator,), + ) + second = _FakeSegment( + segment_index=1, + segment_id="segment-1", + name="place", + calls=(_FakeCompiledCall(1, "place"),), + source_path=("program", "steps", 1), + ) + clock = EnvironmentStepClock(STEP_DT) + sink = BufferedGymCommandSink( + RuntimeCommandFrameEncoder(_QposProvider(torch.zeros(BATCH_SIZE, ROBOT_DOF))), + clock, + ) + runtime = _FakeRuntime(sink, clock, _joint_frame(duration=STEP_DT)) + validator_port = _ValidatorPort(torch.ones(BATCH_SIZE, dtype=torch.bool)) + bridge = AtomicDemoBridge( + _FakeProgram(first, second), + runtime, + sink, + clock, + validator_port=validator_port, + ) + + result = execute_demo_episode(_BridgeExecutorEnv(bridge)) + + assert result.completed + assert len(result.segments) == 2 + assert [segment.success for segment in result.segments] == [True, True] + assert runtime.start_count == 2 + assert validator_port.seen == [first_validator] + + +def test_sequential_segment_analyzes_downstream_calls_but_executes_own_prefix() -> None: + first = _FakeSegment( + calls=(_FakeCompiledCall(0, "pick"),), + ) + second = _FakeSegment( + segment_index=1, + segment_id="segment-1", + name="place", + calls=(_FakeCompiledCall(1, "place"),), + source_path=("program", "steps", 1), + ) + clock = EnvironmentStepClock(STEP_DT) + sink = BufferedGymCommandSink( + RuntimeCommandFrameEncoder(_QposProvider(torch.zeros(BATCH_SIZE, ROBOT_DOF))), + clock, + ) + runtime = _FakeRuntime(sink, clock, _joint_frame(duration=STEP_DT)) + bridge = AtomicDemoBridge(_FakeProgram(first, second), runtime, sink, clock) + segments = iter(bridge.iter_segments()) + + first_demo = next(segments) + tuple(first_demo.actions) + assert first_demo.validator().tolist() == [True, True] + + assert runtime.calls == ("pick", "place") + assert runtime.execution_prefix_lengths == [1] + + tuple(next(segments).actions) + assert runtime.calls == ("place",) + assert runtime.execution_prefix_lengths == [1, 1] + + +def test_parallel_segment_preserves_branches_barrier_and_adopts_state( + monkeypatch: pytest.MonkeyPatch, +) -> None: + left_calls = (_FakeCompiledCall(0, "left-pick"),) + right_calls = ( + _FakeCompiledCall(1, "right-pick"), + _FakeCompiledCall(2, "right-place"), + ) + block = _FakeParallelBlock( + branches=( + _FakeParallelBranch(0, left_calls), + _FakeParallelBranch(1, right_calls), + ) + ) + segment = _FakeSegment( + calls=left_calls + right_calls, + parallel_block=block, + ) + safety_validator = _AcceptParallelSafety() + bridge, runtime, clock = _bridge( + duration=STEP_DT, + segment=segment, + parallel_safety_validator=safety_validator, + ) + captured: dict[str, Any] = {} + fake_parallel = _FakeParallelRuntime(runtime.sink, clock) + + def from_template( + cls: type[ParallelSkillRuntime], + template_runtime: object, + branch_calls: dict[str, tuple[object, ...]], + command_sink: object, + timing_policy: object, + supplied_safety_validator: object, + *, + timeout_steps: int, + failure_policy: str, + workflow_id: str, + branch_paths: dict[str, tuple[object, ...]], + ) -> _FakeParallelRuntime: + del cls + captured.update( + { + "template_runtime": template_runtime, + "branch_calls": branch_calls, + "command_sink": command_sink, + "timing_policy": timing_policy, + "safety_validator": supplied_safety_validator, + "timeout_steps": timeout_steps, + "failure_policy": failure_policy, + "workflow_id": workflow_id, + "branch_paths": branch_paths, + } + ) + return fake_parallel + + monkeypatch.setattr(bridge_module, "SkillRuntime", _FakeRuntime) + monkeypatch.setattr( + ParallelSkillRuntime, + "from_template", + classmethod(from_template), + ) + + demo_segment = next(bridge.iter_segments()) + actions = tuple(demo_segment.actions) + + assert len(actions) == 1 + assert captured["template_runtime"] is runtime + assert captured["command_sink"] is runtime.sink + assert captured["branch_calls"] == { + "branch_0": ("left-pick",), + "branch_1": ("right-pick", "right-place"), + } + assert captured["timing_policy"].step_dt == STEP_DT + assert captured["safety_validator"] is safety_validator + assert captured["timeout_steps"] == 17 + assert captured["failure_policy"] == "fail_fast" + assert captured["workflow_id"].endswith(":parallel_analysis") + assert captured["branch_paths"] == { + "branch_0": segment.source_path, + "branch_1": segment.source_path, + } + assert len(runtime.adopted_states) == 1 + assert demo_segment.failure_policy == "row_independent" + assert demo_segment.metadata["runtime"]["kind"] == "parallel_skill_result" + assert demo_segment.metadata["runtime"]["status"] == "completed" + assert demo_segment.metadata["runtime"]["masks"]["success"] == [True, True] + assert demo_segment.validator().tolist() == [True, True] + + +def test_real_parallel_coordinator_buffers_one_ordered_gym_action_per_step() -> None: + clock = EnvironmentStepClock(STEP_DT) + sink = BufferedGymCommandSink( + RuntimeCommandFrameEncoder(_QposProvider(torch.zeros(BATCH_SIZE, ROBOT_DOF))), + clock, + ) + left_sink = ParallelLaneCommandSink() + right_sink = ParallelLaneCommandSink() + left_runtime = _GridLaneRuntime( + left_sink, + ( + (SkillStatus.RUNNING, _grid_frame("left_arm", 0, 1.0)), + (SkillStatus.COMPLETED, None), + ), + ) + right_runtime = _GridLaneRuntime( + right_sink, + ( + (SkillStatus.RUNNING, _grid_frame("right_arm", 1, 2.0)), + (SkillStatus.RUNNING, _grid_frame("right_arm", 1, 3.0)), + (SkillStatus.COMPLETED, None), + ), + ) + runtime = ParallelSkillRuntime( + ( + ParallelRuntimeBranch( + "left", + (RegisteredSemanticCall("test.left"),), + ResourceClaim(frozenset({"left_arm"}), (0,)), + left_runtime, + left_sink, + ), + ParallelRuntimeBranch( + "right", + (RegisteredSemanticCall("test.right"),), + ResourceClaim(frozenset({"right_arm"}), (1,)), + right_runtime, + right_sink, + ), + ), + sink, + clock, + ParallelTimingPolicy(STEP_DT), + _AcceptParallelSafety(), + timeout_steps=8, + ) + + runtime.start() + first = runtime.step() + assert first.status is SkillStatus.RUNNING + assert sink.pending_count == 1 + first_action = sink.pop() + assert first_action.metadata["bridge_action_kind"] == "runtime_command" + assert torch.equal(first_action.value[:, 0], torch.ones(BATCH_SIZE)) + assert torch.equal(first_action.value[:, 1], torch.full((BATCH_SIZE,), 2.0)) + clock.advance_after_env_step() + + padded = runtime.step() + assert padded.status is SkillStatus.RUNNING + assert sink.pending_count == 1 + padding_action = sink.pop() + assert padding_action.metadata["bridge_action_kind"] == "runtime_safe_hold" + assert torch.equal(padding_action.value, torch.zeros(BATCH_SIZE, ROBOT_DOF)) + lane_steps = (left_runtime.step_count, right_runtime.step_count) + clock.advance_after_env_step() + + deferred = runtime.step() + assert deferred.status is SkillStatus.RUNNING + assert sink.pending_count == 1 + deferred_action = sink.pop() + assert deferred_action.metadata["bridge_action_kind"] == "runtime_command" + assert torch.equal(deferred_action.value[:, 0], torch.zeros(BATCH_SIZE)) + assert torch.equal( + deferred_action.value[:, 1], + torch.full((BATCH_SIZE,), 3.0), + ) + assert (left_runtime.step_count, right_runtime.step_count) == lane_steps + clock.advance_after_env_step() + + completed = runtime.step() + assert completed.status is SkillStatus.COMPLETED + assert sink.pending_count == 1 + terminal_hold = sink.pop() + assert terminal_hold.metadata["bridge_action_kind"] == "runtime_safe_hold" + assert completed.command_count == 2 + clock.advance_after_env_step() + assert clock.step_index == 4 + + +def test_parallel_segment_fails_closed_without_safety_validator() -> None: + block = _FakeParallelBlock( + branches=( + _FakeParallelBranch(0, (_FakeCompiledCall(0, "left"),)), + _FakeParallelBranch(1, (_FakeCompiledCall(1, "right"),)), + ) + ) + segment = _FakeSegment(parallel_block=block) + bridge, runtime, _ = _bridge(duration=STEP_DT, segment=segment) + + with pytest.raises(DemoBridgeError, match="requires an explicit"): + tuple(next(bridge.iter_segments()).actions) + + assert runtime.start_count == 0 diff --git a/tests/gym/envs/expert_program/test_cfg.py b/tests/gym/envs/expert_program/test_cfg.py new file mode 100644 index 000000000..a504f371c --- /dev/null +++ b/tests/gym/envs/expert_program/test_cfg.py @@ -0,0 +1,206 @@ +# ---------------------------------------------------------------------------- +# 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 typed Expert Program configuration values.""" + +from __future__ import annotations + +import math + +import pytest + +from embodichain.lab.gym.envs.expert_program import ( + MAX_REPEAT_COUNT, + CyclicPoseTargetCfg, + ExpertProgramCfg, + ExpertProgramIntegrationCfg, + HandOverCfg, + InvokeCfg, + ObjectNearTargetValidatorCfg, + PickCfg, + PlaceCfg, + PoseCfg, + RegisteredSemanticCallCfg, + RepeatCfg, + SegmentCfg, + SequenceCfg, + TargetRefCfg, + WaitStablePostCfg, +) +from embodichain.utils.configclass import is_configclass + + +def _integration() -> ExpertProgramIntegrationCfg: + """Return one valid provider-free integration selection.""" + return ExpertProgramIntegrationCfg( + robot_profile="auto", + scene_registry="env", + runtime_preset="safe", + ) + + +def _pick_invoke() -> InvokeCfg: + """Return one minimal semantic invocation.""" + return InvokeCfg(call=PickCfg(object="cube")) + + +def test_every_public_schema_value_uses_configclass() -> None: + classes = ( + ExpertProgramCfg, + ExpertProgramIntegrationCfg, + PoseCfg, + TargetRefCfg, + CyclicPoseTargetCfg, + PickCfg, + PlaceCfg, + HandOverCfg, + RegisteredSemanticCallCfg, + WaitStablePostCfg, + ObjectNearTargetValidatorCfg, + InvokeCfg, + SequenceCfg, + RepeatCfg, + SegmentCfg, + ) + + assert all(is_configclass(cls) for cls in classes) + + +def test_call_configs_own_resources_and_registered_payloads() -> None: + resources = {"primary": "left_actor"} + arguments = {"waypoints": [1, {"enabled": True}]} + pick = PickCfg(object="cube", resources=resources) + registered = RegisteredSemanticCallCfg( + call_id="example.inspect", + arguments=arguments, + ) + + resources["primary"] = "right_actor" + arguments["waypoints"][1]["enabled"] = False + + assert pick.resources == {"primary": "left_actor"} + assert registered.arguments == { + "waypoints": (1, {"enabled": True}), + } + + +@pytest.mark.parametrize("count", [False, 0, -1, MAX_REPEAT_COUNT + 1]) +def test_repeat_rejects_non_positive_non_integer_or_excessive_count( + count: object, +) -> None: + with pytest.raises(ValueError, match="count must be an integer"): + RepeatCfg(count=count, body=_pick_invoke()) + + +def test_program_rejects_nested_repeat_expansion_above_static_budget() -> None: + nested = RepeatCfg( + count=MAX_REPEAT_COUNT, + body=RepeatCfg(count=MAX_REPEAT_COUNT, body=_pick_invoke()), + ) + + with pytest.raises(ValueError, match="expands to more than"): + ExpertProgramCfg( + schema_version=1, + program_id="too_large", + integration=_integration(), + targets={}, + program=nested, + ) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({}, "exactly one"), + ( + { + "at": TargetRefCfg(target="drop"), + "on": "tray", + }, + "exactly one", + ), + ], +) +def test_place_requires_exactly_one_typed_destination( + kwargs: dict[str, object], + message: str, +) -> None: + with pytest.raises(ValueError, match=message): + PlaceCfg(object="cube", **kwargs) + + +def test_programmatic_config_rejects_unknown_target_reference() -> None: + program = InvokeCfg( + call=PlaceCfg( + object="cube", + at=TargetRefCfg(target="missing"), + ) + ) + + with pytest.raises(ValueError, match="Unknown target reference 'missing'"): + ExpertProgramCfg( + schema_version=1, + program_id="missing_target", + integration=_integration(), + targets={}, + program=program, + ) + + +@pytest.mark.parametrize( + "arguments", + [ + {"callback": lambda: None}, + {"eval": "1 + 1"}, + {"source": "env.robot.control_parts"}, + {"bad": math.inf}, + ], +) +def test_registered_call_rejects_executable_or_non_declarative_payload( + arguments: dict[str, object], +) -> None: + with pytest.raises((TypeError, ValueError)): + RegisteredSemanticCallCfg( + call_id="example.inspect", + arguments=arguments, + ) + + +def test_pose_rejects_zero_quaternion() -> None: + with pytest.raises(ValueError, match="non-zero magnitude"): + PoseCfg( + position=(0.0, 0.0, 0.0), + quaternion_wxyz=(0.0, 0.0, 0.0, 0.0), + ) + + +def test_segment_owns_post_policy_and_validator_sequences() -> None: + post = [WaitStablePostCfg(entity="cube")] + validators = [ObjectNearTargetValidatorCfg(object="cube", target="drop_pose")] + segment = SegmentCfg( + name="move_cube", + steps=SequenceCfg(items=(_pick_invoke(),)), + post=post, + validators=validators, + ) + + post.clear() + validators.clear() + + assert segment.post == (WaitStablePostCfg(entity="cube"),) + assert segment.validators == ( + ObjectNearTargetValidatorCfg(object="cube", target="drop_pose"), + ) diff --git a/tests/gym/envs/expert_program/test_compiler.py b/tests/gym/envs/expert_program/test_compiler.py new file mode 100644 index 000000000..f3a94abc8 --- /dev/null +++ b/tests/gym/envs/expert_program/test_compiler.py @@ -0,0 +1,549 @@ +# ---------------------------------------------------------------------------- +# 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 provider-free Expert Program compilation and lazy expansion.""" + +from __future__ import annotations + +from itertools import islice + +import pytest +import torch + +from embodichain.lab.gym.envs.expert_program import ( + CyclicPoseTargetCfg, + ExpertProgramCfg, + ExpertProgramCompileError, + ExpertProgramCompiler, + ExpertProgramIntegrationCfg, + HandOverCfg, + InvokeCfg, + MaterializedCompiledProgram, + ObjectNearTargetValidatorCfg, + PickCfg, + PlaceCfg, + PoseCfg, + RegisteredSemanticCallCfg, + RepeatCfg, + SegmentCfg, + SequenceCfg, + TargetRefCfg, + WaitStablePostCfg, +) +from embodichain.lab.sim.atomic_actions import Affordance, EntityState +from embodichain.lab.sim.skills.calls import ( + HandOver, + Pick, + Place, + RegisteredSemanticCall, + SemanticCallSpec, + SemanticPose, +) +from embodichain.lab.sim.skills.scene import ( + SceneAffordanceRef, + SceneArticulationRef, + SceneEntityRegistration, + SceneLinkRef, + SceneObjectRef, + SceneRegistry, +) + + +class _NeverObserveProvider: + """Record and reject every attempted dynamic scene observation.""" + + def __init__(self) -> None: + self.calls = 0 + + def observe( + self, + *, + timestamp: float, + env_ids: torch.Tensor, + ) -> EntityState: + del timestamp, env_ids + self.calls += 1 + raise AssertionError("Expert Program compilation must not observe state.") + + +def _scene_registry() -> tuple[SceneRegistry, _NeverObserveProvider]: + """Return static identities backed by a provider that must stay unused.""" + provider = _NeverObserveProvider() + cube = SceneObjectRef("cube") + tray = SceneObjectRef("tray") + arm = SceneArticulationRef("arm") + registry = SceneRegistry( + ( + SceneEntityRegistration( + ref=cube, + state_provider=provider, + aliases=("sim_cube",), + ), + SceneEntityRegistration(ref=tray, state_provider=provider), + SceneEntityRegistration(ref=arm, state_provider=provider), + SceneEntityRegistration( + ref=SceneLinkRef("arm_tcp"), + state_provider=provider, + parent=arm, + native_name="tcp", + ), + SceneEntityRegistration( + ref=SceneAffordanceRef("cube_grasp"), + aliases=("legacy_grasp",), + parent=cube, + native_name="grasp", + affordance=Affordance(), + relative_pose=torch.eye(4), + ), + SceneEntityRegistration( + ref=SceneAffordanceRef("tray_top"), + parent=tray, + native_name="top", + affordance=Affordance(), + relative_pose=torch.eye(4), + ), + ) + ) + return registry, provider + + +def _integration() -> ExpertProgramIntegrationCfg: + """Return one static integration selection.""" + return ExpertProgramIntegrationCfg( + robot_profile="auto", + scene_registry="env", + runtime_preset="safe", + ) + + +def _pose(x: float, y: float = 0.0, z: float = 0.2) -> PoseCfg: + """Build one target pose with an identity WXYZ quaternion.""" + return PoseCfg( + position=(x, y, z), + quaternion_wxyz=(1.0, 0.0, 0.0, 0.0), + ) + + +def _program( + node: InvokeCfg | SequenceCfg | RepeatCfg | SegmentCfg, + *, + targets: dict[str, CyclicPoseTargetCfg] | None = None, + program_id: str = "test_program", +) -> ExpertProgramCfg: + """Build one valid Version 1 program around a supplied node.""" + return ExpertProgramCfg( + schema_version=1, + program_id=program_id, + integration=_integration(), + program=node, + targets={} if targets is None else targets, + ) + + +def _assert_pose_equal(actual: SemanticPose, expected: SemanticPose) -> None: + """Compare owned pose tensor values.""" + assert torch.allclose(actual.position, expected.position) + assert torch.allclose(actual.quaternion_wxyz, expected.quaternion_wxyz) + + +def _assert_semantic_call_equal( + actual: SemanticCallSpec, + expected: SemanticCallSpec, +) -> None: + """Compare exact semantic call values whose public classes use eq=False.""" + assert type(actual) is type(expected) + assert dict(actual.resources) == dict(expected.resources) + if type(actual) is Pick and type(expected) is Pick: + assert actual.object == expected.object + assert actual.grasp == expected.grasp + elif type(actual) is Place and type(expected) is Place: + assert actual.object == expected.object + assert actual.on == expected.on + assert actual.inside == expected.inside + assert (actual.at is None) == (expected.at is None) + if actual.at is not None and expected.at is not None: + _assert_pose_equal(actual.at, expected.at) + elif type(actual) is HandOver and type(expected) is HandOver: + assert actual.object == expected.object + assert actual.receiver == expected.receiver + assert (actual.final_target is None) == (expected.final_target is None) + if actual.final_target is not None and expected.final_target is not None: + _assert_pose_equal(actual.final_target, expected.final_target) + elif ( + type(actual) is RegisteredSemanticCall + and type(expected) is RegisteredSemanticCall + ): + assert actual.call_id == expected.call_id + assert actual.arguments == expected.arguments + else: # pragma: no cover - exact supported union is exhausted above + raise AssertionError(f"Unsupported call type {type(actual).__name__}.") + + +def test_compiler_matches_direct_python_semantic_calls_and_sequence_order() -> None: + registry, provider = _scene_registry() + target = _pose(0.5, 0.1) + config = _program( + SequenceCfg( + items=( + InvokeCfg( + call=PickCfg( + object="sim_cube", + grasp="legacy_grasp", + resources={"primary": "left_actor"}, + ) + ), + InvokeCfg(call=PlaceCfg(object="sim_cube", on="tray_top")), + InvokeCfg( + call=HandOverCfg( + object="sim_cube", + receiver="right_actor", + final_target=TargetRefCfg(target="handover_pose"), + ) + ), + InvokeCfg( + call=RegisteredSemanticCallCfg( + call_id="example.inspect", + arguments={ + "labels": ["front", "back"], + "options": {"confidence": 0.9}, + }, + ) + ), + ) + ), + targets={"handover_pose": CyclicPoseTargetCfg(values=(target,))}, + ) + + compiled = ExpertProgramCompiler.from_scene_registry(registry).compile(config) + segments = list(compiled) + + expected = ( + Pick( + object=SceneObjectRef("cube"), + grasp=SceneAffordanceRef("cube_grasp"), + resources={"primary": "left_actor"}, + ), + Place( + object=SceneObjectRef("cube"), + on=SceneAffordanceRef("tray_top"), + ), + HandOver( + object=SceneObjectRef("cube"), + receiver="right_actor", + final_target=SemanticPose(target.position, target.quaternion_wxyz), + ), + RegisteredSemanticCall( + call_id="example.inspect", + arguments={ + "labels": ("front", "back"), + "options": {"confidence": 0.9}, + }, + ), + ) + assert len(segments) == len(expected) + assert all(segment.implicit for segment in segments) + assert [segment.segment_index for segment in segments] == list(range(4)) + assert [segment.calls[0].call_index for segment in segments] == list(range(4)) + assert len({segment.segment_id for segment in segments}) == 4 + for segment, expected_call in zip(segments, expected, strict=True): + _assert_semantic_call_equal(segment.calls[0].call, expected_call) + assert provider.calls == 0 + + +def test_repeat_expands_independent_segments_with_cyclic_targets() -> None: + registry, provider = _scene_registry() + poses = (_pose(0.45, -0.2), _pose(0.45, 0.0), _pose(0.45, 0.2)) + body = SegmentCfg( + name="move_cube", + steps=SequenceCfg( + items=( + InvokeCfg(call=PickCfg(object="cube")), + InvokeCfg( + call=PlaceCfg( + object="cube", + at=TargetRefCfg(target="drop_pose"), + ) + ), + ) + ), + post=(WaitStablePostCfg(entity="cube"),), + validators=(ObjectNearTargetValidatorCfg(object="cube", target="drop_pose"),), + ) + config = _program( + RepeatCfg(count=3, body=body), + targets={"drop_pose": CyclicPoseTargetCfg(values=poses)}, + program_id="repeated_cube", + ) + + compiled = ExpertProgramCompiler.from_scene_registry(registry).compile(config) + segments = list(compiled) + second_pass = list(compiled) + + assert [segment.segment_id for segment in segments] == [ + segment.segment_id for segment in second_pass + ] + assert len(segments) == 3 + assert len({segment.segment_id for segment in segments}) == 3 + assert [segment.segment_index for segment in segments] == [0, 1, 2] + assert [call.call_index for segment in segments for call in segment.calls] == list( + range(6) + ) + assert all(not segment.implicit for segment in segments) + assert all(segment is not other for segment, other in zip(segments, second_pass)) + for index, (segment, pose) in enumerate(zip(segments, poses, strict=True)): + assert len(segment.repeat_frames) == 1 + assert segment.repeat_frames[0].path == ("program",) + assert segment.repeat_frames[0].iteration_index == index + assert segment.repeat_frames[0].count == 3 + place = segment.calls[1] + assert type(place.call) is Place + assert place.call.at is not None + _assert_pose_equal( + place.call.at, + SemanticPose(pose.position, pose.quaternion_wxyz), + ) + assert place.target_selections[0].value_index == index + validator = segment.validators[0] + _assert_pose_equal(validator.target_pose, place.call.at) + assert validator.target_selection == place.target_selections[0] + assert segment.post_policies[0].entity == SceneObjectRef("cube") + assert provider.calls == 0 + + +def test_repeat_expansion_is_lazy_and_never_observes_scene_providers() -> None: + registry, provider = _scene_registry() + config = _program( + RepeatCfg( + count=1_000, + body=InvokeCfg(call=PickCfg(object="sim_cube")), + ) + ) + + compiled = ExpertProgramCompiler.from_scene_registry(registry).compile(config) + iterator = iter(compiled) + + assert provider.calls == 0 + first_two = list(islice(iterator, 2)) + assert [segment.segment_index for segment in first_two] == [0, 1] + assert [segment.repeat_frames[0].iteration_index for segment in first_two] == [ + 0, + 1, + ] + assert provider.calls == 0 + + +def test_materialized_program_builds_cross_segment_analysis_windows_provider_free() -> ( + None +): + registry, provider = _scene_registry() + config = _program( + SequenceCfg( + items=( + SegmentCfg( + name="pick", + steps=InvokeCfg(call=PickCfg(object="cube")), + ), + SegmentCfg( + name="place", + steps=InvokeCfg( + call=PlaceCfg( + object="cube", + at=TargetRefCfg(target="drop_pose"), + ) + ), + ), + ) + ), + targets={"drop_pose": CyclicPoseTargetCfg(values=(_pose(0.5),))}, + ) + + materialized = ( + ExpertProgramCompiler.from_scene_registry(registry) + .compile(config) + .materialize() + ) + preflight = materialized.preflight_analyses() + execution = materialized.sequential_execution_analysis(0) + + assert type(materialized) is MaterializedCompiledProgram + assert materialized.segment_count == 2 + assert len(tuple(materialized.iter_segments())) == 2 + assert len(preflight) == 1 + assert preflight[0].kind == "sequential_stretch" + assert [type(call) for call in preflight[0].calls] == [Pick, Place] + assert execution.kind == "sequential_suffix" + assert execution.execution_prefix_length == 1 + assert [type(call) for call in execution.calls] == [Pick, Place] + assert provider.calls == 0 + + +def test_materialization_rechecks_expanded_call_bound_after_config_mutation() -> None: + registry, provider = _scene_registry() + inner = RepeatCfg(count=1, body=InvokeCfg(call=PickCfg(object="cube"))) + config = _program(RepeatCfg(count=1, body=inner)) + assert type(config.program) is RepeatCfg + assert type(config.program.body) is RepeatCfg + config.program.count = 1_000 + config.program.body.count = 1_000 + compiled = ExpertProgramCompiler.from_scene_registry(registry).compile(config) + + with pytest.raises(ExpertProgramCompileError) as error: + compiled.materialize() + + assert error.value.code == "expanded_call_limit" + assert provider.calls == 0 + + +def test_wait_stable_accepts_every_canonical_scene_entity_subtype() -> None: + registry, provider = _scene_registry() + config = _program( + SegmentCfg( + name="link_settle", + steps=InvokeCfg(call=PickCfg(object="cube")), + post=(WaitStablePostCfg(entity="arm_tcp"),), + ) + ) + + segment = next( + iter(ExpertProgramCompiler.from_scene_registry(registry).compile(config)) + ) + + assert segment.post_policies[0].entity == SceneLinkRef("arm_tcp") + assert provider.calls == 0 + + +def test_compiled_program_owns_source_and_each_emitted_mutable_config() -> None: + registry, _ = _scene_registry() + target = CyclicPoseTargetCfg(values=(_pose(0.4), _pose(0.5))) + registered = RegisteredSemanticCallCfg( + call_id="example.inspect", + arguments={"settings": {"enabled": True}}, + ) + repeat = RepeatCfg( + count=2, + body=SegmentCfg( + name="inspect_and_place", + steps=SequenceCfg( + items=( + InvokeCfg(call=registered), + InvokeCfg( + call=PlaceCfg( + object="cube", + at=TargetRefCfg(target="drop_pose"), + ) + ), + ) + ), + post=(WaitStablePostCfg(entity="cube", preset="rigid_object"),), + validators=( + ObjectNearTargetValidatorCfg( + object="cube", + target="drop_pose", + position_tolerance=0.03, + ), + ), + ), + ) + config = _program(repeat, targets={"drop_pose": target}) + compiled = ExpertProgramCompiler.from_scene_registry(registry).compile(config) + + compiled_source = config.program + assert type(compiled_source) is RepeatCfg + source_segment = compiled_source.body + assert type(source_segment) is SegmentCfg + source_steps = source_segment.steps + assert type(source_steps) is SequenceCfg + source_registered = source_steps.items[0].call + assert type(source_registered) is RegisteredSemanticCallCfg + compiled_source.count = 1 + config.targets["drop_pose"].values = (_pose(9.0),) + source_registered.arguments["settings"]["enabled"] = False + source_segment.post[0].preset = "changed" + source_segment.validators[0].position_tolerance = 9.0 + + first_pass = list(compiled) + assert len(first_pass) == 2 + first_registered = first_pass[0].calls[0].call + assert type(first_registered) is RegisteredSemanticCall + assert first_registered.arguments["settings"]["enabled"] is True + first_place = first_pass[0].calls[1].call + assert type(first_place) is Place and first_place.at is not None + assert first_place.at.position[0].item() == pytest.approx(0.4) + assert first_pass[0].post_policies[0].cfg.preset == "rigid_object" + assert first_pass[0].validators[0].cfg.position_tolerance == pytest.approx(0.03) + + first_pass[0].post_policies[0].cfg.preset = "mutated_output" + first_pass[0].validators[0].cfg.position_tolerance = 8.0 + exposed_position = compiled.targets["drop_pose"][0].position + exposed_position[0] = -10.0 + + second_pass = list(compiled) + assert second_pass[0].post_policies[0].cfg.preset == "rigid_object" + assert second_pass[0].validators[0].cfg.position_tolerance == pytest.approx(0.03) + assert compiled.targets["drop_pose"][0].position[0].item() == pytest.approx(0.4) + + +def test_compiler_rejects_nested_segment_at_exact_path() -> None: + registry, _ = _scene_registry() + config = _program( + SegmentCfg( + name="outer", + steps=SegmentCfg( + name="inner", + steps=InvokeCfg(call=PickCfg(object="cube")), + ), + ) + ) + + with pytest.raises(ExpertProgramCompileError) as error: + ExpertProgramCompiler.from_scene_registry(registry).compile(config) + + assert error.value.code == "nested_segment" + assert error.value.path == ("program", "steps") + + +def test_compiler_reports_typed_scene_mismatch_at_reference_site() -> None: + registry, _ = _scene_registry() + config = _program( + InvokeCfg(call=PickCfg(object="tray_top")), + ) + + with pytest.raises(ExpertProgramCompileError) as error: + ExpertProgramCompiler.from_scene_registry(registry).compile(config) + + assert error.value.code == "scene_reference_type_mismatch" + assert error.value.path == ("program", "call", "object") + + +def test_compiler_rechecks_mutated_repeat_and_target_bounds() -> None: + registry, _ = _scene_registry() + repeat = RepeatCfg(count=1, body=InvokeCfg(call=PickCfg(object="cube"))) + target = CyclicPoseTargetCfg(values=(_pose(0.4),)) + config = _program(repeat, targets={"drop_pose": target}) + assert type(config.program) is RepeatCfg + config.program.count = 0 + + with pytest.raises(ExpertProgramCompileError) as repeat_error: + ExpertProgramCompiler.from_scene_registry(registry).compile(config) + assert repeat_error.value.code == "invalid_repeat_count" + assert repeat_error.value.path == ("program", "count") + + config.program.count = 1 + config.targets["drop_pose"].values = () + with pytest.raises(ExpertProgramCompileError) as target_error: + ExpertProgramCompiler.from_scene_registry(registry).compile(config) + assert target_error.value.code == "empty_target_values" + assert target_error.value.path == ("targets", "drop_pose", "values") diff --git a/tests/gym/envs/expert_program/test_completion_metadata.py b/tests/gym/envs/expert_program/test_completion_metadata.py new file mode 100644 index 000000000..1befa0ac6 --- /dev/null +++ b/tests/gym/envs/expert_program/test_completion_metadata.py @@ -0,0 +1,503 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Completion-trace audit across semantic execution and the Gym bridge.""" + +from __future__ import annotations + +from dataclasses import dataclass +import json +from types import SimpleNamespace +from typing import ClassVar +from unittest.mock import Mock + +import torch + +from embodichain.lab.gym.envs.expert_program.bridge import ( + AtomicDemoBridge, + BufferedGymCommandSink, + EnvironmentStepClock, + RuntimeCommandFrameEncoder, +) +from embodichain.lab.sim.atomic_actions import ( + ActionInvocation, + ActionOptions, + ActionPlan, + AtomicAction, + AtomicActionEngine, + EndpointCommand, + EntityState, + JointPositionPayload, + JointPositionTarget, + MotionPolicy, + PlannerDiagnostics, + PlanningContext, + RecoveryPolicy, + ResolvedActionRequest, + RobotObservation, + RuntimeCommandFrame, + SceneSnapshot, + SkillBindingContract, + SkillEndpointRequirement, + SkillResourceSlot, + TaskState, + TimedCommandSequence, +) +from embodichain.lab.sim.skills.calls import RegisteredSemanticCall +from embodichain.lab.sim.skills.compiler import SemanticSkillCompiler +from embodichain.lab.sim.skills.runtime import SkillRuntime, SkillStatus +from embodichain.lab.sim.skills.scene import SceneRegistry + +STEP_DT = 0.02 +BATCH_SIZE = 2 +ROBOT_DOF = 5 +ENV_IDS = torch.tensor([7, 3], dtype=torch.long) +INITIAL_SCENE_VERSION = 41 +REPLANNED_SCENE_VERSION = 42 +INITIAL_COLLISION_REVISIONS = (5, 7) +REPLANNED_COLLISION_REVISIONS = (6, 8) + + +@dataclass(frozen=True, slots=True) +class _TraceGoal: + """Test goal for a deterministic two-phase runtime command sequence.""" + + goal_kind: ClassVar[str] = "completion_trace" + + +class _TraceAction(AtomicAction[_TraceGoal, ActionOptions]): + """Emit named segments and preserve distinct diagnostics on every replan.""" + + skill_id: ClassVar[str] = "completion_trace" + GoalType: ClassVar[type] = _TraceGoal + binding_contract: ClassVar[SkillBindingContract] = SkillBindingContract( + slots=( + SkillResourceSlot( + slot_id="primary", + endpoints=(SkillEndpointRequirement(endpoint_id="motion"),), + ), + ) + ) + + def __init__(self) -> None: + super().__init__() + self.plan_count = 0 + + def _scene_dependencies( + self, + request: ResolvedActionRequest[_TraceGoal, ActionOptions], + ) -> tuple[str, ...]: + del request + return ("trace_target",) + + def _plan( + self, + request: ResolvedActionRequest[_TraceGoal, ActionOptions], + context: PlanningContext, + ) -> ActionPlan: + self.require_goal(request) + generation = self.plan_count + self.plan_count += 1 + target = request.binding.endpoint( + "primary", + "motion", + ).require_target(JointPositionTarget) + frames = tuple( + RuntimeCommandFrame( + commands=( + EndpointCommand( + target=target, + payload=JointPositionPayload( + torch.full( + (context.batch_size, len(target.joint_ids)), + float(generation + phase_index + 1), + device=context.robot.qpos.device, + dtype=context.robot.qpos.dtype, + ) + ), + ), + ), + active_mask=torch.ones( + context.batch_size, + dtype=torch.bool, + device=context.robot.qpos.device, + ), + env_ids=context.env_ids, + hold_duration=torch.full( + (context.batch_size,), + STEP_DT, + device=context.robot.qpos.device, + dtype=context.robot.qpos.dtype, + ), + ) + for phase_index in range(2) + ) + return self.build_command_plan( + request, + context, + success=True, + commands=TimedCommandSequence(frames, context.env_ids), + replannable=True, + diagnostics=PlannerDiagnostics( + backend="completion_trace_planner", + messages=(f"installed generation {generation}",), + metadata={ + "generation": generation, + "quality": {"accepted": True, "score": generation + 0.25}, + }, + ), + segment_lengths={"approach": 1, "commit": 1}, + scene_dependency_monitor_until={"trace_target": 2}, + ) + + +class _TraceObservationProvider: + """Move one scene dependency after the first installed command frame.""" + + def __init__(self, clock: EnvironmentStepClock) -> None: + self.clock = clock + self.calls = 0 + + def observe(self, task_state: TaskState) -> PlanningContext: + self.calls += 1 + replanned_scene = self.calls >= 2 + pose = torch.eye(4).repeat(BATCH_SIZE, 1, 1) + if replanned_scene: + pose[:, 0, 3] = 0.25 + qpos = torch.zeros(BATCH_SIZE, ROBOT_DOF) + timestamp = self.clock.now() + return PlanningContext( + robot=RobotObservation( + timestamp=timestamp, + qpos=qpos, + qvel=torch.zeros_like(qpos), + ), + task=task_state, + scene=SceneSnapshot( + timestamp=timestamp, + version=( + REPLANNED_SCENE_VERSION + if replanned_scene + else INITIAL_SCENE_VERSION + ), + entities={"trace_target": EntityState(pose)}, + collision_world_revision=( + REPLANNED_COLLISION_REVISIONS + if replanned_scene + else INITIAL_COLLISION_REVISIONS + ), + ), + env_ids=ENV_IDS, + ) + + +class _StaticQposProvider: + """Supply full robot state to the bridge's transport encoder.""" + + def current_qpos(self, env_ids: torch.Tensor) -> torch.Tensor: + assert torch.equal(env_ids, ENV_IDS) + return torch.zeros(BATCH_SIZE, ROBOT_DOF) + + +class _UnusedEvidenceCollector: + """Satisfy the runtime port; this action declares no physical effect.""" + + def collect( + self, + spec: object, + *, + timestamp: float, + observation_revision: int, + env_ids: torch.Tensor | None = None, + ) -> dict[str, object]: + del spec, timestamp, observation_revision, env_ids + raise AssertionError("The completion trace must not request effect evidence.") + + +@dataclass(frozen=True, slots=True) +class _TraceWorkflow: + """Minimal analyzed workflow retained by the production runtime.""" + + workflow_id: str + calls: tuple[RegisteredSemanticCall, ...] + + +@dataclass(frozen=True, slots=True) +class _TraceIntegration: + """Production engine and registry exposed through the compiler boundary.""" + + engine: AtomicActionEngine + scene_registry: SceneRegistry + + +@dataclass(frozen=True, slots=True) +class _TraceGroundedCall: + """One grounded invocation with no external effect-verification boundary.""" + + analyzed: object + invocation: ActionInvocation + eligible_mask: torch.Tensor + effect_spec: None = None + effect_monitor: None = None + + +class _TraceCompiler(SemanticSkillCompiler): + """Keep semantic boundaries real while making lowering deterministic.""" + + def __init__(self, engine: AtomicActionEngine) -> None: + self._trace_integration = _TraceIntegration(engine, SceneRegistry()) + + @property + def integration(self) -> _TraceIntegration: + return self._trace_integration + + def analyze( + self, + calls: tuple[RegisteredSemanticCall, ...], + *, + workflow_id: str = "semantic_workflow", + path: tuple[object, ...] = ("workflow",), + ) -> _TraceWorkflow: + del path + return _TraceWorkflow(workflow_id, tuple(calls)) + + def ground( + self, + workflow: _TraceWorkflow, + call_index: int, + context: PlanningContext, + *, + eligible_mask: torch.Tensor | None = None, + revision: int = 0, + path: tuple[object, ...] = ("workflow",), + ) -> _TraceGroundedCall: + del context, path + assert eligible_mask is not None + binding = self.integration.engine.bind_control_parts( + _TraceAction.skill_id, + {"primary": {"motion": "arm"}}, + ) + invocation = ActionInvocation( + skill_id=_TraceAction.skill_id, + goal=_TraceGoal(), + binding=binding, + motion_policy=MotionPolicy( + planner="completion_trace_planner", + sample_count=9, + control_dt=STEP_DT, + ), + recovery_policy=RecoveryPolicy( + max_replans=2, + max_action_retries=1, + action_timeout=1.0, + ), + invocation_id=f"{workflow.workflow_id}:{call_index}", + revision=revision, + ) + analyzed = SimpleNamespace( + bound=SimpleNamespace( + robot_profile=SimpleNamespace(profile_id="completion_trace_robot"), + preset=SimpleNamespace( + preset_id="completion_trace_preset", + schema_version=1, + motion_policy=invocation.motion_policy, + recovery_policy=invocation.recovery_policy, + ), + ) + ) + return _TraceGroundedCall( + analyzed=analyzed, + invocation=invocation, + eligible_mask=eligible_mask.clone(), + ) + + +@dataclass(frozen=True, slots=True) +class _CompiledCall: + """Program-owned semantic call and stable call index.""" + + call_index: int + call: RegisteredSemanticCall + + +@dataclass(frozen=True, slots=True) +class _CompiledSegment: + """One logical program segment consumed by the production bridge.""" + + calls: tuple[_CompiledCall, ...] + segment_index: int = 0 + segment_id: str = "completion-segment" + name: str = "completion-audit" + source_path: tuple[object, ...] = ("program", "steps", 0) + post_policies: tuple[object, ...] = () + validators: tuple[object, ...] = () + parallel_block: None = None + implicit: bool = False + + +@dataclass(frozen=True, slots=True) +class _ProgramAnalysis: + """Sequential look-ahead window selected for one bridge segment.""" + + calls: tuple[RegisteredSemanticCall, ...] + execution_prefix_length: int + + +class _CompiledProgram: + """Single-segment compiled-program port for the completion audit.""" + + schema_version = 2 + program_id = "completion-audit-program" + + def __init__(self, segment: _CompiledSegment) -> None: + self.segment = segment + + def iter_segments(self): + yield self.segment + + def sequential_execution_analysis(self, segment_index: int) -> _ProgramAnalysis: + assert segment_index == self.segment.segment_index + return _ProgramAnalysis( + tuple(compiled.call for compiled in self.segment.calls), + len(self.segment.calls), + ) + + +def _runtime_and_bridge() -> tuple[AtomicDemoBridge, _TraceAction]: + """Assemble real execution/runtime/bridge layers around deterministic ports.""" + robot = Mock() + robot.device = torch.device("cpu") + robot.dof = ROBOT_DOF + robot.control_parts = {"arm": object()} + robot.get_joint_ids.return_value = (1, 3) + robot.get_qpos.return_value = torch.zeros(BATCH_SIZE, ROBOT_DOF) + robot.get_qvel.return_value = torch.zeros(BATCH_SIZE, ROBOT_DOF) + generator = Mock() + generator.robot = robot + generator.device = torch.device("cpu") + generator.planner.cfg.planner_type = "completion_trace_planner" + engine = AtomicActionEngine(generator, load_builtins=False) + action = _TraceAction() + engine.register(action) + + clock = EnvironmentStepClock(STEP_DT) + sink = BufferedGymCommandSink( + RuntimeCommandFrameEncoder(_StaticQposProvider()), + clock, + ) + runtime = SkillRuntime.from_components( + _TraceCompiler(engine), + _TraceObservationProvider(clock), + sink, + _UnusedEvidenceCollector(), + task_state=TaskState.empty(BATCH_SIZE, "cpu"), + clock=clock, + ) + call = RegisteredSemanticCall(call_id="audit.completion_metadata") + segment = _CompiledSegment((_CompiledCall(0, call),)) + bridge = AtomicDemoBridge(_CompiledProgram(segment), runtime, sink, clock) + return bridge, action + + +def test_completion_trace_preserves_every_plan_generation_as_json_metadata() -> None: + """A real scene replan remains complete after SkillResult and bridge snapshots.""" + bridge, action = _runtime_and_bridge() + demo_segment = next(bridge.iter_segments()) + + emitted_actions = tuple(demo_segment.actions) + accepted = demo_segment.validator() + metadata = demo_segment.metadata + + serialized = json.dumps(metadata, allow_nan=False, sort_keys=True) + assert json.loads(serialized) == metadata + assert emitted_actions + assert accepted.tolist() == [True, True] + assert action.plan_count == 2 + + assert metadata["expert_program_schema_version"] == 2 + assert metadata["expert_program_id"] == "completion-audit-program" + assert metadata["program_segment_id"] == "completion-segment" + assert metadata["program_segment_index"] == 0 + assert metadata["program_segment_source_path"] == ["program", "steps", 0] + assert metadata["program_segment_implicit"] is False + assert metadata["semantic_call_indices"] == [0] + assert metadata["post_policy_count"] == 0 + assert metadata["validator_count"] == 0 + assert metadata["parallel"] is False + assert metadata["validation"]["accepted_mask"] == [True, True] + + runtime_trace = metadata["runtime"] + assert runtime_trace["kind"] == "skill_result" + assert runtime_trace["status"] == SkillStatus.COMPLETED.value + call_trace = runtime_trace["calls"][0] + assert call_trace["active_plan_attempt_generation"] == 1 + attempts = call_trace["plan_attempts"] + assert [attempt["attempt_generation"] for attempt in attempts] == [0, 1] + assert [attempt["trigger"] for attempt in attempts] == [ + "action_planned", + "replanned", + ] + assert [attempt["planned_scene_version"] for attempt in attempts] == [ + INITIAL_SCENE_VERSION, + REPLANNED_SCENE_VERSION, + ] + assert [attempt["planned_collision_world_revision"] for attempt in attempts] == [ + list(INITIAL_COLLISION_REVISIONS), + list(REPLANNED_COLLISION_REVISIONS), + ] + assert all( + attempt["scene_dependency_monitor_until"] == {"trace_target": 2} + for attempt in attempts + ) + assert all( + attempt["trajectory_segments"] + == [ + {"name": "approach", "start": 0, "stop": 1, "waypoint_count": 1}, + {"name": "commit", "start": 1, "stop": 2, "waypoint_count": 1}, + ] + for attempt in attempts + ) + assert [attempt["recovery_counters"] for attempt in attempts] == [ + {"action_retries": [0, 0], "replans": [0, 0]}, + {"action_retries": [0, 0], "replans": [1, 1]}, + ] + assert [attempt["planner_diagnostics"] for attempt in attempts] == [ + { + "backend": "completion_trace_planner", + "messages": ["installed generation 0"], + "metadata": { + "generation": 0, + "quality": {"accepted": True, "score": 0.25}, + }, + }, + { + "backend": "completion_trace_planner", + "messages": ["installed generation 1"], + "metadata": { + "generation": 1, + "quality": {"accepted": True, "score": 1.25}, + }, + }, + ] + event_kinds = [event["kind"] for event in call_trace["events"]] + assert runtime_trace["events"] == call_trace["events"] + assert "dynamic_goal_changed" in event_kinds + assert "replanned" in event_kinds + assert event_kinds[-3:] == [ + "trajectory_completed", + "action_completed", + "session_completed", + ] diff --git a/tests/gym/envs/expert_program/test_decoder.py b/tests/gym/envs/expert_program/test_decoder.py new file mode 100644 index 000000000..b73266154 --- /dev/null +++ b/tests/gym/envs/expert_program/test_decoder.py @@ -0,0 +1,530 @@ +# ---------------------------------------------------------------------------- +# 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 strict Expert Program Version 1 decoding.""" + +from __future__ import annotations + +from copy import deepcopy + +import pytest + +from embodichain.lab.gym.envs.expert_program import ( + MAX_REPEAT_COUNT, + ConfigPath, + ExpertProgramDecodeError, + ExpertProgramIntegrationCfg, + ExpertProgramValidationError, + HandOverCfg, + PickCfg, + PlaceCfg, + PostPolicyCfg, + RegisteredSemanticCallCfg, + RepeatCfg, + SceneReferenceRole, + SegmentCfg, + SequenceCfg, + TargetRefCfg, + ValidatorCfg, + decode_expert_program, + render_config_path, +) + + +def _program_data() -> dict[str, object]: + """Return the repeated-cube Version 1 example as plain JSON values.""" + return { + "schema_version": 1, + "program_id": "repeated_cube_pick_place", + "integration": { + "robot_profile": "auto", + "scene_registry": "env", + "runtime_preset": "safe", + }, + "targets": { + "drop_pose": { + "kind": "cyclic_pose", + "values": [ + { + "position": [0.45, -0.20, 0.20], + "quaternion_wxyz": [1.0, 0.0, 0.0, 0.0], + }, + { + "position": [0.45, 0.00, 0.20], + "quaternion_wxyz": [1.0, 0.0, 0.0, 0.0], + }, + { + "position": [0.45, 0.20, 0.20], + "quaternion_wxyz": [1.0, 0.0, 0.0, 0.0], + }, + ], + } + }, + "program": { + "kind": "repeat", + "count": 3, + "body": { + "kind": "segment", + "name": "move_cube", + "steps": { + "kind": "sequence", + "items": [ + { + "kind": "invoke", + "call": {"kind": "pick", "object": "cube"}, + }, + { + "kind": "invoke", + "call": { + "kind": "place", + "object": "cube", + "at": { + "kind": "target_ref", + "target": "drop_pose", + }, + }, + }, + ], + }, + "post": [ + { + "kind": "wait_stable", + "entity": "cube", + "preset": "rigid_object", + } + ], + "validators": [ + { + "kind": "object_near_target", + "object": "cube", + "target": "drop_pose", + "position_tolerance": 0.03, + } + ], + }, + }, + } + + +def _invoke(call: dict[str, object]) -> dict[str, object]: + """Wrap one call mapping in an invoke node.""" + return {"kind": "invoke", "call": call} + + +def test_decoder_builds_owned_repeated_cube_ast() -> None: + data = _program_data() + + config = decode_expert_program(data) + data["program"]["count"] = 99 + data["targets"]["drop_pose"]["values"][0]["position"][0] = -1.0 + + assert type(config.program) is RepeatCfg + assert config.program.count == 3 + assert type(config.program.body) is SegmentCfg + assert type(config.program.body.steps) is SequenceCfg + place = config.program.body.steps.items[1].call + assert type(place) is PlaceCfg + assert place.at == TargetRefCfg(target="drop_pose") + assert config.targets["drop_pose"].values[0].position[0] == pytest.approx(0.45) + + +def test_decoder_supports_every_version_one_semantic_call() -> None: + data = _program_data() + data["program"] = { + "kind": "sequence", + "items": [ + _invoke( + { + "kind": "pick", + "object": "cube", + "grasp": "cube_grasp", + "resources": {"primary": "left_actor"}, + } + ), + _invoke({"kind": "place", "object": "cube", "on": "tray_top"}), + _invoke( + { + "kind": "hand_over", + "object": "cube", + "receiver": "right_actor", + "final_target": { + "kind": "target_ref", + "target": "drop_pose", + }, + } + ), + _invoke( + { + "kind": "registered", + "call_id": "example.inspect", + "schema_version": 1, + "arguments": { + "labels": ["front", "back"], + "options": {"confidence": 0.9}, + }, + } + ), + ], + } + + config = decode_expert_program(data) + assert [type(node.call) for node in config.program.items] == [ + PickCfg, + PlaceCfg, + HandOverCfg, + RegisteredSemanticCallCfg, + ] + handover = config.program.items[2].call + assert handover.resources == {"destination": "right_actor"} + registered = config.program.items[3].call + assert registered.arguments == { + "labels": ("front", "back"), + "options": {"confidence": 0.9}, + } + + +@pytest.mark.parametrize( + ("mutate", "expected_path"), + [ + ( + lambda data: data.update({"unexpected": True}), + "$.unexpected", + ), + ( + lambda data: data["program"]["body"].update({"unexpected": True}), + "$.program.body.unexpected", + ), + ( + lambda data: data["program"]["body"]["steps"]["items"][0]["call"].update( + {"unexpected": True} + ), + "$.program.body.steps.items[0].call.unexpected", + ), + ], +) +def test_decoder_rejects_unknown_fields_with_complete_path( + mutate: object, + expected_path: str, +) -> None: + data = _program_data() + mutate(data) + + with pytest.raises(ExpertProgramDecodeError) as error: + decode_expert_program(data) + + assert error.value.code == "unknown_field" + assert render_config_path(error.value.path) == expected_path + + +def test_decoder_reports_missing_required_field_at_exact_path() -> None: + data = _program_data() + del data["integration"]["runtime_preset"] + + with pytest.raises(ExpertProgramDecodeError) as error: + decode_expert_program(data) + + assert error.value.code == "missing_field" + assert render_config_path(error.value.path) == "$.integration.runtime_preset" + + +@pytest.mark.parametrize( + ("value", "code"), + [ + (None, "missing_discriminator"), + ("parallel", "unknown_discriminator"), + ], +) +def test_decoder_rejects_missing_or_reserved_program_discriminator( + value: str | None, + code: str, +) -> None: + data = _program_data() + if value is None: + del data["program"]["kind"] + else: + data["program"]["kind"] = value + + with pytest.raises(ExpertProgramDecodeError) as error: + decode_expert_program(data) + + assert error.value.code == code + assert error.value.path == ("program", "kind") + + +@pytest.mark.parametrize( + ("mutate", "expected_path"), + [ + ( + lambda data: data["targets"]["drop_pose"].update({"kind": "pose"}), + "$.targets.drop_pose.kind", + ), + ( + lambda data: data["program"]["body"]["steps"]["items"][0]["call"].update( + {"kind": "move"} + ), + "$.program.body.steps.items[0].call.kind", + ), + ( + lambda data: data["program"]["body"]["steps"]["items"][1]["call"][ + "at" + ].update({"kind": "env_ref"}), + "$.program.body.steps.items[1].call.at.kind", + ), + ( + lambda data: data["program"]["body"]["post"][0].update({"kind": "sleep"}), + "$.program.body.post[0].kind", + ), + ( + lambda data: data["program"]["body"]["validators"][0].update( + {"kind": "python"} + ), + "$.program.body.validators[0].kind", + ), + ], +) +def test_every_union_rejects_unknown_discriminator_at_exact_path( + mutate: object, + expected_path: str, +) -> None: + data = _program_data() + mutate(data) + + with pytest.raises(ExpertProgramDecodeError) as error: + decode_expert_program(data) + + assert error.value.code == "unknown_discriminator" + assert render_config_path(error.value.path) == expected_path + + +@pytest.mark.parametrize("schema_version", [False, 0, 3, "1"]) +def test_decoder_rejects_unsupported_top_level_schema_version( + schema_version: object, +) -> None: + data = _program_data() + data["schema_version"] = schema_version + + with pytest.raises(ExpertProgramDecodeError) as error: + decode_expert_program(data) + + assert error.value.code == "unsupported_schema_version" + assert error.value.path == ("schema_version",) + + +def test_decoder_reports_unknown_target_at_reference_site() -> None: + data = _program_data() + data["program"]["body"]["steps"]["items"][1]["call"]["at"]["target"] = "missing" + + with pytest.raises(ExpertProgramDecodeError) as error: + decode_expert_program(data) + + assert error.value.code == "unknown_target" + assert render_config_path(error.value.path) == ( + "$.program.body.steps.items[1].call.at.target" + ) + + +@pytest.mark.parametrize("count", [False, 0, MAX_REPEAT_COUNT + 1]) +def test_decoder_rejects_unbounded_or_invalid_repeat_count(count: object) -> None: + data = _program_data() + data["program"]["count"] = count + + with pytest.raises(ExpertProgramDecodeError) as error: + decode_expert_program(data) + + assert error.value.code == "invalid_repeat_count" + assert render_config_path(error.value.path) == "$.program.count" + + +def test_registered_call_schema_version_error_reports_version_field() -> None: + data = _program_data() + data["program"] = _invoke( + { + "kind": "registered", + "call_id": "example.inspect", + "schema_version": 2, + } + ) + + with pytest.raises(ExpertProgramDecodeError) as error: + decode_expert_program(data) + + assert error.value.code == "invalid_schema_version" + assert render_config_path(error.value.path) == "$.program.call.schema_version" + + +@pytest.mark.parametrize( + ("arguments", "code", "suffix"), + [ + ({"eval": "1 + 1"}, "forbidden_construct", ".arguments.eval"), + ( + {"source": "env.robot.control_parts"}, + "environment_traversal", + ".arguments.source", + ), + ( + {"source": "eval(1 + 1)"}, + "executable_expression", + ".arguments.source", + ), + ({"callback": lambda: None}, "non_declarative_value", ".arguments.callback"), + ({"live": object()}, "non_declarative_value", ".arguments.live"), + ], +) +def test_decoder_rejects_executable_traversal_or_live_registered_payload( + arguments: dict[str, object], + code: str, + suffix: str, +) -> None: + data = _program_data() + data["program"] = _invoke( + { + "kind": "registered", + "call_id": "example.inspect", + "schema_version": 1, + "arguments": arguments, + } + ) + + with pytest.raises(ExpertProgramDecodeError) as error: + decode_expert_program(data) + + assert error.value.code == code + assert render_config_path(error.value.path).endswith(suffix) + + +def test_decoder_rejects_cyclic_input_before_ast_recursion() -> None: + data = _program_data() + cyclic: dict[str, object] = {} + cyclic["self"] = cyclic + data["program"] = _invoke( + { + "kind": "registered", + "call_id": "example.inspect", + "schema_version": 1, + "arguments": cyclic, + } + ) + + with pytest.raises(ExpertProgramDecodeError) as error: + decode_expert_program(data) + + assert error.value.code == "cyclic_input" + + +class _StaticValidationContext: + """Small provider-free reference catalog used by decoder tests.""" + + def __init__( + self, + *, + calls: set[str] | None = None, + scene: set[str] | None = None, + ) -> None: + self.calls = {"pick", "place", "hand_over"} if calls is None else calls + self.scene = {"cube", "cube_grasp", "tray_top"} if scene is None else scene + self.validated_paths: list[ConfigPath] = [] + + def validate_integration( + self, + integration: ExpertProgramIntegrationCfg, + *, + path: ConfigPath, + ) -> None: + if integration.robot_profile != "auto": + raise KeyError(integration.robot_profile) + self.validated_paths.append(path) + + def validate_semantic_call( + self, + call: object, + *, + path: ConfigPath, + ) -> None: + semantic_id = ( + call.call_id if type(call) is RegisteredSemanticCallCfg else call.kind + ) + if semantic_id not in self.calls: + raise KeyError(semantic_id) + self.validated_paths.append(path) + + def validate_scene_reference( + self, + reference: str, + *, + role: SceneReferenceRole, + path: ConfigPath, + ) -> None: + del role + if reference not in self.scene: + raise KeyError(reference) + self.validated_paths.append(path) + + def validate_post_policy( + self, + policy: PostPolicyCfg, + *, + path: ConfigPath, + ) -> None: + if policy.kind != "wait_stable" or policy.preset != "rigid_object": + raise KeyError(policy.preset) + self.validated_paths.append(path) + + def validate_validator( + self, + validator: ValidatorCfg, + *, + path: ConfigPath, + ) -> None: + if validator.kind != "object_near_target": + raise KeyError(validator.kind) + self.validated_paths.append(path) + + +def test_decoder_runs_explicit_provider_free_validation_context() -> None: + data = _program_data() + context = _StaticValidationContext() + + config = decode_expert_program(data, validation_context=context) + + assert config.program_id == "repeated_cube_pick_place" + assert ("integration",) in context.validated_paths + assert ("program", "body", "steps", "items", 0, "call") in (context.validated_paths) + assert ("program", "body", "post", 0, "entity") in (context.validated_paths) + + +def test_validation_context_failure_is_wrapped_at_exact_reference_path() -> None: + data = _program_data() + context = _StaticValidationContext(scene={"cube"}) + data["program"]["body"]["steps"]["items"][0]["call"]["grasp"] = "missing_grasp" + + with pytest.raises(ExpertProgramValidationError) as error: + decode_expert_program(data, validation_context=context) + + assert error.value.code == "reference_validation_failed" + assert render_config_path(error.value.path) == ( + "$.program.body.steps.items[0].call.grasp" + ) + + +def test_decoder_does_not_mutate_caller_input_on_failure() -> None: + data = _program_data() + data["program"]["unexpected"] = True + before = deepcopy(data) + + with pytest.raises(ExpertProgramDecodeError): + decode_expert_program(data) + + assert data == before diff --git a/tests/gym/envs/expert_program/test_environment.py b/tests/gym/envs/expert_program/test_environment.py new file mode 100644 index 000000000..dcca9b241 --- /dev/null +++ b/tests/gym/envs/expert_program/test_environment.py @@ -0,0 +1,946 @@ +# ---------------------------------------------------------------------------- +# 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 reusable environment-backed Expert Program assembly.""" + +from __future__ import annotations + +from collections import Counter +from unittest.mock import Mock + +import pytest +import torch + +from embodichain.lab.gym.envs.expert_program.bridge import ( + AtomicDemoBridge, + DemoBridgeError, + EnvironmentStepClock, + GymPlanningObservationProvider, +) +from embodichain.lab.gym.envs.expert_program.cfg import ( + EXPERT_PROGRAM_SCHEMA_VERSION, + EXPERT_PROGRAM_SCHEMA_VERSION_V2, + BarrierCfg, + CyclicPoseTargetCfg, + ExpertProgramCfg, + ExpertProgramIntegrationCfg, + InvokeCfg, + ObjectNearTargetValidatorCfg, + OperateArticulationCfg, + ParallelCfg, + PickCfg, + PlaceCfg, + PoseCfg, + ProgramNodeCfg, + SegmentCfg, + SequenceCfg, + TargetRefCfg, + WaitStablePostCfg, +) +from embodichain.lab.gym.envs.expert_program.environment import ( + ExpertProgramEnvironmentAdapter, + ExpertProgramEnvironmentMixin, + PlanningObservationPort, +) +from embodichain.lab.sim.atomic_actions import ( + AntipodalAffordance, + ArticulationOperationAffordance, + ArticulationOperationTarget, + AtomicActionEngine, + BATCH_INVERSE_KINEMATICS_CAPABILITY, + CARTESIAN_POSE_CAPABILITY, + ControlPartCommandProfile, + EntityState, + FORWARD_KINEMATICS_CAPABILITY, + GRASP_CAPABILITY, + JOINT_POSITION_CAPABILITY, + MotionPolicy, + PlanningContext, + RobotObservation, + TaskState, +) +from embodichain.lab.sim.skills import ( + ARTICULATION_OPERATION_AFFORDANCE_CAPABILITY, + GRASP_AFFORDANCE_CAPABILITY, + ControlPartEndpoint, + RobotResource, + RobotSkillProfile, + SceneAffordanceRef, + SceneArticulationRef, + SceneCollisionRole, + SceneCollisionWorldMode, + SceneEntityRegistration, + SceneObjectRef, + SceneRegistry, + SkillPolicyPreset, +) +from embodichain.lab.sim.skills.compiler import SemanticSkillCompiler +from embodichain.lab.sim.skills.evidence import EffectEvidenceProvider +from embodichain.lab.sim.skills.integration import SemanticValidationError + +_BATCH_SIZE = 2 +_ROBOT_DOF = 2 +_STEP_DT = 0.02 + + +class _PoseProvider: + """Return a stable owned pose for the fake environment scene.""" + + def __init__(self) -> None: + self._pose = torch.eye(4).repeat(_BATCH_SIZE, 1, 1) + self.calls = 0 + + def observe( + self, + *, + timestamp: float, + env_ids: torch.Tensor, + ) -> EntityState: + """Return rows aligned to the requested environment IDs.""" + del timestamp + self.calls += 1 + return EntityState(self._pose.index_select(0, env_ids)) + + +class _GeometryProvider: + """Return one opaque planner-facing geometry descriptor.""" + + def get_geometry(self) -> object: + return object() + + +def _scene_registry( + *, + dynamic_collision: bool = False, + pose_provider: _PoseProvider | None = None, +) -> SceneRegistry: + """Build an explicitly named object and default grasp affordance.""" + cube = SceneObjectRef("cube") + grasp = SceneAffordanceRef("cube_grasp") + selected_pose_provider = _PoseProvider() if pose_provider is None else pose_provider + return SceneRegistry( + ( + SceneEntityRegistration( + ref=cube, + state_provider=selected_pose_provider, + semantic_type="cube", + default_affordances={GRASP_AFFORDANCE_CAPABILITY: grasp}, + geometry_provider=(_GeometryProvider() if dynamic_collision else None), + collision_role=( + SceneCollisionRole.DYNAMIC + if dynamic_collision + else SceneCollisionRole.NONE + ), + ), + SceneEntityRegistration( + ref=grasp, + parent=cube, + native_name="grasp", + affordance=AntipodalAffordance(), + affordance_capabilities=frozenset({GRASP_AFFORDANCE_CAPABILITY}), + affordance_revision="grasp-v1", + relative_pose=torch.eye(4), + ), + ), + collision_world_mode=( + SceneCollisionWorldMode.PER_ENV if dynamic_collision else None + ), + ) + + +def _robot_profile( + profile_id: str = "fake_robot", + *, + safe_motion_policy: MotionPolicy | None = None, +) -> RobotSkillProfile: + """Build the declarative resource graph used by the fake backend.""" + return RobotSkillProfile( + profile_id=profile_id, + resources={ + "manipulator": RobotResource( + resource_id="manipulator", + endpoints={ + "motion": ControlPartEndpoint( + control_part="arm", + capabilities=frozenset( + { + BATCH_INVERSE_KINEMATICS_CAPABILITY, + CARTESIAN_POSE_CAPABILITY, + FORWARD_KINEMATICS_CAPABILITY, + JOINT_POSITION_CAPABILITY, + } + ), + ), + "grasp": ControlPartEndpoint( + control_part="hand", + capabilities=frozenset({GRASP_CAPABILITY}), + ), + }, + ) + }, + command_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=torch.tensor((0.0,)), + grasp=torch.tensor((1.0,)), + ) + }, + presets={ + "safe": SkillPolicyPreset( + "safe", + motion_policy=safe_motion_policy, + ) + }, + default_preset="safe", + ) + + +def _parallel_articulation_scene_registry() -> SceneRegistry: + """Build one drawer whose exact joint key is statically discoverable.""" + drawer = SceneArticulationRef("drawer") + handle = SceneAffordanceRef("drawer_handle") + return SceneRegistry( + ( + SceneEntityRegistration( + ref=drawer, + state_provider=_PoseProvider(), + semantic_type="drawer", + default_affordances={ + ARTICULATION_OPERATION_AFFORDANCE_CAPABILITY: handle + }, + ), + SceneEntityRegistration( + ref=handle, + state_provider=_PoseProvider(), + parent=drawer, + native_name="handle", + affordance=ArticulationOperationAffordance( + joint_id="drawer_slide", + operation_axis=torch.tensor((1.0, 0.0, 0.0)), + semantic_targets={ + "open": ArticulationOperationTarget( + target_position=0.4, + displacement=0.35, + ) + }, + ), + affordance_capabilities=frozenset( + {ARTICULATION_OPERATION_AFFORDANCE_CAPABILITY} + ), + affordance_revision="drawer-operation-v1", + ), + ) + ) + + +def _parallel_articulation_profile() -> RobotSkillProfile: + """Build two physically disjoint resources that can address one drawer.""" + + def resource(resource_id: str) -> RobotResource: + return RobotResource( + resource_id=resource_id, + endpoints={ + "motion": ControlPartEndpoint( + control_part=f"{resource_id}_arm", + capabilities=frozenset( + {CARTESIAN_POSE_CAPABILITY, JOINT_POSITION_CAPABILITY} + ), + ), + "interaction": ControlPartEndpoint( + control_part=f"{resource_id}_hand", + capabilities=frozenset({GRASP_CAPABILITY}), + ), + }, + ) + + return RobotSkillProfile( + profile_id="parallel_articulation_robot", + resources={ + "left": resource("left"), + "right": resource("right"), + }, + command_profiles={ + hand: ControlPartCommandProfile.joint_positions( + open=torch.tensor((0.0,)), + grasp=torch.tensor((1.0,)), + ) + for hand in ("left_hand", "right_hand") + }, + presets={"safe": SkillPolicyPreset("safe")}, + default_preset="safe", + ) + + +def _parallel_articulation_engine( + profile: RobotSkillProfile, +) -> AtomicActionEngine: + """Build the disjoint four-control-part engine used only for preflight.""" + robot = Mock() + robot.device = torch.device("cpu") + robot.dof = 4 + robot.control_parts = { + "left_arm": object(), + "left_hand": object(), + "right_arm": object(), + "right_hand": object(), + } + robot.get_qpos.return_value = torch.zeros(_BATCH_SIZE, robot.dof) + robot.get_qvel.return_value = torch.zeros(_BATCH_SIZE, robot.dof) + joint_ids = { + "left_arm": [0], + "left_hand": [1], + "right_arm": [2], + "right_hand": [3], + } + robot.get_joint_ids.side_effect = lambda name: joint_ids[name] + robot.get_solver.return_value = object() + generator = Mock() + generator.robot = robot + generator.device = torch.device("cpu") + generator.planner.cfg.planner_type = "fake_planner" + return AtomicActionEngine(generator, skill_profile=profile) + + +def _engine(profile: RobotSkillProfile) -> AtomicActionEngine: + """Build a CPU-only engine around a minimal typed robot surface.""" + robot = Mock() + robot.device = torch.device("cpu") + robot.dof = _ROBOT_DOF + robot.control_parts = {"arm": object(), "hand": object()} + robot.get_qpos.return_value = torch.zeros(_BATCH_SIZE, _ROBOT_DOF) + robot.get_qvel.return_value = torch.zeros(_BATCH_SIZE, _ROBOT_DOF) + robot.get_joint_ids.side_effect = lambda name: {"arm": [0], "hand": [1]}[name] + robot.get_solver.return_value = object() + generator = Mock() + generator.robot = robot + generator.device = torch.device("cpu") + generator.planner.cfg.planner_type = "fake_planner" + return AtomicActionEngine(generator, skill_profile=profile) + + +class _FakeEnvironmentFactory: + """Count every explicit factory boundary used by the production adapter.""" + + scene_registry_id = "fake_scene" + robot_profile_id = "fake_robot" + + def __init__(self, *, returned_profile_id: str = "fake_robot") -> None: + self.returned_profile_id = returned_profile_id + self.calls: Counter[str] = Counter() + self.observation_samples = 0 + + def create_scene_registry(self) -> SceneRegistry: + """Create a fresh live registry.""" + self.calls["scene"] += 1 + return _scene_registry() + + def create_robot_skill_profile(self) -> RobotSkillProfile: + """Create the configured robot profile.""" + self.calls["profile"] += 1 + return _robot_profile(self.returned_profile_id) + + def create_atomic_action_engine( + self, + profile: RobotSkillProfile, + ) -> AtomicActionEngine: + """Create an engine for exactly the supplied profile.""" + self.calls["engine"] += 1 + return _engine(profile) + + def create_planning_observation_provider( + self, + *, + scene_registry: SceneRegistry, + engine: AtomicActionEngine, + clock: EnvironmentStepClock, + ) -> GymPlanningObservationProvider: + """Create a callback-backed Gym observation port.""" + self.calls["observation"] += 1 + scene_provider = scene_registry.make_scene_provider(batch_size=_BATCH_SIZE) + + def capture(task_state: TaskState) -> PlanningContext: + self.observation_samples += 1 + env_ids = torch.arange(_BATCH_SIZE, dtype=torch.long) + timestamp = clock.now() + return PlanningContext( + robot=RobotObservation( + timestamp=timestamp, + qpos=engine.robot.get_qpos(), + qvel=engine.robot.get_qvel(), + ), + task=task_state, + scene=scene_provider.snapshot( + timestamp=timestamp, + env_ids=env_ids, + ), + env_ids=env_ids, + ) + + return GymPlanningObservationProvider(capture) + + def create_effect_evidence_providers( + self, + *, + scene_registry: SceneRegistry, + engine: AtomicActionEngine, + observation_provider: PlanningObservationPort, + ) -> tuple[EffectEvidenceProvider, ...]: + """Return the fake environment's explicit evidence-provider set.""" + del scene_registry, engine, observation_provider + self.calls["evidence"] += 1 + return () + + +class _ParallelArticulationFactory(_FakeEnvironmentFactory): + """Expose two robot resources and one shared articulation write target.""" + + scene_registry_id = "parallel_articulation_scene" + robot_profile_id = "parallel_articulation_robot" + + def create_scene_registry(self) -> SceneRegistry: + self.calls["scene"] += 1 + return _parallel_articulation_scene_registry() + + def create_robot_skill_profile(self) -> RobotSkillProfile: + self.calls["profile"] += 1 + return _parallel_articulation_profile() + + def create_atomic_action_engine( + self, + profile: RobotSkillProfile, + ) -> AtomicActionEngine: + self.calls["engine"] += 1 + return _parallel_articulation_engine(profile) + + +class _DynamicCollisionFactory(_FakeEnvironmentFactory): + """Expose a safe dynamic scene backed by an unsupported planner.""" + + def __init__(self) -> None: + super().__init__() + self.pose_provider = _PoseProvider() + self.last_engine: AtomicActionEngine | None = None + + def create_scene_registry(self) -> SceneRegistry: + self.calls["scene"] += 1 + return _scene_registry( + dynamic_collision=True, + pose_provider=self.pose_provider, + ) + + def create_robot_skill_profile(self) -> RobotSkillProfile: + self.calls["profile"] += 1 + return _robot_profile( + safe_motion_policy=MotionPolicy(strategy="motion_gen"), + ) + + def create_atomic_action_engine( + self, + profile: RobotSkillProfile, + ) -> AtomicActionEngine: + self.calls["engine"] += 1 + engine = _engine(profile) + engine.motion_generator.supports_dynamic_collision_world = False + self.last_engine = engine + return engine + + +class _FakeDeclarativeEnvironment(ExpertProgramEnvironmentMixin): + """Environment surface requiring no task-level motion implementation.""" + + def __init__(self, adapter: ExpertProgramEnvironmentAdapter) -> None: + self._adapter = adapter + + @property + def expert_program_adapter(self) -> ExpertProgramEnvironmentAdapter: + """Return the reusable environment adapter.""" + return self._adapter + + +class _AcceptParallelSafety: + """Accept test-only merged commands after static preflight succeeds.""" + + def validate( + self, + *, + branch_frames: object, + merged_frame: object, + ) -> None: + del branch_frames, merged_frame + + +class _PresetCheckingPostPolicyPort: + """Pure test port that rejects policies outside its preset table.""" + + def __init__(self, preset_ids: tuple[str, ...]) -> None: + self._preset_ids = frozenset(preset_ids) + self.validated_presets: list[str] = [] + + def validate_policy(self, policy: object, *, segment: object) -> None: + del segment + cfg = getattr(policy, "cfg") + preset = getattr(cfg, "preset") + self.validated_presets.append(preset) + if preset not in self._preset_ids: + raise KeyError(f"Unknown settle preset {preset!r}.") + + def actions( + self, + policy: object, + *, + segment: object, + active_mask: torch.Tensor, + ) -> tuple[torch.Tensor, ...]: + del policy, segment, active_mask + raise AssertionError("Preflight must not request post-policy actions.") + + +def _program( + *, + robot_profile: str = "fake_robot", + scene_registry: str = "fake_scene", + runtime_preset: str = "safe", + schema_version: int = EXPERT_PROGRAM_SCHEMA_VERSION, + node: ProgramNodeCfg | None = None, + targets: dict[str, CyclicPoseTargetCfg] | None = None, +) -> ExpertProgramCfg: + """Build one minimal declarative pick program.""" + return ExpertProgramCfg( + schema_version=schema_version, + program_id="fake_pick", + integration=ExpertProgramIntegrationCfg( + robot_profile=robot_profile, + scene_registry=scene_registry, + runtime_preset=runtime_preset, + ), + program=(InvokeCfg(call=PickCfg(object="cube")) if node is None else node), + targets={} if targets is None else targets, + ) + + +def _program_with_later_parallel_conflict() -> ExpertProgramCfg: + """Build an early sequential call followed by conflicting branch claims.""" + return _program( + schema_version=EXPERT_PROGRAM_SCHEMA_VERSION_V2, + node=SequenceCfg( + items=( + InvokeCfg(call=PickCfg(object="cube")), + ParallelCfg( + branches=( + InvokeCfg(call=PickCfg(object="cube")), + InvokeCfg(call=PickCfg(object="cube")), + ), + barrier=BarrierCfg(name="conflicting_join"), + ), + ) + ), + ) + + +def _program_with_later_segment_hooks( + *, + post: tuple[WaitStablePostCfg, ...] = (), + validators: tuple[ObjectNearTargetValidatorCfg, ...] = (), +) -> ExpertProgramCfg: + """Build a valid pick/place flow whose hooks live on the later segment.""" + return _program( + node=SequenceCfg( + items=( + SegmentCfg( + name="pick", + steps=InvokeCfg(call=PickCfg(object="cube")), + ), + SegmentCfg( + name="place", + steps=InvokeCfg( + call=PlaceCfg( + object="cube", + at=TargetRefCfg(target="drop"), + ) + ), + post=post, + validators=validators, + ), + ) + ), + targets={ + "drop": CyclicPoseTargetCfg( + values=( + PoseCfg( + position=(0.4, 0.1, 0.2), + quaternion_wxyz=(1.0, 0.0, 0.0, 0.0), + ), + ) + ) + }, + ) + + +def _parallel_articulation_program() -> ExpertProgramCfg: + """Operate one joint from disjoint resources in two parallel branches.""" + return ExpertProgramCfg( + schema_version=EXPERT_PROGRAM_SCHEMA_VERSION_V2, + program_id="conflicting_drawer_operations", + integration=ExpertProgramIntegrationCfg( + robot_profile="parallel_articulation_robot", + scene_registry="parallel_articulation_scene", + runtime_preset="safe", + ), + targets={}, + program=ParallelCfg( + branches=tuple( + InvokeCfg( + call=OperateArticulationCfg( + articulation="drawer", + target="open", + resources={"primary": resource_id}, + ) + ) + for resource_id in ("left", "right") + ), + barrier=BarrierCfg(name="drawer_join"), + ), + ) + + +def test_mixin_compiles_and_assembles_bridge_without_task_motion_code() -> None: + """One adapter property implements both EmbodiedEnv integration hooks.""" + factory = _FakeEnvironmentFactory() + adapter = ExpertProgramEnvironmentAdapter(factory, step_dt=_STEP_DT) + env = _FakeDeclarativeEnvironment(adapter) + + compiled = env.compile_expert_program(_program()) + + assert factory.calls == Counter(scene=1) + bridge = env.create_expert_program_bridge(compiled) + assert isinstance(bridge, AtomicDemoBridge) + assert factory.calls == Counter( + scene=2, + profile=1, + engine=1, + observation=1, + evidence=1, + ) + segment_iterator = bridge.iter_segments() + segment = next(segment_iterator) + assert segment.name == "invoke:pick" + assert segment.failure_policy == "row_independent" + segment_iterator.close() + + +def test_later_sequential_resource_error_fails_before_observation_or_action() -> None: + factory = _FakeEnvironmentFactory() + adapter = ExpertProgramEnvironmentAdapter(factory, step_dt=_STEP_DT) + compiled = adapter.compile( + _program( + node=SequenceCfg( + items=( + InvokeCfg(call=PickCfg(object="cube")), + InvokeCfg( + call=PickCfg( + object="cube", + resources={"primary": "missing"}, + ) + ), + ) + ) + ) + ) + + with pytest.raises(SemanticValidationError) as error: + adapter.create_bridge(compiled) + + assert error.value.diagnostic.code == "unknown_resource" + assert factory.calls == Counter(scene=2, profile=1, engine=1) + assert factory.observation_samples == 0 + + +def test_later_post_policy_requires_port_before_runtime_assembly() -> None: + """A later hook cannot defer its missing-port error until segment execution.""" + factory = _FakeEnvironmentFactory() + adapter = ExpertProgramEnvironmentAdapter(factory, step_dt=_STEP_DT) + compiled = adapter.compile( + _program_with_later_segment_hooks( + post=(WaitStablePostCfg(entity="cube", preset="fast"),), + ) + ) + + with pytest.raises(DemoBridgeError, match="SegmentPostPolicyPort"): + adapter.create_bridge(compiled) + + assert factory.calls == Counter(scene=1) + assert factory.observation_samples == 0 + + +def test_later_validator_requires_port_before_runtime_assembly() -> None: + """A later validator must have an installed pure-validation boundary.""" + factory = _FakeEnvironmentFactory() + adapter = ExpertProgramEnvironmentAdapter(factory, step_dt=_STEP_DT) + compiled = adapter.compile( + _program_with_later_segment_hooks( + validators=( + ObjectNearTargetValidatorCfg( + object="cube", + target="drop", + ), + ), + ) + ) + + with pytest.raises(DemoBridgeError, match="SegmentValidatorPort"): + adapter.create_bridge(compiled) + + assert factory.calls == Counter(scene=1) + assert factory.observation_samples == 0 + + +def test_later_unknown_settle_preset_fails_during_pure_preflight() -> None: + """Every declared preset is checked before semantic or live runtime assembly.""" + factory = _FakeEnvironmentFactory() + post_policy_port = _PresetCheckingPostPolicyPort(("fast",)) + adapter = ExpertProgramEnvironmentAdapter( + factory, + step_dt=_STEP_DT, + post_policy_port=post_policy_port, + ) + compiled = adapter.compile( + _program_with_later_segment_hooks( + post=(WaitStablePostCfg(entity="cube", preset="missing"),), + ) + ) + + with pytest.raises(KeyError, match="Unknown settle preset 'missing'"): + adapter.create_bridge(compiled) + + assert post_policy_port.validated_presets == ["missing"] + assert factory.calls == Counter(scene=1) + assert factory.observation_samples == 0 + + +def test_preflight_preserves_pick_target_lookahead_across_explicit_segments( + monkeypatch: pytest.MonkeyPatch, +) -> None: + factory = _FakeEnvironmentFactory() + adapter = ExpertProgramEnvironmentAdapter(factory, step_dt=_STEP_DT) + config = _program( + node=SequenceCfg( + items=( + SegmentCfg( + name="pick", + steps=InvokeCfg(call=PickCfg(object="cube")), + ), + SegmentCfg( + name="place", + steps=InvokeCfg( + call=PlaceCfg( + object="cube", + at=TargetRefCfg(target="drop"), + ) + ), + ), + ) + ), + targets={ + "drop": CyclicPoseTargetCfg( + values=( + PoseCfg( + position=(0.4, 0.1, 0.2), + quaternion_wxyz=(1.0, 0.0, 0.0, 0.0), + ), + ) + ) + }, + ) + workflows: list[object] = [] + original_analyze = SemanticSkillCompiler.analyze + + def record_analyze( + compiler: SemanticSkillCompiler, + calls: object, + **kwargs: object, + ) -> object: + workflow = original_analyze(compiler, calls, **kwargs) + workflows.append(workflow) + return workflow + + monkeypatch.setattr(SemanticSkillCompiler, "analyze", record_analyze) + + adapter.create_bridge(adapter.compile(config)) + + assert len(workflows) == 1 + workflow = workflows[0] + assert len(workflow.calls) == 2 # type: ignore[attr-defined] + downstream = workflow.calls[0].downstream_object_targets # type: ignore[attr-defined] + assert len(downstream) == 1 + assert downstream[0].pose is not None + torch.testing.assert_close( + downstream[0].pose.position, + torch.tensor((0.4, 0.1, 0.2)), + ) + assert factory.observation_samples == 0 + + +def test_parallel_program_requires_safety_validator_before_first_action() -> None: + factory = _FakeEnvironmentFactory() + adapter = ExpertProgramEnvironmentAdapter(factory, step_dt=_STEP_DT) + compiled = adapter.compile(_program_with_later_parallel_conflict()) + + with pytest.raises(ValueError, match="ParallelCommandSafetyValidator"): + adapter.create_bridge(compiled) + + assert factory.observation_samples == 0 + + +def test_later_parallel_claim_conflict_fails_during_whole_program_preflight() -> None: + factory = _FakeEnvironmentFactory() + adapter = ExpertProgramEnvironmentAdapter( + factory, + step_dt=_STEP_DT, + parallel_safety_validator=_AcceptParallelSafety(), + ) + compiled = adapter.compile(_program_with_later_parallel_conflict()) + + with pytest.raises(ValueError, match="overlapping resource claims"): + adapter.create_bridge(compiled) + + assert factory.observation_samples == 0 + + +def test_parallel_symbolic_write_conflict_fails_before_observation_or_action() -> None: + factory = _ParallelArticulationFactory() + adapter = ExpertProgramEnvironmentAdapter( + factory, + step_dt=_STEP_DT, + parallel_safety_validator=_AcceptParallelSafety(), + ) + compiled = adapter.compile(_parallel_articulation_program()) + materialized = compiled.materialize() + parallel_block = tuple(materialized.iter_segments())[0].parallel_block + assert parallel_block is not None + expected_path = parallel_block.branches[1].source_path + + with pytest.raises(SemanticValidationError) as error: + adapter.create_bridge(compiled) + + diagnostic = error.value.diagnostic + assert diagnostic.code == "parallel_symbolic_write_conflict" + assert diagnostic.path == expected_path + assert "articulation_joint['drawer', 'drawer_slide']" in diagnostic.message + assert factory.observation_samples == 0 + + +def test_runtime_assembly_shares_exact_bound_components() -> None: + """Compiler, runtime, clock, sink, scene, and profile form one ownership graph.""" + factory = _FakeEnvironmentFactory() + adapter = ExpertProgramEnvironmentAdapter(factory, step_dt=_STEP_DT) + + assembly = adapter.assemble_runtime(_program().integration) + + assert assembly.compiler.integration.scene_registry is assembly.scene_registry + assert assembly.compiler.integration.manifest is assembly.manifest + assert assembly.compiler.integration.engine is assembly.engine + assert assembly.compiler.integration.manifest.runtime_preset == "safe" + assert assembly.compiler.integration.robot_profile.engine is assembly.engine + assert assembly.runtime.compiler is assembly.compiler + assert assembly.runtime.clock is assembly.clock + assert assembly.command_sink.clock is assembly.clock + assert assembly.clock.step_dt == pytest.approx(_STEP_DT) + assert assembly.evidence_collector.registry.providers == {} + + +def test_safe_dynamic_collision_fails_before_observation_planning_or_command() -> None: + factory = _DynamicCollisionFactory() + adapter = ExpertProgramEnvironmentAdapter(factory, step_dt=_STEP_DT) + compiled = adapter.compile(_program(runtime_preset="safe")) + + with pytest.raises(SemanticValidationError) as error: + adapter.create_bridge(compiled) + + diagnostic = error.value.diagnostic + assert diagnostic.code == "safe_dynamic_collision_unsupported" + assert diagnostic.path == ( + "integration", + "robot_profile", + "presets", + "safe", + "motion_policy", + "dynamic_collision_mode", + ) + assert factory.calls == Counter(scene=2, profile=1, engine=1) + assert factory.pose_provider.calls == 0 + assert factory.observation_samples == 0 + assert factory.last_engine is not None + factory.last_engine.motion_generator.generate.assert_not_called() + + +@pytest.mark.parametrize( + ("field", "value", "match"), + ( + ("robot_profile", "other_robot", "selects robot_profile"), + ("scene_registry", "other_scene", "selects scene_registry"), + ), +) +def test_integration_id_mismatch_fails_before_live_factory_access( + field: str, + value: str, + match: str, +) -> None: + """Static selection drift never reaches simulation or motion factories.""" + factory = _FakeEnvironmentFactory() + adapter = ExpertProgramEnvironmentAdapter(factory, step_dt=_STEP_DT) + options = {field: value} + + with pytest.raises(ValueError, match=match): + adapter.compile(_program(**options)) + + assert factory.calls == Counter() + + +def test_robot_profile_factory_drift_fails_before_engine_creation() -> None: + """The declared profile ID must match the concrete factory output.""" + factory = _FakeEnvironmentFactory(returned_profile_id="different_robot") + adapter = ExpertProgramEnvironmentAdapter(factory, step_dt=_STEP_DT) + + with pytest.raises(ValueError, match="profile declaration drifted"): + adapter.assemble_runtime(_program().integration) + + assert factory.calls == Counter(scene=1, profile=1) + + +def test_factory_selection_declaration_drift_fails_before_live_access() -> None: + """A mutable factory cannot silently change IDs after adapter creation.""" + factory = _FakeEnvironmentFactory() + adapter = ExpertProgramEnvironmentAdapter(factory, step_dt=_STEP_DT) + factory.scene_registry_id = "changed_scene" + + with pytest.raises(ValueError, match="scene registry declaration drifted"): + adapter.compile(_program()) + + assert factory.calls == Counter() + + +def test_unknown_runtime_preset_fails_during_manifest_assembly() -> None: + """Runtime preset names are validated against the selected robot profile.""" + factory = _FakeEnvironmentFactory() + adapter = ExpertProgramEnvironmentAdapter(factory, step_dt=_STEP_DT) + + with pytest.raises(ValueError, match="Unknown runtime preset"): + adapter.assemble_runtime(_program(runtime_preset="unregistered").integration) + + +def test_factory_protocol_is_required() -> None: + """Loose objects cannot enter the production assembly boundary.""" + with pytest.raises(TypeError, match="ExpertProgramEnvironmentFactory"): + ExpertProgramEnvironmentAdapter(object(), step_dt=_STEP_DT) diff --git a/tests/gym/envs/expert_program/test_loader.py b/tests/gym/envs/expert_program/test_loader.py new file mode 100644 index 000000000..4ef5c4bb6 --- /dev/null +++ b/tests/gym/envs/expert_program/test_loader.py @@ -0,0 +1,252 @@ +# ---------------------------------------------------------------------------- +# 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 strict serialized Expert Program loading.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +import yaml + +from embodichain.lab.gym.envs.expert_program import ( + ConfigPath, + ExpertProgramDecodeError, + ExpertProgramValidationError, + InvokeCfg, + PickCfg, + load_expert_program, + loads_expert_program_json, + parse_expert_program_json, +) + + +def _program_data(*, schema_version: int = 1) -> dict[str, object]: + """Return one minimal complete Expert Program JSON value.""" + program: dict[str, object] = { + "kind": "invoke", + "call": {"kind": "pick", "object": "cube"}, + } + if schema_version == 2: + program = { + "kind": "parallel", + "branches": [ + program, + { + "kind": "invoke", + "call": {"kind": "pick", "object": "other_cube"}, + }, + ], + "barrier": { + "kind": "barrier", + "name": "both_picked", + "timeout_steps": 20, + "failure_policy": "fail_fast", + }, + } + return { + "schema_version": schema_version, + "program_id": "loader_pick", + "integration": { + "robot_profile": "test_robot", + "scene_registry": "test_scene", + "runtime_preset": "safe", + }, + "targets": {}, + "program": program, + } + + +def _program_json() -> str: + """Serialize the minimal program using standards-compliant JSON.""" + return json.dumps(_program_data()) + + +class _RejectingValidationContext: + """Reject integration references after recording their exact path.""" + + def __init__(self) -> None: + self.integration_paths: list[ConfigPath] = [] + + def validate_integration( + self, + integration: object, + *, + path: ConfigPath, + ) -> None: + del integration + self.integration_paths.append(path) + raise KeyError("unavailable integration") + + def validate_semantic_call(self, call: object, *, path: ConfigPath) -> None: + del call, path + + def validate_scene_reference( + self, + reference: str, + *, + role: str, + path: ConfigPath, + ) -> None: + del reference, role, path + + def validate_post_policy(self, policy: object, *, path: ConfigPath) -> None: + del policy, path + + def validate_validator(self, validator: object, *, path: ConfigPath) -> None: + del validator, path + + +def test_parse_expert_program_json_preserves_predecode_mapping() -> None: + value = parse_expert_program_json('{"host_integration_pending": [true, null, 3.5]}') + + assert value == {"host_integration_pending": [True, None, 3.5]} + + +@pytest.mark.parametrize("response", ["[]", "null", '"program"']) +def test_parse_expert_program_json_requires_top_level_mapping(response: str) -> None: + with pytest.raises(ExpertProgramDecodeError) as error: + parse_expert_program_json(response) + + assert error.value.code == "expected_mapping" + assert error.value.path == () + + +def test_loads_expert_program_json_decodes_one_plain_document() -> None: + config = loads_expert_program_json(f"\n{_program_json()}\t") + + assert type(config.program) is InvokeCfg + assert type(config.program.call) is PickCfg + assert config.program.call.object == "cube" + + +@pytest.mark.parametrize("suffix", [".json", ".yaml"]) +@pytest.mark.parametrize("schema_version", [1, 2]) +def test_load_expert_program_forwards_validation_context_for_each_format( + tmp_path: Path, + suffix: str, + schema_version: int, +) -> None: + data = _program_data(schema_version=schema_version) + serialized = json.dumps(data) if suffix == ".json" else yaml.safe_dump(data) + path = tmp_path / f"program{suffix}" + path.write_text(serialized, encoding="utf-8") + context = _RejectingValidationContext() + + with pytest.raises(ExpertProgramValidationError) as error: + load_expert_program(path, validation_context=context) + + assert error.value.code == "reference_validation_failed" + assert error.value.path == ("integration",) + assert context.integration_paths == [("integration",)] + + +def test_loads_expert_program_json_rejects_nested_duplicate_keys() -> None: + duplicate = _program_json().replace( + '"object": "cube"', + '"object": "cube", "object": "other"', + ) + + with pytest.raises(ExpertProgramDecodeError) as error: + loads_expert_program_json(duplicate) + + assert error.value.code == "duplicate_json_key" + + +@pytest.mark.parametrize( + "invalid_response", + [ + "```json\n{}\n```", + f"{_program_json()} trailing text", + f"{_program_json()} {_program_json()}", + ], +) +def test_loads_expert_program_json_requires_one_unfenced_document( + invalid_response: str, +) -> None: + with pytest.raises(ExpertProgramDecodeError) as error: + loads_expert_program_json(invalid_response) + + assert error.value.code == "invalid_json" + + +@pytest.mark.parametrize("number", ["NaN", "Infinity", "-Infinity", "1e400"]) +def test_loads_expert_program_json_rejects_non_finite_numbers(number: str) -> None: + response = f'{{"value": {number}}}' + + with pytest.raises(ExpertProgramDecodeError) as error: + loads_expert_program_json(response) + + assert error.value.code == "non_finite_number" + + +def test_loads_expert_program_json_enforces_utf8_byte_limit() -> None: + response = _program_json() + too_small = len(response.encode("utf-8")) - 1 + + with pytest.raises(ExpertProgramDecodeError) as error: + loads_expert_program_json(response, max_bytes=too_small) + + assert error.value.code == "input_too_large" + + +def test_loads_expert_program_json_normalizes_invalid_utf8_text() -> None: + response = "\ud800" + + with pytest.raises(ExpertProgramDecodeError) as error: + loads_expert_program_json(response) + + assert error.value.code == "invalid_utf8" + + +def test_loads_expert_program_json_rejects_escaped_unpaired_surrogate() -> None: + response = _program_json().replace("loader_pick", r"\ud800") + + with pytest.raises(ExpertProgramDecodeError) as error: + loads_expert_program_json(response) + + assert error.value.code == "invalid_utf8" + + +def test_loads_expert_program_json_accepts_escaped_surrogate_pair() -> None: + response = _program_json().replace("loader_pick", r"\ud83d\ude00") + + config = loads_expert_program_json(response) + + assert config.program_id == "😀" + + +def test_loads_expert_program_json_normalizes_oversized_integer() -> None: + data = _program_data() + data["targets"] = { + "goal": { + "kind": "cyclic_pose", + "values": [ + { + "position": [10**400, 0, 0], + "quaternion_wxyz": [1, 0, 0, 0], + } + ], + } + } + + with pytest.raises(ExpertProgramDecodeError) as error: + loads_expert_program_json(json.dumps(data)) + + assert error.value.code == "invalid_value" + assert error.value.path == ("targets", "goal", "values", 0) diff --git a/tests/gym/envs/expert_program/test_parallel_compiler.py b/tests/gym/envs/expert_program/test_parallel_compiler.py new file mode 100644 index 000000000..84c678c3e --- /dev/null +++ b/tests/gym/envs/expert_program/test_parallel_compiler.py @@ -0,0 +1,221 @@ +# ---------------------------------------------------------------------------- +# 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 provider-free schema-v2 parallel program compilation.""" + +from __future__ import annotations + +import pytest +import torch + +from embodichain.lab.gym.envs.expert_program import ( + BarrierCfg, + ExpertProgramCfg, + ExpertProgramCompileError, + ExpertProgramCompiler, + ExpertProgramIntegrationCfg, + InvokeCfg, + ParallelCfg, + PickCfg, + RepeatCfg, + SegmentCfg, + SequenceCfg, +) +from embodichain.lab.sim.atomic_actions import EntityState +from embodichain.lab.sim.skills.calls import Pick +from embodichain.lab.sim.skills.scene import ( + SceneEntityRegistration, + SceneObjectRef, + SceneRegistry, +) + + +class _NeverObserveProvider: + """Reject dynamic observation during provider-free compilation.""" + + def observe( + self, + *, + timestamp: float, + env_ids: torch.Tensor, + ) -> EntityState: + del timestamp, env_ids + raise AssertionError("Compilation must not observe the scene.") + + +def _compiler() -> ExpertProgramCompiler: + provider = _NeverObserveProvider() + registry = SceneRegistry( + tuple( + SceneEntityRegistration( + ref=SceneObjectRef(entity_id), + state_provider=provider, + ) + for entity_id in ("left_cube", "right_cube") + ) + ) + return ExpertProgramCompiler.from_scene_registry(registry) + + +def _integration() -> ExpertProgramIntegrationCfg: + return ExpertProgramIntegrationCfg( + robot_profile="dual_arm", + scene_registry="scene", + runtime_preset="safe", + ) + + +def _parallel() -> ParallelCfg: + return ParallelCfg( + branches=( + SequenceCfg( + items=( + InvokeCfg(call=PickCfg(object="left_cube")), + InvokeCfg(call=PickCfg(object="left_cube")), + ) + ), + RepeatCfg( + count=2, + body=InvokeCfg(call=PickCfg(object="right_cube")), + ), + ), + barrier=BarrierCfg( + name="both_arms_done", + timeout_steps=240, + failure_policy="fail_fast", + ), + ) + + +def _config(program: ParallelCfg | SegmentCfg | SequenceCfg) -> ExpertProgramCfg: + return ExpertProgramCfg( + schema_version=2, + program_id="parallel_pick", + integration=_integration(), + program=program, + targets={}, + ) + + +def test_parallel_compiles_independent_ordered_lanes_and_explicit_join() -> None: + segment = tuple(_compiler().compile(_config(_parallel())))[0] + + assert segment.implicit + assert segment.name == "parallel:both_arms_done" + assert segment.parallel_block is not None + block = segment.parallel_block + assert block.barrier.name == "both_arms_done" + assert block.barrier.timeout_steps == 240 + assert block.barrier.failure_policy == "fail_fast" + assert tuple(branch.branch_index for branch in block.branches) == (0, 1) + assert tuple(len(branch.calls) for branch in block.branches) == (2, 2) + assert tuple(call.call_index for call in segment.calls) == (0, 1, 2, 3) + assert tuple(call.segment_call_index for call in segment.calls) == (0, 1, 2, 3) + assert segment.calls == tuple( + call for branch in block.branches for call in branch.calls + ) + assert all(type(call.call) is Pick for call in segment.calls) + assert tuple(call.call.object.entity_id for call in block.branches[0].calls) == ( + "left_cube", + "left_cube", + ) + assert tuple(call.call.object.entity_id for call in block.branches[1].calls) == ( + "right_cube", + "right_cube", + ) + assert tuple( + frame.iteration_index + for call in block.branches[1].calls + for frame in call.repeat_frames + ) == (0, 1) + + +def test_segment_may_wrap_one_parallel_block() -> None: + segment = tuple( + _compiler().compile(_config(SegmentCfg(name="dual_pick", steps=_parallel()))) + )[0] + + assert not segment.implicit + assert segment.name == "dual_pick" + assert segment.parallel_block is not None + assert len(segment.calls) == 4 + + +def test_materialized_analysis_stops_sequential_lookahead_at_parallel_barriers() -> ( + None +): + config = _config( + SequenceCfg( + items=( + InvokeCfg(call=PickCfg(object="left_cube")), + _parallel(), + InvokeCfg(call=PickCfg(object="left_cube")), + ) + ) + ) + + program = _compiler().compile(config).materialize() + analyses = program.preflight_analyses() + + assert [analysis.kind for analysis in analyses] == [ + "sequential_stretch", + "parallel_branch", + "parallel_branch", + "sequential_stretch", + ] + assert [analysis.segment_indices for analysis in analyses] == [ + (0,), + (1,), + (1,), + (2,), + ] + assert program.sequential_execution_analysis(0).segment_indices == (0,) + assert program.sequential_execution_analysis(2).segment_indices == (2,) + with pytest.raises(ValueError, match="Parallel segments"): + program.sequential_execution_analysis(1) + + +def test_parallel_branch_rejects_segment_owned_lifecycle() -> None: + invoke = InvokeCfg(call=PickCfg(object="left_cube")) + parallel = ParallelCfg( + branches=( + SegmentCfg(name="branch", steps=invoke), + invoke, + ), + barrier=BarrierCfg(name="join"), + ) + + with pytest.raises(ValueError, match="wrap the Parallel node in one Segment"): + _config(parallel) + + +def test_segment_rejects_mixed_sequential_and_parallel_tree() -> None: + invoke = InvokeCfg(call=PickCfg(object="left_cube")) + config = _config( + SegmentCfg( + name="ambiguous_boundary", + steps=SequenceCfg(items=(invoke, _parallel())), + ) + ) + + with pytest.raises( + ExpertProgramCompileError, + match="either a call-only program or one direct Parallel", + ): + _compiler().compile(config) + + +__all__: list[str] = [] diff --git a/tests/gym/envs/expert_program/test_parallel_schema.py b/tests/gym/envs/expert_program/test_parallel_schema.py new file mode 100644 index 000000000..b34ffcb69 --- /dev/null +++ b/tests/gym/envs/expert_program/test_parallel_schema.py @@ -0,0 +1,137 @@ +# ---------------------------------------------------------------------------- +# 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 Expert Program schema Version 2 parallel nodes.""" + +from __future__ import annotations + +import pytest + +from embodichain.lab.gym.envs.expert_program.cfg import ( + BarrierCfg, + ExpertProgramCfg, + ExpertProgramIntegrationCfg, + InvokeCfg, + ParallelCfg, + PickCfg, +) +from embodichain.lab.gym.envs.expert_program.decoder import ( + ExpertProgramDecodeError, + decode_expert_program, +) + + +def _payload(*, schema_version: int = 2) -> dict[str, object]: + return { + "schema_version": schema_version, + "program_id": "parallel_pick", + "integration": { + "robot_profile": "dual_arm", + "scene_registry": "scene", + "runtime_preset": "safe", + }, + "targets": {}, + "program": { + "kind": "parallel", + "branches": [ + { + "kind": "invoke", + "call": {"kind": "pick", "object": "left_cube"}, + }, + { + "kind": "invoke", + "call": {"kind": "pick", "object": "right_cube"}, + }, + ], + "barrier": { + "kind": "barrier", + "name": "both_picked", + "timeout_steps": 200, + "failure_policy": "fail_fast", + }, + }, + } + + +def test_decode_schema_v2_parallel_with_explicit_barrier() -> None: + config = decode_expert_program(_payload()) + + assert config.schema_version == 2 + assert type(config.program) is ParallelCfg + assert len(config.program.branches) == 2 + assert config.program.barrier == BarrierCfg( + name="both_picked", + timeout_steps=200, + failure_policy="fail_fast", + ) + + +def test_schema_v1_rejects_parallel_discriminator() -> None: + with pytest.raises(ExpertProgramDecodeError) as error: + decode_expert_program(_payload(schema_version=1)) + + assert error.value.code == "unknown_discriminator" + assert error.value.path == ("program", "kind") + + +def test_parallel_requires_two_branches_and_explicit_barrier() -> None: + payload = _payload() + program = payload["program"] + assert type(program) is dict + program["branches"] = program["branches"][:1] # type: ignore[index] + with pytest.raises(ExpertProgramDecodeError) as error: + decode_expert_program(payload) + assert error.value.code == "parallel_branch_count" + + payload = _payload() + program = payload["program"] + assert type(program) is dict + del program["barrier"] + with pytest.raises(ExpertProgramDecodeError) as error: + decode_expert_program(payload) + assert error.value.code == "missing_field" + assert error.value.path == ("program", "barrier") + + +def test_barrier_is_not_valid_as_a_standalone_program() -> None: + with pytest.raises(ValueError, match="only be owned by Parallel"): + ExpertProgramCfg( + schema_version=2, + program_id="invalid_barrier", + integration=ExpertProgramIntegrationCfg( + robot_profile="profile", + scene_registry="scene", + runtime_preset="safe", + ), + targets={}, + program=BarrierCfg(name="orphan"), + ) + + +def test_parallel_cfg_rejects_nested_parallel() -> None: + invoke = InvokeCfg(call=PickCfg(object="cube")) + nested = ParallelCfg( + branches=(invoke, invoke), + barrier=BarrierCfg(name="inner"), + ) + with pytest.raises(ValueError, match="Nested Parallel"): + ParallelCfg( + branches=(nested, invoke), + barrier=BarrierCfg(name="outer"), + ) + + +__all__: list[str] = [] diff --git a/tests/gym/envs/expert_program/test_simulation.py b/tests/gym/envs/expert_program/test_simulation.py new file mode 100644 index 000000000..652df4702 --- /dev/null +++ b/tests/gym/envs/expert_program/test_simulation.py @@ -0,0 +1,436 @@ +# ---------------------------------------------------------------------------- +# 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 explicit Expert Program simulation bindings.""" + +from __future__ import annotations + +from dataclasses import dataclass, replace + +import pytest +import torch + +from embodichain.lab.gym.envs.expert_program.simulation import ( + AntipodalGraspAffordanceBinding, + ArticulationOperationAffordanceBinding, + ArticulationOperationTargetBinding, + ControlPartCommandPreset, + ControlPartEndpointBinding, + ControlPartResourceBinding, + RobotResourceBinding, + SimulationArticulationBinding, + SimulationArticulationLinkBinding, + SimulationRigidObjectBinding, + SimulationResourceEndpointBinding, + SimulationRobotResourceBinding, + SimulationRobotSkillProfileBinding, + SimulationSceneBinding, +) +from embodichain.lab.sim.atomic_actions import ( + AntipodalAffordance, + ArticulationOperationAffordance, + CARTESIAN_POSE_CAPABILITY, + GRASP_CAPABILITY, +) +from embodichain.lab.sim.skills import ( + ARTICULATION_OPERATION_AFFORDANCE_CAPABILITY, + GRASP_AFFORDANCE_CAPABILITY, + SceneAffordanceRef, + SceneArticulationRef, + SceneDynamics, + SceneLinkRef, + SceneObjectRef, + SkillPolicyPreset, +) +from embodichain.lab.sim.skills.profiles import ResourceEndpoint + +_BATCH_SIZE = 2 +_OPEN_TARGET = 0.42 +_OPEN_DISPLACEMENT = 0.4 + + +class _RigidObject: + """Minimal selected rigid object with a batched triangle mesh.""" + + def __init__(self) -> None: + self.pose = torch.eye(4).repeat(_BATCH_SIZE, 1, 1) + self.vertices = torch.tensor( + ( + (0.0, 0.0, 0.0), + (1.0, 0.0, 0.0), + (0.0, 1.0, 0.0), + ), + dtype=torch.float32, + ).repeat(_BATCH_SIZE, 1, 1) + self.triangles = torch.tensor( + (((0, 1, 2),),), + dtype=torch.int32, + ).repeat(_BATCH_SIZE, 1, 1) + + def get_local_pose(self, *, to_matrix: bool) -> torch.Tensor: + assert to_matrix is True + return self.pose + + def get_vertices( + self, + env_ids: list[int], + *, + scale: bool, + ) -> torch.Tensor: + assert scale is True + return self.vertices[env_ids] + + def get_triangles(self, env_ids: list[int]) -> torch.Tensor: + return self.triangles[env_ids] + + +class _Articulation: + """Minimal articulation exposing exact joint and link lookup surfaces.""" + + joint_names = ("drawer_slide",) + link_names = ("drawer_handle",) + + def __init__(self) -> None: + self.pose = torch.eye(4).repeat(_BATCH_SIZE, 1, 1) + self.qpos = torch.tensor(((0.1,), (0.2,)), dtype=torch.float32) + self.link_pose = torch.eye(4).repeat(_BATCH_SIZE, 1, 1) + self.link_pose[:, 0, 3] = torch.tensor((0.3, 0.4)) + + def get_local_pose(self, *, to_matrix: bool) -> torch.Tensor: + assert to_matrix is True + return self.pose + + def get_qpos(self, *, target: bool) -> torch.Tensor: + assert target is False + return self.qpos + + def get_link_pose( + self, + link_name: str, + *, + env_ids: list[int], + to_matrix: bool, + ) -> torch.Tensor: + assert link_name == "drawer_handle" + assert to_matrix is True + return self.link_pose[env_ids] + + +class _Simulation: + """Explicit native-UID lookup fixture.""" + + def __init__(self) -> None: + self.rigid_object = _RigidObject() + self.articulation = _Articulation() + + def get_rigid_object(self, uid: str) -> _RigidObject | None: + return self.rigid_object if uid == "native_cube" else None + + def get_articulation(self, uid: str) -> _Articulation | None: + return self.articulation if uid == "native_drawer" else None + + +class _Robot: + """Minimal robot control-part lookup fixture.""" + + control_parts = {"arm": object(), "hand": object()} + + def get_joint_ids(self, *, name: str) -> list[int]: + return {"arm": [0, 1], "hand": [2]}[name] + + +@dataclass(frozen=True, slots=True) +class _MobileEndpoint(ResourceEndpoint): + """Test-only non-joint endpoint declaration.""" + + controller_id: str + + +def _scene_binding() -> SimulationSceneBinding: + """Build one cube-and-drawer binding using only typed declarations.""" + return SimulationSceneBinding( + registry_id="tabletop", + rigid_objects=( + SimulationRigidObjectBinding( + entity_id="cube", + simulation_uid="native_cube", + aliases=("perceived_cube",), + dynamics=SceneDynamics.DYNAMIC, + semantic_type="cube", + default_grasp_affordance="cube_grasp", + ), + ), + articulations=( + SimulationArticulationBinding( + entity_id="drawer", + simulation_uid="native_drawer", + semantic_type="drawer", + default_operation_affordance="drawer_handle_operation", + ), + ), + links=( + SimulationArticulationLinkBinding( + entity_id="drawer_handle_link", + articulation_id="drawer", + native_link_name="drawer_handle", + ), + ), + antipodal_grasps=( + AntipodalGraspAffordanceBinding( + entity_id="cube_grasp", + object_id="cube", + native_name="mesh_antipodal", + revision="cube-grasp-v1", + ), + ), + articulation_operations=( + ArticulationOperationAffordanceBinding( + entity_id="drawer_handle_operation", + articulation_id="drawer", + link_id="drawer_handle_link", + joint_id="drawer_slide", + revision="drawer-operation-v1", + semantic_targets={ + "open": ArticulationOperationTargetBinding( + target_position=_OPEN_TARGET, + displacement=_OPEN_DISPLACEMENT, + ) + }, + ), + ), + ) + + +def _profile_binding() -> SimulationRobotSkillProfileBinding: + """Build one manipulation profile declaration.""" + return SimulationRobotSkillProfileBinding( + profile_id="test_robot", + resources=( + ControlPartResourceBinding( + resource_id="manipulator", + endpoints=( + ControlPartEndpointBinding( + endpoint_id="motion", + control_part="arm", + capabilities=frozenset({CARTESIAN_POSE_CAPABILITY}), + ), + ControlPartEndpointBinding( + endpoint_id="grasp", + control_part="hand", + capabilities=frozenset({GRASP_CAPABILITY}), + command_preset="parallel_gripper", + ), + ), + ), + ), + command_presets=( + ControlPartCommandPreset( + preset_id="parallel_gripper", + control_part="hand", + commands={"open": (0.0,), "grasp": (1.0,)}, + ), + ), + defaults={"pick_up": {"primary": "manipulator"}}, + presets=(SkillPolicyPreset("safe"),), + default_preset="safe", + ) + + +def test_scene_binding_builds_existing_registry_contracts() -> None: + simulation = _Simulation() + + registry = _scene_binding().build(simulation) # type: ignore[arg-type] + snapshot = registry.make_scene_provider().snapshot( + timestamp=0.0, + env_ids=torch.tensor((0, 1), dtype=torch.long), + ) + + assert registry.resolve("perceived_cube") == SceneObjectRef("cube") + assert registry.resolve("drawer") == SceneArticulationRef("drawer") + assert registry.resolve("drawer_handle_link") == SceneLinkRef("drawer_handle_link") + grasp_ref = registry.resolve_affordance( + "cube", + capability=GRASP_AFFORDANCE_CAPABILITY, + ) + assert grasp_ref == SceneAffordanceRef("cube_grasp") + grasp = registry.lookup(grasp_ref).affordance + assert isinstance(grasp, AntipodalAffordance) + assert grasp.mesh_vertices is not None and grasp.mesh_vertices.shape == (3, 3) + operation_ref = registry.resolve_affordance( + "drawer", + capability=ARTICULATION_OPERATION_AFFORDANCE_CAPABILITY, + ) + operation = registry.lookup(operation_ref).affordance + assert isinstance(operation, ArticulationOperationAffordance) + assert operation.joint_id == "drawer_slide" + assert operation.semantic_targets["open"].target_position == pytest.approx( + _OPEN_TARGET + ) + assert torch.equal( + snapshot.articulation_joints[("drawer", "drawer_slide")].position, + simulation.articulation.qpos, + ) + assert torch.equal( + snapshot.entities["drawer_handle_operation"].pose, + simulation.articulation.link_pose, + ) + + +def test_scene_binding_fails_closed_on_missing_native_entity() -> None: + binding = _scene_binding() + missing = replace( + binding.rigid_objects[0], + simulation_uid="missing_cube", + ) + + with pytest.raises(KeyError, match="missing_cube"): + replace(binding, rigid_objects=(missing,)).build( # type: ignore[arg-type] + _Simulation() + ) + + +def test_scene_binding_fails_closed_on_missing_native_link() -> None: + binding = _scene_binding() + missing = replace(binding.links[0], native_link_name="missing_handle") + + with pytest.raises(KeyError, match="missing_handle"): + replace(binding, links=(missing,)).build( # type: ignore[arg-type] + _Simulation() + ) + + +def test_scene_binding_fails_closed_on_missing_native_joint() -> None: + binding = _scene_binding() + missing = replace( + binding.articulation_operations[0], + joint_id="missing_joint", + ) + + with pytest.raises(KeyError, match="missing_joint"): + replace(binding, articulation_operations=(missing,)).build( # type: ignore[arg-type] + _Simulation() + ) + + +def test_robot_profile_binding_builds_existing_profile_contracts() -> None: + profile = _profile_binding().build(_Robot()) # type: ignore[arg-type] + + resource = profile.resources["manipulator"] + motion = resource.endpoints["motion"] + grasp = resource.endpoints["grasp"] + assert motion.control_part == "arm" + assert motion.capabilities == frozenset({CARTESIAN_POSE_CAPABILITY}) + assert grasp.control_part == "hand" + assert grasp.command_profile == "parallel_gripper" + command = profile.command_profiles["parallel_gripper"].commands["grasp"] + assert torch.equal(command.positions, torch.tensor((1.0,))) + assert profile.defaults["pick_up"].resources == {"primary": "manipulator"} + + +def test_generic_resource_binding_owns_arbitrary_typed_endpoint() -> None: + endpoint = _MobileEndpoint( + controller_id="base_controller", + capabilities=frozenset({"motion.base.velocity"}), + ) + binding = RobotResourceBinding( + resource_id="mobile_base", + endpoints={"motion": endpoint}, + ) + + assert isinstance(binding, SimulationRobotResourceBinding) + resource = binding.build(object()) # type: ignore[arg-type] + built_endpoint = resource.endpoints["motion"] + + assert isinstance(built_endpoint, _MobileEndpoint) + assert built_endpoint is not endpoint + assert built_endpoint.controller_id == "base_controller" + assert resource.members == () + + +def test_control_part_endpoint_binding_implements_public_build_protocol() -> None: + binding = ControlPartEndpointBinding( + endpoint_id="motion", + control_part="arm", + capabilities=frozenset({CARTESIAN_POSE_CAPABILITY}), + ) + + assert isinstance(binding, SimulationResourceEndpointBinding) + + +def test_whole_body_control_part_remains_supported_and_strict() -> None: + class WholeBodyRobot: + control_parts = {"whole_body": object()} + + def get_joint_ids(self, *, name: str) -> list[int]: + assert name == "whole_body" + return [0, 1, 2, 3] + + binding = SimulationRobotSkillProfileBinding( + profile_id="whole_body_robot", + resources=( + ControlPartResourceBinding( + resource_id="body", + endpoints=( + ControlPartEndpointBinding( + endpoint_id="motion", + control_part="whole_body", + capabilities=frozenset({"motion.whole_body"}), + ), + ), + ), + ), + ) + + profile = binding.build(WholeBodyRobot()) # type: ignore[arg-type] + endpoint = profile.resources["body"].endpoints["motion"] + + assert endpoint.control_part == "whole_body" + assert endpoint.capabilities == frozenset({"motion.whole_body"}) + + +def test_robot_profile_binding_fails_closed_on_missing_control_part() -> None: + binding = _profile_binding() + resource = binding.resources[0] + missing_endpoint = replace( + resource.endpoints[0], + control_part="missing_arm", + ) + + with pytest.raises(KeyError, match="missing_arm"): + replace( + binding, + resources=( + replace( + resource, + endpoints=(missing_endpoint, resource.endpoints[1]), + ), + ), + ).build( + _Robot() + ) # type: ignore[arg-type] + + +def test_robot_profile_binding_rejects_wrong_command_width() -> None: + binding = _profile_binding() + invalid = replace( + binding.command_presets[0], + commands={"open": (0.0, 0.0), "grasp": (1.0, 1.0)}, + ) + + with pytest.raises(ValueError, match="has 2 positions.*has 1 joints"): + replace(binding, command_presets=(invalid,)).build( # type: ignore[arg-type] + _Robot() + ) diff --git a/tests/gym/envs/expert_program/test_simulation_environment.py b/tests/gym/envs/expert_program/test_simulation_environment.py new file mode 100644 index 000000000..0b8c5c186 --- /dev/null +++ b/tests/gym/envs/expert_program/test_simulation_environment.py @@ -0,0 +1,1893 @@ +# ---------------------------------------------------------------------------- +# 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 reusable simulation-backed Expert Program assembly.""" + +from __future__ import annotations + +import ast +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, fields, is_dataclass +import inspect +import textwrap +from types import MethodType, SimpleNamespace +from typing import Any, ClassVar +from unittest.mock import MagicMock + +import pytest +import torch + +from embodichain.lab.gym.envs.expert_program import ( + AntipodalGraspAffordanceBinding, + ControlCommandStateEvidenceTracker, + ControlPartCommandPreset, + ControlPartEndpointBinding, + ControlPartResourceBinding, + ExpertProgramCompiler, + ExpertProgramCfg, + ExpertProgramIntegrationCfg, + ExpertProgramEnvironmentAdapter, + ExpertProgramRuntimeAssembly, + HandOverCfg, + InvokeCfg, + RobotResourceBinding, + SharedTickSceneProvider, + SimulationExpertProgramFactory, + SimulationPlanningObservationProvider, + SimulationRigidObjectBinding, + SimulationRobotSkillProfileBinding, + SimulationSceneBinding, + create_simulation_expert_program_adapter, + decode_expert_program, +) +from embodichain.lab.gym.envs.expert_program.bridge import EnvironmentStepClock +from embodichain.lab.sim.atomic_actions import ( + ActionInvocation, + Affordance, + BATCH_INVERSE_KINEMATICS_CAPABILITY, + CARTESIAN_POSE_CAPABILITY, + CommandAcknowledgement, + EntityState, + FORWARD_KINEMATICS_CAPABILITY, + GRASP_CAPABILITY, + HeldObjectState, + MotionPolicy, + ObservedArticulationJointState, + PlanningContext, + StateDelta, + TaskState, +) +from embodichain.lab.sim.atomic_actions.runner import ExecutionRunnerCfg +from embodichain.lab.sim.atomic_actions.bindings import JointPositionTarget +from embodichain.lab.sim.atomic_actions.bindings import RuntimeEndpointTarget +from embodichain.lab.sim.atomic_actions.control import ControlPartCommandProfile +from embodichain.lab.sim.atomic_actions.runtime_commands import ( + EndpointCommand, + JointPositionPayload, + RuntimeCommandFrame, +) +from embodichain.lab.sim.planners import MotionGenerator +from embodichain.lab.sim.skills import ( + AtomicSkills, + BoundSemanticCall, + EndpointResolution, + HandOver, + HandOverPoseProvider, + HandOverPoseTargets, + Pick, + Place, + RelationTargetGrounder, + ResourceEndpoint, + ResourceEndpointAdapter, + SceneArticulationRef, + SceneEntityRegistration, + SceneRegistry, + SemanticCallSpec, + SemanticObjectTarget, + SemanticPose, + SemanticRelationTarget, + SemanticValidationError, + SkillPolicyPreset, +) +from embodichain.lab.sim.skills.effects import ( + CONSTRAINT_EFFECT_CHANNEL, + CONTROL_PART_EVIDENCE_PROVIDER_ID, + CONTROL_PART_EVIDENCE_PROVIDER_REVISION, + BinaryEffectClause, + BinaryEvidenceKind, + ControlPartEvidenceAddress, + EffectEvidenceSourceRef, + HeldObjectRelation, + HeldObjectStateExpectation, +) +from embodichain.lab.sim.skills.evidence import ( + BinaryEffectEvidenceQuery, + BinaryEffectEvidenceBatch, + BinaryEffectObservation, + EffectEvidenceCollectionContext, + PoseRelationEvidenceBatch, +) +from embodichain.lab.sim.skills.runtime import ( + SkillEffectTrace, + SkillResult, + SkillRuntime, + SkillStatus, +) +from embodichain.lab.sim.skills.scene import SceneObjectRef + +_BATCH_SIZE = 3 +_ROBOT_DOF = 2 +_STEP_DT = 0.04 +_TRACKER_ENV_IDS = torch.tensor((7, 3, 11), dtype=torch.long) +_HAND_OPEN_POSITION = 0.0 +_HAND_GRASP_POSITION = 0.8 +_HAND_INTERMEDIATE_POSITION = 0.4 +_DUAL_ROBOT_DOF = 4 +_RELEASE_SEPARATION = 0.2 +_DIRECT_PLACE_TARGET = SemanticPose( + position=(0.0, 0.0, 0.0), + quaternion_wxyz=(1.0, 0.0, 0.0, 0.0), +) +_QUICKSTART_MAX_LINES = 15 + + +def _command_state_tracker() -> ControlCommandStateEvidenceTracker: + """Build a three-row tracker for one semantic gripper profile.""" + profile = ControlPartCommandProfile.joint_positions( + open=torch.tensor((_HAND_OPEN_POSITION,)), + grasp=torch.tensor((_HAND_GRASP_POSITION,)), + ) + return ControlCommandStateEvidenceTracker( + {"hand": profile}, + _TRACKER_ENV_IDS, + ) + + +def _hand_command_frame( + *, + env_ids: tuple[int, ...], + positions: tuple[float, ...], + active: tuple[bool, ...] | None = None, +) -> RuntimeCommandFrame: + """Build one row-addressed semantic hand command frame.""" + batch_size = len(env_ids) + if len(positions) != batch_size: + raise ValueError("positions must have one value per environment ID.") + if active is None: + active = (True,) * batch_size + return RuntimeCommandFrame( + commands=( + EndpointCommand( + target=JointPositionTarget("hand", (1,)), + payload=JointPositionPayload( + torch.tensor(positions, dtype=torch.float32).unsqueeze(1) + ), + ), + ), + active_mask=torch.tensor(active, dtype=torch.bool), + env_ids=torch.tensor(env_ids, dtype=torch.long), + hold_duration=torch.full((batch_size,), _STEP_DT), + ) + + +def _hand_state_observation( + tracker: ControlCommandStateEvidenceTracker, + *env_ids: int, +) -> BinaryEffectObservation: + """Observe command-state evidence in an explicit stable-ID order.""" + expectation = HeldObjectStateExpectation( + expectation_id="held-cube", + relation=HeldObjectRelation.ATTACHED, + object_id="cube", + slot_id="primary", + resource_id="manipulator", + task_state_key="held-cube", + ) + query = BinaryEffectEvidenceQuery( + BinaryEffectClause( + clause_id="hand-constraint", + expectation_id=expectation.expectation_id, + source=EffectEvidenceSourceRef( + CONTROL_PART_EVIDENCE_PROVIDER_ID, + CONTROL_PART_EVIDENCE_PROVIDER_REVISION, + ControlPartEvidenceAddress("hand", CONSTRAINT_EFFECT_CHANNEL), + ), + evidence_kind=BinaryEvidenceKind.CONSTRAINT, + expected=True, + ), + expectation, + ) + context = EffectEvidenceCollectionContext( + timestamp=0.0, + observation_revision=0, + env_ids=torch.tensor(env_ids, dtype=torch.long), + ) + return tracker.observe(query, context) + + +class _CountingEntityProvider: + """Return row-addressed poses and record every native acquisition.""" + + def __init__(self) -> None: + self.calls: list[tuple[float, torch.Tensor]] = [] + + def observe(self, *, timestamp: float, env_ids: torch.Tensor) -> EntityState: + """Return one distinct x translation for each environment ID.""" + self.calls.append((timestamp, env_ids.clone())) + pose = torch.eye(4).repeat(env_ids.numel(), 1, 1) + pose[:, 0, 3] = env_ids.to(dtype=pose.dtype) + return EntityState(pose) + + +class _CountingJointProvider: + """Return row-addressed articulation state and record acquisitions.""" + + def __init__(self) -> None: + self.calls: list[tuple[float, torch.Tensor]] = [] + + def observe_joints( + self, + *, + timestamp: float, + env_ids: torch.Tensor, + ) -> dict[str, ObservedArticulationJointState]: + """Return one scalar joint position per environment ID.""" + self.calls.append((timestamp, env_ids.clone())) + position = env_ids.to(dtype=torch.float32).unsqueeze(1) + return { + "slide": ObservedArticulationJointState( + position, + torch.ones(env_ids.numel(), dtype=torch.bool), + ) + } + + +def _shared_scene_provider() -> tuple[ + SharedTickSceneProvider, + _CountingEntityProvider, + _CountingJointProvider, +]: + """Build one full-batch registry provider with observable acquisitions.""" + entity_provider = _CountingEntityProvider() + joint_provider = _CountingJointProvider() + registry = SceneRegistry( + ( + SceneEntityRegistration( + ref=SceneArticulationRef("drawer"), + state_provider=entity_provider, + joint_state_provider=joint_provider, + ), + ) + ) + delegate = registry.make_scene_provider(batch_size=_BATCH_SIZE) + full_env_ids = torch.arange(_BATCH_SIZE, dtype=torch.long) + return ( + SharedTickSceneProvider(delegate, full_env_ids), + entity_provider, + joint_provider, + ) + + +def test_shared_tick_scene_provider_projects_partial_rows_without_resampling() -> None: + """Planning full batch and evidence subsets share one native acquisition.""" + provider, entity_provider, joint_provider = _shared_scene_provider() + full_env_ids = torch.arange(_BATCH_SIZE, dtype=torch.long) + + full = provider.snapshot(timestamp=0.0, env_ids=full_env_ids) + subset = provider.snapshot( + timestamp=0.0, + env_ids=torch.tensor((2, 0), dtype=torch.long), + ) + + assert len(entity_provider.calls) == 1 + assert len(joint_provider.calls) == 1 + assert torch.equal(entity_provider.calls[0][1], full_env_ids) + assert full.entities["drawer"].pose[:, 0, 3].tolist() == [0.0, 1.0, 2.0] + assert subset.entities["drawer"].pose[:, 0, 3].tolist() == [2.0, 0.0] + joint = subset.articulation_joints[("drawer", "slide")] + assert joint.position[:, 0].tolist() == [2.0, 0.0] + assert joint.valid_mask is not None and joint.valid_mask.tolist() == [True, True] + assert subset.collision_world_revision == (0, 0) + + +def test_shared_tick_scene_provider_captures_full_batch_when_subset_arrives_first() -> ( + None +): + """A partial first consumer cannot poison the delegate's stable batch.""" + provider, entity_provider, joint_provider = _shared_scene_provider() + requested = torch.tensor((1,), dtype=torch.long) + + first = provider.snapshot(timestamp=0.0, env_ids=requested) + second = provider.snapshot( + timestamp=0.0, + env_ids=torch.tensor((2, 1), dtype=torch.long), + ) + + expected_full = torch.arange(_BATCH_SIZE, dtype=torch.long) + assert torch.equal(entity_provider.calls[0][1], expected_full) + assert torch.equal(joint_provider.calls[0][1], expected_full) + assert len(entity_provider.calls) == 1 + assert first.entities["drawer"].pose[:, 0, 3].tolist() == [1.0] + assert second.entities["drawer"].pose[:, 0, 3].tolist() == [2.0, 1.0] + + +def test_shared_tick_scene_provider_rejects_unknown_or_regressing_rows() -> None: + """Unknown correlations and time regressions fail before native sampling.""" + provider, entity_provider, _ = _shared_scene_provider() + provider.snapshot( + timestamp=0.5, + env_ids=torch.tensor((0, 2), dtype=torch.long), + ) + + with pytest.raises(ValueError, match="absent from full_env_ids"): + provider.snapshot( + timestamp=0.5, + env_ids=torch.tensor((3,), dtype=torch.long), + ) + with pytest.raises(ValueError, match="monotonic"): + provider.snapshot( + timestamp=0.4, + env_ids=torch.tensor((0,), dtype=torch.long), + ) + + assert len(entity_provider.calls) == 1 + + +def test_command_state_tracker_correlates_open_and_grasp_across_subsets() -> None: + """Stable IDs, not subset row positions, own accepted gripper state.""" + tracker = _command_state_tracker() + + tracker.accepted( + _hand_command_frame( + env_ids=(3, 7), + positions=(_HAND_GRASP_POSITION, _HAND_OPEN_POSITION), + ) + ) + observation = _hand_state_observation(tracker, 7, 11, 3) + + assert tracker.tracked_control_parts == ("hand",) + assert observation.values.tolist() == [False, False, True] + assert observation.valid is not None + assert observation.valid.tolist() == [True, False, True] + assert observation.acquisition_errors[0] is None + assert observation.acquisition_errors[1] is not None + assert observation.acquisition_errors[2] is None + + +def test_command_state_tracker_preserves_intermediate_and_inactive_rows() -> None: + """Unrecognized targets and inactive rows cannot overwrite prior evidence.""" + tracker = _command_state_tracker() + tracker.accepted( + _hand_command_frame( + env_ids=(7, 3), + positions=(_HAND_OPEN_POSITION, _HAND_GRASP_POSITION), + ) + ) + + tracker.accepted( + _hand_command_frame( + env_ids=(3, 7), + positions=( + _HAND_INTERMEDIATE_POSITION, + _HAND_INTERMEDIATE_POSITION, + ), + ) + ) + after_intermediate = _hand_state_observation(tracker, 3, 7) + assert after_intermediate.values.tolist() == [True, False] + assert after_intermediate.valid is not None + assert after_intermediate.valid.tolist() == [True, True] + + tracker.accepted( + _hand_command_frame( + env_ids=(3, 7), + positions=(_HAND_OPEN_POSITION, _HAND_GRASP_POSITION), + active=(False, True), + ) + ) + after_inactive_row = _hand_state_observation(tracker, 3, 7) + assert after_inactive_row.values.tolist() == [True, True] + assert after_inactive_row.valid is not None + assert after_inactive_row.valid.tolist() == [True, True] + + +def test_command_state_tracker_cancel_invalidates_target_state() -> None: + """Cancelling a hand destination invalidates every correlated hand row.""" + tracker = _command_state_tracker() + frame = _hand_command_frame( + env_ids=(7, 3), + positions=(_HAND_OPEN_POSITION, _HAND_GRASP_POSITION), + ) + tracker.accepted(frame) + + tracker.cancelled(frame.targets) + observation = _hand_state_observation(tracker, 3, 7) + + assert observation.values.tolist() == [False, False] + assert observation.valid is not None + assert observation.valid.tolist() == [False, False] + assert all(error is not None for error in observation.acquisition_errors) + + +def test_command_state_tracker_discard_invalidates_all_state() -> None: + """A fail-closed sink discard removes every accepted row state.""" + tracker = _command_state_tracker() + tracker.accepted( + _hand_command_frame( + env_ids=(11, 3), + positions=(_HAND_GRASP_POSITION, _HAND_OPEN_POSITION), + ) + ) + + tracker.discarded() + observation = _hand_state_observation(tracker, 11, 3) + + assert observation.values.tolist() == [False, False] + assert observation.valid is not None + assert observation.valid.tolist() == [False, False] + + +def test_command_state_tracker_rejects_unknown_environment_ids() -> None: + """Unknown correlation IDs fail before tracker state can be mutated or read.""" + tracker = _command_state_tracker() + + with pytest.raises(ValueError, match="absent from tracker env_ids"): + tracker.accepted( + _hand_command_frame( + env_ids=(99,), + positions=(_HAND_GRASP_POSITION,), + ) + ) + with pytest.raises(ValueError, match="absent from tracker env_ids"): + _hand_state_observation(tracker, 99) + + observation = _hand_state_observation(tracker, 7, 3, 11) + assert observation.valid is not None + assert observation.valid.tolist() == [False, False, False] + + +class _Robot: + """Minimal typed robot surface used by the production factory.""" + + uid = "robot" + device = torch.device("cpu") + dof = _ROBOT_DOF + control_parts = {"arm": ("joint_0",)} + + def __init__(self) -> None: + self.qpos = torch.zeros(_BATCH_SIZE, _ROBOT_DOF) + + def get_qpos( + self, + name: str | None = None, + target: bool = False, + ) -> torch.Tensor: + """Return full or control-part positions.""" + del target + return self.qpos if name is None else self.qpos[:, :1] + + def get_qvel( + self, + name: str | None = None, + target: bool = False, + ) -> torch.Tensor: + """Return zero measured velocities.""" + return torch.zeros_like(self.get_qpos(name=name, target=target)) + + def get_qf(self, name: str | None = None) -> torch.Tensor: + """Return zero measured effort.""" + return torch.zeros_like(self.get_qpos(name=name)) + + def get_joint_ids(self, name: str) -> list[int]: + """Resolve the only declared control part.""" + if name != "arm": + raise KeyError(name) + return [0] + + def get_solver(self, name: str) -> object: + """Return a configured solver marker for Cartesian capability.""" + if name != "arm": + raise KeyError(name) + return object() + + def compute_fk( + self, + qpos: torch.Tensor, + name: str | None = None, + env_ids: list[int] | None = None, + to_matrix: bool = False, + ) -> torch.Tensor: + """Return identity endpoint poses for evidence adapter validation.""" + del name, env_ids + if not to_matrix: + raise ValueError("Tests require matrix FK output.") + return torch.eye(4).repeat(qpos.shape[0], 1, 1) + + +class _EvidenceRobot(_Robot): + """Joint-backed arm and hand with a mutable measured endpoint pose.""" + + control_parts = { + "arm": ("joint_0",), + "hand": ("joint_1",), + } + + def __init__(self) -> None: + super().__init__() + self.endpoint_pose = torch.eye(4).repeat(_BATCH_SIZE, 1, 1) + + def get_qpos( + self, + name: str | None = None, + target: bool = False, + ) -> torch.Tensor: + """Return the full state or the selected control-part state.""" + del target + if name is None: + return self.qpos + joint_id = self.get_joint_ids(name)[0] + return self.qpos[:, joint_id : joint_id + 1] + + def get_joint_ids(self, name: str) -> list[int]: + """Resolve the disjoint arm and hand joints.""" + if name == "arm": + return [0] + if name == "hand": + return [1] + raise KeyError(name) + + def get_solver(self, name: str) -> object: + """Return the configured arm solver marker.""" + if name != "arm": + raise KeyError(name) + return object() + + def compute_fk( + self, + qpos: torch.Tensor, + name: str | None = None, + env_ids: list[int] | None = None, + to_matrix: bool = False, + ) -> torch.Tensor: + """Return the live arm endpoint pose for requested simulator rows.""" + del qpos + if name != "arm" or not to_matrix: + raise ValueError("Evidence FK requires the arm matrix pose.") + rows = list(range(_BATCH_SIZE)) if env_ids is None else env_ids + return self.endpoint_pose[rows].clone() + + +class _DualRobot(_Robot): + """Four-part dual-arm robot used for provider-aware helper preflight.""" + + uid = "dual_robot" + dof = _DUAL_ROBOT_DOF + control_parts = { + "left_arm": ("left_arm_joint",), + "left_hand": ("left_hand_joint",), + "right_arm": ("right_arm_joint",), + "right_hand": ("right_hand_joint",), + } + _joint_ids = { + "left_arm": (0,), + "left_hand": (1,), + "right_arm": (2,), + "right_hand": (3,), + } + + def __init__(self) -> None: + self.qpos = torch.zeros(_BATCH_SIZE, self.dof) + + def get_qpos( + self, + name: str | None = None, + target: bool = False, + ) -> torch.Tensor: + """Return full state or the selected one-joint control part.""" + del target + if name is None: + return self.qpos + return self.qpos[:, list(self._joint_ids[name])] + + def get_joint_ids(self, name: str) -> list[int]: + """Resolve one disjoint arm or hand joint.""" + return list(self._joint_ids[name]) + + def get_solver(self, name: str) -> object: + """Return configured solver markers for both motion endpoints.""" + if name not in {"left_arm", "right_arm"}: + raise KeyError(name) + return object() + + def compute_fk( + self, + qpos: torch.Tensor, + name: str | None = None, + env_ids: list[int] | None = None, + to_matrix: bool = False, + ) -> torch.Tensor: + """Return identity arm endpoint poses for runtime assembly checks.""" + del env_ids + if name not in {"left_arm", "right_arm"} or not to_matrix: + raise ValueError("Dual-arm evidence requires an arm matrix pose.") + return torch.eye(4).repeat(qpos.shape[0], 1, 1) + + +class _RigidObject: + """Mutable batched rigid object with the mesh surface required by binding.""" + + def __init__(self) -> None: + self.pose = torch.eye(4).repeat(_BATCH_SIZE, 1, 1) + + def get_local_pose(self, *, to_matrix: bool = False) -> torch.Tensor: + """Return the current measured object pose.""" + if not to_matrix: + raise ValueError("Tests require matrix object poses.") + return self.pose.clone() + + def get_vertices( + self, + *, + env_ids: list[int], + scale: bool = True, + ) -> torch.Tensor: + """Return one minimal triangular mesh per requested row.""" + del scale + vertices = torch.tensor(((0.0, 0.0, 0.0), (0.04, 0.0, 0.0), (0.0, 0.04, 0.0))) + return vertices.unsqueeze(0).repeat(len(env_ids), 1, 1) + + def get_triangles(self, *, env_ids: list[int]) -> torch.Tensor: + """Return one valid triangle per requested row.""" + return ( + torch.tensor(((0, 1, 2),), dtype=torch.long) + .unsqueeze(0) + .repeat(len(env_ids), 1, 1) + ) + + +class _ForwardedRelationGrounder(RelationTargetGrounder): + """Sentinel relation grounder installed only to prove helper forwarding.""" + + capability: ClassVar[str] = "test.place_relation" + affordance_type: ClassVar[type[Affordance]] = Affordance + affordance_revision: ClassVar[str] = "test-v1" + + def ground( + self, + relation: SemanticRelationTarget, + *, + affordance: Affordance, + context: PlanningContext, + ) -> torch.Tensor: + """Return a direct identity target when explicitly exercised.""" + del relation, affordance, context + return torch.eye(4) + + +class _ForwardedHandOverPoseProvider(HandOverPoseProvider): + """Sentinel embodiment provider installed only through the standard helper.""" + + provider_id: ClassVar[str] = "test.handover_pose" + + def __init__(self) -> None: + self.calls = 0 + + def resolve( + self, + call: HandOver, + *, + context: PlanningContext, + bound: BoundSemanticCall, + ) -> HandOverPoseTargets: + """Return owned direct targets without embedding task-side motion code.""" + del call, context, bound + self.calls += 1 + pose = SemanticPose( + position=(0.0, 0.0, 0.5), + quaternion_wxyz=(1.0, 0.0, 0.0, 0.0), + ) + return HandOverPoseTargets( + middle=SemanticObjectTarget(pose=pose), + final=SemanticObjectTarget(pose=pose), + ) + + +@dataclass(frozen=True, slots=True) +class _MobileEndpoint(ResourceEndpoint): + """Non-joint endpoint used by the standard simulation factory test.""" + + controller_id: str + + +@dataclass(frozen=True, slots=True) +class _MobileTarget(RuntimeEndpointTarget): + """Runtime destination for the test mobile controller.""" + + controller_id: str + + @property + def transport_id(self) -> str: + """Return the matching test Gym transport ID.""" + return "test.mobile_velocity" + + @property + def target_id(self) -> str: + """Return the selected controller ID.""" + return self.controller_id + + +class _MobileEndpointAdapter(ResourceEndpointAdapter): + """Resolve a mobile endpoint without consulting robot control parts.""" + + adapter_id: ClassVar[str] = "test.mobile_velocity" + endpoint_type: ClassVar[type[ResourceEndpoint]] = _MobileEndpoint + + def resolve( + self, + endpoint: ResourceEndpoint, + *, + engine: Any, + ) -> EndpointResolution: + """Resolve one exclusive controller claim.""" + del engine + if not isinstance(endpoint, _MobileEndpoint): + raise TypeError("_MobileEndpointAdapter requires _MobileEndpoint.") + return EndpointResolution( + runtime_target=_MobileTarget(endpoint.controller_id), + claim_tokens=frozenset({f"controller:{endpoint.controller_id}"}), + ) + + +class _MobileTransportEncoder: + """Minimal Gym encoder registered for the custom mobile target.""" + + @property + def transport_id(self) -> str: + """Return the custom mobile transport ID.""" + return "test.mobile_velocity" + + def encode( + self, + command: EndpointCommand, + *, + base_action: Any, + active_mask: torch.Tensor, + ) -> Any: + """Preserve the base action in this assembly-only test transport.""" + del command, active_mask + return base_action.clone() + + def hold( + self, + targets: tuple[RuntimeEndpointTarget, ...], + *, + base_action: Any, + context: Any, + ) -> Any: + """Preserve the base action for a mobile safe hold.""" + del targets, context + return base_action.clone() + + +class _MobileRobot: + """Full-state robot fixture with no control-parts or joint-ID surface.""" + + uid = "mobile_robot" + device = torch.device("cpu") + dof = _ROBOT_DOF + + def __init__(self) -> None: + self.qpos = torch.zeros(_BATCH_SIZE, _ROBOT_DOF) + + def get_qpos(self) -> torch.Tensor: + """Return the full controller hold state.""" + return self.qpos + + def get_qvel(self) -> torch.Tensor: + """Return the full measured velocity state.""" + return torch.zeros_like(self.qpos) + + def get_qf(self) -> torch.Tensor: + """Return the full measured effort state.""" + return torch.zeros_like(self.qpos) + + +class _Simulation: + """Minimal simulation registry for one exact robot.""" + + def __init__( + self, + robot: _Robot, + rigid_objects: dict[str, _RigidObject] | None = None, + ) -> None: + self.robot = robot + self.rigid_objects = {} if rigid_objects is None else dict(rigid_objects) + + def get_robot(self, uid: str) -> _Robot | None: + """Resolve the selected robot UID.""" + return self.robot if uid == self.robot.uid else None + + def get_rigid_object(self, uid: str) -> _RigidObject | None: + """Resolve one explicitly registered rigid-object UID.""" + return self.rigid_objects.get(uid) + + +def _profile_binding() -> SimulationRobotSkillProfileBinding: + """Build one motion-only profile with an intentionally wrong cadence.""" + return SimulationRobotSkillProfileBinding( + profile_id="robot_profile", + resources=( + ControlPartResourceBinding( + resource_id="manipulator", + endpoints=( + ControlPartEndpointBinding( + endpoint_id="motion", + control_part="arm", + capabilities=frozenset({CARTESIAN_POSE_CAPABILITY}), + ), + ), + ), + ), + presets=( + SkillPolicyPreset( + "safe", + motion_policy=MotionPolicy(control_dt=0.01), + ), + ), + default_preset="safe", + ) + + +def _handover_profile_binding() -> SimulationRobotSkillProfileBinding: + """Declare two disjoint manipulators and one selected pose provider ID.""" + motion_capabilities = frozenset( + { + CARTESIAN_POSE_CAPABILITY, + FORWARD_KINEMATICS_CAPABILITY, + } + ) + resources = tuple( + ControlPartResourceBinding( + resource_id=side, + endpoints=( + ControlPartEndpointBinding( + endpoint_id="motion", + control_part=f"{side}_arm", + capabilities=motion_capabilities, + ), + ControlPartEndpointBinding( + endpoint_id="grasp", + control_part=f"{side}_hand", + capabilities=frozenset({GRASP_CAPABILITY}), + command_preset=f"{side}_hand_commands", + ), + ), + ) + for side in ("left", "right") + ) + command_presets = tuple( + ControlPartCommandPreset( + preset_id=f"{side}_hand_commands", + control_part=f"{side}_hand", + commands={ + "open": (_HAND_OPEN_POSITION,), + "grasp": (_HAND_GRASP_POSITION,), + }, + ) + for side in ("left", "right") + ) + return SimulationRobotSkillProfileBinding( + profile_id="handover_profile", + resources=resources, + command_presets=command_presets, + defaults={ + "hand_over": {"source": "left", "destination": "right"}, + }, + presets=(SkillPolicyPreset("safe"),), + default_preset="safe", + grounding_providers={ + "hand_over": _ForwardedHandOverPoseProvider.provider_id, + }, + ) + + +def _handover_helper_inputs() -> tuple[ + SimpleNamespace, + SimulationSceneBinding, + SimulationRobotSkillProfileBinding, +]: + """Build standard-helper inputs for one provider-aware HandOver program.""" + robot = _DualRobot() + cube = _RigidObject() + simulation = _Simulation(robot, {"cube_native": cube}) + environment = SimpleNamespace( + sim=simulation, + robot=robot, + step_dt=_STEP_DT, + ) + scene_binding = SimulationSceneBinding( + registry_id="handover_scene", + rigid_objects=( + SimulationRigidObjectBinding( + entity_id="cube", + simulation_uid="cube_native", + default_grasp_affordance="cube_grasp", + ), + ), + antipodal_grasps=( + AntipodalGraspAffordanceBinding( + entity_id="cube_grasp", + object_id="cube", + native_name="body", + revision="1", + ), + ), + ) + return environment, scene_binding, _handover_profile_binding() + + +def _handover_program() -> ExpertProgramCfg: + """Build one external-held-state HandOver call for static preflight.""" + return ExpertProgramCfg( + schema_version=1, + program_id="handover_preflight", + integration=ExpertProgramIntegrationCfg( + robot_profile="handover_profile", + scene_registry="handover_scene", + runtime_preset="safe", + ), + program=InvokeCfg(call=HandOverCfg(object="cube")), + ) + + +def _motion_generator(robot: _Robot) -> MotionGenerator: + """Build a type-checkable motion-generator test double.""" + generator = MagicMock(spec=MotionGenerator) + generator.robot = robot + generator.device = robot.device + generator.planner = SimpleNamespace(cfg=SimpleNamespace(planner_type="test")) + generator.dynamic_collision_entity_ids = () + generator.collision_world_entity_ids = () + generator.supports_dynamic_collision_world = False + generator.collision_world_batch_mode = None + return generator + + +def _factory() -> tuple[SimulationExpertProgramFactory, _Robot]: + """Create one production factory around CPU-only test doubles.""" + robot = _Robot() + simulation = _Simulation(robot) + return ( + SimulationExpertProgramFactory( + simulation, # type: ignore[arg-type] + robot, # type: ignore[arg-type] + SimulationSceneBinding(registry_id="scene"), + _profile_binding(), + step_dt=_STEP_DT, + motion_generator_factory=lambda: _motion_generator(robot), + ), + robot, + ) + + +def _evidence_profile_binding() -> SimulationRobotSkillProfileBinding: + """Declare one manipulation resource with exact open/grasp semantics.""" + motion_capabilities = frozenset( + { + BATCH_INVERSE_KINEMATICS_CAPABILITY, + CARTESIAN_POSE_CAPABILITY, + FORWARD_KINEMATICS_CAPABILITY, + } + ) + return SimulationRobotSkillProfileBinding( + profile_id="evidence_profile", + resources=( + ControlPartResourceBinding( + resource_id="manipulator", + endpoints=( + ControlPartEndpointBinding( + endpoint_id="motion", + control_part="arm", + capabilities=motion_capabilities, + ), + ControlPartEndpointBinding( + endpoint_id="grasp", + control_part="hand", + capabilities=frozenset({GRASP_CAPABILITY}), + command_preset="hand_commands", + ), + ), + ), + ), + command_presets=( + ControlPartCommandPreset( + preset_id="hand_commands", + control_part="hand", + commands={ + "open": (_HAND_OPEN_POSITION,), + "grasp": (_HAND_GRASP_POSITION,), + }, + ), + ), + defaults={ + "pick_up": {"primary": "manipulator"}, + "place": {"primary": "manipulator"}, + }, + presets=(SkillPolicyPreset("evidence"),), + default_preset="evidence", + ) + + +def _pick_evidence_plan(action: Any, request: Any, context: Any) -> Any: + """Build one grasp frame and an identity object-to-endpoint expectation.""" + goal = action.require_goal(request) + trajectory = context.robot.qpos.unsqueeze(1).clone() + trajectory[:, 0, 1] = _HAND_GRASP_POSITION + relation = torch.eye(4).repeat(context.batch_size, 1, 1) + held = HeldObjectState( + semantics=goal.semantics, + object_to_eef=relation, + grasp_xpos=relation, + ) + return action.build_plan( + request, + context, + success=torch.ones(context.batch_size, dtype=torch.bool), + trajectory=trajectory, + expected_effects=StateDelta( + held_object_updates={"manipulator": held}, + ), + replannable=False, + scene_dependency_monitor_until={"cube": 0}, + ) + + +def _place_evidence_plan(action: Any, request: Any, context: Any) -> Any: + """Build one open frame and the matching held-object removal delta.""" + trajectory = context.robot.qpos.unsqueeze(1).clone() + trajectory[:, 0, 1] = _HAND_OPEN_POSITION + return action.build_plan( + request, + context, + success=torch.ones(context.batch_size, dtype=torch.bool), + trajectory=trajectory, + expected_effects=StateDelta( + held_object_updates={"manipulator": None}, + ), + replannable=False, + ) + + +def _evidence_runtime() -> tuple[ + ExpertProgramRuntimeAssembly, + _EvidenceRobot, + _RigidObject, +]: + """Assemble the production Pick/Place evidence chain on CPU fixtures.""" + robot = _EvidenceRobot() + cube = _RigidObject() + simulation = _Simulation(robot, {"cube_native": cube}) + scene_binding = SimulationSceneBinding( + registry_id="evidence_scene", + rigid_objects=( + SimulationRigidObjectBinding( + entity_id="cube", + simulation_uid="cube_native", + default_grasp_affordance="cube_grasp", + ), + ), + antipodal_grasps=( + AntipodalGraspAffordanceBinding( + entity_id="cube_grasp", + object_id="cube", + native_name="body", + revision="1", + ), + ), + ) + factory = SimulationExpertProgramFactory( + simulation, # type: ignore[arg-type] + robot, # type: ignore[arg-type] + scene_binding, + _evidence_profile_binding(), + step_dt=_STEP_DT, + motion_generator_factory=lambda: _motion_generator(robot), + ) + adapter = factory.create_adapter( + runner_cfg=ExecutionRunnerCfg( + minimum_cycle_time=0.0, + hold_on_completion=False, + ) + ) + assembly = adapter.assemble_runtime( + ExpertProgramIntegrationCfg( + robot_profile="evidence_profile", + scene_registry="evidence_scene", + runtime_preset="evidence", + ) + ) + pick_action = assembly.engine.actions["pick_up"] + place_action = assembly.engine.actions["place"] + pick_action._plan = MethodType(_pick_evidence_plan, pick_action) + place_action._plan = MethodType(_place_evidence_plan, place_action) + return assembly, robot, cube + + +def _consume_buffered_action( + assembly: ExpertProgramRuntimeAssembly, + robot: _EvidenceRobot, +) -> None: + """Apply one accepted Gym action and advance the authoritative clock.""" + processed = assembly.command_sink.pop() + if not isinstance(processed.value, torch.Tensor): + raise TypeError("Joint-backed evidence actions must be tensors.") + robot.qpos = processed.value.clone() + assembly.clock.advance_after_env_step() + + +def _accept_hand_command( + assembly: ExpertProgramRuntimeAssembly, + robot: _EvidenceRobot, + position: float, +) -> None: + """Accept and consume one semantic hand command through the Gym sink.""" + assert assembly.command_sink.pending_count == 0 + frame = _hand_command_frame( + env_ids=tuple(range(_BATCH_SIZE)), + positions=(position,) * _BATCH_SIZE, + ) + acknowledgement = assembly.command_sink.send(frame, timeout=1.0) + assert acknowledgement.accepted + _consume_buffered_action(assembly, robot) + + +def _sample_effect( + assembly: ExpertProgramRuntimeAssembly, + robot: _EvidenceRobot, + *, + expected_trace_count: int, + advance_clock: bool = True, +) -> tuple[Any, SkillEffectTrace]: + """Advance one fresh environment tick and return its production trace.""" + if advance_clock: + assembly.clock.advance_after_env_step() + result = assembly.runtime.step() + assert len(result.effects) == expected_trace_count + while assembly.command_sink.pending_count: + _consume_buffered_action(assembly, robot) + return result, result.effects[-1] + + +class _SynchronousEvidenceClock: + """Advance the fixture's simulation clock during standalone facade waits.""" + + def __init__(self, clock: EnvironmentStepClock) -> None: + self._clock = clock + + def now(self) -> float: + """Return the simulation fixture's authoritative time.""" + return self._clock.now() + + def sleep(self, duration: float) -> None: + """Advance the exact number of fixture ticks requested by the runner.""" + steps = self._clock.steps_for_duration(duration) + if steps: + self._clock.advance_after_env_step(steps) + + +class _ImmediateEvidenceCommandSink: + """Apply accepted endpoint frames immediately for standalone CPU execution.""" + + def __init__( + self, + assembly: ExpertProgramRuntimeAssembly, + robot: _EvidenceRobot, + cube: _RigidObject, + ) -> None: + self._encoder = assembly.command_encoder + self._observer = assembly.accepted_command_observer + self._clock = assembly.clock + self._robot = robot + self._cube = cube + + def send( + self, + command: RuntimeCommandFrame, + *, + timeout: float, + ) -> CommandAcknowledgement: + """Apply one command and publish its accepted semantic hand state.""" + assert timeout > 0.0 + action = self._encoder.encode(command) + if not isinstance(action, torch.Tensor): + raise TypeError("The CPU quickstart fixture requires tensor actions.") + self._robot.qpos = action.clone() + if self._observer is None: + raise RuntimeError("The evidence fixture requires an accepted observer.") + self._observer.accepted(command.snapshot()) + if torch.allclose( + action[:, 1], + torch.full_like(action[:, 1], _HAND_OPEN_POSITION), + ): + self._cube.pose[:, 0, 3] = _RELEASE_SEPARATION + self._clock.advance_after_env_step() + return CommandAcknowledgement.accepted_ack() + + def hold( + self, + targets: tuple[RuntimeEndpointTarget, ...], + context: PlanningContext, + *, + timeout: float, + ) -> CommandAcknowledgement: + """Apply the encoder's observed-position hold immediately.""" + assert timeout > 0.0 + action = self._encoder.encode_hold(targets, context) + if not isinstance(action, torch.Tensor): + raise TypeError("The CPU quickstart fixture requires tensor actions.") + self._robot.qpos = action.clone() + return CommandAcknowledgement.accepted_ack() + + def cancel( + self, + targets: tuple[RuntimeEndpointTarget, ...], + *, + timeout: float, + ) -> CommandAcknowledgement: + """Clear accepted command evidence for cancelled destinations.""" + assert timeout > 0.0 + if self._observer is not None: + self._observer.cancelled(targets) + return CommandAcknowledgement.accepted_ack() + + +class _QuickstartRuntimeProvider: + """Explicit provider used by the public ``AtomicSkills.from_env`` path.""" + + def __init__(self, runtime: SkillRuntime) -> None: + self._runtime = runtime + self.presets: list[str] = [] + + def create_skill_runtime(self, *, preset: str) -> SkillRuntime: + """Return the configured canonical runtime and record preset selection.""" + self.presets.append(preset) + return self._runtime + + +def _quickstart_runtime_provider() -> _QuickstartRuntimeProvider: + """Build a synchronous provider from the shared production CPU fixture.""" + assembly, robot, cube = _evidence_runtime() + runtime = SkillRuntime.from_components( + assembly.compiler, + assembly.observation_provider, + _ImmediateEvidenceCommandSink(assembly, robot, cube), + assembly.evidence_collector, + clock=_SynchronousEvidenceClock(assembly.clock), + runner_cfg=ExecutionRunnerCfg( + minimum_cycle_time=0.0, + hold_on_completion=False, + ), + ) + return _QuickstartRuntimeProvider(runtime) + + +def _documented_pick_place_quickstart( + runtime_provider: _QuickstartRuntimeProvider, +) -> SkillResult: + """Run the application-facing quickstart, excluding scene construction.""" + skills = AtomicSkills.from_env(runtime_provider, preset="evidence") + cube = skills.scene.object("cube") + return skills.run( + Pick(object=cube), + Place(object=cube, at=_DIRECT_PLACE_TARGET), + ) + + +def _python_pick_place_calls() -> tuple[SemanticCallSpec, ...]: + """Return the application-facing calls used by both acceptance paths.""" + cube = SceneObjectRef("cube") + return ( + Pick(object=cube), + Place(object=cube, at=_DIRECT_PLACE_TARGET), + ) + + +def _decoded_pick_place_calls( + registry: SceneRegistry, +) -> tuple[SemanticCallSpec, ...]: + """Decode and provider-free compile the config equivalent of Python calls.""" + config = decode_expert_program( + { + "schema_version": 1, + "program_id": "pick_place_equivalence", + "integration": { + "robot_profile": "evidence_profile", + "scene_registry": "evidence_scene", + "runtime_preset": "evidence", + }, + "targets": { + "place_target": { + "kind": "cyclic_pose", + "values": [ + { + "position": _DIRECT_PLACE_TARGET.position.tolist(), + "quaternion_wxyz": ( + _DIRECT_PLACE_TARGET.quaternion_wxyz.tolist() + ), + } + ], + } + }, + "program": { + "kind": "sequence", + "items": [ + { + "kind": "invoke", + "call": {"kind": "pick", "object": "cube"}, + }, + { + "kind": "invoke", + "call": { + "kind": "place", + "object": "cube", + "at": { + "kind": "target_ref", + "target": "place_target", + }, + }, + }, + ], + }, + } + ) + program = ExpertProgramCompiler.from_scene_registry(registry).compile(config) + return tuple( + compiled_call.call for segment in program for compiled_call in segment.calls + ) + + +def _capture_grounded_invocations( + monkeypatch: pytest.MonkeyPatch, + assembly: ExpertProgramRuntimeAssembly, +) -> list[ActionInvocation[Any, Any]]: + """Record the production compiler's final lowering without replacing it.""" + invocations: list[ActionInvocation[Any, Any]] = [] + ground = assembly.compiler.ground + + def recording_ground(*args: Any, **kwargs: Any) -> Any: + grounded = ground(*args, **kwargs) + invocations.append(grounded.invocation) + return grounded + + monkeypatch.setattr(assembly.compiler, "ground", recording_ground) + return invocations + + +def _run_evidence_pick_place( + assembly: ExpertProgramRuntimeAssembly, + robot: _EvidenceRobot, + cube: _RigidObject, + calls: tuple[SemanticCallSpec, ...], +) -> tuple[SkillResult, HeldObjectState]: + """Drive one happy-path workflow through accepted commands and live evidence.""" + result = assembly.runtime.start(calls, workflow_id="pick_place_equivalence") + verified_pick: HeldObjectState | None = None + for _ in range(32): + while assembly.command_sink.pending_count: + _consume_buffered_action(assembly, robot) + if result.terminal: + break + if result.current_call_index == 1: + if verified_pick is None: + verified_pick = result.task_state.get_held_object("manipulator") + cube.pose[:, 0, 3] = _RELEASE_SEPARATION + assembly.clock.advance_after_env_step() + result = assembly.runtime.step() + + assert result.status is SkillStatus.COMPLETED + assert verified_pick is not None + assert result.task_state.get_held_object("manipulator") is None + return result, verified_pick + + +def _assert_typed_equivalent( + actual: object, + expected: object, + *, + path: str = "value", +) -> None: + """Compare nested typed compiler output, including owned tensor values.""" + assert type(actual) is type(expected), path + if isinstance(actual, torch.Tensor): + assert isinstance(expected, torch.Tensor) + torch.testing.assert_close(actual, expected) + return + if isinstance(actual, Mapping): + assert isinstance(expected, Mapping) + assert tuple(actual) == tuple(expected) + for key in actual: + _assert_typed_equivalent( + actual[key], + expected[key], + path=f"{path}[{key!r}]", + ) + return + if isinstance(actual, Sequence) and not isinstance(actual, (str, bytes)): + assert isinstance(expected, Sequence) + assert len(actual) == len(expected) + for index, (actual_item, expected_item) in enumerate( + zip(actual, expected, strict=True) + ): + _assert_typed_equivalent( + actual_item, + expected_item, + path=f"{path}[{index}]", + ) + return + if is_dataclass(actual) and not isinstance(actual, type): + assert is_dataclass(expected) and not isinstance(expected, type) + for data_field in fields(actual): + _assert_typed_equivalent( + getattr(actual, data_field.name), + getattr(expected, data_field.name), + path=f"{path}.{data_field.name}", + ) + return + assert actual == expected, path + + +def _assert_invocation_equivalent( + actual: ActionInvocation[Any, Any], + expected: ActionInvocation[Any, Any], +) -> None: + """Compare semantic lowering while ignoring engine-instance owner UUIDs.""" + assert actual.skill_id == expected.skill_id + assert actual.invocation_id == expected.invocation_id + assert actual.revision == expected.revision + _assert_typed_equivalent(actual.goal, expected.goal, path="invocation.goal") + _assert_typed_equivalent( + actual.binding.endpoints, + expected.binding.endpoints, + path="invocation.binding.endpoints", + ) + _assert_typed_equivalent( + actual.motion_policy, + expected.motion_policy, + path="invocation.motion_policy", + ) + _assert_typed_equivalent( + actual.recovery_policy, + expected.recovery_policy, + path="invocation.recovery_policy", + ) + _assert_typed_equivalent( + actual.skill_options, + expected.skill_options, + path="invocation.skill_options", + ) + _assert_typed_equivalent( + actual.control_overrides, + expected.control_overrides, + path="invocation.control_overrides", + ) + + +def test_simulation_factory_aligns_every_motion_policy_to_gym_step() -> None: + """The environment cadence replaces unrelated preset fallback timing.""" + factory, _ = _factory() + + profile = factory.create_robot_skill_profile() + + assert profile.presets["safe"].motion_policy.control_dt == pytest.approx(_STEP_DT) + + +def test_decoded_program_and_python_calls_share_invocations_and_results( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Both frontends reach equivalent core invocations and verified state.""" + python_assembly, python_robot, python_cube = _evidence_runtime() + config_assembly, config_robot, config_cube = _evidence_runtime() + python_invocations = _capture_grounded_invocations(monkeypatch, python_assembly) + config_invocations = _capture_grounded_invocations(monkeypatch, config_assembly) + + python_result, python_held = _run_evidence_pick_place( + python_assembly, + python_robot, + python_cube, + _python_pick_place_calls(), + ) + config_result, config_held = _run_evidence_pick_place( + config_assembly, + config_robot, + config_cube, + _decoded_pick_place_calls(config_assembly.scene_registry), + ) + + assert len(python_invocations) == len(config_invocations) == 2 + for python_invocation, config_invocation in zip( + python_invocations, + config_invocations, + strict=True, + ): + _assert_invocation_equivalent(python_invocation, config_invocation) + _assert_typed_equivalent(python_held, config_held) + _assert_typed_equivalent(python_result, config_result) + + +def test_atomic_skills_from_env_runs_documented_pick_place_quickstart() -> None: + """The small public facade executes without exposing core motion plumbing.""" + provider = _quickstart_runtime_provider() + + result = _documented_pick_place_quickstart(provider) + + source = textwrap.dedent(inspect.getsource(_documented_pick_place_quickstart)) + function = ast.parse(source).body[0] + assert isinstance(function, ast.FunctionDef) + executable = function.body[1:] # Exclude the helper's docstring. + assert executable[-1].end_lineno is not None + assert executable[-1].end_lineno - executable[0].lineno + 1 <= ( + _QUICKSTART_MAX_LINES + ) + identifiers = { + identifier + for node in ast.walk(function) + for identifier in ( + node.id if isinstance(node, ast.Name) else None, + node.attr if isinstance(node, ast.Attribute) else None, + ) + if identifier is not None + } + assert identifiers.isdisjoint( + { + "qpos", + "matrix", + "planner", + "session", + "MotionGenerator", + "PlanningContext", + "ExecutionSession", + } + ) + assert provider.presets == ["evidence"] + assert result.status is SkillStatus.COMPLETED + assert result.success_mask.tolist() == [True] * _BATCH_SIZE + assert [call.semantic_id for call in result.calls] == ["pick", "place"] + assert result.task_state.get_held_object("manipulator") is None + + +def test_simulation_factory_builds_shared_observation_and_evidence_ports() -> None: + """Observation and both built-in evidence providers share one scene source.""" + factory, robot = _factory() + registry = factory.create_scene_registry() + profile = factory.create_robot_skill_profile() + engine = factory.create_atomic_action_engine(profile) + clock = EnvironmentStepClock(_STEP_DT) + + observation = factory.create_planning_observation_provider( + scene_registry=registry, + engine=engine, + clock=clock, + ) + assert type(observation) is SimulationPlanningObservationProvider + context = observation.observe(TaskState.empty(_BATCH_SIZE, robot.device)) + providers = tuple( + factory.create_effect_evidence_providers( + scene_registry=registry, + engine=engine, + observation_provider=observation, + ) + ) + accepted_command_observer = factory.create_accepted_runtime_command_observer( + scene_registry=registry, + engine=engine, + observation_provider=observation, + ) + + assert context.robot.timestamp == pytest.approx(0.0) + assert torch.equal(observation.current_qpos(context.env_ids), robot.qpos) + assert accepted_command_observer is observation.command_state_tracker + assert len(providers) == 2 + assert all( + getattr(provider, "_scene_provider") is observation.scene_provider + for provider in providers + ) + + +def test_simulation_factory_returns_exact_environment_adapter() -> None: + """The convenience path remains compatible with the exact-type mixin check.""" + factory, _ = _factory() + + adapter = factory.create_adapter() + + assert type(adapter) is ExpertProgramEnvironmentAdapter + assert adapter.step_dt == pytest.approx(_STEP_DT) + assert factory.segment_policy_port is not None + + +def test_simulation_helper_forwards_semantic_grounding_extensions() -> None: + """Both explicit grounding seams reach the runtime compiler unchanged.""" + robot = _Robot() + environment = SimpleNamespace( + sim=_Simulation(robot), + robot=robot, + step_dt=_STEP_DT, + ) + relation_grounder = _ForwardedRelationGrounder() + handover_provider = _ForwardedHandOverPoseProvider() + adapter = create_simulation_expert_program_adapter( + environment, # type: ignore[arg-type] + scene_binding=SimulationSceneBinding(registry_id="scene"), + robot_profile_binding=_profile_binding(), + motion_generator_factory=lambda: _motion_generator(robot), + relation_grounders=(relation_grounder,), + handover_pose_providers=(handover_provider,), + ) + + assembly = adapter.assemble_runtime( + ExpertProgramIntegrationCfg( + robot_profile="robot_profile", + scene_registry="scene", + runtime_preset="safe", + ) + ) + + assert tuple(assembly.compiler.relation_grounders.values()) == (relation_grounder,) + assert tuple(assembly.compiler.handover_pose_providers.values()) == ( + handover_provider, + ) + + +def test_simulation_helper_handover_preflight_is_fail_closed_by_default() -> None: + """Selecting a provider ID does not infer or auto-install an implementation.""" + environment, scene_binding, profile_binding = _handover_helper_inputs() + robot = environment.robot + adapter = create_simulation_expert_program_adapter( + environment, # type: ignore[arg-type] + scene_binding=scene_binding, + robot_profile_binding=profile_binding, + motion_generator_factory=lambda: _motion_generator(robot), + ) + compiled = adapter.compile(_handover_program()) + + with pytest.raises(SemanticValidationError) as error: + adapter.create_bridge(compiled) + + assert error.value.diagnostic.code == "handover_grounding_provider_not_installed" + + +def test_simulation_helper_forwards_handover_provider_to_preflight() -> None: + """An explicitly supplied embodiment provider satisfies standard preflight.""" + environment, scene_binding, profile_binding = _handover_helper_inputs() + robot = environment.robot + provider = _ForwardedHandOverPoseProvider() + adapter = create_simulation_expert_program_adapter( + environment, # type: ignore[arg-type] + scene_binding=scene_binding, + robot_profile_binding=profile_binding, + motion_generator_factory=lambda: _motion_generator(robot), + handover_pose_providers=(provider,), + ) + + bridge = adapter.create_bridge(adapter.compile(_handover_program())) + + assert bridge is not None + assert provider.calls == 0 + + +def test_simulation_helper_assembles_mobile_endpoint_and_transport_without_joints() -> ( + None +): + """The one-line factory path supports a custom non-joint controller.""" + robot = _MobileRobot() + simulation = _Simulation(robot) # type: ignore[arg-type] + profile_binding = SimulationRobotSkillProfileBinding( + profile_id="mobile_profile", + resources=( + RobotResourceBinding( + resource_id="mobile_base", + endpoints={ + "motion": _MobileEndpoint( + controller_id="base_velocity", + capabilities=frozenset({"motion.base.velocity"}), + ) + }, + ), + ), + presets=(SkillPolicyPreset("runtime"),), + default_preset="runtime", + ) + environment = SimpleNamespace( + sim=simulation, + robot=robot, + step_dt=_STEP_DT, + ) + + adapter = create_simulation_expert_program_adapter( + environment, # type: ignore[arg-type] + scene_binding=SimulationSceneBinding(registry_id="mobile_scene"), + robot_profile_binding=profile_binding, + motion_generator_factory=lambda: _motion_generator(robot), # type: ignore[arg-type] + endpoint_adapters={_MobileEndpoint: _MobileEndpointAdapter()}, + runtime_transports=(_MobileTransportEncoder(),), + ) + assembly = adapter.assemble_runtime( + ExpertProgramIntegrationCfg( + robot_profile="mobile_profile", + scene_registry="mobile_scene", + runtime_preset="runtime", + ) + ) + + endpoint = assembly.robot_profile.resources["mobile_base"].endpoints["motion"] + assert isinstance(endpoint, _MobileEndpoint) + assert "test.mobile_velocity" in assembly.command_encoder.transport_ids + assert assembly.engine.skill_profile is not None + resolved = assembly.engine.skill_profile.resources["mobile_base"] + assert isinstance(resolved.endpoints["motion"].runtime_target, _MobileTarget) + assert resolved.claim.claim_tokens == frozenset({"controller:base_velocity"}) + + +def test_pick_place_effects_require_accepted_hand_state_and_live_pose() -> None: + """Production Pick/Place evidence stays conjunctive through runtime traces.""" + assembly, robot, cube = _evidence_runtime() + assert type(assembly.accepted_command_observer) is ( + ControlCommandStateEvidenceTracker + ) + cube.pose[:, 0, 3] = 0.2 + result = assembly.runtime.start( + ( + Pick(object=SceneObjectRef("cube")), + Place( + object=SceneObjectRef("cube"), + at=SemanticPose( + position=(0.0, 0.0, 0.0), + quaternion_wxyz=(1.0, 0.0, 0.0, 0.0), + ), + ), + ), + workflow_id="production_evidence_chain", + ) + assert result.status is SkillStatus.RUNNING + + result = assembly.runtime.step() + assert assembly.command_sink.pending_count == 1 + assert len(result.effects) == 0 + _consume_buffered_action(assembly, robot) + result = assembly.runtime.step() + assert len(result.effects) == 0 + while assembly.command_sink.pending_count: + _consume_buffered_action(assembly, robot) + + result, pick_pose_missing = _sample_effect( + assembly, + robot, + expected_trace_count=1, + ) + pick_pose = pick_pose_missing.evidence["destination.pose"] + pick_constraint = pick_pose_missing.evidence["destination.constraint"] + assert type(pick_pose) is PoseRelationEvidenceBatch + assert type(pick_constraint) is BinaryEffectEvidenceBatch + assert pick_pose.object_to_endpoint[:, 0, 3].tolist() == pytest.approx( + [-0.2] * _BATCH_SIZE + ) + assert pick_constraint.values.tolist() == [True] * _BATCH_SIZE + assert pick_constraint.valid.tolist() == [True] * _BATCH_SIZE + assert not pick_pose_missing.success_mask.any() + + cube.pose = torch.eye(4).repeat(_BATCH_SIZE, 1, 1) + assembly.command_sink.discard_pending() + result, pick_command_missing = _sample_effect( + assembly, + robot, + expected_trace_count=2, + ) + pick_pose = pick_command_missing.evidence["destination.pose"] + pick_constraint = pick_command_missing.evidence["destination.constraint"] + torch.testing.assert_close( + pick_pose.object_to_endpoint, + torch.eye(4).repeat(_BATCH_SIZE, 1, 1), + ) + assert pick_constraint.valid.tolist() == [False] * _BATCH_SIZE + assert not pick_command_missing.success_mask.any() + + _accept_hand_command(assembly, robot, _HAND_GRASP_POSITION) + result, pick_first_complete_sample = _sample_effect( + assembly, + robot, + expected_trace_count=3, + advance_clock=False, + ) + assert not pick_first_complete_sample.success_mask.any() + result, pick_success = _sample_effect( + assembly, + robot, + expected_trace_count=4, + ) + assert pick_success.call_index == 0 + assert pick_success.effect_spec.semantic_id == "pick" + assert pick_success.success_mask.tolist() == [True] * _BATCH_SIZE + assert ( + pick_success.evidence["destination.constraint"].values.tolist() + == [True] * _BATCH_SIZE + ) + assert result.task_state.get_held_object("manipulator") is not None + assert result.current_call_index == 1 + + result = assembly.runtime.step() + assert assembly.command_sink.pending_count == 1 + assert len(result.effects) == 4 + _consume_buffered_action(assembly, robot) + result = assembly.runtime.step() + assert len(result.effects) == 4 + while assembly.command_sink.pending_count: + _consume_buffered_action(assembly, robot) + + result, place_pose_missing = _sample_effect( + assembly, + robot, + expected_trace_count=5, + ) + place_pose = place_pose_missing.evidence["source.pose"] + place_constraint = place_pose_missing.evidence["source.constraint"] + assert type(place_pose) is PoseRelationEvidenceBatch + assert type(place_constraint) is BinaryEffectEvidenceBatch + torch.testing.assert_close( + place_pose.object_to_endpoint, + torch.eye(4).repeat(_BATCH_SIZE, 1, 1), + ) + assert place_constraint.values.tolist() == [False] * _BATCH_SIZE + assert place_constraint.valid.tolist() == [True] * _BATCH_SIZE + assert not place_pose_missing.success_mask.any() + + cube.pose[:, 0, 3] = 0.2 + assembly.command_sink.discard_pending() + result, place_command_missing = _sample_effect( + assembly, + robot, + expected_trace_count=6, + ) + place_pose = place_command_missing.evidence["source.pose"] + place_constraint = place_command_missing.evidence["source.constraint"] + assert place_pose.object_to_endpoint[:, 0, 3].tolist() == pytest.approx( + [-0.2] * _BATCH_SIZE + ) + assert place_constraint.valid.tolist() == [False] * _BATCH_SIZE + assert not place_command_missing.success_mask.any() + + _accept_hand_command(assembly, robot, _HAND_OPEN_POSITION) + result, place_first_complete_sample = _sample_effect( + assembly, + robot, + expected_trace_count=7, + advance_clock=False, + ) + assert not place_first_complete_sample.success_mask.any() + result, place_success = _sample_effect( + assembly, + robot, + expected_trace_count=8, + ) + assert result.status is SkillStatus.COMPLETED + assert place_success.call_index == 1 + assert place_success.effect_spec.semantic_id == "place" + assert place_success.success_mask.tolist() == [True] * _BATCH_SIZE + assert ( + place_success.evidence["source.constraint"].values.tolist() + == [False] * _BATCH_SIZE + ) + assert result.task_state.get_held_object("manipulator") is None + assert [len(call.effects) for call in result.calls] == [4, 4] + assert assembly.command_sink.accepted_action_count >= 4 diff --git a/tests/gym/envs/expert_program/test_simulation_policies.py b/tests/gym/envs/expert_program/test_simulation_policies.py new file mode 100644 index 000000000..51b3bea13 --- /dev/null +++ b/tests/gym/envs/expert_program/test_simulation_policies.py @@ -0,0 +1,451 @@ +# ---------------------------------------------------------------------------- +# 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 explicit simulation-backed segment policies.""" + +from __future__ import annotations + +import json +from types import SimpleNamespace + +import pytest +import torch + +from embodichain.lab.gym.envs.expert_program import ( + ExpertProgramCompiler, + SimulationRigidObjectBinding, + SimulationSceneBinding, + decode_expert_program, +) +from embodichain.lab.gym.envs.expert_program.bridge import ( + SegmentPostPolicyMetadataPort, + SegmentPostPolicyPort, + SegmentPostPolicyResultPort, + SegmentValidatorMetadataPort, + SegmentValidatorPort, +) +from embodichain.lab.gym.envs.expert_program.simulation_policies import ( + SimulationSegmentPolicyPort, +) +from embodichain.lab.gym.envs.settling import DynamicSettleMonitorCfg +from embodichain.lab.sim.atomic_actions import EntityState +from embodichain.lab.sim.skills.scene import ( + SceneEntityRegistration, + SceneObjectRef, + SceneRegistry, +) + + +class _StaticStateProvider: + """Provide an inert object state for provider-free compilation.""" + + def observe( + self, + *, + timestamp: float, + env_ids: torch.Tensor, + ) -> EntityState: + del timestamp + return EntityState( + torch.eye(4, device=env_ids.device).expand(env_ids.numel(), -1, -1) + ) + + +class _RigidObject: + """Small live rigid-object double with mutable velocities and poses.""" + + def __init__(self, positions: torch.Tensor) -> None: + batch_size = positions.shape[0] + self.is_non_dynamic = False + self.pose_reads = 0 + self.body_data = SimpleNamespace( + lin_vel=torch.zeros(batch_size, 3), + ang_vel=torch.zeros(batch_size, 3), + ) + self._pose = torch.eye(4).expand(batch_size, -1, -1).clone() + self._pose[:, :3, 3] = positions + + def get_local_pose(self, *, to_matrix: bool) -> torch.Tensor: + assert to_matrix + self.pose_reads += 1 + return self._pose.clone() + + +class _Robot: + """Full-qpos source used by post-policy hold actions.""" + + def __init__(self, qpos: torch.Tensor) -> None: + self.qpos = qpos + self.qpos_reads = 0 + + def get_qpos(self) -> torch.Tensor: + self.qpos_reads += 1 + return self.qpos.clone() + + +class _Simulation: + """Resolve only one explicitly selected native rigid object.""" + + def __init__(self, entity: _RigidObject) -> None: + self.entity = entity + + def get_rigid_object(self, uid: str) -> _RigidObject | None: + return self.entity if uid == "native_cube" else None + + def get_articulation(self, uid: str) -> None: + del uid + return None + + +def _compiled_segment(*, settle_preset: str = "fast"): + """Compile one segment containing both supported policy types.""" + payload = { + "schema_version": 1, + "program_id": "policy_test", + "integration": { + "robot_profile": "test_robot", + "scene_registry": "test_scene", + "runtime_preset": "safe", + }, + "targets": { + "drop": { + "kind": "cyclic_pose", + "values": [ + { + "position": [0.0, 0.0, 0.0], + "quaternion_wxyz": [1.0, 0.0, 0.0, 0.0], + } + ], + } + }, + "program": { + "kind": "segment", + "name": "place", + "steps": { + "kind": "invoke", + "call": { + "kind": "place", + "object": "cube", + "at": {"kind": "target_ref", "target": "drop"}, + }, + }, + "post": [ + { + "kind": "wait_stable", + "entity": "cube", + "preset": settle_preset, + } + ], + "validators": [ + { + "kind": "object_near_target", + "object": "cube", + "target": "drop", + "position_tolerance": 0.05, + } + ], + }, + } + registry = SceneRegistry( + ( + SceneEntityRegistration( + ref=SceneObjectRef("cube"), + state_provider=_StaticStateProvider(), + ), + ) + ) + compiled = ExpertProgramCompiler.from_scene_registry(registry).compile( + decode_expert_program(payload) + ) + return next(compiled.iter_segments()) + + +def _port( + positions: torch.Tensor, + *, + preset: DynamicSettleMonitorCfg | None = None, +) -> tuple[SimulationSegmentPolicyPort, _RigidObject, _Robot]: + """Build one policy port and expose its mutable test doubles.""" + entity = _RigidObject(positions) + robot = _Robot(torch.tensor([[1.0, 2.0], [3.0, 4.0]])) + port = SimulationSegmentPolicyPort( + _Simulation(entity), + robot, + SimulationSceneBinding( + registry_id="test_scene", + rigid_objects=( + SimulationRigidObjectBinding( + entity_id="cube", + simulation_uid="native_cube", + ), + ), + ), + settle_presets={ + "fast": preset + or DynamicSettleMonitorCfg( + min_steps=0, + max_steps=3, + check_interval_steps=1, + required_stable_checks=2, + ) + }, + ) + return port, entity, robot + + +def test_port_implements_both_bridge_policy_protocols() -> None: + """One shared instance serves post-policy and validator boundaries.""" + port, _, _ = _port(torch.zeros(2, 3)) + + assert isinstance(port, SegmentPostPolicyPort) + assert isinstance(port, SegmentPostPolicyMetadataPort) + assert isinstance(port, SegmentPostPolicyResultPort) + assert isinstance(port, SegmentValidatorPort) + assert isinstance(port, SegmentValidatorMetadataPort) + assert port.settle_preset_ids == ("fast",) + + +def test_pure_preflight_validates_hooks_without_reading_live_state() -> None: + """Static hook validation emits no hold and samples no pose or qpos.""" + segment = _compiled_segment() + port, entity, robot = _port(torch.zeros(2, 3)) + + port.validate_policy(segment.post_policies[0], segment=segment) + port.validate_validator(segment.validators[0], segment=segment) + + assert robot.qpos_reads == 1 + assert entity.pose_reads == 0 + + +def test_pure_preflight_rejects_unknown_settle_preset_without_observation() -> None: + """An unknown preset fails before policy iteration can sample live state.""" + segment = _compiled_segment(settle_preset="missing") + port, entity, robot = _port(torch.zeros(2, 3)) + + with pytest.raises(KeyError, match="Unknown settle preset 'missing'"): + port.validate_policy(segment.post_policies[0], segment=segment) + + assert robot.qpos_reads == 1 + assert entity.pose_reads == 0 + + +def test_wait_stable_yields_fresh_full_qpos_holds_through_gym() -> None: + """Settling observes only after each yielded hold has been consumed.""" + segment = _compiled_segment() + port, _, robot = _port(torch.zeros(2, 3)) + actions = port.actions( + segment.post_policies[0], + segment=segment, + active_mask=torch.ones(2, dtype=torch.bool), + ) + + first = next(actions) + assert torch.equal(first, robot.qpos) + first.fill_(99.0) + with pytest.raises(StopIteration): + next(actions) + assert torch.equal(robot.qpos, torch.tensor([[1.0, 2.0], [3.0, 4.0]])) + + metadata = port.post_policy_metadata( + segment.post_policies[0], + segment=segment, + ) + assert metadata["status"] == "settled" + assert metadata["preset"] == "fast" + assert metadata["thresholds"] == { + "linear_velocity": 0.03, + "angular_velocity": 0.2, + "min_steps": 0, + "max_steps": 3, + "check_interval_steps": 1, + "required_stable_checks": 2, + } + assert metadata["state"]["elapsed_steps"] == 1 + assert metadata["state"]["settled_mask"] == [True, True] + assert metadata["state"]["timeout_mask"] == [False, False] + assert metadata["state"]["max_linear_speed"] == [0.0, 0.0] + assert port.post_policy_result( + segment.post_policies[0], + segment=segment, + ).tolist() == [True, True] + + +def test_wait_stable_returns_row_local_timeout_result_and_metadata() -> None: + """A moving row times out without failing a settled peer or the batch.""" + segment = _compiled_segment() + port, entity, _ = _port(torch.zeros(2, 3)) + entity.body_data.lin_vel[1, 0] = 1.0 + actions = port.actions( + segment.post_policies[0], + segment=segment, + active_mask=torch.ones(2, dtype=torch.bool), + ) + + assert sum(1 for _ in (next(actions), next(actions), next(actions))) == 3 + with pytest.raises(StopIteration): + next(actions) + + metadata = port.post_policy_metadata( + segment.post_policies[0], + segment=segment, + ) + assert metadata["status"] == "timed_out" + assert metadata["state"]["elapsed_steps"] == 3 + assert metadata["state"]["settled_mask"] == [True, False] + assert metadata["state"]["timeout_mask"] == [False, True] + assert metadata["state"]["max_linear_speed"] == [0.0, 1.0] + assert port.post_policy_result( + segment.post_policies[0], + segment=segment, + ).tolist() == [True, False] + + +def test_in_progress_settling_metadata_uses_json_null_for_unchecked_speeds() -> None: + segment = _compiled_segment() + port, _, _ = _port( + torch.zeros(2, 3), + preset=DynamicSettleMonitorCfg( + min_steps=2, + max_steps=4, + check_interval_steps=1, + required_stable_checks=1, + ), + ) + actions = port.actions( + segment.post_policies[0], + segment=segment, + active_mask=torch.ones(2, dtype=torch.bool), + ) + + next(actions) + metadata = port.post_policy_metadata( + segment.post_policies[0], + segment=segment, + ) + actions.close() + + assert metadata["status"] == "running" + assert metadata["state"]["max_linear_speed"] == [None, None] + assert metadata["state"]["max_angular_speed"] == [None, None] + json.dumps(metadata, allow_nan=False, sort_keys=True) + + +def test_wait_stable_excludes_inactive_moving_row_from_completion() -> None: + """A failed runtime row cannot block or pass a later settling policy.""" + segment = _compiled_segment() + port, entity, _ = _port(torch.zeros(2, 3)) + entity.body_data.lin_vel[1, 0] = 1.0 + active_mask = torch.tensor([True, False]) + actions = port.actions( + segment.post_policies[0], + segment=segment, + active_mask=active_mask, + ) + + assert sum(1 for _ in actions) == 1 + + metadata = port.post_policy_metadata( + segment.post_policies[0], + segment=segment, + ) + assert metadata["status"] == "settled" + assert metadata["active_mask"] == [True, False] + assert metadata["state"]["active_mask"] == [True, False] + assert metadata["state"]["settled_mask"] == [True, False] + assert metadata["state"]["timeout_mask"] == [False, False] + assert metadata["state"]["max_linear_speed"] == [0.0, None] + assert port.post_policy_result( + segment.post_policies[0], + segment=segment, + ).tolist() == [True, False] + + +def test_wait_stable_skips_when_no_rows_remain_active() -> None: + """An empty active cohort completes without an environment hold.""" + segment = _compiled_segment() + port, _, _ = _port(torch.zeros(2, 3)) + + actions = port.actions( + segment.post_policies[0], + segment=segment, + active_mask=torch.zeros(2, dtype=torch.bool), + ) + + assert tuple(actions) == () + metadata = port.post_policy_metadata( + segment.post_policies[0], + segment=segment, + ) + assert metadata["status"] == "skipped" + assert metadata["state"]["settled_mask"] == [False, False] + assert metadata["state"]["timeout_mask"] == [False, False] + assert port.post_policy_result( + segment.post_policies[0], + segment=segment, + ).tolist() == [False, False] + + +def test_object_near_target_validates_rows_independently() -> None: + """The validator compares explicit native object poses row by row.""" + segment = _compiled_segment() + port, _, _ = _port(torch.tensor([[0.01, 0.0, 0.0], [0.20, 0.0, 0.0]])) + + result = port.validate(segment.validators[0], segment=segment) + + assert result.dtype == torch.bool + assert result.tolist() == [True, False] + metadata = port.validator_metadata(segment.validators[0], segment=segment) + assert metadata["kind"] == "object_near_target" + assert metadata["object_id"] == "cube" + assert metadata["target_id"] == "drop" + assert metadata["position_tolerance"] == 0.05 + assert metadata["position_error"] == pytest.approx([0.01, 0.20]) + assert metadata["accepted_mask"] == [True, False] + + +def test_policy_port_rejects_unbound_native_entities_and_foreign_members() -> None: + """Bindings and compiled segment ownership are exact fail-closed boundaries.""" + binding = SimulationSceneBinding( + registry_id="test_scene", + rigid_objects=( + SimulationRigidObjectBinding( + entity_id="missing", + simulation_uid="unknown", + ), + ), + ) + robot = _Robot(torch.zeros(2, 2)) + with pytest.raises(KeyError, match="unknown"): + SimulationSegmentPolicyPort( + _Simulation(_RigidObject(torch.zeros(2, 3))), + robot, + binding, + ) + + segment = _compiled_segment() + other = _compiled_segment() + port, _, _ = _port(torch.zeros(2, 3)) + with pytest.raises(ValueError, match="does not belong"): + tuple( + port.actions( + other.post_policies[0], + segment=segment, + active_mask=torch.ones(2, dtype=torch.bool), + ) + ) + + +__all__: list[str] = [] diff --git a/tests/gym/envs/test_demo.py b/tests/gym/envs/test_demo.py index 1cbfcb29e..b0ef3a740 100644 --- a/tests/gym/envs/test_demo.py +++ b/tests/gym/envs/test_demo.py @@ -20,15 +20,99 @@ import threading from typing import Any +from unittest.mock import Mock import pytest import torch from tensordict import TensorDict -from embodichain.lab.gym.envs.demo import DemoSegment, execute_demo_episode +from embodichain.lab.gym.envs.demo import ( + DemoSegment, + DemoSegmentResult, + ProcessedEnvAction, + execute_demo_episode, +) from embodichain.lab.gym.envs.embodied_env import EmbodiedEnv +def test_processed_env_action_owns_value_and_metadata() -> None: + value = torch.tensor([[1.0, 2.0]]) + metadata = {"semantic_id": "pick", "segments": ["approach"]} + + action = ProcessedEnvAction(value=value, metadata=metadata) + value.zero_() + metadata["segments"].append("close") + snapshot = action.snapshot() + + assert action.value.tolist() == [[1.0, 2.0]] + assert dict(action.metadata) == { + "semantic_id": "pick", + "segments": ["approach"], + } + assert snapshot is not action + assert snapshot.value is not action.value + + +def test_demo_segment_result_owns_json_safe_lifecycle_metadata() -> None: + metadata = { + "runtime": {"status": "completed"}, + "validation": {"accepted_mask": [True, False]}, + } + result = DemoSegmentResult( + segment_id=0, + name="place", + start_step=0, + end_step=2, + success=False, + metadata=metadata, + ) + + metadata["runtime"]["status"] = "mutated" + exported = result.to_metadata() + exported["metadata"]["validation"]["accepted_mask"][0] = False + + assert result.metadata["runtime"]["status"] == "completed" + assert result.metadata["validation"]["accepted_mask"] == [True, False] + + +def test_demo_segment_result_rejects_non_json_metadata() -> None: + with pytest.raises(TypeError, match="non-JSON value Tensor"): + DemoSegmentResult( + segment_id=0, + name="place", + start_step=0, + end_step=1, + success=True, + metadata={"mask": torch.tensor([True])}, + ) + + +def test_embodied_env_skips_preprocessing_for_processed_action() -> None: + env = object.__new__(EmbodiedEnv) + env._num_envs = 2 + env._traj_buffer = None + env.action_manager = Mock() + env._demo_no_auto_reset = False + action = ProcessedEnvAction(value=torch.ones(2, 3)) + + normalized = env._normalize_demo_action(action) + processed = env._preprocess_action(normalized) + + assert isinstance(normalized, ProcessedEnvAction) + assert normalized is not action + assert torch.equal(processed, action.value) + env.action_manager.process_action.assert_not_called() + + +def test_embodied_env_validates_processed_action_batch_size() -> None: + env = object.__new__(EmbodiedEnv) + env._num_envs = 2 + action = ProcessedEnvAction(value=torch.ones(1, 3)) + + with pytest.raises(ValueError, match="batch size"): + env._normalize_demo_action(action) + + class _SegmentedEnv: """Small environment stub that supports lazy two-segment planning.""" @@ -94,6 +178,142 @@ def test_execute_demo_episode_runs_lazy_segments_as_one_episode() -> None: assert not env._demo_no_auto_reset +class _LifecycleMetadataEnv: + """Populate one shared metadata mapping at lazy lifecycle boundaries.""" + + def __init__(self) -> None: + self.num_envs = 1 + self.lifecycle = {"runtime": None, "validation": None} + + def create_demo_segments(self): + def actions(): + yield 1 + self.lifecycle["runtime"] = {"status": "completed"} + + def validate() -> bool: + self.lifecycle["validation"] = {"accepted_mask": [True]} + return True + + return ( + DemoSegment( + actions=actions(), + name="lifecycle", + metadata=self.lifecycle, + validator=validate, + ), + ) + + def step(self, action: int): + del action + return ( + None, + torch.zeros(1), + torch.zeros(1, dtype=torch.bool), + torch.zeros(1, dtype=torch.bool), + {"success": torch.tensor([True])}, + ) + + def is_task_success(self) -> torch.Tensor: + return torch.tensor([True]) + + +class _EmptySuccessfulSegmentEnv: + """Expose an empty ordinary segment whose callbacks otherwise succeed.""" + + num_envs = 1 + + def __init__(self) -> None: + self.validator_calls = 0 + self.step_calls = 0 + + def create_demo_segments(self): + return ( + DemoSegment( + actions=(), + name="empty", + validator=self._validate, + ), + ) + + def _validate(self) -> bool: + self.validator_calls += 1 + return True + + def step(self, action: object): + del action + self.step_calls += 1 + raise AssertionError("An empty segment must not call env.step().") + + @staticmethod + def is_task_success() -> torch.Tensor: + return torch.tensor([True]) + + +def test_execute_demo_episode_snapshots_finalized_lifecycle_metadata() -> None: + env = _LifecycleMetadataEnv() + + result = execute_demo_episode(env) + env.lifecycle["runtime"]["status"] = "mutated" + + assert result.segments[0].metadata == { + "runtime": {"status": "completed"}, + "validation": {"accepted_mask": [True]}, + } + + +def test_empty_ordinary_segment_keeps_existing_empty_segment_guard() -> None: + env = _EmptySuccessfulSegmentEnv() + + result = execute_demo_episode(env) + + assert env.step_calls == 0 + assert env.validator_calls == 0 + assert not result.completed + assert result.terminal_reason == "empty_segment" + assert result.segments[0].failure_reason == "empty_segment" + + +class _GeneratorFailureEnv: + """Raise between lazy actions and expose an emergency hold callback.""" + + def __init__(self) -> None: + self.num_envs = 1 + self.actions: list[int] = [] + self.abort_calls: list[tuple[str, bool]] = [] + + def create_demo_segments(self): + def actions(): + yield 1 + raise ValueError("planner stream failed") + + def abort(reason: str, *, last_action_consumed: bool): + self.abort_calls.append((reason, last_action_consumed)) + yield 0 + + return (DemoSegment(actions=actions(), abort_actions=abort),) + + def step(self, action: int): + self.actions.append(action) + return ( + None, + torch.zeros(1), + torch.zeros(1, dtype=torch.bool), + torch.zeros(1, dtype=torch.bool), + {}, + ) + + +def test_action_generator_failure_safe_stops_before_propagating() -> None: + env = _GeneratorFailureEnv() + + with pytest.raises(RuntimeError, match="action generation") as error: + execute_demo_episode(env) + + assert isinstance(error.value.__cause__, ValueError) + assert env.actions == [1, 0] + assert env.abort_calls == [("action_generation_failed", True)] + + class _TerminatingEnv(_SegmentedEnv): def create_demo_segments(self): return (DemoSegment(actions=(1, 2, 3), name="pick"),) @@ -213,6 +433,30 @@ def test_vector_failure_aborts_peer_and_preserves_per_env_reason() -> None: assert result.lengths == (2, 2) +class _RowIndependentFailureEnv(_VectorFailureEnv): + def create_demo_segments(self): + return ( + DemoSegment( + actions=(1, 2, 3), + name="shared", + failure_policy="row_independent", + ), + ) + + +def test_row_independent_failure_freezes_only_failed_environment() -> None: + env = _RowIndependentFailureEnv() + + result = execute_demo_episode(env) + + assert env.actions == [1, 2, 3] + assert env.masked_actions == [(3, (False, True))] + assert result.completed_by_env == (False, True) + assert result.terminal_reasons == ("failure", "success") + assert result.success == (False, True) + assert result.lengths == (2, 3) + + class _ValidatedSegmentEnv(_SegmentedEnv): def __init__(self, validation: bool) -> None: super().__init__() @@ -388,6 +632,35 @@ def test_validator_batch_abort_has_consistent_peer_status() -> None: assert result.segments[0].failure_reason == "segment_validation_failed" +class _RowIndependentValidatorEnv(_VectorValidatorEnv): + def create_demo_segments(self): + return ( + DemoSegment( + actions=(1,), + name="validated", + validator=lambda: torch.tensor([True, False]), + failure_policy="row_independent", + ), + ) + + +def test_row_independent_validator_keeps_accepted_peer_active() -> None: + result = execute_demo_episode(_RowIndependentValidatorEnv()) + + assert result.segments[0].successes == (True, False) + assert result.segments[0].failure_reasons == ( + None, + "segment_validation_failed", + ) + assert result.completed_by_env == (True, False) + assert result.terminal_reasons == ("success", "segment_validation_failed") + + +def test_demo_segment_rejects_unknown_failure_policy() -> None: + with pytest.raises(ValueError, match="failure_policy"): + DemoSegment(actions=(1,), failure_policy="continue") + + class _CancellationEnv(_ValidatedSegmentEnv): def __init__(self) -> None: super().__init__(validation=True) diff --git a/tests/gym/envs/test_embodied_env_expert_program.py b/tests/gym/envs/test_embodied_env_expert_program.py new file mode 100644 index 000000000..fc7892b9f --- /dev/null +++ b/tests/gym/envs/test_embodied_env_expert_program.py @@ -0,0 +1,92 @@ +# ---------------------------------------------------------------------------- +# 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 explicit Expert Program environment integration hooks.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from embodichain.lab.gym.envs.demo import DemoSegment +from embodichain.lab.gym.envs.embodied_env import EmbodiedEnv, EmbodiedEnvCfg + + +class _FakeBridge: + """Minimal bridge protocol used by the environment adapter test.""" + + def __init__(self, segment: DemoSegment) -> None: + self._segment = segment + self.iteration_count = 0 + + def iter_segments(self): + """Yield the configured segment lazily.""" + self.iteration_count += 1 + yield self._segment + + +class _DeclarativeEnv(EmbodiedEnv): + """Environment stub with explicit compiler and bridge factories.""" + + def compile_expert_program(self, program): + self.compiled_input = program + return self.compiled_program + + def create_expert_program_bridge(self, program): + self.bridge_input = program + return self.bridge + + +def _uninitialized_env(cls: type[EmbodiedEnv], expert_program: object) -> EmbodiedEnv: + """Create an environment instance without starting simulation.""" + env = object.__new__(cls) + env.cfg = SimpleNamespace(expert_program=expert_program) + return env + + +def test_embodied_env_cfg_disables_expert_program_by_default() -> None: + """Declarative execution remains an explicit opt-in configuration.""" + cfg = EmbodiedEnvCfg() + + assert cfg.expert_program is None + + +def test_create_demo_segments_uses_explicit_compiler_and_bridge_hooks() -> None: + """Configured programs flow through provider and runtime factories lazily.""" + program = object() + compiled_program = object() + expected_segment = DemoSegment(actions=(), name="declarative") + bridge = _FakeBridge(expected_segment) + env = _uninitialized_env(_DeclarativeEnv, program) + env.compiled_program = compiled_program + env.bridge = bridge + + segments = env.create_demo_segments(debug_mode=True) + + assert bridge.iteration_count == 0 + assert tuple(segments) == (expected_segment,) + assert bridge.iteration_count == 1 + assert env.compiled_input is program + assert env.bridge_input is compiled_program + + +def test_configured_program_requires_explicit_scene_provider_hook() -> None: + """The base environment never guesses a live scene provider.""" + env = _uninitialized_env(EmbodiedEnv, object()) + + with pytest.raises(NotImplementedError, match="explicit scene resolver"): + env.create_demo_segments() diff --git a/tests/gym/envs/test_settling.py b/tests/gym/envs/test_settling.py new file mode 100644 index 000000000..d1476b842 --- /dev/null +++ b/tests/gym/envs/test_settling.py @@ -0,0 +1,145 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import pytest +import torch + +from embodichain.lab.gym.envs.settling import ( + DynamicSettleMonitor, + DynamicSettleMonitorCfg, + DynamicSettleSample, +) + + +def _sample( + linear: tuple[float, ...], angular: tuple[float, ...] +) -> DynamicSettleSample: + return DynamicSettleSample( + entity_id="cube", + linear_speed=torch.tensor(linear, dtype=torch.float32).unsqueeze(1), + angular_speed=torch.tensor(angular, dtype=torch.float32).unsqueeze(1), + ) + + +def test_settle_monitor_tracks_rows_independently_and_owns_metadata() -> None: + env_ids = torch.tensor([4, 9], dtype=torch.long) + monitor = DynamicSettleMonitor( + DynamicSettleMonitorCfg( + min_steps=1, + max_steps=5, + check_interval_steps=1, + required_stable_checks=2, + ), + env_ids, + ) + + first = monitor.observe((_sample((0.0, 1.0), (0.0, 1.0)),), elapsed_steps=1) + second = monitor.observe((_sample((0.0, 1.0), (0.0, 1.0)),), elapsed_steps=2) + third = monitor.observe((_sample((0.0, 0.0), (0.0, 0.0)),), elapsed_steps=3) + final = monitor.observe((_sample((0.0, 0.0), (0.0, 0.0)),), elapsed_steps=4) + + assert first.stable_counts.tolist() == [1, 0] + assert second.settled_mask.tolist() == [True, False] + assert third.stable_counts.tolist() == [2, 1] + assert final.settled_mask.tolist() == [True, True] + assert final.timeout_mask.tolist() == [False, False] + assert final.complete is True + metadata = final.to_metadata() + assert metadata["env_ids"] == [4, 9] + assert metadata["settled_mask"] == [True, True] + + env_ids[0] = 100 + assert monitor.env_ids.tolist() == [4, 9] + + +def test_settle_monitor_duplicate_observation_is_idempotent() -> None: + monitor = DynamicSettleMonitor( + DynamicSettleMonitorCfg( + min_steps=0, + max_steps=3, + check_interval_steps=1, + required_stable_checks=2, + ), + torch.tensor([0], dtype=torch.long), + ) + + first = monitor.observe((_sample((0.0,), (0.0,)),), elapsed_steps=0) + duplicate = monitor.observe((_sample((0.0,), (0.0,)),), elapsed_steps=0) + second = monitor.observe((_sample((0.0,), (0.0,)),), elapsed_steps=1) + + assert first.checked is True + assert duplicate.checked is False + assert duplicate.stable_counts.tolist() == [1] + assert second.settled_mask.tolist() == [True] + assert second.observation_count == 2 + + +def test_settle_monitor_marks_only_unresolved_rows_timed_out() -> None: + monitor = DynamicSettleMonitor( + DynamicSettleMonitorCfg( + min_steps=0, + max_steps=1, + check_interval_steps=1, + required_stable_checks=1, + ), + torch.tensor([0, 1], dtype=torch.long), + ) + + state = monitor.observe((_sample((0.0, 1.0), (0.0, 1.0)),), elapsed_steps=1) + + assert state.settled_mask.tolist() == [True, False] + assert state.timeout_mask.tolist() == [False, True] + assert state.complete is True + + +@pytest.mark.parametrize( + ("kwargs", "match"), + ( + ({"min_steps": -1}, "min_steps"), + ({"max_steps": 1, "min_steps": 2}, "max_steps"), + ({"check_interval_steps": 0}, "check_interval_steps"), + ({"linear_velocity_threshold": float("nan")}, "linear_velocity_threshold"), + ( + {"min_steps": 0, "max_steps": 0, "required_stable_checks": 2}, + "cannot be reached", + ), + ), +) +def test_settle_monitor_cfg_rejects_invalid_values( + kwargs: dict[str, object], match: str +) -> None: + with pytest.raises((TypeError, ValueError), match=match): + DynamicSettleMonitorCfg(**kwargs) + + +def test_settle_monitor_rejects_regressing_steps_and_incomplete_samples() -> None: + monitor = DynamicSettleMonitor( + DynamicSettleMonitorCfg( + min_steps=0, + max_steps=2, + required_stable_checks=1, + ), + torch.tensor([0], dtype=torch.long), + ) + sample = _sample((1.0,), (1.0,)) + monitor.observe((sample,), elapsed_steps=1) + + with pytest.raises(ValueError, match="monotonic"): + monitor.observe((sample,), elapsed_steps=0) + with pytest.raises(ValueError, match="contain DynamicSettleSample"): + monitor.observe((), elapsed_steps=2) diff --git a/tests/gym/utils/test_gym_utils.py b/tests/gym/utils/test_gym_utils.py index 12d46874b..db3119281 100644 --- a/tests/gym/utils/test_gym_utils.py +++ b/tests/gym/utils/test_gym_utils.py @@ -509,6 +509,37 @@ def test_different_max_episode_steps(): class TestConfigToCfgFromFile: + @staticmethod + def _minimal_gym_config() -> dict[str, object]: + """Return a minimal config that reaches the generic parser.""" + return { + "id": "EmbodiedEnv-v1", + "env": {}, + "robot": { + "class_type": "URRobot", + "robot_type": "ur5", + "uid": "TestUR5", + }, + } + + @staticmethod + def _expert_program_payload() -> dict[str, object]: + """Return one minimal strict Expert Program payload.""" + return { + "schema_version": 1, + "program_id": "configured_pick", + "integration": { + "robot_profile": "default_robot", + "scene_registry": "default_scene", + "runtime_preset": "default_runtime", + }, + "targets": {}, + "program": { + "kind": "invoke", + "call": {"kind": "pick", "object": "cube"}, + }, + } + def test_robot_class_type_preserves_ur_variant(self): config = { "id": "EmbodiedEnv-v1", @@ -532,6 +563,107 @@ def test_robot_class_type_preserves_ur_variant(self): "uid": "TestUR5", } + def test_expert_program_path_is_resolved_from_gym_config_source( + self, + tmp_path, + ) -> None: + """A serialized program path is relative to its Gym config file.""" + gym_dir = tmp_path / "gym" / "task" + program_dir = tmp_path / "expert_program" + gym_dir.mkdir(parents=True) + program_dir.mkdir() + gym_path = gym_dir / "gym_config.json" + program_path = program_dir / "program.yaml" + save_config(program_path, self._expert_program_payload()) + config = self._minimal_gym_config() + config["expert_program_path"] = "../../expert_program/program.yaml" + + cfg = config_to_cfg( + config, + manager_modules=DEFAULT_MANAGER_MODULES, + source_path=gym_path, + ) + + assert cfg.expert_program.program_id == "configured_pick" + assert cfg.expert_program.integration.scene_registry == "default_scene" + + def test_build_env_cfg_loads_source_relative_expert_program( + self, + tmp_path, + ) -> None: + """The normal file launcher attaches the decoded program before init.""" + gym_dir = tmp_path / "gym" + program_dir = tmp_path / "programs" + gym_dir.mkdir() + program_dir.mkdir() + gym_path = gym_dir / "gym_config.json" + program_path = program_dir / "program.json" + save_config(program_path, self._expert_program_payload()) + config = self._minimal_gym_config() + config["expert_program_path"] = "../programs/program.json" + save_config(gym_path, config) + args = argparse.Namespace( + gym_config=str(gym_path), + num_envs=1, + device="cpu", + headless=True, + renderer=None, + gpu_id=0, + arena_space=2.0, + max_episodes=None, + filter_visual_rand=False, + filter_dataset_saving=False, + preview=False, + action_config=None, + ) + + cfg, _, _ = build_env_cfg_from_args(args) + + assert cfg.expert_program.program_id == "configured_pick" + + def test_config_to_cfg_uses_cwd_without_source_path( + self, + tmp_path, + monkeypatch, + ) -> None: + """Dictionary-only callers retain explicit current-directory semantics.""" + program_path = tmp_path / "program.yaml" + save_config(program_path, self._expert_program_payload()) + monkeypatch.chdir(tmp_path) + config = self._minimal_gym_config() + config["expert_program_path"] = "program.yaml" + + cfg = config_to_cfg(config, manager_modules=DEFAULT_MANAGER_MODULES) + + assert cfg.expert_program.program_id == "configured_pick" + + @pytest.mark.parametrize("value", [None, True, 1, {}, "", " program.yaml"]) + def test_expert_program_path_rejects_ambiguous_values( + self, + value, + ) -> None: + """The path field never accepts coercion, null, or outer whitespace.""" + config = self._minimal_gym_config() + config["expert_program_path"] = value + + with pytest.raises((TypeError, ValueError), match="expert_program_path"): + config_to_cfg(config, manager_modules=DEFAULT_MANAGER_MODULES) + + def test_expert_program_path_missing_file_fails_before_environment_init( + self, + tmp_path, + ) -> None: + """A configured program must exist when the Gym config is decoded.""" + config = self._minimal_gym_config() + config["expert_program_path"] = "missing.yaml" + + with pytest.raises(FileNotFoundError, match="missing.yaml"): + config_to_cfg( + config, + manager_modules=DEFAULT_MANAGER_MODULES, + source_path=tmp_path / "gym_config.json", + ) + def test_yaml_gym_config_parses_to_cfg(self, tmp_path): config = { "id": "EmbodiedEnv-v1", diff --git a/tests/lab/scripts/test_run_env.py b/tests/lab/scripts/test_run_env.py index 0c495a4f5..788a89f9b 100644 --- a/tests/lab/scripts/test_run_env.py +++ b/tests/lab/scripts/test_run_env.py @@ -16,6 +16,7 @@ from __future__ import annotations +import json from types import SimpleNamespace from unittest.mock import MagicMock @@ -27,6 +28,7 @@ from embodichain.lab.scripts import run_env from embodichain.lab.scripts.run_env import ( _create_parser, + _load_expert_program, _run_replay_control_loop, generate_function, ) @@ -40,6 +42,24 @@ VISER_POLL_INTERVAL = 0.05 +def _expert_program_payload() -> dict[str, object]: + """Return one minimal strict Expert Program payload.""" + return { + "schema_version": 1, + "program_id": "cli_pick", + "integration": { + "robot_profile": "default_robot", + "scene_registry": "default_scene", + "runtime_preset": "default_runtime", + }, + "targets": {}, + "program": { + "kind": "invoke", + "call": {"kind": "pick", "object": "cube"}, + }, + } + + class _LegacyProgressEnv: num_envs = 1 @@ -123,6 +143,86 @@ def test_run_env_preserves_configured_viser_image_fps() -> None: assert merged["visualization"]["sensor_image_fps"] == configured_fps +def test_run_env_parser_accepts_expert_program_path() -> None: + """The declarative program is an explicit, opt-in CLI input.""" + program_path = "program.yaml" + + args = _create_parser().parse_args( + ["--gym_config", GYM_CONFIG_PATH, "--expert-program", program_path] + ) + + assert args.expert_program == program_path + + +def test_run_env_parser_accepts_debug_trace_mode() -> None: + """Failed Expert Program attempts can expose their structured trace.""" + args = _create_parser().parse_args( + ["--gym_config", GYM_CONFIG_PATH, "--debug-mode"] + ) + + assert args.debug_mode is True + + +@pytest.mark.parametrize("suffix", [".json", ".yaml", ".yml"]) +def test_load_expert_program_safely_decodes_supported_files( + tmp_path, + suffix: str, +) -> None: + """JSON and safe YAML inputs share the same strict schema decoder.""" + path = tmp_path / f"program{suffix}" + payload = _expert_program_payload() + if suffix == ".json": + serialized = json.dumps(payload) + else: + import yaml + + serialized = yaml.safe_dump(payload) + path.write_text(serialized, encoding="utf-8") + + program = _load_expert_program(path) + + assert program.program_id == "cli_pick" + assert program.integration.scene_registry == "default_scene" + + +@pytest.mark.parametrize( + ("filename", "serialized", "message"), + [ + ( + "program.json", + '{"schema_version": 1, "schema_version": 1}', + "Duplicate JSON key", + ), + ( + "program.yaml", + "schema_version: 1\nschema_version: 1\n", + "found duplicate key", + ), + ], +) +def test_load_expert_program_rejects_duplicate_mapping_keys( + tmp_path, + filename: str, + serialized: str, + message: str, +) -> None: + """Ambiguous duplicate keys are rejected before schema decoding.""" + path = tmp_path / filename + path.write_text(serialized, encoding="utf-8") + + with pytest.raises(ValueError, match=message): + _load_expert_program(path) + + +def test_load_expert_program_rejects_unsupported_file_extension(tmp_path) -> None: + """Only explicit JSON and YAML file formats are accepted.""" + path = tmp_path / "program.toml" + path.write_text("schema_version = 1", encoding="utf-8") + + with pytest.raises(ValueError, match=".json, .yaml, or .yml"): + _load_expert_program(path) + + def test_replay_restores_wrapper_state_without_closing_caller_env(monkeypatch) -> None: """Replay leaves the environment close to its CLI owner.""" env = MagicMock() @@ -230,6 +330,34 @@ def test_generate_function_discards_retry_then_commits_once(monkeypatch) -> None assert env.reset_options == [{"save_data": False}, None] +def test_generate_function_logs_failed_trace_in_debug_mode(monkeypatch) -> None: + """Debug retries expose the owned structured episode trace.""" + env = _ResetTrackingEnv() + result = _episode_result(success=False, reason="segment_validation_failed") + warnings: list[str] = [] + monkeypatch.setattr( + "embodichain.lab.scripts.run_env.execute_demo_episode", + lambda *args, **kwargs: result, + ) + monkeypatch.setattr( + "embodichain.lab.scripts.run_env.log_warning", + warnings.append, + ) + + generated = generate_function( + env, + max_attempts=1, + reset_before=False, + debug_mode=True, + ) + + assert not generated + debug_trace = next( + message for message in warnings if "Failed demo trace" in message + ) + assert '"terminal_reason":"segment_validation_failed"' in debug_trace + + def test_generate_function_commits_failed_episode_when_configured(monkeypatch) -> None: """A recorded task failure is a persisted result when explicitly enabled.""" env = _ResetTrackingEnv(save_failed_episodes=True) @@ -421,6 +549,49 @@ def test_cli_aborts_before_closing_environment_once(monkeypatch) -> None: assert env.events == [abort_event, abort_event, ("close", None)] +def test_cli_injects_decoded_expert_program_before_environment_creation( + monkeypatch, +) -> None: + """The CLI attaches the strict program config to the environment config.""" + env = _LifecycleTrackingEnv() + env_cfg = SimpleNamespace(expert_program=None) + decoded_program = object() + args = SimpleNamespace( + replay=False, + replay_mode="kinematic", + preview=True, + expert_program="program.yaml", + ) + parser = MagicMock() + parser.parse_args.return_value = args + make = MagicMock(return_value=env) + + monkeypatch.setattr(run_env, "_create_parser", lambda: parser) + monkeypatch.setattr(run_env, "discover_task_packages", lambda: None) + monkeypatch.setattr(run_env, "execute_init_hooks", lambda: None) + monkeypatch.setattr( + run_env, + "build_env_cfg_from_args", + lambda parsed_args: (env_cfg, {"id": GYM_ID}, {}), + ) + monkeypatch.setattr( + run_env, + "_load_expert_program", + MagicMock(return_value=decoded_program), + ) + monkeypatch.setattr(run_env.gymnasium, "make", make) + monkeypatch.setattr(run_env, "main", lambda *args, **kwargs: None) + monkeypatch.setattr( + "embodichain.lab.sim.sim_manager.SimulationManager.flush_cleanup_queue", + lambda: None, + ) + + run_env.cli([]) + + assert env_cfg.expert_program is decoded_program + make.assert_called_once_with(id=GYM_ID, cfg=env_cfg) + + def test_close_durability_failure_is_not_swallowed() -> None: """A failed recorder barrier makes the runner fail after aborting pending data.""" env = _LifecycleTrackingEnv() diff --git a/tests/utils/test_config_paths.py b/tests/utils/test_config_paths.py new file mode 100644 index 000000000..b05c9fe04 --- /dev/null +++ b/tests/utils/test_config_paths.py @@ -0,0 +1,68 @@ +# ---------------------------------------------------------------------------- +# 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 stable configuration-path resolution.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from embodichain.utils import resolve_config_path as exported_resolve_config_path +from embodichain.utils.config_paths import resolve_config_path + + +def test_resolve_config_path_preserves_existing_path(tmp_path: Path) -> None: + config_path = tmp_path / "config.yaml" + config_path.write_text("id: Test-v0\n", encoding="utf-8") + + assert resolve_config_path(config_path) == config_path + + +def test_resolve_config_path_is_exported_from_utils_package() -> None: + assert exported_resolve_config_path is resolve_config_path + + +def test_resolve_config_path_preserves_ordinary_relative_path( + tmp_path: Path, + monkeypatch, +) -> None: + monkeypatch.chdir(tmp_path) + + assert resolve_config_path("local/config.yaml") == Path("local/config.yaml") + + +def test_resolve_config_path_redirects_packaged_task_config( + tmp_path: Path, + monkeypatch, +) -> None: + monkeypatch.chdir(tmp_path) + + resolved = resolve_config_path("embodichain_tasks/configs/gym/cobotmagic.json") + + assert resolved.is_file() + assert resolved.name == "cobotmagic.json" + + +def test_resolve_config_path_rejects_packaged_path_escape( + tmp_path: Path, + monkeypatch, +) -> None: + monkeypatch.chdir(tmp_path) + + with pytest.raises(ValueError, match="stay within the package"): + resolve_config_path("embodichain_tasks/configs/../VERSION")