diff --git a/embodichain/gen_sim/action_engine/__init__.py b/embodichain/gen_sim/action_engine/__init__.py new file mode 100644 index 000000000..cd941d35b --- /dev/null +++ b/embodichain/gen_sim/action_engine/__init__.py @@ -0,0 +1,50 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Capability-driven planning and live execution for generated simulations. + +Action Engine deliberately exposes a small public surface. Natural-language +goals become a typed TaskSpec, deterministic planning lowers them into a +coordinate-free SeedGraph, and the runtime grounds that graph only against live +simulator state. +""" + +from __future__ import annotations + +from .unbound import ( + UNBOUND_ACTION_PLAN_SCHEMA, + UnboundActionPlan, + build_unbound_action_plan, + validate_unbound_action_plan, +) + +from .protocol import ( + ACTION_ENGINE_CONFIG_SCHEMA, + ACTION_ENGINE_ENV_ID, + EXECUTION_PROGRAM_SCHEMA, + TASK_AGENT_SCHEMA, +) + +__all__ = [ + "ACTION_ENGINE_CONFIG_SCHEMA", + "ACTION_ENGINE_ENV_ID", + "EXECUTION_PROGRAM_SCHEMA", + "TASK_AGENT_SCHEMA", + "UNBOUND_ACTION_PLAN_SCHEMA", + "UnboundActionPlan", + "build_unbound_action_plan", + "validate_unbound_action_plan", +] diff --git a/embodichain/gen_sim/action_engine/agent.py b/embodichain/gen_sim/action_engine/agent.py new file mode 100644 index 000000000..55def0bec --- /dev/null +++ b/embodichain/gen_sim/action_engine/agent.py @@ -0,0 +1,762 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Grounded-plan compilation and compact execution reporting.""" + +from __future__ import annotations + +from collections.abc import Collection, Mapping, Sequence +from copy import deepcopy +from dataclasses import asdict, is_dataclass +from datetime import datetime, timezone +import hashlib +import json +from pathlib import Path +from types import SimpleNamespace +from typing import Any, Callable, TypeAlias + +import numpy as np +import torch + +from embodichain.gen_sim.action_engine.capabilities import ( + AtomicCapabilityRegistry, + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.domain import ( + seed_graph_hash, + validate_seed_graph, +) +from embodichain.gen_sim.action_engine.planning.linker import ( + validate_persisted_contracts, +) +from embodichain.gen_sim.action_engine.runtime import ( + ExecutionProgram, + ExecutionReport, + ExecutionResult, + ProgramExecutor, + build_execution_provenance, + load_execution_program, + validate_execution_report, + write_execution_report, +) +from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph +from embodichain.gen_sim.action_engine.unbound import ( + ActionCapabilityError, + UnboundActionPlan, + build_unbound_action_plan, + validate_unbound_action_plan, +) + +__all__ = ["ActionAgent", "ActionGraph"] + +ExecutorFactory = Callable[..., ProgramExecutor] +ActionGraph: TypeAlias = dict[str, Any] + + +class ActionAgent: + """Compile, preflight, execute, and report one grounded task plan.""" + + def __init__( + self, + *, + registry: AtomicCapabilityRegistry | None = None, + executor_factory: ExecutorFactory = ProgramExecutor, + ) -> None: + self.registry = registry or build_atomic_capability_registry() + self.executor_factory = executor_factory + + def plan(self, grounded_plan: Mapping[str, Any]) -> ActionGraph: + """Compile a validated GroundedTaskPlan to the public SeedGraph v3.""" + plan = _validate_grounded_plan(grounded_plan) + task_spec = _mapping(plan.get("task_spec"), "GroundedTaskPlan.task_spec") + bindings = _role_binding_map(plan.get("role_bindings")) + graph = instantiate_seed_graph( + task_spec, + bindings, + registry=self.registry, + ) + known_uids = _known_uids( + plan.get("scene_manifest"), + bindings=bindings, + ) + known_uids.add("table") + graph = validate_seed_graph( + graph, + known_objects=known_uids or None, + known_actions=self.registry.names(), + executable_actions=self.registry.executable_names(), + require_executable=False, + ) + validate_persisted_contracts(graph, self.registry) + return graph + + def draft(self, candidate: Mapping[str, Any]) -> UnboundActionPlan: + """Create an Action-owned draft before final scene UID binding. + + Args: + candidate: One validated Task Engine candidate. + + Returns: + A scene-independent action plan whose selectors contain no UIDs. + """ + plan = build_unbound_action_plan(candidate) + names = getattr(self.registry, "names", None) + executable_names = getattr(self.registry, "executable_names", None) + if callable(names): + missing = sorted(set(plan["required_actions"]) - set(names())) + if missing: + raise ActionCapabilityError( + "Required AtomicAction is not registered: " + ", ".join(missing) + ) + if callable(executable_names): + unavailable = sorted( + set(plan["required_actions"]) - set(executable_names()) + ) + if unavailable: + raise ActionCapabilityError( + "Required AtomicAction is not executable: " + ", ".join(unavailable) + ) + return plan + + def bind_and_plan( + self, + unbound_plan: Mapping[str, Any], + grounded_plan: Mapping[str, Any], + ) -> ActionGraph: + """Bind one audited unbound plan through a final GroundedTaskPlan. + + The grounded task draft must reproduce the exact unbound IR. This + prevents the final planner from silently reinterpreting a candidate + after Scene Engine work has run concurrently. + """ + unbound = validate_unbound_action_plan(unbound_plan) + grounded = _validate_grounded_plan(grounded_plan) + expected = build_unbound_action_plan( + { + "candidate_id": grounded["selected_candidate_id"], + "draft": grounded["task_draft"], + } + ) + if unbound != expected: + raise ValueError( + "UnboundActionPlan does not match the final GroundedTaskPlan." + ) + return self.plan(grounded) + + def preflight( + self, + action_graph: Mapping[str, Any] | str | Path, + *, + scene_manifest: Mapping[str, Any] | None = None, + known_uids: Collection[str] | None = None, + ) -> ExecutionProgram: + """Reject invalid and planning-only graphs before simulator motion.""" + known = set(str(uid) for uid in (known_uids or ()) if str(uid)) + known.update(_known_uids(scene_manifest)) + if isinstance(action_graph, Mapping): + metadata = action_graph.get("metadata", {}) + if isinstance(metadata, Mapping): + bindings = metadata.get("role_bindings", {}) + if isinstance(bindings, Mapping): + known.update(str(uid) for uid in bindings.values() if str(uid)) + if known: + known.add("table") + return load_execution_program( + action_graph, + known_objects=known or None, + registry=self.registry, + require_executable=True, + ) + + def execute( + self, + action_graph: Mapping[str, Any] | str | Path, + env: Any, + *, + grounded_plan: Mapping[str, Any] | None = None, + scene_manifest: Mapping[str, Any] | None = None, + known_uids: Collection[str] | None = None, + run_id: str | None = None, + episode_index: int = 0, + episode_seed: int | None = None, + runtime_arguments: Mapping[str, Any] | None = None, + executor_kwargs: Mapping[str, Any] | None = None, + ) -> ExecutionReport: + """Preflight and execute a graph, converting all outcomes to a report.""" + task_id = _task_id(grounded_plan, action_graph) + plan_hash = _plan_hash(grounded_plan) + graph_hash = _action_graph_hash(action_graph) + effective_run_id = run_id or _new_run_id() + provenance = build_execution_provenance( + episode_seed=episode_seed, + runtime_arguments=runtime_arguments, + ) + effective_manifest = scene_manifest + if effective_manifest is None and grounded_plan is not None: + value = grounded_plan.get("scene_manifest") + if isinstance(value, Mapping): + effective_manifest = value + + try: + program = self.preflight( + action_graph, + scene_manifest=effective_manifest, + known_uids=known_uids, + ) + except (TypeError, ValueError, OSError) as exc: + return self._empty_report( + env, + task_id=task_id, + plan_hash=plan_hash, + graph_hash=graph_hash, + status="rejected", + run_id=effective_run_id, + episode_index=episode_index, + provenance=provenance, + error=_error_message(exc), + ) + + kwargs = dict(executor_kwargs or {}) + kwargs.setdefault("capability_registry", self.registry) + try: + executor = self.executor_factory(program, env, **kwargs) + result = executor.run( + run_id=effective_run_id, + episode_index=episode_index, + ) + except Exception as exc: + return self._empty_report( + env, + task_id=task_id, + plan_hash=plan_hash, + graph_hash=graph_hash, + status="aborted", + run_id=effective_run_id, + episode_index=episode_index, + provenance=provenance, + error=_error_message(exc), + ) + if not isinstance(result, ExecutionResult): + return self._empty_report( + env, + task_id=task_id, + plan_hash=plan_hash, + graph_hash=graph_hash, + status="aborted", + run_id=effective_run_id, + episode_index=episode_index, + provenance=provenance, + error="TypeError: ProgramExecutor.run must return ExecutionResult.", + ) + try: + return self.report_execution_result( + result, + action_graph=action_graph, + grounded_plan=grounded_plan, + run_id=effective_run_id, + episode_index=episode_index, + episode_seed=episode_seed, + runtime_arguments=runtime_arguments, + ) + except (TypeError, ValueError, OverflowError) as exc: + return self._empty_report( + env, + task_id=task_id, + plan_hash=plan_hash, + graph_hash=graph_hash, + status="aborted", + run_id=effective_run_id, + episode_index=episode_index, + provenance=provenance, + error=_error_message(exc), + ) + + def run( + self, + grounded_plan: Mapping[str, Any], + env: Any, + *, + run_id: str | None = None, + episode_index: int = 0, + episode_seed: int | None = None, + runtime_arguments: Mapping[str, Any] | None = None, + executor_kwargs: Mapping[str, Any] | None = None, + ) -> ExecutionReport: + """Compile and execute one GroundedTaskPlan through the full pipeline.""" + effective_run_id = run_id or _new_run_id() + try: + graph = self.plan(grounded_plan) + except (TypeError, ValueError, OSError) as exc: + return self._empty_report( + env, + task_id=_task_id(grounded_plan, {}), + plan_hash=_plan_hash(grounded_plan), + graph_hash=_document_hash({}), + status="rejected", + run_id=effective_run_id, + episode_index=episode_index, + provenance=build_execution_provenance( + episode_seed=episode_seed, + runtime_arguments=runtime_arguments, + ), + error=_error_message(exc), + ) + return self.execute( + graph, + env, + grounded_plan=grounded_plan, + run_id=effective_run_id, + episode_index=episode_index, + episode_seed=episode_seed, + runtime_arguments=runtime_arguments, + executor_kwargs=executor_kwargs, + ) + + def report_execution_result( + self, + result: ExecutionResult, + *, + action_graph: Mapping[str, Any] | str | Path, + grounded_plan: Mapping[str, Any] | None = None, + run_id: str | None = None, + episode_index: int = 0, + episode_seed: int | None = None, + runtime_arguments: Mapping[str, Any] | None = None, + ) -> ExecutionReport: + """Convert a result already executed by the legacy runner to a report.""" + if not isinstance(result, ExecutionResult): + raise TypeError("result must be an ExecutionResult.") + return self._result_report( + result, + task_id=_task_id(grounded_plan, action_graph), + plan_hash=_plan_hash(grounded_plan), + graph_hash=_action_graph_hash(action_graph), + run_id=run_id or _new_run_id(), + episode_index=episode_index, + provenance=build_execution_provenance( + episode_seed=episode_seed, + runtime_arguments=runtime_arguments, + ), + ) + + def rejection_report( + self, + action_graph: Mapping[str, Any] | str | Path, + error: BaseException | str, + *, + grounded_plan: Mapping[str, Any] | None = None, + environment_count: int = 1, + run_id: str | None = None, + episode_index: int = 0, + episode_seed: int | None = None, + runtime_arguments: Mapping[str, Any] | None = None, + ) -> ExecutionReport: + """Build a zero-action report for a preflight rejection.""" + message = error if isinstance(error, str) else _error_message(error) + return self._empty_report( + SimpleNamespace(num_envs=max(1, int(environment_count))), + task_id=_task_id(grounded_plan, action_graph), + plan_hash=_plan_hash(grounded_plan), + graph_hash=_action_graph_hash(action_graph), + status="rejected", + run_id=run_id or _new_run_id(), + episode_index=episode_index, + provenance=build_execution_provenance( + episode_seed=episode_seed, + runtime_arguments=runtime_arguments, + ), + error=str(message), + ) + + def abortion_report( + self, + action_graph: Mapping[str, Any] | str | Path, + error: BaseException | str, + *, + grounded_plan: Mapping[str, Any] | None = None, + environment_count: int = 1, + run_id: str | None = None, + episode_index: int = 0, + episode_seed: int | None = None, + runtime_arguments: Mapping[str, Any] | None = None, + ) -> ExecutionReport: + """Build a zero-action report for an unexpected runtime exception.""" + message = error if isinstance(error, str) else _error_message(error) + return self._empty_report( + SimpleNamespace(num_envs=max(1, int(environment_count))), + task_id=_task_id(grounded_plan, action_graph), + plan_hash=_plan_hash(grounded_plan), + graph_hash=_action_graph_hash(action_graph), + status="aborted", + run_id=run_id or _new_run_id(), + episode_index=episode_index, + provenance=build_execution_provenance( + episode_seed=episode_seed, + runtime_arguments=runtime_arguments, + ), + error=str(message), + ) + + def _result_report( + self, + result: ExecutionResult, + *, + task_id: str, + plan_hash: str, + graph_hash: str, + run_id: str, + episode_index: int, + provenance: Mapping[str, Any], + ) -> ExecutionReport: + _persist_executed_trajectory(result) + success = _bool_vector(result.success) + semantics = { + str(step_id): _bool_vector(mask) + for step_id, mask in result.semantic_success.items() + } + failures = tuple(_json_safe(item) for item in result.failure_events) + revisions = tuple(_json_safe(item) for item in result.runtime_revisions) + action_count = len(result.actions) + environments = tuple( + { + "env_id": str(env_id), + "success": value, + "semantic_success": { + step_id: values[env_id] + for step_id, values in semantics.items() + if env_id < len(values) + }, + "action_count": action_count, + "retry_count": _retry_count_for_env(result, env_id), + "recovery_count": _revision_count_for_env( + revisions, env_id, kind="insert_recovery" + ), + "revision_count": _revision_count_for_env(revisions, env_id), + "failures": _events_for_env(failures, env_id), + } + for env_id, value in enumerate(success) + ) + report = ExecutionReport( + task_id=task_id, + plan_hash=plan_hash, + action_graph_hash=graph_hash, + status="succeeded" if all(success) else "failed", + run_id=run_id, + episode_id=str(episode_index), + provenance=deepcopy(dict(provenance)), + environments=environments, + action_count=action_count, + retry_count=int(result.retry_count), + recovery_count=int(result.recovery_count), + revision_count=int(result.revision_count), + failure_events=failures, + graph_revisions=revisions, + record_dir=result.record_dir, + error=None, + ) + validated = _validated_report(report) + _publish_execution_report(validated) + return validated + + def _empty_report( + self, + env: Any, + *, + task_id: str, + plan_hash: str, + graph_hash: str, + status: str, + run_id: str, + episode_index: int, + provenance: Mapping[str, Any], + error: str, + ) -> ExecutionReport: + count = _environment_count(env) + report = ExecutionReport( + task_id=task_id, + plan_hash=plan_hash, + action_graph_hash=graph_hash, + status=status, + run_id=run_id, + episode_id=str(episode_index), + provenance=deepcopy(dict(provenance)), + environments=tuple( + { + "env_id": str(env_id), + "success": False, + "semantic_success": {}, + "action_count": 0, + "retry_count": 0, + "recovery_count": 0, + "revision_count": 0, + "failures": [], + } + for env_id in range(count) + ), + error=error, + ) + return _validated_report(report) + + +def _validate_grounded_plan(value: Mapping[str, Any]) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise TypeError("GroundedTaskPlan must be a mapping.") + # GroundedTaskPlan is a cross-engine protocol owned by Task Engine. + # Import lazily so Action Engine remains importable without initializing + # the coordinator or Scene Adapter. + try: + from embodichain.gen_sim.task_engine.orchestration.contracts import ( + validate_grounded_task_plan, + ) + except (ImportError, AttributeError): + return deepcopy(dict(value)) + return validate_grounded_task_plan(value) + + +def _validated_report(report: ExecutionReport) -> ExecutionReport: + payload = report.as_mapping() + validate_execution_report(payload) + return report + + +def _persist_executed_trajectory(result: ExecutionResult) -> None: + """Persist every emitted control tensor beside the runtime graph audit.""" + if not result.record_dir: + return + root = Path(result.record_dir).expanduser().resolve() + root.mkdir(parents=True, exist_ok=True) + actions = [action.detach().cpu() for action in result.actions] + temporary = root / ".executed_trajectory.pt.tmp" + destination = root / "executed_trajectory.pt" + torch.save({"actions": actions}, temporary) + temporary.replace(destination) + manifest = { + "schema_version": "action_engine_executed_trajectory/v1", + "path": destination.name, + "action_count": len(actions), + "actions": [ + { + "index": index, + "shape": list(action.shape), + "dtype": str(action.dtype), + } + for index, action in enumerate(actions) + ], + } + manifest_path = root / "executed_trajectory.json" + manifest_path.write_text( + json.dumps(manifest, ensure_ascii=False, indent=2, allow_nan=False) + "\n", + encoding="utf-8", + ) + + +def _mapping(value: Any, label: str) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise ValueError(f"{label} must be a mapping.") + return deepcopy(dict(value)) + + +def _role_binding_map(value: Any) -> dict[str, str]: + source = _mapping(value, "GroundedTaskPlan.role_bindings") + nested = source.get("role_bindings") + if isinstance(nested, Mapping): + source = dict(nested) + result = {str(role): str(uid) for role, uid in source.items()} + if not result or any(not role or not uid for role, uid in result.items()): + raise ValueError("GroundedTaskPlan role bindings must not be empty.") + return result + + +def _known_uids( + manifest: Any, + *, + bindings: Mapping[str, str] | None = None, +) -> set[str]: + result = {str(uid) for uid in (bindings or {}).values() if str(uid)} + if not isinstance(manifest, Mapping): + return result + objects = manifest.get("objects", ()) + if isinstance(objects, Sequence) and not isinstance( + objects, (str, bytes, bytearray) + ): + for item in objects: + if isinstance(item, Mapping): + uid = item.get("uid", item.get("runtime_uid")) + if isinstance(uid, str) and uid: + result.add(uid) + return result + + +def _task_id( + plan: Mapping[str, Any] | None, + graph: Mapping[str, Any] | str | Path, +) -> str: + if isinstance(plan, Mapping): + value = plan.get("task_id") + if isinstance(value, str) and value: + return value + if isinstance(graph, Mapping): + value = graph.get("task_id") + if isinstance(value, str) and value: + return value + return "unknown_task" + + +def _plan_hash(plan: Mapping[str, Any] | None) -> str: + if isinstance(plan, Mapping): + hashes = plan.get("hashes", {}) + if isinstance(hashes, Mapping): + value = hashes.get("plan") + if isinstance(value, str) and value: + return value + return _safe_document_hash(plan) + return _document_hash({}) + + +def _action_graph_hash(value: Mapping[str, Any] | str | Path) -> str: + if isinstance(value, Mapping): + try: + return seed_graph_hash(value) + except (TypeError, ValueError): + return _safe_document_hash(value) + path = Path(value).expanduser() + try: + loaded = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return hashlib.sha256(str(path).encode("utf-8")).hexdigest() + return ( + _action_graph_hash(loaded) + if isinstance(loaded, Mapping) + else _safe_document_hash(loaded) + ) + + +def _safe_document_hash(value: Any) -> str: + try: + return _document_hash(value) + except (TypeError, ValueError, OverflowError): + return hashlib.sha256(repr(value).encode("utf-8")).hexdigest() + + +def _document_hash(value: Any) -> str: + payload = json.dumps( + _json_safe(value), + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def _bool_vector(value: Any) -> list[bool]: + if isinstance(value, torch.Tensor): + return [bool(item) for item in value.detach().cpu().reshape(-1).tolist()] + if isinstance(value, np.ndarray): + return [bool(item) for item in value.reshape(-1).tolist()] + if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + return [bool(item) for item in value] + return [bool(value)] + + +def _events_for_env( + events: Sequence[Mapping[str, Any]], env_id: int +) -> list[dict[str, Any]]: + result = [] + for event in events: + env_ids = event.get("env_ids") + if isinstance(env_ids, Sequence) and not isinstance( + env_ids, (str, bytes, bytearray) + ): + if env_id not in env_ids: + continue + item = deepcopy(dict(event)) + item["env_ids"] = [env_id] + result.append(item) + else: + result.append(deepcopy(dict(event))) + return result + + +def _revision_count_for_env( + revisions: Sequence[Mapping[str, Any]], + env_id: int, + *, + kind: str | None = None, +) -> int: + count = 0 + for revision in revisions: + if kind is not None and revision.get("kind") != kind: + continue + active = revision.get("active_env_ids") + if ( + isinstance(active, Sequence) + and not isinstance(active, (str, bytes, bytearray)) + and env_id not in active + ): + continue + count += 1 + return count + + +def _retry_count_for_env(result: ExecutionResult, env_id: int) -> int: + counts = result.retry_counts + if env_id < len(counts): + return int(counts[env_id]) + return int(result.retry_count) + + +def _environment_count(env: Any) -> int: + value = getattr(env, "num_envs", 1) + try: + return max(1, int(value)) + except (TypeError, ValueError): + return 1 + + +def _json_safe(value: Any) -> Any: + if isinstance(value, torch.Tensor): + return value.detach().cpu().tolist() + if isinstance(value, np.ndarray): + return value.tolist() + if isinstance(value, np.generic): + return value.item() + if isinstance(value, Path): + return value.as_posix() + if is_dataclass(value): + return _json_safe(asdict(value)) + if isinstance(value, Mapping): + return {str(key): _json_safe(item) for key, item in value.items()} + if isinstance(value, (list, tuple, set, frozenset)): + return [_json_safe(item) for item in value] + if value is None or isinstance(value, (str, int, float, bool)): + return value + return str(value) + + +def _new_run_id() -> str: + return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ") + + +def _error_message(exc: BaseException) -> str: + return f"{type(exc).__name__}: {exc}" + + +def _publish_execution_report(report: ExecutionReport) -> None: + """Atomically publish the compact report beside runtime episode records.""" + if not report.record_dir: + return + write_execution_report(report.record_dir, report) diff --git a/embodichain/gen_sim/action_engine/cli/run_agent.py b/embodichain/gen_sim/action_engine/cli/run_agent.py new file mode 100644 index 000000000..13cffe627 --- /dev/null +++ b/embodichain/gen_sim/action_engine/cli/run_agent.py @@ -0,0 +1,1620 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Run a generated Action Engine configuration.""" + +from __future__ import annotations + +import argparse +from collections.abc import Callable, Mapping +from copy import deepcopy +from dataclasses import dataclass, replace +from datetime import datetime, timezone +import json +import multiprocessing as mp +import os +from pathlib import Path +import shutil +from types import SimpleNamespace +from typing import Any + +import gymnasium +import numpy as np +import torch + +from embodichain.gen_sim.action_engine.config import generation_defaults +from embodichain.gen_sim.action_engine.environment import ( # noqa: F401 + ACTION_ENGINE_ENV_ID, +) +from embodichain.gen_sim.action_engine.runtime import ( + ExecutionReport, + load_agent_execution_program, + write_execution_report, +) +from embodichain.lab.gym.utils.gym_utils import ( + add_env_launcher_args_to_parser, + build_env_cfg_from_args, +) +from embodichain.utils import set_seed +from embodichain.utils.logger import log_info, log_warning +from embodichain.utils.utility import load_config + +__all__ = ["build_parser", "cli"] + +_DEFAULT_MAX_EPISODES = int(generation_defaults()["task"]["max_episodes"]) + + +def build_parser() -> argparse.ArgumentParser: + """Build the command-line parser used by generated demo commands.""" + parser = argparse.ArgumentParser(description="Execute an Action Engine task agent.") + add_env_launcher_args_to_parser(parser) + parser.add_argument("--task_name", required=True, help="Generated task name.") + parser.add_argument( + "--agent_config", + required=True, + help="Path to action_engine_config_v2 JSON.", + ) + parser.add_argument( + "--regenerate", + action="store_true", + help="Rebuild SeedGraph from TaskSpec in memory before execution.", + ) + parser.add_argument( + "--show-physical-collision", + action="store_true", + help="Show physical collision geometry after every reset.", + ) + parser.add_argument( + "--seed", + type=int, + default=None, + help="Base random seed; episode N uses seed + N.", + ) + parser.add_argument( + "--runtime-backend", + choices=("independent",), + default="independent", + help="Execution backend. Action Engine owns the production runtime.", + ) + parser.add_argument( + "--vlm-model", + default=None, + help="Optional runtime override for A/B visual facts and online planning.", + ) + parser.add_argument( + "--task-engine-report", + action="store_true", + help=argparse.SUPPRESS, + ) + return parser + + +def _validate_gym_id(config: dict[str, Any]) -> None: + if config.get("id") != ACTION_ENGINE_ENV_ID: + raise ValueError( + f"Gym config id must be {ACTION_ENGINE_ENV_ID!r}, " + f"got {config.get('id')!r}." + ) + + +def _validate_run_contract( + gym_config: dict[str, Any], + agent_config: dict[str, Any], + task_name: str, +) -> None: + """Validate the small cross-artifact contract before simulator startup.""" + configured_task = agent_config.get("task_name") + if configured_task != task_name: + raise ValueError( + f"--task_name {task_name!r} does not match agent_config task " + f"{configured_task!r}." + ) + extension = gym_config.get("env", {}).get("extensions", {}).get("action_engine", {}) + if extension.get("task_name") != task_name: + raise ValueError("Gym and agent configs describe different tasks.") + gym_hash = extension.get("seed_task_graph_hash") + agent_hash = agent_config.get("seed_task_graph_hash") + if not isinstance(agent_hash, str) or not agent_hash or gym_hash != agent_hash: + raise ValueError("Gym and agent configs have different program hashes.") + agent_mode = str(agent_config.get("planning_mode", "offline")) + gym_mode = str(extension.get("planning_mode", "offline")) + if agent_mode != gym_mode: + raise ValueError( + f"Gym and agent configs have different planning modes: " + f"gym={gym_mode!r}, agent={agent_mode!r}." + ) + + +def cli() -> int | None: + """Launch the environment and execute all configured episodes.""" + np.set_printoptions(precision=5, suppress=True) + torch.set_printoptions(precision=5, sci_mode=False) + args = build_parser().parse_args() + if args.seed is not None: + set_seed(args.seed) + env_cfg, gym_config, _ = build_env_cfg_from_args(args) + if args.seed is not None: + env_cfg.seed = args.seed + _validate_gym_id(gym_config) + agent_config = load_config(args.agent_config) + if not isinstance(agent_config, dict): + raise ValueError("agent_config must contain a JSON object.") + _validate_run_contract(gym_config, agent_config, args.task_name) + planning_mode = str(agent_config.get("planning_mode", "offline")) + if planning_mode == "ab": + _run_ab( + args, + env_cfg=env_cfg, + gym_config=gym_config, + agent_config=agent_config, + ) + return 0 if args.task_engine_report else None + if planning_mode != "offline": + raise ValueError(f"Unsupported Action Engine planning_mode {planning_mode!r}.") + execution_program = load_agent_execution_program( + agent_config, + agent_config_path=args.agent_config, + regenerate=bool(args.regenerate), + ) + grounded_plan = _load_grounded_task_plan(args.agent_config) + action_reporter = None + if grounded_plan is not None: + from embodichain.gen_sim.action_engine.agent import ActionAgent + + action_reporter = ActionAgent() + + run_id = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ") + episodes = int(gym_config.get("max_episodes", _DEFAULT_MAX_EPISODES)) + runtime_arguments = { + "agent_config": str(Path(args.agent_config).expanduser().resolve()), + "base_seed": args.seed, + "gym_config": str(Path(args.gym_config).expanduser().resolve()), + "max_episodes": episodes, + "planning_mode": planning_mode, + "regenerate": bool(args.regenerate), + "runtime_backend": str(args.runtime_backend), + "task_name": str(args.task_name), + } + any_failed = False + task_engine_reports: list[ExecutionReport] = [] + episode_index = 0 + episode_seed = None + seed_graph = getattr(execution_program, "seed_graph", None) + env = None + try: + env = gymnasium.make( + id=gym_config["id"], + cfg=env_cfg, + agent_config=agent_config, + agent_config_path=args.agent_config, + task_name=args.task_name, + runtime_backend=args.runtime_backend, + ) + for episode_index in range(episodes): + episode_seed = None if args.seed is None else int(args.seed) + episode_index + env.reset(seed=episode_seed) + if args.show_physical_collision: + _show_physical_collision(env) + execute = env.get_wrapper_attr("create_demo_action_list") + result = execute( + regenerate=bool(args.regenerate), + runtime_run_id=run_id, + episode_index=episode_index, + ) + if not getattr(result, "already_executed", False): + raise RuntimeError( + "Action Engine env returned an offline action sequence." + ) + success = torch.as_tensor( + getattr(result, "runtime_success"), + dtype=torch.bool, + ) + any_failed = any_failed or not bool(success.all()) + log_info( + "Action Engine episode " + f"{episode_index}: {int(success.sum())}/{success.numel()} " + "environments succeeded.", + color="green", + ) + record_dir = getattr(result, "runtime_graph_output_dir", None) + if record_dir: + log_info(f"Runtime records: {record_dir}", color="green") + if action_reporter is not None and isinstance(seed_graph, Mapping): + report = action_reporter.report_execution_result( + result, + action_graph=seed_graph, + grounded_plan=grounded_plan, + run_id=run_id, + episode_index=episode_index, + episode_seed=episode_seed, + runtime_arguments=runtime_arguments, + ) + task_engine_reports.append(report) + _publish_task_engine_report( + args.agent_config, + report, + enabled=bool(args.task_engine_report), + ) + log_info( + "Execution report: " + f"status={report.status}, actions={report.action_count}", + color="green" if report.status == "succeeded" else "yellow", + ) + # EmbodiedEnv publishes the just-finished rollout during reset. Flush + # the final episode as well; otherwise only episodes followed by a next + # iteration reach the configured dataset recorder. + env.reset(options={"final": True}) + except KeyboardInterrupt: + log_warning("Action Engine run interrupted by user.") + return 130 if args.task_engine_report else None + except Exception as exc: + if action_reporter is not None and isinstance(seed_graph, Mapping): + report = action_reporter.abortion_report( + seed_graph, + exc, + grounded_plan=grounded_plan, + environment_count=_runtime_environment_count(env), + run_id=run_id, + episode_index=episode_index, + episode_seed=episode_seed, + runtime_arguments=runtime_arguments, + ) + write_execution_report(Path(args.agent_config).resolve().parent, report) + if args.task_engine_report: + log_warning(f"Action Engine execution aborted: {type(exc).__name__}: {exc}") + return 3 + raise + finally: + close = getattr(env, "close", None) if env is not None else None + if callable(close): + close() + if not args.task_engine_report: + return None + return _task_engine_exit_code(any_failed, task_engine_reports) + + +def _publish_task_engine_report( + agent_config_path: str | Path, + report: ExecutionReport, + *, + enabled: bool, +) -> Path | None: + """Mirror one normal execution report into its Task Engine bundle.""" + if not enabled: + return None + return write_execution_report(Path(agent_config_path).resolve().parent, report) + + +def _task_engine_exit_code( + any_failed: bool, + reports: list[ExecutionReport], +) -> int: + """Return a report-authoritative exit code for Task Engine execution.""" + return int( + bool(any_failed) or any(report.status != "succeeded" for report in reports) + ) + + +def _load_grounded_task_plan(agent_config_path: str | Path) -> dict[str, Any] | None: + """Load the optional Task Engine hand-off beside an agent config.""" + path = ( + Path(agent_config_path).expanduser().resolve().parent + / "grounded_task_plan.json" + ) + if not path.is_file(): + return None + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError(f"Unable to read GroundedTaskPlan at {path}: {exc}") from exc + if not isinstance(value, Mapping): + raise ValueError("grounded_task_plan.json must contain a JSON object.") + from embodichain.gen_sim.task_engine.orchestration.contracts import ( + validate_grounded_task_plan, + ) + + return validate_grounded_task_plan(value) + + +def _runtime_environment_count(env: Any) -> int: + value = getattr(getattr(env, "unwrapped", env), "num_envs", 1) + try: + return max(1, int(value)) + except (TypeError, ValueError): + return 1 + + +class _BranchExecutor: + def __init__( + self, + graph: dict[str, Any], + env: gymnasium.Env, + *, + record_root: Path, + ) -> None: + self.graph = graph + self.env = env + self.record_root = record_root + + def preflight(self) -> bool: + """Compile and capability-check the branch without sending motion.""" + route = getattr(self.env.unwrapped, "action_engine_ab_route", None) + if route in {"offline", "online"} and self.graph.get("planner_route") != route: + raise ValueError( + f"A/B branch route {route!r} cannot execute graph route " + f"{self.graph.get('planner_route')!r}." + ) + try: + preflight = self.env.get_wrapper_attr("preflight_seed_graph") + except AttributeError: + preflight = None + if callable(preflight): + value = preflight(self.graph) + return value is not False + # Older generated environments expose only execute_seed_graph. The + # loader is still a useful structural/capability preflight and does + # not step the simulator. + from embodichain.gen_sim.action_engine.runtime import load_execution_program + + source = self.env.unwrapped.agent_config.get("source") + if source is None: + source = {} + if not isinstance(source, dict): + raise ValueError("agent_config.source must be a mapping when provided.") + uid_map = source.get("uid_map", {}) + if not isinstance(uid_map, dict): + raise ValueError( + "agent_config.source.uid_map must be a mapping when provided." + ) + known_objects = {str(uid) for uid in uid_map.values()} + load_execution_program(self.graph, known_objects=known_objects or None) + return True + + def run(self, *, run_id: str, episode_index: int) -> Any: + execute = self.env.get_wrapper_attr("execute_seed_graph") + return execute( + self.graph, + runtime_run_id=run_id, + episode_index=episode_index, + record_root=self.record_root.as_posix(), + ) + + +@dataclass(frozen=True) +class _ABWorkerConfig: + """Serializable startup contract for one process-isolated A/B branch.""" + + route: str + gym_config: dict[str, Any] + env_options: dict[str, Any] + gym_id: str + agent_config: dict[str, Any] + agent_config_path: str + task_name: str + runtime_backend: str + seed: int + camera_uids: tuple[str, ...] + staging_dir: str + + +class _ABBranchWorker: + """Small RPC proxy for one simulator process. + + DexSim entities resolve through a process-global default world. Keeping + each branch in a separate process is therefore a correctness requirement, + not merely a way to parallelize A/B execution. + """ + + _STARTUP_TIMEOUT_SECONDS = 300.0 + _COMMAND_TIMEOUT_SECONDS = 1800.0 + _SHUTDOWN_TIMEOUT_SECONDS = 30.0 + + def __init__(self, config: _ABWorkerConfig) -> None: + self.action_engine_ab_route = config.route + self._config = config + self._closed = False + self._context = mp.get_context("spawn") + self._connection, child_connection = self._context.Pipe(duplex=True) + self._process = self._context.Process( + target=_ab_worker_main, + args=(child_connection, config), + name=f"action-engine-ab-{config.route}", + ) + try: + self._process.start() + except BaseException: + child_connection.close() + self._connection.close() + raise + child_connection.close() + try: + startup = self._receive( + "startup", timeout_seconds=self._STARTUP_TIMEOUT_SECONDS + ) + except Exception: + self.close() + raise + if not isinstance(startup, dict): + self.close() + raise RuntimeError( + f"A/B {config.route} worker returned an invalid startup payload." + ) + snapshot = startup.get("snapshot") + if not isinstance(snapshot, dict): + self.close() + raise RuntimeError( + f"A/B {config.route} worker did not return its reset snapshot." + ) + self.startup_snapshot = snapshot + self.startup_observation = startup.get("observation") + + def snapshot(self) -> dict[str, Any]: + value = self._request("snapshot") + if not isinstance(value, dict): + raise RuntimeError( + f"A/B {self.action_engine_ab_route} worker returned an invalid snapshot." + ) + return value + + def preflight(self, graph: dict[str, Any]) -> bool: + value = self._request("preflight", graph=graph) + return value is not False + + def run( + self, + graph: dict[str, Any], + *, + run_id: str, + episode_index: int, + record_root: Path, + ) -> Any: + value = self._request( + "run", + graph=graph, + run_id=run_id, + episode_index=int(episode_index), + record_root=record_root.as_posix(), + ) + if not isinstance(value, dict): + raise RuntimeError( + f"A/B {self.action_engine_ab_route} worker returned an invalid result." + ) + return _execution_result_from_wire(value) + + def finalize(self, branch_dir: Path, *, episode_index: int) -> list[str]: + value = self._request( + "finalize", + branch_dir=branch_dir.as_posix(), + episode_index=int(episode_index), + ) + if not isinstance(value, list) or not all( + isinstance(path, str) and path for path in value + ): + raise RuntimeError( + f"A/B {self.action_engine_ab_route} worker returned invalid video paths." + ) + return value + + def close(self) -> None: + """Ask the worker to clean up, then force-stop only if it is stuck.""" + if self._closed: + return + self._closed = True + try: + if self._process.is_alive(): + try: + self._connection.send({"op": "shutdown"}) + self._receive( + "shutdown", timeout_seconds=self._SHUTDOWN_TIMEOUT_SECONDS + ) + except Exception: + # The process is still joined/terminated below. Cleanup + # errors cannot justify leaking a simulator child. + pass + finally: + try: + self._connection.close() + finally: + self._process.join(timeout=self._SHUTDOWN_TIMEOUT_SECONDS) + if self._process.is_alive(): + self._process.terminate() + self._process.join(timeout=self._SHUTDOWN_TIMEOUT_SECONDS) + + def _request(self, operation: str, **payload: Any) -> Any: + if self._closed: + raise RuntimeError( + f"A/B {self.action_engine_ab_route} worker is already closed." + ) + try: + self._connection.send({"op": operation, **payload}) + except (BrokenPipeError, EOFError, OSError) as exc: + raise RuntimeError( + f"A/B {self.action_engine_ab_route} worker could not receive " + f"{operation!r}." + ) from exc + return self._receive(operation, timeout_seconds=self._COMMAND_TIMEOUT_SECONDS) + + def _receive(self, operation: str, *, timeout_seconds: float) -> Any: + try: + ready = self._connection.poll(timeout_seconds) + except (EOFError, OSError) as exc: + raise RuntimeError( + f"A/B {self.action_engine_ab_route} worker closed during {operation}." + ) from exc + if not ready: + exit_code = self._process.exitcode + if exit_code is not None: + raise RuntimeError( + f"A/B {self.action_engine_ab_route} worker exited with code " + f"{exit_code} during {operation}." + ) + raise TimeoutError( + f"A/B {self.action_engine_ab_route} worker timed out during " + f"{operation} after {timeout_seconds:.0f}s." + ) + try: + response = self._connection.recv() + except (EOFError, OSError) as exc: + raise RuntimeError( + f"A/B {self.action_engine_ab_route} worker closed during {operation}." + ) from exc + if not isinstance(response, dict) or "ok" not in response: + raise RuntimeError( + f"A/B {self.action_engine_ab_route} worker returned a malformed " + f"response during {operation}." + ) + if response["ok"] is True: + return response.get("value") + message = response.get("error") + if not isinstance(message, str) or not message: + message = "unknown worker error" + raise RuntimeError( + f"A/B {self.action_engine_ab_route} worker failed during {operation}: " + f"{message}" + ) + + +class _SerializedABBranch: + """Run one branch in fresh isolated workers when two worlds do not fit. + + Each worker still owns a separate DexSim process and is reset from the + exact same seed. The proxy only serializes their GPU residency: it probes + a reset for planning, starts a fresh worker for preflight, then starts one + more fresh worker for execution. Every startup digest must match the + planning reset before an RPC is allowed to progress. + """ + + def __init__( + self, + config: _ABWorkerConfig, + *, + startup_snapshot: Mapping[str, Any], + startup_observation: Any, + expected_initial_state_digest: str, + worker_factory: Callable[[_ABWorkerConfig], Any] | None = None, + ) -> None: + self.action_engine_ab_route = config.route + self.startup_snapshot = deepcopy(dict(startup_snapshot)) + self.startup_observation = startup_observation + self._config = config + self._expected_initial_state_digest = expected_initial_state_digest + self._worker_factory = worker_factory or _ABBranchWorker + self._active_worker: Any | None = None + self._closed = False + + def snapshot(self) -> dict[str, Any]: + """Return the verified reset snapshot without rehydrating a GPU world.""" + return deepcopy(self.startup_snapshot) + + def preflight(self, graph: dict[str, Any]) -> bool: + worker = self._start_worker("preflight") + try: + return worker.preflight(graph) + finally: + worker.close() + + def run( + self, + graph: dict[str, Any], + *, + run_id: str, + episode_index: int, + record_root: Path, + ) -> Any: + if self._active_worker is not None: + raise RuntimeError( + f"A/B {self.action_engine_ab_route} execution worker is already active." + ) + worker = self._start_worker("execute") + self._active_worker = worker + return worker.run( + graph, + run_id=run_id, + episode_index=episode_index, + record_root=record_root, + ) + + def finalize(self, branch_dir: Path, *, episode_index: int) -> list[str]: + worker = self._active_worker + if worker is None: + raise RuntimeError( + f"A/B {self.action_engine_ab_route} has no execution worker to finalize." + ) + try: + return worker.finalize(branch_dir, episode_index=episode_index) + finally: + self._active_worker = None + worker.close() + + def close(self) -> None: + self._closed = True + worker = self._active_worker + self._active_worker = None + if worker is not None: + worker.close() + + def _start_worker(self, phase: str) -> Any: + if self._closed: + raise RuntimeError( + f"A/B {self.action_engine_ab_route} serialized branch is closed." + ) + worker = self._worker_factory(_ab_phase_worker_config(self._config, phase)) + try: + from embodichain.gen_sim.action_engine.evaluation import state_digest + + snapshot = worker.startup_snapshot + actual_digest = state_digest(snapshot) + if actual_digest != self._expected_initial_state_digest: + raise RuntimeError( + "Strict A/B serialized reset mismatch before " + f"{phase}: route={self.action_engine_ab_route}, " + f"expected={self._expected_initial_state_digest}, " + f"actual={actual_digest}." + ) + return worker + except BaseException: + worker.close() + raise + + +def _ab_phase_worker_config(config: _ABWorkerConfig, phase: str) -> _ABWorkerConfig: + """Give serial lifecycle phases distinct recorder and dataset roots.""" + if not phase: + return config + staging_dir = Path(config.staging_dir) + return replace( + config, + staging_dir=(staging_dir.parent / phase / staging_dir.name).as_posix(), + ) + + +def _prepare_ab_branches( + configs: Mapping[str, _ABWorkerConfig], + *, + worker_factory: Callable[[_ABWorkerConfig], Any] = _ABBranchWorker, + prefer_serial: bool | None = None, + gpu_id: int | None = None, +) -> tuple[dict[str, Any], dict[str, dict[str, Any]]]: + """Start concurrent worlds, with a digest-checked serialized fallback. + + A renderer can consume several GiB per DexSim process. On smaller GPUs, + starting the second isolated branch may fail before any action is sent. In + that case keeping one world resident is not a semantic requirement, while + the reset digest is; use fresh one-at-a-time workers instead. + """ + if prefer_serial is None: + prefer_serial = _prefer_serial_ab_startup(gpu_id=gpu_id) + if prefer_serial: + log_warning( + "A/B GPU capacity is below the concurrent-world budget; " + "using serialized isolated workers with reset-digest checks." + ) + return _prepare_serial_ab_branches(configs, worker_factory=worker_factory) + + workers: dict[str, Any] = {} + try: + for route in ("offline", "online"): + workers[route] = worker_factory(configs[route]) + except Exception as error: + for worker in workers.values(): + worker.close() + if not _is_gpu_memory_error(error): + raise + log_warning( + "A/B concurrent simulator startup exhausted GPU memory; " + "using serialized isolated workers with reset-digest checks." + ) + return _prepare_serial_ab_branches(configs, worker_factory=worker_factory) + return ( + workers, + {route: worker.startup_snapshot for route, worker in workers.items()}, + ) + + +def _prepare_serial_ab_branches( + configs: Mapping[str, _ABWorkerConfig], + *, + worker_factory: Callable[[_ABWorkerConfig], Any], +) -> tuple[dict[str, Any], dict[str, dict[str, Any]]]: + """Probe one branch at a time and return lazy serialized branch proxies.""" + + snapshots: dict[str, dict[str, Any]] = {} + observations: dict[str, Any] = {} + for route in ("offline", "online"): + worker = worker_factory(_ab_phase_worker_config(configs[route], "probe")) + try: + snapshots[route] = worker.startup_snapshot + observations[route] = worker.startup_observation + finally: + worker.close() + from embodichain.gen_sim.action_engine.evaluation import state_digest + + expected_digest = state_digest(snapshots["offline"]) + return ( + { + route: _SerializedABBranch( + configs[route], + startup_snapshot=snapshots[route], + startup_observation=observations[route], + expected_initial_state_digest=expected_digest, + worker_factory=worker_factory, + ) + for route in ("offline", "online") + }, + snapshots, + ) + + +def _is_gpu_memory_error(error: BaseException) -> bool: + """Recognize the process-startup failures where serialization is safe.""" + message = str(error).lower() + memory_markers = ( + "out of memory", + "out_of_memory", + "out_of_device_memory", + "outofmemory", + "resource exhausted", + ) + return any(marker in message for marker in memory_markers) and ( + "cuda" in message + or "gpu" in message + or "vulkan" in message + or "device" in message + ) + + +def _prefer_serial_ab_startup(*, gpu_id: int | None = None) -> bool: + """Avoid a known OOM trial on GPUs too small for two renderer worlds.""" + if not torch.cuda.is_available(): + return False + try: + device = torch.device(f"cuda:{int(gpu_id)}" if gpu_id is not None else "cuda") + free, _ = torch.cuda.mem_get_info(device=device) + except (RuntimeError, ValueError): + return False + # One hybrid DexSim world with the four VLM cameras can occupy roughly + # 11--13 GiB on the supported RTX setup. Reserve 24 GiB for two worlds; + # larger cards still attempt concurrent startup and retain the OOM fallback + # for unusually heavy scenes. + return int(free) < 24 * 1024**3 + + +class _RemoteBranchExecutor: + """Executor adapter which keeps simulator calls inside the branch worker.""" + + def __init__( + self, + graph: dict[str, Any], + worker: Any, + *, + record_root: Path, + ) -> None: + self.graph = graph + self.worker = worker + self.record_root = record_root + + def preflight(self) -> bool: + return self.worker.preflight(self.graph) + + def run(self, *, run_id: str, episode_index: int) -> Any: + return self.worker.run( + self.graph, + run_id=run_id, + episode_index=episode_index, + record_root=self.record_root, + ) + + +def _ab_worker_main(connection: Any, config: _ABWorkerConfig) -> None: + """Create and drive exactly one real environment in a child process.""" + # SimulationManager otherwise exits the whole worker with os._exit(0) + # during environment cleanup, bypassing the artifact/RPC shutdown contract. + os.environ["EMBODICHAIN_SIM_EXIT_PROCESS"] = "0" + env: gymnasium.Env | None = None + try: + from embodichain.lab.gym.utils.gym_utils import ( + config_to_cfg, + get_manager_modules, + ) + + # ``config_to_cfg`` creates local component-config classes which are + # intentionally not picklable. Send the merged JSON contract over IPC + # and reconstruct it inside each worker instead of pickling ``env_cfg``. + branch_cfg = config_to_cfg( + deepcopy(config.gym_config), + manager_modules=get_manager_modules(), + ) + _apply_ab_env_options(branch_cfg, config.env_options) + branch_cfg.seed = int(config.seed) + _configure_ab_branch_cfg( + branch_cfg, + staging_dir=Path(config.staging_dir), + dataset_dir=Path(config.staging_dir).parent / ".dataset", + ) + set_seed(int(config.seed)) + env = gymnasium.make( + id=config.gym_id, + cfg=branch_cfg, + agent_config=deepcopy(config.agent_config), + agent_config_path=config.agent_config_path, + task_name=config.task_name, + runtime_backend=config.runtime_backend, + ) + setattr(env.unwrapped, "action_engine_ab_route", config.route) + env.reset(seed=int(config.seed)) + startup: dict[str, Any] = { + "snapshot": _snapshot_environment(env, list(config.camera_uids)), + } + if config.route == "online": + from embodichain.gen_sim.action_engine.planning import ( + collect_scene_observation, + ) + + startup["observation"] = collect_scene_observation( + env.unwrapped, + camera_uids=config.camera_uids, + env_id=0, + ) + # The recorder normally receives its first frame from an interval + # event during ``env.step``. Capture one reset-time, no-motion frame so + # an execution branch that fails before its first action still has a + # valid video artifact after the mandatory final reset. + _capture_ab_initial_frame(env) + _worker_send(connection, ok=True, value=startup) + while True: + try: + request = connection.recv() + except EOFError: + break + if not isinstance(request, dict): + raise ValueError("A/B worker request must be a mapping.") + operation = request.get("op") + try: + if operation == "snapshot": + value = _snapshot_environment(env, list(config.camera_uids)) + elif operation == "preflight": + graph = _worker_graph(request.get("graph")) + value = _BranchExecutor( + graph, + env, + record_root=Path(config.staging_dir).parent / "runtime", + ).preflight() + elif operation == "run": + graph = _worker_graph(request.get("graph")) + run_id = request.get("run_id") + record_root = request.get("record_root") + if not isinstance(run_id, str) or not run_id: + raise ValueError( + "A/B worker run_id must be a non-empty string." + ) + if not isinstance(record_root, str) or not record_root: + raise ValueError( + "A/B worker record_root must be a non-empty path string." + ) + result = _BranchExecutor( + graph, + env, + record_root=Path(record_root).expanduser().resolve(), + ).run( + run_id=run_id, + episode_index=int(request.get("episode_index", 0)), + ) + value = _execution_result_to_wire(result) + elif operation == "finalize": + branch_dir = request.get("branch_dir") + if not isinstance(branch_dir, str) or not branch_dir: + raise ValueError( + "A/B worker branch_dir must be a non-empty path string." + ) + value = _finalize_ab_branch_video( + env, + staging_dir=Path(config.staging_dir), + branch_dir=Path(branch_dir).expanduser().resolve(), + ) + elif operation == "shutdown": + _worker_send(connection, ok=True, value=True) + break + else: + raise ValueError(f"Unknown A/B worker operation {operation!r}.") + except BaseException as exc: + _worker_send(connection, ok=False, error=_worker_error(exc)) + except BaseException as exc: + _worker_send(connection, ok=False, error=_worker_error(exc)) + finally: + if env is not None: + try: + env.close() + except BaseException: + pass + try: + from embodichain.lab.sim import SimulationManager + + SimulationManager.flush_cleanup_queue() + except BaseException: + pass + try: + connection.close() + except OSError: + pass + + +def _worker_graph(value: Any) -> dict[str, Any]: + if not isinstance(value, dict): + raise ValueError("A/B worker SeedGraph must be a JSON object.") + return value + + +def _worker_send( + connection: Any, + *, + ok: bool, + value: Any = None, + error: str | None = None, +) -> None: + try: + payload: dict[str, Any] = {"ok": bool(ok)} + if ok: + payload["value"] = value + else: + payload["error"] = error or "unknown worker error" + connection.send(payload) + except (BrokenPipeError, EOFError, OSError): + pass + + +def _worker_error(error: BaseException) -> str: + return f"{type(error).__name__}: {error}" + + +def _capture_ab_initial_frame(env: gymnasium.Env) -> None: + """Append one audience-camera frame without advancing simulation state.""" + base = env.unwrapped + manager = getattr(base, "event_manager", None) + mode_cfgs = getattr(manager, "_mode_functor_cfgs", {}) + candidates: list[tuple[Any, dict[str, Any]]] = [] + for configured in mode_cfgs.values(): + for functor_cfg in configured: + functor = _config_member(functor_cfg, "func") + class_name = getattr(type(functor), "__name__", "") + if not callable(functor) or class_name not in { + "record_camera_data", + "record_camera_data_async", + }: + continue + params = _config_member(functor_cfg, "params") or {} + if not isinstance(params, Mapping): + raise ValueError("A/B record_camera params must be a mapping.") + params = dict(params) + if params.get("name") == "record_cam_audience_view": + candidates.insert(0, (functor, params)) + else: + candidates.append((functor, params)) + if candidates: + # Prefer the explicitly generated audience recorder. A single + # unnamed legacy recorder remains a compatible fallback; selecting + # among multiple non-audience recorders would silently produce the + # wrong camera view, so fail instead. + if len(candidates) > 1 and candidates[0][1].get("name") != ( + "record_cam_audience_view" + ): + raise RuntimeError( + "A/B environment has multiple camera recorders but none is " + "named 'record_cam_audience_view'." + ) + functor, params = candidates[0] + functor(base, None, **params) + return + raise RuntimeError( + "A/B environment must define a record_camera_data audience recorder." + ) + + +def _execution_result_to_wire(result: Any) -> dict[str, Any]: + """Strip simulator-owned state from an execution result before IPC.""" + + actions = [_wire_tensor(action) for action in list(getattr(result, "actions", ()))] + success = _wire_tensor(getattr(result, "success", False)) + return { + "actions": actions, + "success": success, + "record_dir": getattr(result, "record_dir", None), + "already_executed": bool(getattr(result, "already_executed", True)), + "retry_count": int(getattr(result, "retry_count", 0)), + "recovery_count": int(getattr(result, "recovery_count", 0)), + "revision_count": int(getattr(result, "revision_count", 0)), + "failure_events": list(getattr(result, "failure_events", ())), + "runtime_revisions": list(getattr(result, "runtime_revisions", ())), + } + + +def _wire_tensor(value: Any) -> Any: + if isinstance(value, torch.Tensor): + return value.detach().cpu() + return value + + +def _execution_result_from_wire(value: dict[str, Any]) -> SimpleNamespace: + required = { + "actions", + "success", + "retry_count", + "recovery_count", + "revision_count", + "failure_events", + "runtime_revisions", + } + missing = sorted(required - set(value)) + if missing: + raise RuntimeError(f"A/B worker result is missing fields: {missing}.") + return SimpleNamespace(**value) + + +def _run_ab( + args: argparse.Namespace, + *, + env_cfg: Any, + gym_config: dict[str, Any], + agent_config: dict[str, Any], +) -> None: + """Plan and execute strict offline/online branches for every episode.""" + from embodichain.gen_sim.action_engine.evaluation import run_strict_ab, state_digest + from embodichain.gen_sim.action_engine.planning import ( + plan_candidates_parallel, + plan_online_seed_graph, + ) + from embodichain.gen_sim.action_engine.generation import VLM_CAMERA_UIDS + + config_path = Path(args.agent_config).expanduser().resolve() + task_path = _resolve_artifact_path( + agent_config, + config_path, + "task_spec", + "task_spec_path", + ) + task_spec = _read_json(task_path, "TaskSpec") + reference_program = load_agent_execution_program( + agent_config, + agent_config_path=config_path, + regenerate=bool(getattr(args, "regenerate", False)), + require_executable=False, + ) + if reference_program.seed_graph is None: + raise ValueError("A/B execution requires an immutable offline SeedGraph.") + reference_graph = reference_program.seed_graph + source = agent_config.get("source") + if not isinstance(source, dict): + source = {} + uid_map = source.get("uid_map", {}) + if not isinstance(uid_map, dict): + raise ValueError("agent_config.source.uid_map must be a mapping when provided.") + known_objects = {str(uid) for uid in uid_map.values() if str(uid)} + online_config = agent_config.get("online_planning", {}) + if online_config is None: + online_config = {} + if not isinstance(online_config, dict): + raise ValueError("agent_config.online_planning must be a mapping.") + camera_uids = online_config.get("camera_uids") or agent_config.get( + "vlm_camera_uids", [] + ) + if camera_uids != list(VLM_CAMERA_UIDS): + raise ValueError( + "A/B execution requires the canonical VLM cameras " + f"{list(VLM_CAMERA_UIDS)}." + ) + vlm_model = ( + getattr(args, "vlm_model", None) + or online_config.get("vlm_model") + or agent_config.get("vlm_model") + ) + robot_profile = str(agent_config.get("robot_profile", "dual_ur10")) + base_seed = 0 if args.seed is None else int(args.seed) + if args.seed is None: + set_seed(base_seed) + run_id = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ") + output_root = config_path.parent / "ab_runs" / run_id + episodes = int(gym_config.get("max_episodes", _DEFAULT_MAX_EPISODES)) + summaries = [] + env_options = _ab_env_options(env_cfg) + + for episode_index in range(episodes): + episode_seed = base_seed + episode_index + episode_root = output_root / f"episode_{episode_index:04d}" + branch_envs: dict[str, Any] = {} + ownership_transferred = False + try: + # The online VLM observes the same reset that its branch executes. + # Each branch owns one simulator process because DexSim entities + # resolve through a process-global default world. + worker_gym_config = _ab_runtime_gym_config(gym_config, env_cfg) + worker_configs = { + route: _ABWorkerConfig( + route=route, + gym_config=worker_gym_config, + env_options=deepcopy(env_options), + gym_id=str(gym_config["id"]), + agent_config=agent_config, + agent_config_path=config_path.as_posix(), + task_name=args.task_name, + runtime_backend=getattr(args, "runtime_backend", "independent"), + seed=episode_seed, + camera_uids=tuple(str(uid) for uid in camera_uids), + staging_dir=(episode_root / ".work" / route / "video").as_posix(), + ) + for route in ("offline", "online") + } + branch_envs, snapshots = _prepare_ab_branches( + worker_configs, + gpu_id=getattr(getattr(env_cfg, "sim_cfg", None), "gpu_id", None), + ) + planning_digest = state_digest(snapshots["offline"]) + if planning_digest != state_digest(snapshots["online"]): + raise RuntimeError( + "Strict A/B initial state mismatch before planning: " + f"offline={planning_digest}, " + f"online={state_digest(snapshots['online'])}." + ) + observation = branch_envs["online"].startup_observation + if observation is None: + raise RuntimeError( + "Online A/B worker did not return scene observation." + ) + visual_facts: dict[str, Any] = {} + + def offline_planner(*, task_spec: dict[str, Any]) -> dict[str, Any]: + # Generation has already materialized the fixed recipe from the + # shared TaskSpec. Reuse that immutable artifact verbatim so A/B + # also supports legacy v2 bundles whose graph metadata predates + # ``role_bindings``. + del task_spec + return deepcopy(reference_graph) + + def online_planner(*, task_spec: dict[str, Any]) -> dict[str, Any]: + graph, facts = plan_online_seed_graph( + task_spec, + observation, + vlm_model=vlm_model, + robot_profile=robot_profile, + ) + visual_facts.update(facts) + return graph + + candidates = plan_candidates_parallel( + task_spec, + offline_planner=offline_planner, + online_planner=online_planner, + known_objects=known_objects or None, + robot_profile=robot_profile, + ) + if candidates.offline != reference_graph: + from embodichain.gen_sim.action_engine.domain import seed_graph_hash + + if seed_graph_hash(candidates.offline) != seed_graph_hash( + reference_graph + ): + raise RuntimeError( + "A/B offline recipe no longer matches the generated " + "reference graph." + ) + + # Visual evidence is an online-planning artifact, not an execution + # artifact. Persist it as soon as both candidates have passed their + # static checks so it remains auditable even if a later preflight, + # runtime action, or video flush fails. + _write_json(episode_root / "online" / "visual_facts.json", visual_facts) + + # Rendering and external planning must be side-effect free. Take + # a second full snapshot immediately before executor preflight so + # an accidental simulation advance cannot be hidden behind the + # reset-time digest used for visual planning. + snapshots = { + route: worker.snapshot() for route, worker in branch_envs.items() + } + execution_digests = { + route: state_digest(snapshot) for route, snapshot in snapshots.items() + } + if ( + execution_digests["offline"] != planning_digest + or execution_digests["online"] != planning_digest + ): + raise RuntimeError( + "Strict A/B initial state changed during visual planning: " + f"offline={execution_digests['offline']}, " + f"online={execution_digests['online']}, " + f"expected={planning_digest}." + ) + + def executor_factory(graph: dict[str, Any], worker: Any) -> Any: + route = worker.action_engine_ab_route + if route not in branch_envs: + raise ValueError(f"Unknown A/B worker route {route!r}.") + return _RemoteBranchExecutor( + graph, + worker, + record_root=episode_root / ".work" / route / "runtime", + ) + + def branch_finalizer(**kwargs: Any) -> list[str]: + worker = kwargs["env"] + branch_dir = Path(kwargs["branch_dir"]) + return worker.finalize( + branch_dir, + episode_index=int(kwargs.get("episode_index", episode_index)), + ) + + result = run_strict_ab( + task_spec, + candidates.offline, + candidates.online, + executor_factory=executor_factory, + snapshot_reader=lambda env: _snapshot_environment(env, camera_uids), + output_dir=episode_root, + seed=episode_seed, + shared_config={ + "robot_profile": robot_profile, + "camera_uids": camera_uids, + "vlm_model": vlm_model, + "strict_state_digest": True, + }, + planning_metrics=candidates.planning_metrics, + known_objects=known_objects or None, + expected_initial_state_digest=planning_digest, + branch_finalizer=branch_finalizer, + episode_index=episode_index, + strict_state_digest=True, + prepared_environments=branch_envs, + prepared_snapshots=snapshots, + require_branch_videos=True, + ) + # run_strict_ab owns and closes prepared workers on both its normal + # and exceptional execution paths. Do not claim ownership until + # it has entered/returned from that cleanup boundary; this also + # closes workers when graph validation fails before its try/finally. + ownership_transferred = True + finally: + if not ownership_transferred: + for worker in branch_envs.values(): + worker.close() + summaries.append( + { + "episode_index": episode_index, + "seed": episode_seed, + "comparison": result.comparison_path.as_posix(), + "initial_state_digest": result.initial_state_digest, + } + ) + log_info( + "Action Engine A/B episode " + f"{episode_index}: offline=" + f"{result.comparison['branches']['offline']['success_rate']:.3f}, " + f"online={result.comparison['branches']['online']['success_rate']:.3f}.", + color="green", + ) + + summary_path = output_root / "run_summary.json" + _write_json( + summary_path, + { + "schema_version": "action_engine_ab_run_v1", + "task_id": args.task_name, + "run_id": run_id, + "episodes": summaries, + }, + ) + log_info(f"A/B comparison artifacts: {output_root}", color="green") + + +def _configure_ab_branch_cfg( + env_cfg: Any, + *, + staging_dir: Path, + dataset_dir: Path, +) -> None: + """Give one worker exclusive recorder paths before env construction.""" + staging_dir.mkdir(parents=True, exist_ok=True) + dataset_dir.mkdir(parents=True, exist_ok=True) + events = _config_member(env_cfg, "events") + recorder = _config_member(events, "record_camera") + if recorder is None: + raise ValueError("A/B environment config must define record_camera.") + _set_config_param(recorder, "save_path", staging_dir.as_posix()) + + # Dataset output is not part of the A/B contract, but leaving the + # generated path shared would still let two workers overwrite each other. + dataset = _config_member(env_cfg, "dataset") + if dataset is not None: + for name in ("lerobot", "record", "dataset"): + term = _config_member(dataset, name) + if term is not None: + _set_config_param(term, "save_path", dataset_dir.as_posix()) + + +def _ab_runtime_gym_config( + gym_config: Mapping[str, Any], env_cfg: Any +) -> dict[str, Any]: + """Carry launcher-resolved simulation settings into spawned workers.""" + result = deepcopy(dict(gym_config)) + sim_cfg = getattr(env_cfg, "sim_cfg", None) + if sim_cfg is None: + return result + result.update( + { + "device": str(getattr(sim_cfg, "sim_device", "cpu")), + "gpu_id": int(getattr(sim_cfg, "gpu_id", 0)), + "headless": bool(getattr(sim_cfg, "headless", False)), + "arena_space": float(getattr(sim_cfg, "arena_space", 5.0)), + "num_envs": int(getattr(sim_cfg, "num_envs", result.get("num_envs", 1))), + } + ) + render_cfg = getattr(sim_cfg, "render_cfg", None) + renderer = getattr(render_cfg, "renderer", None) + if renderer is not None: + result["renderer"] = str(renderer) + return result + + +def _ab_env_options(env_cfg: Any) -> dict[str, Any]: + """Extract the non-JSON flags applied after gym config parsing.""" + profiler = getattr(env_cfg, "profiler", None) + return { + "filter_visual_rand": bool(getattr(env_cfg, "filter_visual_rand", False)), + "filter_dataset_saving": bool(getattr(env_cfg, "filter_dataset_saving", False)), + "record_trajectory": bool(getattr(env_cfg, "record_trajectory", False)), + "trajectory_save_dir": getattr(env_cfg, "trajectory_save_dir", None), + "profile": bool(getattr(profiler, "enable_time", False)), + "profile_output": getattr(profiler, "output_path", None), + } + + +def _apply_ab_env_options(env_cfg: Any, options: Mapping[str, Any]) -> None: + """Apply launcher flags after reconstructing a worker's config.""" + env_cfg.filter_visual_rand = bool(options.get("filter_visual_rand", False)) + env_cfg.filter_dataset_saving = bool(options.get("filter_dataset_saving", False)) + env_cfg.record_trajectory = bool(options.get("record_trajectory", False)) + trajectory_dir = options.get("trajectory_save_dir") + if trajectory_dir: + env_cfg.trajectory_save_dir = str(trajectory_dir) + if bool(options.get("profile", False)): + from embodichain.lab.gym.utils.profiler import EnvProfilerCfg + + env_cfg.profiler = EnvProfilerCfg( + enable_time=True, + output_path=options.get("profile_output"), + ) + + +def _config_member(value: Any, name: str) -> Any: + if isinstance(value, dict): + return value.get(name) + return getattr(value, name, None) if value is not None else None + + +def _set_config_param(term: Any, name: str, value: Any) -> None: + params = _config_member(term, "params") + if params is None: + params = {} + if isinstance(term, dict): + term["params"] = params + else: + setattr(term, "params", params) + if not isinstance(params, dict): + raise ValueError( + f"A/B recorder params must be a mapping, got {type(params)!r}." + ) + params[name] = value + + +def _finalize_ab_branch_video( + env: gymnasium.Env, + *, + staging_dir: Path, + branch_dir: Path, +) -> list[str]: + """Flush the final episode and publish exactly this worker's video.""" + before = { + path: (path.stat().st_mtime_ns, path.stat().st_size) + for path in staging_dir.glob("episode_*_record_cam_audience_view.mp4") + if path.is_file() + } + env.reset(options={"final": True}) + video = _publish_branch_video(staging_dir, branch_dir, before=before) + return [video.as_posix()] + + +def _snapshot_environment( + env: gymnasium.Env, + camera_uids: list[str], +) -> dict[str, Any]: + base = env.unwrapped + sim = base.sim + object_poses = {} + for uid in sim.get_rigid_object_uid_list(): + entity = sim.get_rigid_object(uid) + if entity is not None: + object_poses[str(uid)] = _snapshot_tensor( + entity.get_local_pose(to_matrix=True) + ) + articulation_state = {} + for uid in getattr(sim, "get_articulation_uid_list", lambda: [])(): + entity = sim.get_articulation(uid) + if entity is None: + continue + articulation_state[str(uid)] = { + "pose": _snapshot_tensor(entity.get_local_pose(to_matrix=True)), + "qpos": _snapshot_tensor(entity.get_qpos()), + } + camera_calibration = {} + available_sensor_uids = getattr(sim, "get_sensor_uid_list", lambda: [])() + snapshot_camera_uids = sorted( + {str(uid) for uid in [*camera_uids, *available_sensor_uids] if str(uid)} + ) + for uid in snapshot_camera_uids: + sensor = sim.get_sensor(uid) + if sensor is None: + raise ValueError(f"A/B snapshot cannot find camera {uid!r}.") + camera_calibration[uid] = { + "intrinsics": _snapshot_sensor_value(sensor, "get_intrinsics"), + "extrinsics": _snapshot_sensor_value( + sensor, "get_arena_pose", to_matrix=True + ), + } + return { + "robot_qpos": _snapshot_tensor(base.robot.get_qpos()), + "object_poses": object_poses, + "articulation_state": articulation_state, + "camera_calibration": camera_calibration, + } + + +def _snapshot_tensor(value: Any) -> torch.Tensor: + """Normalize simulator values for deterministic digesting.""" + tensor = torch.as_tensor(value) + return tensor.detach().cpu().contiguous() + + +def _snapshot_sensor_value( + sensor: Any, + method_name: str, + **kwargs: Any, +) -> torch.Tensor: + method = getattr(sensor, method_name, None) + if not callable(method): + raise ValueError(f"A/B snapshot sensor lacks {method_name}().") + try: + value = method(**kwargs) + except TypeError: + value = method() + return _snapshot_tensor(value) + + +def _publish_branch_video( + staging_dir: Path, + branch_dir: Path, + *, + before: dict[Path, tuple[int, int]] | None = None, +) -> Path: + candidates = sorted( + ( + path + for path in staging_dir.glob("episode_*_record_cam_audience_view.mp4") + if path.is_file() + and ( + before is None + or path not in before + or (path.stat().st_mtime_ns, path.stat().st_size) != before[path] + ) + ), + key=lambda path: path.stat().st_mtime_ns, + ) + if not candidates: + raise RuntimeError(f"No completed A/B audience video found in {staging_dir}.") + source = candidates[-1] + destination = branch_dir / "video.mp4" + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source, destination) + if destination.stat().st_size == 0: + raise RuntimeError(f"A/B audience video is empty: {destination}.") + return destination + + +def _resolve_artifact_path( + config: dict[str, Any], + config_path: Path, + *keys: str, +) -> Path: + for key in keys: + value = config.get(key) + if value is None: + continue + if not isinstance(value, str) or not value: + raise ValueError(f"agent_config.{key} must be a non-empty path string.") + path = Path(value).expanduser() + return ( + path.resolve() + if path.is_absolute() + else (config_path.parent / path).resolve() + ) + joined = " or ".join(f"agent_config.{key}" for key in keys) + raise ValueError(f"A/B execution requires {joined}.") + + +def _read_json(path: Path, label: str) -> dict[str, Any]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError(f"Unable to read {label} at {path}: {exc}") from exc + if not isinstance(value, dict): + raise ValueError(f"{label} must contain a JSON object.") + return value + + +def _write_json(path: Path, value: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(value, ensure_ascii=False, indent=2, allow_nan=False) + "\n", + encoding="utf-8", + ) + + +def _show_physical_collision(env: gymnasium.Env) -> None: + """Enable physical-shape visualization for all supported scene assets.""" + sim = env.get_wrapper_attr("sim") + uids: list[str] = [] + for getter_name in ( + "get_rigid_object_uid_list", + "get_rigid_object_group_uid_list", + "get_articulation_uid_list", + ): + getter = getattr(sim, getter_name, None) + if callable(getter): + uids.extend(getter()) + visible = 0 + for uid in uids: + asset = sim.get_asset(uid) + if asset is None or not hasattr(asset, "set_physical_visible"): + continue + try: + asset.set_physical_visible( + visible=True, + rgba=[1.0, 0.15, 0.1, 0.35], + ) + visible += 1 + except Exception as exc: + log_warning(f"Unable to show collision geometry for {uid!r}: {exc}") + log_info(f"Physical collision geometry visible for {visible} assets.") + + +if __name__ == "__main__": + raise SystemExit(cli()) diff --git a/embodichain/gen_sim/action_engine/environment/__init__.py b/embodichain/gen_sim/action_engine/environment/__init__.py new file mode 100644 index 000000000..269a419c1 --- /dev/null +++ b/embodichain/gen_sim/action_engine/environment/__init__.py @@ -0,0 +1,23 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Tracked Action Engine environment package.""" + +from __future__ import annotations + +from .agent_env import ACTION_ENGINE_ENV_ID, ActionEngineEnv + +__all__ = ["ACTION_ENGINE_ENV_ID", "ActionEngineEnv"] diff --git a/embodichain/gen_sim/action_engine/environment/agent_env.py b/embodichain/gen_sim/action_engine/environment/agent_env.py new file mode 100644 index 000000000..ef53c8226 --- /dev/null +++ b/embodichain/gen_sim/action_engine/environment/agent_env.py @@ -0,0 +1,490 @@ +# ---------------------------------------------------------------------------- +# 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 environment that executes Action Engine programs against live state.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import torch + +from embodichain.gen_sim.action_engine.config import ( + RuntimePolicyCfg, + generation_defaults, + resolve_agent_runtime_policy, +) +from embodichain.gen_sim.action_engine.capabilities import ( + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.domain import validate_seed_graph +from embodichain.gen_sim.action_engine.protocol import ACTION_ENGINE_ENV_ID +from embodichain.gen_sim.action_engine.runtime import ( + ProgramExecutor, + evaluate_predicate, + load_agent_execution_program, + load_execution_program, +) +from embodichain.gen_sim.action_engine.runtime.solver_compat import ( + install_action_engine_solver_compat, + repair_action_engine_ur5_solver_cfg, +) +from embodichain.gen_sim.action_engine.runtime.motion_policy import ( + resolve_motion_policy, +) +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env + +__all__ = ["ACTION_ENGINE_ENV_ID", "ActionEngineEnv"] + +_MAX_EPISODE_STEPS = int(generation_defaults()["task"]["max_episode_steps"]) + + +@register_env(ACTION_ENGINE_ENV_ID, max_episode_steps=_MAX_EPISODE_STEPS) +class ActionEngineEnv(EmbodiedEnv): + """EmbodiedEnv adapter for in-memory compiled execution programs.""" + + def __init__(self, cfg: EmbodiedEnvCfg | None = None, **kwargs: Any) -> None: + agent_config = kwargs.pop("agent_config", None) + task_name = kwargs.pop("task_name", None) + agent_config_path = kwargs.pop("agent_config_path", None) + runtime_backend = kwargs.pop("runtime_backend", "independent") + runtime_policy = kwargs.pop("runtime_policy", None) + if not isinstance(agent_config, Mapping): + raise ValueError("ActionEngineEnv requires an agent_config mapping.") + if not isinstance(task_name, str) or not task_name: + raise ValueError("ActionEngineEnv requires a non-empty task_name.") + if not isinstance(agent_config_path, str) or not agent_config_path: + raise ValueError("ActionEngineEnv requires agent_config_path.") + self.agent_config = dict(agent_config) + self.agent_config_path = agent_config_path + self.task_name = task_name + if runtime_policy is None: + runtime_policy = resolve_agent_runtime_policy(self.agent_config) + if not isinstance(runtime_policy, RuntimePolicyCfg): + raise TypeError("ActionEngineEnv runtime_policy must be RuntimePolicyCfg.") + self.runtime_policy = runtime_policy + if runtime_backend != "independent": + raise ValueError( + "ActionEngineEnv only supports its independent runtime, got " + f"{runtime_backend!r}." + ) + self.runtime_backend = str(runtime_backend) + self.last_execution: Any | None = None + self._runtime_state_ready = False + repair_action_engine_ur5_solver_cfg(getattr(cfg, "robot", None)) + super().__init__(cfg, **kwargs) + install_action_engine_solver_compat(self.robot) + if bool(getattr(self, "ignore_terminations_during_agent", False)): + # Atomic trajectories execute online through env.step(). Prevent a + # transient task signal from resetting an environment mid-program. + self.cfg.ignore_terminations = True + self._capture_runtime_state() + + def reset( + self, + seed: int | None = None, + options: dict[str, Any] | None = None, + ) -> tuple[Any, dict[str, Any]]: + self._runtime_state_ready = False + observation, info = super().reset(seed=seed, options=options) + self.last_execution = None + self._capture_runtime_state() + return observation, info + + def _capture_runtime_state(self) -> None: + """Capture reset-relative robot and object state used by symbolic bindings.""" + self.init_qpos = self.robot.get_qpos().clone() + self._agent_arm_slots = self._resolve_arm_slots() + for side in ("left", "right"): + self._initialize_arm(side, self._agent_arm_slots.get(side)) + + default_open = getattr(self, "gripper_open_state", (0.04, 0.04)) + default_close = getattr(self, "gripper_close_state", (0.0, 0.0)) + self.open_state = torch.as_tensor( + getattr(self, "agent_open_state", default_open), + dtype=self.init_qpos.dtype, + device=self.init_qpos.device, + ).flatten() + self.close_state = torch.as_tensor( + getattr(self, "agent_close_state", default_close), + dtype=self.init_qpos.dtype, + device=self.init_qpos.device, + ).flatten() + self.left_arm_current_gripper_state = self._hand_qpos("left") + self.right_arm_current_gripper_state = self._hand_qpos("right") + self.update_obj_info() + self.agent_initial_object_poses = { + uid: item["pose"].clone() for uid, item in self.obj_info.items() + } + self.agent_initial_object_heights = { + uid: item["height"].clone() for uid, item in self.obj_info.items() + } + self._runtime_state_ready = True + + def _resolve_arm_slots(self) -> dict[str, dict[str, str | None] | None]: + configured = getattr(self, "agent_arm_slots", None) + if isinstance(configured, Mapping): + result: dict[str, dict[str, str | None] | None] = { + "left": None, + "right": None, + } + for side in result: + value = configured.get(side) + if isinstance(value, str): + result[side] = {"arm": value, "eef": None} + elif isinstance(value, Mapping): + result[side] = { + "arm": value.get("arm", value.get("arm_control_part")), + "eef": value.get( + "eef", + value.get("hand", value.get("eef_control_part")), + ), + } + return result + parts = getattr(self.robot, "control_parts", {}) or {} + if "left_arm" in parts or "right_arm" in parts: + return { + "left": {"arm": "left_arm", "eef": "left_eef"}, + "right": {"arm": "right_arm", "eef": "right_eef"}, + } + if "arm" in parts: + side = str(getattr(self, "agent_single_arm_slot", "right")) + result = {"left": None, "right": None} + result[side] = {"arm": "arm", "eef": "hand"} + return result + raise ValueError("Robot exposes no arm control part for Action Engine.") + + def _initialize_arm( + self, + side: str, + slot: dict[str, str | None] | None, + ) -> None: + arm = None if slot is None else slot.get("arm") + eef = None if slot is None else slot.get("eef") + arm_ids = self._control_part_ids(arm) + eef_ids = self._control_part_ids(eef) + setattr(self, f"{side}_arm_joints", arm_ids) + setattr(self, f"{side}_eef_joints", eef_ids) + arm_qpos = self.init_qpos[:, arm_ids] + setattr(self, f"{side}_arm_init_qpos", arm_qpos.clone()) + setattr(self, f"{side}_arm_current_qpos", arm_qpos.clone()) + if arm is None or not arm_ids: + setattr(self, f"{side}_arm_init_xpos", None) + setattr(self, f"{side}_arm_current_xpos", None) + return + xpos = self.robot.compute_fk(arm_qpos, name=arm, to_matrix=True) + setattr(self, f"{side}_arm_init_xpos", xpos.clone()) + setattr(self, f"{side}_arm_current_xpos", xpos.clone()) + + def _control_part_ids(self, name: str | None) -> list[int]: + if name is None: + return [] + parts = getattr(self.robot, "control_parts", {}) or {} + if name not in parts: + return [] + return list(self.robot.get_joint_ids(name=name)) + + def _hand_qpos(self, side: str) -> torch.Tensor: + ids = list(getattr(self, f"{side}_eef_joints", ())) + return self.init_qpos[:, ids].clone() + + def get_agent_arm_control_part(self, is_left: bool) -> str: + value = self._agent_arm_slots["left" if is_left else "right"] + arm = None if value is None else value.get("arm") + if not isinstance(arm, str) or not arm: + raise ValueError(f"{'left' if is_left else 'right'} arm is not configured.") + return arm + + def get_agent_eef_control_part(self, is_left: bool) -> str | None: + value = self._agent_arm_slots["left" if is_left else "right"] + eef = None if value is None else value.get("eef") + return str(eef) if eef else None + + def get_current_qpos_agent(self) -> tuple[torch.Tensor, torch.Tensor]: + qpos = self.robot.get_qpos() + return tuple( + qpos[:, list(getattr(self, f"{side}_arm_joints", ()))].clone() + for side in ("left", "right") + ) + + def set_current_qpos_agent( + self, + arm_qpos: torch.Tensor, + is_left: bool, + ) -> None: + side = "left" if is_left else "right" + setattr(self, f"{side}_arm_current_qpos", arm_qpos) + + def get_current_xpos_agent( + self, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + qpos = self.robot.get_qpos() + result = [] + for side in ("left", "right"): + slot = self._agent_arm_slots.get(side) + arm = None if slot is None else slot.get("arm") + arm_ids = list(getattr(self, f"{side}_arm_joints", ())) + if not arm or not arm_ids: + result.append(None) + continue + result.append( + self.robot.compute_fk( + qpos[:, arm_ids], + name=arm, + to_matrix=True, + ) + ) + return result[0], result[1] + + def set_current_xpos_agent( + self, + arm_xpos: torch.Tensor, + is_left: bool, + ) -> None: + side = "left" if is_left else "right" + setattr(self, f"{side}_arm_current_xpos", arm_xpos) + + def get_current_gripper_state_agent( + self, + ) -> tuple[torch.Tensor, torch.Tensor]: + qpos = self.robot.get_qpos() + return tuple( + qpos[:, list(getattr(self, f"{side}_eef_joints", ()))].clone() + for side in ("left", "right") + ) + + def set_current_gripper_state_agent( + self, + arm_gripper_state: torch.Tensor, + is_left: bool, + ) -> None: + side = "left" if is_left else "right" + setattr(self, f"{side}_arm_current_gripper_state", arm_gripper_state) + + def get_arm_fk(self, qpos: torch.Tensor, is_left: bool) -> torch.Tensor: + return self.robot.compute_fk( + name=self.get_agent_arm_control_part(is_left), + qpos=torch.as_tensor(qpos, device=self.robot.device), + to_matrix=True, + ) + + def sync_agent_state_from_qpos(self, qpos: torch.Tensor) -> None: + """Keep arm-selection seeds synchronized with the command sent to sim.""" + qpos = torch.as_tensor( + qpos, + dtype=self.init_qpos.dtype, + device=self.init_qpos.device, + ) + for side in ("left", "right"): + arm_ids = list(getattr(self, f"{side}_arm_joints", ())) + hand_ids = list(getattr(self, f"{side}_eef_joints", ())) + arm_qpos = qpos[:, arm_ids] + setattr(self, f"{side}_arm_current_qpos", arm_qpos.clone()) + slot = self._agent_arm_slots.get(side) + arm = None if slot is None else slot.get("arm") + if arm and arm_ids: + xpos = self.robot.compute_fk(arm_qpos, name=arm, to_matrix=True) + setattr(self, f"{side}_arm_current_xpos", xpos) + setattr( + self, + f"{side}_arm_current_gripper_state", + qpos[:, hand_ids].clone(), + ) + + def get_arm_ik( + self, + target_xpos: torch.Tensor, + is_left: bool, + qpos_seed: torch.Tensor | None = None, + env_ids: list[int] | None = None, + ) -> tuple[bool, torch.Tensor]: + success, qpos = self.robot.compute_ik( + name=self.get_agent_arm_control_part(is_left), + pose=target_xpos, + joint_seed=qpos_seed, + env_ids=env_ids, + ) + success_value = ( + bool(torch.as_tensor(success).all().item()) + if isinstance(success, torch.Tensor) + else bool(success) + ) + return success_value, qpos + + def update_obj_info(self) -> None: + info = getattr(self, "obj_info", {}) + for uid in self.sim.get_rigid_object_uid_list(): + entity = self.sim.get_rigid_object(uid) + if entity is None: + continue + pose = entity.get_local_pose(to_matrix=True) + info[uid] = {"pose": pose, "height": pose[:, 2, 3]} + self.obj_info = info + + def create_demo_action_list( + self, + regenerate: bool = False, + **kwargs: Any, + ) -> Any: + """Compile in memory when requested, then execute the program online.""" + program = load_agent_execution_program( + self.agent_config, + agent_config_path=self.agent_config_path, + regenerate=regenerate, + ) + executor = ProgramExecutor( + program, + self, + max_transitions=( + int(self.action_engine_max_transitions) + if hasattr(self, "action_engine_max_transitions") + else None + ), + settle_steps=( + int(self.action_engine_settle_steps) + if hasattr(self, "action_engine_settle_steps") + else None + ), + record_runtime=bool(getattr(self, "action_engine_record_runtime", True)), + record_root=getattr(self, "action_engine_record_root", None), + runtime_policy=self.runtime_policy, + ) + self.last_execution = executor.run( + run_id=kwargs.get("runtime_run_id"), + episode_index=int(kwargs.get("episode_index", 0)), + ) + return self.last_execution + + def execute_seed_graph( + self, + seed_graph: Mapping[str, Any], + *, + runtime_run_id: str, + episode_index: int, + record_root: str | None = None, + ) -> Any: + """Execute one already validated branch graph without rewriting config.""" + program = self.preflight_seed_graph(seed_graph) + route = getattr(self, "action_engine_ab_route", None) + graph_route = seed_graph.get("planner_route") + if route in {"offline", "online"} and graph_route != route: + raise ValueError( + f"A/B branch route {route!r} cannot execute graph route " + f"{graph_route!r}." + ) + executor = ProgramExecutor( + program, + self, + max_transitions=( + int(self.action_engine_max_transitions) + if hasattr(self, "action_engine_max_transitions") + else None + ), + settle_steps=( + int(self.action_engine_settle_steps) + if hasattr(self, "action_engine_settle_steps") + else None + ), + record_runtime=bool(getattr(self, "action_engine_record_runtime", True)), + record_root=record_root, + runtime_policy=self.runtime_policy, + ) + self.last_execution = executor.run( + run_id=runtime_run_id, + episode_index=episode_index, + ) + return self.last_execution + + def preflight_seed_graph(self, seed_graph: Mapping[str, Any]) -> Any: + """Validate/compile one branch graph without stepping the simulator. + + This hook is intentionally separate from :meth:`execute_seed_graph` so + strict A/B can preflight both branches before either executor sends a + command to the robot. + """ + source = self.agent_config.get("source", {}) + if not isinstance(source, Mapping): + source = {} + uid_map = source.get("uid_map", {}) + if not isinstance(uid_map, Mapping): + uid_map = {} + known_objects = {str(uid) for uid in uid_map.values() if str(uid)} + registry = build_atomic_capability_registry() + graph = validate_seed_graph( + seed_graph, + known_objects=known_objects or None, + known_actions=registry.names(), + executable_actions=registry.executable_names(), + require_executable=True, + ) + for node in graph["nodes"]: + registry.validate_binding(node) + resolve_motion_policy( + str( + self.agent_config.get( + "robot_profile", + getattr(self, "agent_robot_profile", "dual_ur10"), + ) + ), + node["motion_policy"], + ) + return load_execution_program( + graph, + known_objects=known_objects or None, + registry=registry, + ) + + def _normalize_demo_action_list(self, action_list: Any) -> Any: + """Preserve metadata on action streams that already ran online. + + ``EmbodiedEnv`` normally rebuilds returned sequences after validating + their action width. Rebuilding an ``ExecutionResult`` would discard its + success masks and runtime-record location, and its commands have + already been sent to the simulator, so no replay normalization is + needed. + """ + if getattr(action_list, "already_executed", False): + return action_list + return super()._normalize_demo_action_list(action_list) + + def is_task_success(self, **_: Any) -> torch.Tensor: + configured = getattr(self, "agent_success", None) + if isinstance(configured, Mapping): + return evaluate_predicate(self, configured) + if self.last_execution is not None: + return torch.as_tensor( + getattr( + self.last_execution, + "runtime_success", + getattr(self.last_execution, "success", False), + ), + dtype=torch.bool, + device=self.device, + ) + return torch.zeros( + int(self.num_envs), + dtype=torch.bool, + device=self.device, + ) + + def compute_task_state( + self, + **_: Any, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + success = self.is_task_success() + return success, torch.zeros_like(success), {} diff --git a/embodichain/gen_sim/action_engine/evaluation/__init__.py b/embodichain/gen_sim/action_engine/evaluation/__init__.py new file mode 100644 index 000000000..2d2c15190 --- /dev/null +++ b/embodichain/gen_sim/action_engine/evaluation/__init__.py @@ -0,0 +1,27 @@ +# ---------------------------------------------------------------------------- +# 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 offline/online comparison utilities.""" + +from __future__ import annotations + +from .ab import ABExecutionResult, run_strict_ab, state_digest + +__all__ = [ + "ABExecutionResult", + "run_strict_ab", + "state_digest", +] diff --git a/embodichain/gen_sim/action_engine/evaluation/ab.py b/embodichain/gen_sim/action_engine/evaluation/ab.py new file mode 100644 index 000000000..1c95a61ea --- /dev/null +++ b/embodichain/gen_sim/action_engine/evaluation/ab.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. +# ---------------------------------------------------------------------------- + +"""Execute offline and online SeedGraphs from strictly identical resets.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from dataclasses import dataclass +import hashlib +import json +from pathlib import Path +from time import perf_counter +from typing import Any + +import numpy as np +import torch + +from embodichain.gen_sim.action_engine.domain import ( + seed_graph_hash, + validate_seed_graph, + validate_task_spec, +) +from embodichain.gen_sim.action_engine.capabilities import ( + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.protocol import ( + COMPARISON_FILENAME, + EXECUTION_PROGRAM_FILENAME, +) +from embodichain.gen_sim.action_engine.runtime.motion_policy import ( + resolve_motion_policy, +) + +__all__ = ["ABExecutionResult", "run_strict_ab", "state_digest"] + +EnvFactory = Callable[..., Any] +ExecutorFactory = Callable[[Mapping[str, Any], Any], Any] +SnapshotReader = Callable[[Any], Mapping[str, Any]] +SuccessEvaluator = Callable[..., Any] +BranchFinalizer = Callable[..., list[str]] + +_FULL_SNAPSHOT_KEYS = frozenset( + {"robot_qpos", "object_poses", "articulation_state", "camera_calibration"} +) +_PRIVATE_OR_LIVE_KEYS = frozenset( + { + "absolute_position", + "coordinates", + "grasp_pose", + "joint_positions", + "live_pose", + "live_transform", + "oracle", + "pose", + "positions", + "qpos", + "target_pose", + "trajectory", + "waypoints", + "xpos", + } +) + + +@dataclass(frozen=True) +class ABExecutionResult: + """Paths and summaries from one strict A/B run.""" + + comparison_path: Path + offline_dir: Path + online_dir: Path + initial_state_digest: str + comparison: dict[str, Any] + + +def state_digest(snapshot: Mapping[str, Any]) -> str: + """Hash nested tensors/arrays/mappings without lossy JSON conversion.""" + digest = hashlib.sha256() + _update_digest(digest, snapshot) + return digest.hexdigest() + + +def run_strict_ab( + task_spec: Mapping[str, Any], + offline_graph: Mapping[str, Any], + online_graph: Mapping[str, Any], + *, + env_factory: EnvFactory | None = None, + executor_factory: ExecutorFactory, + snapshot_reader: SnapshotReader, + output_dir: str | Path, + seed: int, + shared_config: Mapping[str, Any] | None = None, + planning_metrics: Mapping[str, Mapping[str, Any]] | None = None, + success_evaluator: SuccessEvaluator | None = None, + known_objects: set[str] | None = None, + expected_initial_state_digest: str | None = None, + branch_finalizer: BranchFinalizer | None = None, + episode_index: int = 0, + strict_state_digest: bool | None = None, + prepared_environments: Mapping[str, Any] | None = None, + prepared_snapshots: Mapping[str, Mapping[str, Any]] | None = None, + require_branch_videos: bool = False, +) -> ABExecutionResult: + """Run both planners in isolated environments after exact state checks. + + Callers that need visual observations before planning may supply two already + reset environments and their snapshots. The environments remain owned by + this function once supplied and are closed on every exit path. + + Set ``require_branch_videos`` for production A/B runs. In that mode each + finalizer must publish one non-empty ``video.mp4`` in its branch directory. + """ + supplied_environments = ( + tuple(prepared_environments.values()) + if prepared_environments is not None + else () + ) + try: + task, config, offline, online = _validate_ab_inputs( + task_spec, + offline_graph, + online_graph, + shared_config=shared_config, + success_evaluator=success_evaluator, + known_objects=known_objects, + ) + except BaseException: + # Prepared environments are already live before graph validation. The + # caller transfers ownership at function entry, including invalid-input + # paths that return before the normal environment scope below. + _close_environments(supplied_environments) + raise + + metrics = dict(planning_metrics or {}) + environments: dict[str, Any] = {} + try: + routes = ("offline", "online") + if prepared_environments is not None: + # Keep every supplied object in ``environments`` until after shape + # validation so the finally block closes extras on this error path. + environments = dict(prepared_environments) + if set(environments) != set(routes): + raise ValueError( + "prepared_environments must contain exactly offline and online." + ) + environments = {route: prepared_environments[route] for route in routes} + else: + if not callable(env_factory): + raise TypeError( + "env_factory is required when prepared_environments is not supplied." + ) + for route in routes: + environments[route] = env_factory( + route=route, + seed=int(seed), + config=config, + ) + if id(environments["offline"]) == id(environments["online"]): + raise RuntimeError( + "Strict A/B requires two isolated environment instances; " + "env_factory returned the same object twice." + ) + for route, env in environments.items(): + marker = getattr(env, "action_engine_ab_route", None) + if marker is not None and str(marker) != route: + raise RuntimeError( + f"A/B environment route marker {marker!r} does not match {route!r}." + ) + snapshots: dict[str, Mapping[str, Any]] = {} + if prepared_snapshots is not None and prepared_environments is None: + raise ValueError( + "prepared_snapshots requires prepared_environments so the state " + "being compared is unambiguous." + ) + if prepared_snapshots is not None: + if set(prepared_snapshots) != set(routes): + raise ValueError( + "prepared_snapshots must contain exactly offline and online." + ) + snapshots = {route: prepared_snapshots[route] for route in routes} + digests = {} + for route, env in environments.items(): + if prepared_snapshots is None: + env.reset(seed=int(seed)) + snapshots[route] = snapshot_reader(env) + _validate_snapshot(snapshots[route], route=route, require_full=False) + if strict_state_digest is None: + strict_state_digest = bool(config.get("strict_state_digest", False)) + # Automatically enforce the expanded contract whenever a caller + # supplies any of the new state components, while retaining the + # two-field v1 test helper compatibility. + strict_state_digest = strict_state_digest or any( + set(snapshot) & (_FULL_SNAPSHOT_KEYS - {"robot_qpos", "object_poses"}) + for snapshot in snapshots.values() + ) + if strict_state_digest: + for route, snapshot in snapshots.items(): + _validate_snapshot(snapshot, route=route, require_full=True) + for route, snapshot in snapshots.items(): + digests[route] = state_digest(snapshots[route]) + if digests["offline"] != digests["online"]: + raise RuntimeError( + "Strict A/B initial state mismatch: " + f"offline={digests['offline']}, online={digests['online']}." + ) + if ( + expected_initial_state_digest is not None + and digests["offline"] != expected_initial_state_digest + ): + raise RuntimeError( + "Strict A/B execution state does not match the online-planning " + f"snapshot: planning={expected_initial_state_digest}, " + f"execution={digests['offline']}." + ) + + root = Path(output_dir).expanduser().resolve() + branch_dirs = {route: root / route for route in environments} + for branch_dir in branch_dirs.values(): + branch_dir.mkdir(parents=True, exist_ok=True) + # Construct and preflight both executors before invoking either run. + # A route-specific executor may perform capability/robot checks that + # cannot be expressed in the serializable SeedGraph validator. + executors: dict[str, Any] = {} + for route, graph in (("offline", offline), ("online", online)): + _write_json(branch_dirs[route] / EXECUTION_PROGRAM_FILENAME, graph) + executors[route] = executor_factory(graph, environments[route]) + preflight_errors: dict[str, Exception] = {} + for route, executor in executors.items(): + preflight = getattr(executor, "preflight", None) + if not callable(preflight): + preflight = getattr(executor, "validate", None) + if not callable(preflight): + continue + try: + outcome = _call_preflight( + preflight, + route=route, + graph=(offline if route == "offline" else online), + env=environments[route], + ) + if outcome is not None: + try: + preflight_ok = bool(outcome) + except (TypeError, ValueError, RuntimeError) as exc: + raise RuntimeError( + f"{route} executor preflight returned a non-scalar result." + ) from exc + if not preflight_ok: + raise RuntimeError( + f"{route} executor preflight returned false." + ) + except Exception as exc: + preflight_errors[route] = exc + if preflight_errors: + detail = "; ".join( + f"{route}: {type(error).__name__}: {error}" + for route, error in sorted(preflight_errors.items()) + ) + raise RuntimeError( + "Strict A/B preflight failed; no branch was allowed to move. " + detail + ) + results = {} + finalization_errors: dict[str, Exception] = {} + for route, graph in (("offline", offline), ("online", online)): + result = None + started = perf_counter() + try: + executor = executors[route] + result = executor.run( + run_id=f"ab-{seed}-{route}", + episode_index=episode_index, + ) + elapsed = perf_counter() - started + success_override = ( + success_evaluator( + task_spec=task, + graph=graph, + env=environments[route], + result=result, + route=route, + ) + if success_evaluator is not None + else None + ) + results[route] = _result_summary( + result, + elapsed, + graph, + metrics.get(route, {}), + success_override=success_override, + ) + except Exception as exc: + elapsed = perf_counter() - started + results[route] = _error_result_summary( + exc, + elapsed, + graph, + metrics.get(route, {}), + ) + try: + raw_video_paths = ( + branch_finalizer( + route=route, + env=environments[route], + result=result, + branch_dir=branch_dirs[route], + episode_index=episode_index, + ) + if branch_finalizer is not None + else list(getattr(result, "video_paths", ())) + ) + video_paths = [str(path) for path in raw_video_paths] + if require_branch_videos: + _validate_branch_video_paths( + video_paths, + route=route, + branch_dir=branch_dirs[route], + ) + except Exception as exc: + video_paths = [] + results[route]["video_error"] = f"{type(exc).__name__}: {exc}" + finalization_errors[route] = exc + results[route]["video_paths"] = video_paths + results[route]["initial_state_digest"] = digests[route] + results[route]["seed_graph_hash"] = seed_graph_hash(graph) + _write_json( + branch_dirs[route] / "runtime_revisions.json", + { + "schema_version": "action_engine_runtime_revisions_v1", + "task_id": task["task_id"], + "route": route, + "revisions": list(getattr(result, "runtime_revisions", ())), + }, + ) + _write_json(branch_dirs[route] / "result.json", results[route]) + + comparison = { + "schema_version": "action_engine_ab_comparison_v1", + "task_id": task["task_id"], + "seed": int(seed), + "shared_config": config, + "initial_state_digest": digests["offline"], + "initial_state_digests": dict(digests), + "strict_state_digest": bool(strict_state_digest), + "graph_hashes": { + "offline": seed_graph_hash(offline), + "online": seed_graph_hash(online), + }, + "branches": { + "offline": { + **results["offline"], + }, + "online": { + **results["online"], + }, + }, + "graph_difference": _graph_difference(offline, online), + "video_finalization_errors": { + route: f"{type(error).__name__}: {error}" + for route, error in sorted(finalization_errors.items()) + }, + } + comparison_path = root / COMPARISON_FILENAME + _write_json(comparison_path, comparison) + if finalization_errors: + detail = "; ".join( + f"{route}: {type(error).__name__}: {error}" + for route, error in sorted(finalization_errors.items()) + ) + first_error = next(iter(finalization_errors.values())) + raise RuntimeError( + "Strict A/B branch video finalization failed; comparison report " + "was written with artifact errors. " + detail + ) from first_error + return ABExecutionResult( + comparison_path, + branch_dirs["offline"], + branch_dirs["online"], + digests["offline"], + comparison, + ) + finally: + _close_environments(environments.values()) + + +def _validate_ab_inputs( + task_spec: Mapping[str, Any], + offline_graph: Mapping[str, Any], + online_graph: Mapping[str, Any], + *, + shared_config: Mapping[str, Any] | None, + success_evaluator: SuccessEvaluator | None, + known_objects: set[str] | None, +) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any], dict[str, Any]]: + """Validate every serializable input before either branch may move.""" + task = validate_task_spec(task_spec) + if task["level"] == "L4" and not callable(success_evaluator): + raise ValueError( + "Strict L4 A/B requires a path-independent private-oracle " + "success_evaluator." + ) + config = dict(shared_config or {}) + capabilities = build_atomic_capability_registry() + offline = validate_seed_graph( + offline_graph, + known_objects=known_objects, + known_actions=capabilities.names(), + executable_actions=capabilities.executable_names(), + require_executable=True, + ) + online = validate_seed_graph( + online_graph, + known_objects=known_objects, + known_actions=capabilities.names(), + executable_actions=capabilities.executable_names(), + require_executable=True, + ) + robot_profile = str(config.get("robot_profile", "dual_ur10")) + from embodichain.gen_sim.action_engine.planning.linker import ( + validate_persisted_contracts, + ) + + for graph in (offline, online): + if graph["capability_catalog_hash"] != capabilities.catalog_hash(): + raise ValueError("A/B SeedGraph capability catalog does not match runtime.") + validate_persisted_contracts(graph, capabilities) + for node in graph["nodes"]: + capabilities.validate_binding(node) + resolve_motion_policy( + robot_profile, + node["atomic_action"], + node["motion_policy"], + ) + _reject_private_or_live_fields(graph, "A/B SeedGraph") + if offline["task_id"] != task["task_id"] or online["task_id"] != task["task_id"]: + raise ValueError("A/B graphs and TaskSpec must have the same task_id.") + for route, graph in (("offline", offline), ("online", online)): + if ( + graph["level"] != task["level"] + or graph["reasoning_type"] != task["reasoning_type"] + ): + raise ValueError( + f"A/B {route} SeedGraph level/reasoning does not match TaskSpec." + ) + _validate_task_group_coverage(task, graph, route=route) + if offline["planner_route"] != "offline" or online["planner_route"] != "online": + raise ValueError( + "Strict A/B requires explicit offline and online graph routes." + ) + return task, config, offline, online + + +def _validate_branch_video_paths( + video_paths: list[str], *, route: str, branch_dir: Path +) -> None: + """Require the normalized video artifact used by strict production A/B.""" + expected = (branch_dir / "video.mp4").resolve() + if len(video_paths) != 1: + raise RuntimeError( + f"Strict A/B {route} branch must publish exactly one video.mp4." + ) + published = Path(video_paths[0]).expanduser().resolve() + if published != expected: + raise RuntimeError( + f"Strict A/B {route} video must be published as {expected.as_posix()}." + ) + if not expected.is_file() or expected.stat().st_size <= 0: + raise RuntimeError(f"Strict A/B {route} video.mp4 is missing or empty.") + + +def _close_environments(environments: Any) -> None: + """Best-effort close every distinct supplied environment exactly once.""" + seen: set[int] = set() + for env in environments: + if id(env) in seen: + continue + seen.add(id(env)) + close = getattr(env, "close", None) + if not callable(close): + continue + try: + close() + except Exception: + # Preserve the validation/execution failure that triggered cleanup, + # but continue closing the other independent branch. + continue + + +def _result_summary( + result: Any, + elapsed: float, + graph: Mapping[str, Any], + planning_metrics: Mapping[str, Any], + *, + success_override: Any | None, +) -> dict[str, Any]: + success = torch.as_tensor( + ( + getattr(result, "success", False) + if success_override is None + else success_override + ), + dtype=torch.bool, + ) + actions = list(getattr(result, "actions", ())) + retries = int(getattr(result, "retry_count", 0)) + recoveries = int(getattr(result, "recovery_count", 0)) + revisions = int(getattr(result, "revision_count", 0)) + metadata = graph.get("metadata", {}) + if not isinstance(metadata, Mapping): + metadata = {} + return { + "route": str(graph.get("planner_route", "")), + "seed_graph_hash": seed_graph_hash(graph), + "planning_seconds": float( + planning_metrics.get( + "planning_seconds", + metadata.get("planning_latency_seconds", 0.0), + ) + ), + "execution_seconds": float(elapsed), + "vlm_call_count": int( + planning_metrics.get( + "vlm_call_count", + metadata.get("vlm_call_count", 0), + ) + ), + "success": success.tolist(), + "success_source": ( + "runtime_postconditions" if success_override is None else "private_oracle" + ), + "success_rate": float(success.float().mean()) if success.numel() else 0.0, + "action_command_count": len(actions), + "path_length": _path_length(actions), + "retry_count": retries, + "recovery_count": recoveries, + "revision_count": revisions, + "failure_events": list(getattr(result, "failure_events", ())), + "ik_failure_count": sum( + item.get("failure_type") in {"plan_failed", "search_exhausted"} + for item in getattr(result, "failure_events", ()) + ), + "record_dir": getattr(result, "record_dir", None), + "video_paths": list(getattr(result, "video_paths", ())), + } + + +def _error_result_summary( + error: Exception, + elapsed: float, + graph: Mapping[str, Any], + planning_metrics: Mapping[str, Any], +) -> dict[str, Any]: + metadata = graph.get("metadata", {}) + if not isinstance(metadata, Mapping): + metadata = {} + return { + "route": str(graph.get("planner_route", "")), + "seed_graph_hash": seed_graph_hash(graph), + "planning_seconds": float( + planning_metrics.get( + "planning_seconds", + metadata.get("planning_latency_seconds", 0.0), + ) + ), + "execution_seconds": float(elapsed), + "vlm_call_count": int( + planning_metrics.get("vlm_call_count", metadata.get("vlm_call_count", 0)) + ), + "success": [False], + "success_source": "runtime_exception", + "success_rate": 0.0, + "action_command_count": 0, + "path_length": 0.0, + "retry_count": 0, + "recovery_count": 0, + "revision_count": 0, + "failure_events": [], + "ik_failure_count": 0, + "record_dir": None, + "video_paths": [], + "error": f"{type(error).__name__}: {error}", + } + + +def _path_length(actions: list[Any]) -> float: + if len(actions) < 2: + return 0.0 + tensors = [torch.as_tensor(action, dtype=torch.float32) for action in actions] + return float( + sum( + torch.linalg.vector_norm(current - previous, dim=-1).sum() + for previous, current in zip(tensors, tensors[1:]) + ) + ) + + +def _graph_difference( + offline: Mapping[str, Any], online: Mapping[str, Any] +) -> dict[str, Any]: + offline_nodes = {str(node["id"]): node for node in offline["nodes"]} + online_nodes = {str(node["id"]): node for node in online["nodes"]} + offline_actions = [node["atomic_action"] for node in offline["nodes"]] + online_actions = [node["atomic_action"] for node in online["nodes"]] + common_ids = sorted(set(offline_nodes) & set(online_nodes)) + node_changes = [] + for node_id in common_ids: + left = offline_nodes[node_id] + right = online_nodes[node_id] + changed_fields = sorted( + key for key in set(left) | set(right) if left.get(key) != right.get(key) + ) + if changed_fields: + node_changes.append({"id": node_id, "changed_fields": changed_fields}) + offline_groups = {str(group["id"]): group for group in offline["task_groups"]} + online_groups = {str(group["id"]): group for group in online["task_groups"]} + common_group_ids = sorted(set(offline_groups) & set(online_groups)) + group_changes = [] + for group_id in common_group_ids: + left = offline_groups[group_id] + right = online_groups[group_id] + changed_fields = sorted( + key for key in set(left) | set(right) if left.get(key) != right.get(key) + ) + if changed_fields: + group_changes.append({"id": group_id, "changed_fields": changed_fields}) + atomic_action_difference = { + "offline": offline_actions, + "online": online_actions, + "same_sequence": offline_actions == online_actions, + "added_node_ids": sorted(set(online_nodes) - set(offline_nodes)), + "removed_node_ids": sorted(set(offline_nodes) - set(online_nodes)), + "changed_nodes": node_changes, + } + task_group_difference = { + "offline_ids": sorted(offline_groups), + "online_ids": sorted(online_groups), + "added_ids": sorted(set(online_groups) - set(offline_groups)), + "removed_ids": sorted(set(offline_groups) - set(online_groups)), + "same_ids": set(offline_groups) == set(online_groups), + "changed_groups": group_changes, + } + return { + "offline_node_count": len(offline_actions), + "online_node_count": len(online_actions), + "same_action_sequence": offline_actions == online_actions, + "offline_actions": offline_actions, + "online_actions": online_actions, + "added_node_ids": sorted(set(online_nodes) - set(offline_nodes)), + "removed_node_ids": sorted(set(offline_nodes) - set(online_nodes)), + "changed_nodes": node_changes, + "offline_task_group_ids": sorted(offline_groups), + "online_task_group_ids": sorted(online_groups), + "added_task_group_ids": sorted(set(online_groups) - set(offline_groups)), + "removed_task_group_ids": sorted(set(offline_groups) - set(online_groups)), + "same_task_group_ids": set(offline_groups) == set(online_groups), + "changed_task_groups": group_changes, + "atomic_action_difference": atomic_action_difference, + "task_group_difference": task_group_difference, + } + + +def _validate_task_group_coverage( + task: Mapping[str, Any], graph: Mapping[str, Any], *, route: str +) -> None: + """Check explicit TaskSpec instances are neither dropped nor duplicated.""" + if task.get("level") == "L4": + return + expected = { + str(item["id"]) + for item in task.get("task_instances", ()) + if isinstance(item, Mapping) + } + actual = {str(group["id"]) for group in graph.get("task_groups", ())} + if expected != actual: + raise ValueError( + f"A/B {route} TaskGroup coverage mismatch; " + f"missing={sorted(expected - actual)}, " + f"unexpected={sorted(actual - expected)}." + ) + + +def _validate_snapshot( + snapshot: Mapping[str, Any], *, route: str, require_full: bool = False +) -> None: + if not isinstance(snapshot, Mapping): + raise TypeError(f"A/B {route} snapshot must be a mapping.") + required = {"robot_qpos", "object_poses"} + if require_full: + required = set(_FULL_SNAPSHOT_KEYS) + missing = required - set(snapshot) + if missing: + raise ValueError( + f"A/B {route} snapshot is missing required state {sorted(missing)}." + ) + qpos = torch.as_tensor(snapshot["robot_qpos"]) + if qpos.numel() == 0 or not bool(torch.isfinite(qpos).all()): + raise ValueError(f"A/B {route} robot_qpos must be finite and non-empty.") + object_poses = snapshot["object_poses"] + if not isinstance(object_poses, Mapping) or not object_poses: + raise ValueError(f"A/B {route} object_poses must be a non-empty mapping.") + for uid, pose in object_poses.items(): + tensor = torch.as_tensor(pose) + if not isinstance(uid, str) or not uid or tensor.numel() == 0: + raise ValueError(f"A/B {route} object_poses contains an invalid entry.") + if not bool(torch.isfinite(tensor).all()): + raise ValueError(f"A/B {route} pose for {uid!r} must be finite.") + if "articulation_state" in snapshot: + articulation_state = snapshot["articulation_state"] + if not isinstance(articulation_state, Mapping): + raise ValueError(f"A/B {route} articulation_state must be a mapping.") + for uid, state in articulation_state.items(): + if not isinstance(uid, str) or not uid: + raise ValueError( + f"A/B {route} articulation_state contains invalid UID." + ) + if not isinstance(state, Mapping): + raise ValueError( + f"A/B {route} articulation state for {uid!r} must be a mapping." + ) + if not state: + raise ValueError( + f"A/B {route} articulation state for {uid!r} is empty." + ) + for name, value in state.items(): + tensor = torch.as_tensor(value) + if tensor.numel() == 0 or not bool(torch.isfinite(tensor).all()): + raise ValueError( + f"A/B {route} articulation {uid!r}.{name} must be finite." + ) + if "camera_calibration" in snapshot: + calibrations = snapshot["camera_calibration"] + if not isinstance(calibrations, Mapping): + raise ValueError(f"A/B {route} camera_calibration must be a mapping.") + if require_full and not calibrations: + raise ValueError(f"A/B {route} camera_calibration must not be empty.") + for uid, calibration in calibrations.items(): + if not isinstance(uid, str) or not uid: + raise ValueError( + f"A/B {route} camera_calibration contains invalid UID." + ) + if not isinstance(calibration, Mapping): + raise ValueError( + f"A/B {route} calibration for {uid!r} must be a mapping." + ) + for name in ("intrinsics", "extrinsics"): + if name not in calibration: + raise ValueError( + f"A/B {route} calibration for {uid!r} is missing {name}." + ) + tensor = torch.as_tensor(calibration[name]) + if tensor.numel() == 0 or not bool(torch.isfinite(tensor).all()): + raise ValueError( + f"A/B {route} calibration {uid!r}.{name} must be finite." + ) + + +def _update_digest(digest: Any, value: Any) -> None: + if isinstance(value, Mapping): + digest.update(b"mapping{") + for key in sorted(value, key=str): + _update_digest(digest, str(key)) + _update_digest(digest, value[key]) + digest.update(b"}") + return + if isinstance(value, (list, tuple)): + digest.update(b"sequence[") + for item in value: + _update_digest(digest, item) + digest.update(b"]") + return + if isinstance(value, torch.Tensor): + value = value.detach().cpu().contiguous().numpy() + if isinstance(value, np.ndarray): + digest.update(str(value.dtype).encode("ascii")) + digest.update(str(tuple(value.shape)).encode("ascii")) + digest.update(value.tobytes(order="C")) + return + digest.update(type(value).__name__.encode("ascii")) + digest.update(repr(value).encode("utf-8")) + + +def _reject_private_or_live_fields(value: Any, context: str) -> None: + """Reject oracle/grounded fields before either branch can execute.""" + if isinstance(value, Mapping): + for key, child in value.items(): + if str(key).lower() in _PRIVATE_OR_LIVE_KEYS: + raise ValueError(f"{context} contains private/live field {key!r}.") + _reject_private_or_live_fields(child, f"{context}.{key}") + elif isinstance(value, (list, tuple)): + for index, child in enumerate(value): + _reject_private_or_live_fields(child, f"{context}[{index}]") + + +def _call_preflight( + callback: Callable[..., Any], *, route: str, graph: Mapping[str, Any], env: Any +) -> Any: + """Call executor preflight hooks across the small supported API variants.""" + try: + return callback() + except TypeError as first_error: + # Third-party branch executors often expose contextual keyword-only + # arguments. Retry only for an argument-binding TypeError; if the + # callback itself raised TypeError, preserve that original failure. + try: + return callback(route=route, graph=graph, env=env) + except TypeError: + raise first_error + + +def _write_json(path: Path, value: Mapping[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(value, ensure_ascii=False, indent=2, allow_nan=False) + "\n", + encoding="utf-8", + ) diff --git a/embodichain/gen_sim/action_engine/evaluation/e1_e2_scene_action.py b/embodichain/gen_sim/action_engine/evaluation/e1_e2_scene_action.py new file mode 100644 index 000000000..b4b256806 --- /dev/null +++ b/embodichain/gen_sim/action_engine/evaluation/e1_e2_scene_action.py @@ -0,0 +1,448 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Measure deterministic E1/E2 scene feasibility and graph compilation. + +This CPU benchmark exercises the contract path before simulator motion. It +checks that Scene Engine v1 output adapts successfully, required capabilities +are executable, and E1/E2 compile to action graphs containing pickup, +held-object motion, and placement. + +Run this module with ``--iterations 100`` for the default benchmark. +""" + +from __future__ import annotations + +import argparse +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from time import perf_counter +import tracemalloc +from types import SimpleNamespace +from typing import Any + +from embodichain.gen_sim.action_engine.capabilities import ( + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.domain import ( + TASK_CONTRACTS, + task_success_type, + validate_scene_requirements, + validate_task_spec, +) +from embodichain.gen_sim.action_engine.protocol import ( + SCENE_REQUIREMENTS_SCHEMA, + TASK_SPEC_SCHEMA, +) +from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph +from embodichain.gen_sim.task_engine.scene import ( + FeasibilityBroker, + SceneEngineV1Adapter, +) + +__all__ = ["BenchmarkResult", "run_benchmark"] + + +@dataclass(frozen=True) +class BenchmarkResult: + """One scenario's contract latency, memory, and correctness metrics.""" + + scenario: str + iterations: int + elapsed_seconds: float + peak_bytes: int + success_count: int + feasibility_status: str + action_count: int + unknown_checks: int + runtime_probe_checks: int + + @property + def success_rate(self) -> float: + """Return successful iterations divided by all iterations.""" + return self.success_count / self.iterations + + @property + def mean_milliseconds(self) -> float: + """Return mean contract latency in milliseconds.""" + return self.elapsed_seconds * 1000.0 / self.iterations + + +def run_benchmark( + *, + iterations: int = 100, + output_dir: str | Path = "outputs/benchmarks", +) -> tuple[tuple[BenchmarkResult, ...], Path]: + """Run E1/E2 contract regressions and write one Markdown report.""" + if ( + isinstance(iterations, bool) + or not isinstance(iterations, int) + or iterations < 1 + ): + raise ValueError("iterations must be a positive integer.") + results = tuple( + _benchmark_scenario(task_type, iterations=iterations) + for task_type in ("E1", "E2") + ) + report = _write_report(results, output_dir=Path(output_dir)) + return results, report + + +def _benchmark_scenario(task_type: str, *, iterations: int) -> BenchmarkResult: + task, requirements = _generated_task(task_type) + bindings = { + item["role_id"]: f"{task_type.lower()}_{item['role_id']}" + for item in requirements["objects"] + } + manifest = _static_manifest(task_type, requirements, bindings) + candidate, reference_bindings = _candidate( + task_type, + task, + requirements, + bindings, + ) + registry = build_atomic_capability_registry() + broker = FeasibilityBroker() + task_actions = { + name: contract.core_actions for name, contract in TASK_CONTRACTS.items() + } + success_count = 0 + last_report: dict[str, Any] = {} + last_graph: dict[str, Any] = {} + + tracemalloc.start() + start = perf_counter() + try: + for _ in range(iterations): + last_report = broker.assess( + candidate, + reference_bindings, + manifest, + capability_catalog=registry.catalog(), + task_actions=task_actions, + ) + last_graph = instantiate_seed_graph(task, bindings, registry=registry) + actions = { + str(node.get("atomic_action")) + for node in last_graph["nodes"] + if node.get("atomic_action") + } + if ( + last_report["status"] != "contradicted" + and {"PickUp", "MoveHeldObject", "Place"} <= actions + ): + success_count += 1 + finally: + elapsed = perf_counter() - start + _, peak_bytes = tracemalloc.get_traced_memory() + tracemalloc.stop() + + action_count = sum( + bool(node.get("atomic_action")) for node in last_graph.get("nodes", ()) + ) + return BenchmarkResult( + scenario=task_type, + iterations=iterations, + elapsed_seconds=elapsed, + peak_bytes=peak_bytes, + success_count=success_count, + feasibility_status=str(last_report.get("status", "unknown")), + action_count=action_count, + unknown_checks=int(last_report.get("summary", {}).get("unknown", 0)), + runtime_probe_checks=int( + last_report.get("summary", {}).get("runtime_probe", 0) + ), + ) + + +def _generated_task(task_type: str) -> tuple[dict[str, Any], dict[str, Any]]: + if task_type not in {"E1", "E2"}: + raise ValueError("This benchmark supports only E1 and E2 fixtures.") + params: dict[str, Any] = {"object_role": "object"} + initial_state = {} + if task_type == "E1": + params.update({"target_role": "target", "relation": "inside"}) + else: + params.update( + { + "orientation_goal": "upright", + "support_role": "table", + "upright_local_axis": "long_axis", + } + ) + initial_state = {"orientation": "fallen"} + task_id = f"benchmark-{task_type.lower()}" + task = validate_task_spec( + { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": task_id, + "level": "L1", + "instruction": "benchmark-instruction", + "reasoning_type": "none", + "task_instances": [ + { + "id": "task_01", + "task_type": task_type, + "params": params, + "depends_on": [], + "role": "primary", + } + ], + "success": { + "type": task_success_type(task_type, params), + "task_instance_id": "task_01", + }, + "oracle": {}, + "metadata": {"benchmark_fixture": True}, + } + ) + objects = [ + { + "role_id": "object", + "category": "can", + "count": 1, + "affordances": sorted(TASK_CONTRACTS[task_type].scene_affordances), + "initial_state": initial_state, + "attributes": {}, + } + ] + if task_type == "E1": + objects.append( + { + "role_id": "target", + "category": "container", + "count": 1, + "affordances": ["container", "support_surface"], + "initial_state": {}, + "attributes": {}, + } + ) + requirements = validate_scene_requirements( + { + "schema_version": SCENE_REQUIREMENTS_SCHEMA, + "task_id": task_id, + "objects": objects, + "cameras": [], + "spatial_constraints": [ + {"type": "reachable", "roles": "all_interaction_objects"} + ], + "distractor_count": 0, + "metadata": {"benchmark_fixture": True}, + } + ) + return task, requirements + + +def _static_manifest( + task_type: str, + requirements: dict[str, Any], + bindings: dict[str, str], +) -> dict[str, Any]: + planner_objects = [] + runtime_objects = [] + for index, requirement in enumerate(requirements["objects"]): + role_id = str(requirement["role_id"]) + uid = bindings[role_id] + role = "rigid_object" + planner_objects.append( + { + "uid": uid, + "source_uid": f"{uid}_0", + "role": role, + "name": uid, + "description": f"Synthetic {task_type} benchmark object.", + "category": str(requirement["category"]), + "color": requirement.get("attributes", {}).get("color"), + "shape": {"shape_type": "Mesh", "fpath": f"/{uid}.glb"}, + "init_pos": [0.15 * index, 0.0, 0.7], + "init_rot": [90.0, 0.0, 0.0] if task_type == "E2" else [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + "affordances": list(requirement["affordances"]), + "initial_state": dict(requirement["initial_state"]), + "attributes": dict(requirement["attributes"]), + } + ) + runtime_objects.append( + { + "uid": uid, + "shape": {"shape_type": "Mesh", "fpath": f"/{uid}.glb"}, + "attrs": {"mass": 0.1}, + "body_type": "dynamic", + } + ) + prepared = SimpleNamespace( + source_config_path=Path("/synthetic/scene_config.json"), + planner_objects=tuple(planner_objects), + background=(), + rigid_objects=tuple(runtime_objects), + articulations=(), + asset_hashes={ + uid: uid.encode().hex().ljust(64, "0")[:64] for uid in bindings.values() + }, + ) + return SceneEngineV1Adapter().adapt_prepared_scene( + prepared, + source_format="benchmark", + robot_profile="dual_franka", + ) + + +def _candidate( + task_type: str, + task: dict[str, Any], + requirements: dict[str, Any], + bindings: dict[str, str], +) -> tuple[dict[str, Any], dict[str, list[str]]]: + references = [] + reference_bindings = {} + for index, requirement in enumerate(requirements["objects"]): + role_id = str(requirement["role_id"]) + role = "object" if index == 0 else "target" + reference_id = f"task_01.{role}" + references.append( + { + "reference_id": reference_id, + "role": role, + "source_structure": "rigid_object", + "affordances": list(requirement["affordances"]), + "initial_state": dict(requirement["initial_state"]), + "attributes": dict(requirement["attributes"]), + } + ) + reference_bindings[reference_id] = [bindings[role_id]] + return ( + { + "candidate_id": "candidate_01", + "draft": { + "task_id": task["task_id"], + "steps": [{"id": "task_01", "task_type": task_type}], + }, + "scene_request": {"references": references}, + }, + reference_bindings, + ) + + +def _write_report( + results: tuple[BenchmarkResult, ...], + *, + output_dir: Path, +) -> Path: + output_dir.mkdir(parents=True, exist_ok=True) + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + path = output_dir / f"e1_e2_scene_action_{timestamp}.md" + performance_rows = [ + { + "Scenario": item.scenario, + "Iterations": item.iterations, + "Total ms": f"{item.elapsed_seconds * 1000.0:.3f}", + "Mean ms": f"{item.mean_milliseconds:.3f}", + "Peak KiB": f"{item.peak_bytes / 1024.0:.1f}", + } + for item in results + ] + metric_rows = [ + { + "Scenario": item.scenario, + "Success rate": f"{item.success_rate:.3f}", + "Feasibility": item.feasibility_status, + "Actions": item.action_count, + "Unknown checks": item.unknown_checks, + "Runtime probes": item.runtime_probe_checks, + } + for item in results + ] + leaderboard_rows = [ + { + "Rank": rank, + "Scenario": item.scenario, + "Success rate": f"{item.success_rate:.3f}", + "Mean ms": f"{item.mean_milliseconds:.3f}", + } + for rank, item in enumerate( + sorted( + results, + key=lambda value: (-value.success_rate, value.mean_milliseconds), + ), + start=1, + ) + ] + lines = [ + "# E1/E2 Scene-Action Contract Benchmark", + "", + f"Generated at: {datetime.now().isoformat(timespec='seconds')}", + "", + "## Time & Memory", + "", + *_table(performance_rows), + "", + "## Success & Other Metrics", + "", + *_table(metric_rows), + "", + "## Leaderboard", + "", + *_table(leaderboard_rows), + "", + "## Notes", + "", + "- This benchmark covers deterministic contracts and graph compilation, not GPU motion execution.", + ] + path.write_text("\n".join(lines) + "\n", encoding="utf-8") + return path + + +def _table(rows: list[dict[str, object]]) -> list[str]: + headers = list(rows[0]) + return [ + "| " + " | ".join(headers) + " |", + "| " + " | ".join("---" for _ in headers) + " |", + *[ + "| " + " | ".join(str(row[header]) for header in headers) + " |" + for row in rows + ], + ] + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Benchmark E1/E2 scene-action contract stability." + ) + parser.add_argument("--iterations", type=int, default=100) + parser.add_argument("--output-dir", type=Path, default=Path("outputs/benchmarks")) + return parser + + +def main() -> int: + """Run from the command line and print the generated report path.""" + args = _build_parser().parse_args() + results, report = run_benchmark( + iterations=args.iterations, + output_dir=args.output_dir, + ) + for result in results: + print( + f"{result.scenario}: success={result.success_rate:.3f}, " + f"mean={result.mean_milliseconds:.3f} ms, " + f"peak={result.peak_bytes / 1024.0:.1f} KiB" + ) + print(f"Report: {report}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/embodichain/gen_sim/action_engine/runtime/__init__.py b/embodichain/gen_sim/action_engine/runtime/__init__.py index a91770d6e..64774425f 100644 --- a/embodichain/gen_sim/action_engine/runtime/__init__.py +++ b/embodichain/gen_sim/action_engine/runtime/__init__.py @@ -18,15 +18,49 @@ from __future__ import annotations +from .executor import ProgramExecutor +from .dynamic import DynamicRecoveryController, RecoveryDirective from .loader import load_agent_execution_program, load_execution_program +from .recovery import ( + FAILURE_TYPES, + GraphRevision, + RetryDecision, + RuntimeGraph, + build_upright_recovery, + classify_failure, +) from .models import ExecutionProgram, ExecutionReport, ExecutionResult +from .reporting import ( + EXECUTION_REPORT_FILENAME, + EXECUTION_REPORT_SCHEMA, + build_execution_provenance, + validate_execution_report, + write_execution_report, +) +from .predicates import PREDICATE_TYPES, evaluate_predicate from .state import ExecutionState __all__ = [ "ExecutionProgram", "ExecutionState", + "DynamicRecoveryController", + "PREDICATE_TYPES", "ExecutionResult", "ExecutionReport", + "EXECUTION_REPORT_FILENAME", + "EXECUTION_REPORT_SCHEMA", + "FAILURE_TYPES", + "GraphRevision", + "ProgramExecutor", + "RetryDecision", + "RecoveryDirective", + "RuntimeGraph", + "build_execution_provenance", + "build_upright_recovery", + "classify_failure", + "evaluate_predicate", "load_agent_execution_program", "load_execution_program", + "validate_execution_report", + "write_execution_report", ] diff --git a/embodichain/gen_sim/action_engine/runtime/dynamic.py b/embodichain/gen_sim/action_engine/runtime/dynamic.py new file mode 100644 index 000000000..616e14c80 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/dynamic.py @@ -0,0 +1,164 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Route-explicit recovery and suffix-replanning coordinator.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +from .recovery import RuntimeGraph + +__all__ = ["DynamicRecoveryController", "RecoveryDirective"] + +Replanner = Callable[..., Mapping[str, Any]] + + +@dataclass(frozen=True) +class RecoveryDirective: + """A graph revision that must execute before suffix replanning.""" + + failure_type: str + failed_node_id: str + recovery_group_id: str | None + graph: dict[str, Any] + requires_recovery_execution: bool + active_env_ids: tuple[int, ...] + + +class DynamicRecoveryController: + """Keep offline and online dynamic replanning as separately testable modes.""" + + def __init__( + self, + runtime_graph: RuntimeGraph, + *, + mode: str, + offline_replanner: Replanner | None = None, + online_replanner: Replanner | None = None, + ) -> None: + if mode not in {"offline_dynamic", "online_dynamic"}: + raise ValueError("Dynamic mode must be offline_dynamic or online_dynamic.") + selected = offline_replanner if mode == "offline_dynamic" else online_replanner + if not callable(selected): + raise ValueError(f"{mode} requires its matching replanner callback.") + self.runtime_graph = runtime_graph + self.mode = mode + self._replanner = selected + + def handle_failure( + self, + *, + failed_node_id: str, + failure_type: str, + active_env_ids: Sequence[int] | None = None, + ) -> RecoveryDirective: + """Insert recovery when known; otherwise request immediate full replanning.""" + env_ids = tuple( + sorted( + set( + range(self.runtime_graph.num_envs) + if active_env_ids is None + else (int(env_id) for env_id in active_env_ids) + ) + ) + ) + if not env_ids or env_ids[0] < 0 or env_ids[-1] >= self.runtime_graph.num_envs: + raise ValueError( + "Recovery active_env_ids are outside the environment range." + ) + if failure_type == "object_fallen": + graph = self.runtime_graph.insert_default_recovery( + failed_node_id=failed_node_id, + failure_type=failure_type, + active_env_ids=env_ids, + ) + group_id = self.runtime_graph.revisions[-1].inserted_group_ids[0] + return RecoveryDirective( + failure_type, + failed_node_id, + group_id, + graph, + True, + self.runtime_graph.revisions[-1].active_env_ids, + ) + return RecoveryDirective( + failure_type, + failed_node_id, + None, + self.runtime_graph.graph, + False, + env_ids, + ) + + def handle_execution_result(self, result: Any) -> RecoveryDirective: + """Create a directive from the first actionable runtime failure event.""" + events = getattr(result, "failure_events", None) + if not isinstance(events, Sequence) or not events: + raise ValueError("Execution result contains no recoverable failure event.") + event = next( + ( + item + for item in events + if isinstance(item, Mapping) and bool(item.get("fatal", True)) + ), + None, + ) + if event is None: + raise ValueError( + "Execution result contains no fatal recoverable failure event." + ) + if not isinstance(event, Mapping): + raise ValueError("Execution failure events must be mappings.") + node_id = event.get("node_id") + failure_type = event.get("failure_type") + env_ids = event.get("env_ids", ()) + if not isinstance(node_id, str) or not node_id: + raise ValueError("Dynamic recovery requires a v3 SeedGraph node_id.") + if not isinstance(failure_type, str): + raise ValueError("Execution failure event requires failure_type.") + return self.handle_failure( + failed_node_id=node_id, + failure_type=failure_type, + active_env_ids=env_ids, + ) + + def replan( + self, + directive: RecoveryDirective, + *, + completed_group_ids: Sequence[str], + recovery_succeeded: bool, + observations: Mapping[str, Any] | None = None, + ) -> dict[str, Any]: + """Replace only the unfinished suffix after recovery or escalation.""" + if directive.requires_recovery_execution and not recovery_succeeded: + reason = f"{directive.failure_type}:recovery_failed" + else: + reason = f"{directive.failure_type}:state_restored" + replacement = self._replanner( + graph=self.runtime_graph.graph, + completed_group_ids=tuple(completed_group_ids), + failure_type=directive.failure_type, + observations=dict(observations or {}), + ) + return self.runtime_graph.replace_unfinished_suffix( + replacement, + completed_group_ids=completed_group_ids, + reason=reason, + ) diff --git a/embodichain/gen_sim/action_engine/runtime/executor.py b/embodichain/gen_sim/action_engine/runtime/executor.py new file mode 100644 index 000000000..690a82787 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/executor.py @@ -0,0 +1,4331 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Closed-loop executor for action-engine execution-program DAGs.""" + +from __future__ import annotations + +from collections.abc import Iterator, Mapping, Sequence +from contextlib import contextmanager +from copy import deepcopy +from dataclasses import dataclass, field, replace +import logging +from threading import RLock +from typing import Any + +import numpy as np +import torch + +from embodichain.gen_sim.action_engine.config import ( + ArmSelectionPolicyCfg, + RuntimePolicyCfg, + default_runtime_policy, + runtime_policy_hash, +) +from embodichain.gen_sim.action_engine.domain import normalize_placement_relation +from embodichain.gen_sim.action_engine.orientation import ( + AlignAxisConstraint, + MatchRotationConstraint, + compile_orientation_constraint, +) +from embodichain.lab.sim.atomic_actions import ( + HeldObjectState, + SceneProvider, + StateDelta, +) +from embodichain.utils import logger as project_logger +from embodichain.utils.logger import log_info, log_warning + +from .actions import AtomicActionAdapter +from .frames import DIRECTIONAL_RELATIONS, robot_frame_axes +from .grounding import ActionGrounder, LiveArrangementPlan, LivePlacementPlan +from .models import ( + ActionOutcome, + ExecutionEdge, + ExecutionProgram, + ExecutionResult, + GroundedAction, + SemanticStep, +) +from .predicates import evaluate_predicate +from .recording import RuntimeRecorder +from .recovery import RuntimeGraph +from .robot_parts import arm_control_part +from .state import ExecutionState + +__all__ = ["ProgramExecutor"] + + +@dataclass +class _Candidate: + feasible: torch.Tensor + cost: torch.Tensor + plans: dict[str, tuple[GroundedAction, ActionOutcome]] + score_components: dict[str, torch.Tensor] = field(default_factory=dict) + warnings: tuple[str, ...] = () + blockers: tuple[dict[str, Any], ...] = () + + +@dataclass +class _EdgeResult: + actions: list[torch.Tensor] + failed: torch.Tensor + grounded: list[GroundedAction] + planner_traces: list[dict[str, Any]] = field(default_factory=list) + executed: torch.Tensor | None = None + + +@dataclass(frozen=True) +class _SupportRelation: + support_uid: str + semantic_step_id: str + + +@dataclass +class _PlacementRecoveryResult: + failed: torch.Tensor + succeeded: torch.Tensor + observed: torch.Tensor + actions: list[torch.Tensor] + failure_events: list[dict[str, Any]] = field(default_factory=list) + covered_failures: torch.Tensor | None = None + + +def _score_arm_candidate( + *, + arm: str, + motion_cost: torch.Tensor, + source_pose: torch.Tensor, + target_pose: torch.Tensor | None, + workspace_center_xy: torch.Tensor, + workspace_half_width: torch.Tensor, + robot_lateral_axis: torch.Tensor, + policy: ArmSelectionPolicyCfg, +) -> dict[str, torch.Tensor]: + """Combine motion length with soft, table-normalized cross-zone costs.""" + arm_sign = 1.0 if arm == "left_arm" else -1.0 + deadband = workspace_half_width * float(policy.crossing_deadband_ratio) + + def crossing(pose: torch.Tensor | None, weight: float) -> torch.Tensor: + if pose is None: + return torch.zeros_like(motion_cost) + lateral = torch.sum( + (pose[:, :2, 3] - workspace_center_xy) * robot_lateral_axis, + dim=1, + ) + wrong_side_depth = torch.clamp( + -arm_sign * lateral - deadband, + min=0.0, + ) + return weight * torch.square(wrong_side_depth / workspace_half_width) + + normalized_motion = motion_cost / float(policy.motion_cost_scale) + pickup_penalty = crossing(source_pose, float(policy.pickup_crossing_weight)) + placement_penalty = crossing( + target_pose, + float(policy.placement_crossing_weight), + ) + return { + "motion_cost": motion_cost, + "normalized_motion_cost": normalized_motion, + "pickup_crossing_penalty": pickup_penalty, + "placement_crossing_penalty": placement_penalty, + "total_cost": normalized_motion + pickup_penalty + placement_penalty, + } + + +_SPECULATIVE_LOG_LOCK = RLock() + + +@contextmanager +def _capture_speculative_warnings() -> Iterator[list[str]]: + """Temporarily capture project warnings without changing its log level.""" + messages: list[str] = [] + collector = logging.Handler(level=logging.WARNING) + collector.emit = lambda record: messages.append(record.getMessage()) + logger = project_logger.logger + with _SPECULATIVE_LOG_LOCK: + handlers = list(logger.handlers) + propagate = logger.propagate + try: + logger.handlers[:] = [collector] + logger.propagate = False + yield messages + finally: + logger.handlers[:] = handlers + logger.propagate = propagate + + +class ProgramExecutor: + """Schedule, ground, plan, execute, and verify one immutable program.""" + + def __init__( + self, + program: ExecutionProgram, + env: Any, + *, + max_transitions: int | None = None, + settle_steps: int | None = None, + record_runtime: bool = True, + record_root: str | None = None, + runtime_policy: RuntimePolicyCfg | None = None, + capability_registry: Any | None = None, + scene_provider: SceneProvider | None = None, + ) -> None: + self.program = program + self.env = env + self.record_runtime = bool(record_runtime) + self.record_root = record_root + if runtime_policy is None: + profile = str(getattr(env, "agent_robot_profile", "dual_ur10")) + runtime_policy = default_runtime_policy(profile) + if not isinstance(runtime_policy, RuntimePolicyCfg): + raise TypeError("ProgramExecutor runtime_policy must be RuntimePolicyCfg.") + self.runtime_policy = runtime_policy + self.capability_registry = capability_registry + self.env.runtime_policy = runtime_policy + execution = runtime_policy.execution + self.max_transitions = int( + execution["max_transitions"] if max_transitions is None else max_transitions + ) + self.settle_steps = int( + execution["semantic_step_settle_steps"] + if settle_steps is None + else settle_steps + ) + self.max_retries_per_action = int(execution["max_retries_per_action"]) + self.support_stability_samples = int(execution["support_stability_samples"]) + self.support_stability_interval_steps = int( + execution["support_stability_interval_steps"] + ) + self.support_linear_velocity_tolerance = float( + execution["support_linear_velocity_tolerance"] + ) + self.support_angular_velocity_tolerance = float( + execution["support_angular_velocity_tolerance"] + ) + self.placement_recovery_attempts = int( + runtime_policy.grounding["placement"]["recovery_attempts"] + ) + self.runtime_graph = ( + RuntimeGraph( + program.seed_graph, + num_envs=int(env.num_envs), + max_retries=self.max_retries_per_action, + max_revisions=int(execution["max_graph_revisions"]), + max_recovery_actions=int(execution["max_recovery_actions"]), + registry=capability_registry, + ) + if program.seed_graph is not None + else None + ) + self.retry_count = 0 + self.edges = {edge.id: edge for edge in program.edges} + self.steps = {step.id: step for step in program.semantic_steps} + self.step_by_edge = { + edge_id: step + for step in program.semantic_steps + for edge_id in step.edge_ids + } + missing = set(self.edges) - set(self.step_by_edge) + if missing: + raise ValueError( + "Every execution edge must belong to one semantic step; missing " + f"{sorted(missing)}." + ) + self._completion_only_dependencies = self._completion_only_dependency_edges() + self.group_by_step = { + str(step_id): group + for group in program.allocation_groups + for step_id in group.get("semantic_step_ids", ()) + } + arrangement_steps = [ + step + for step in program.semantic_steps + if step.operator in {"arrange_line", "place_in_line"} + ] + arrangement_groups: dict[str, list[SemanticStep]] = {} + for step in arrangement_steps: + arrangement_groups.setdefault(step.parent_step_id, []).append(step) + arrangement_policy = runtime_policy.grounding["arrangement"] + plans = [ + LiveArrangementPlan( + env, + steps, + slot_margin=float(arrangement_policy["slot_margin"]), + minimum_spacing=float(arrangement_policy["minimum_spacing"]), + clearance=float(arrangement_policy["layout_clearance"]), + row_search_step=float(arrangement_policy["row_search_step"]), + row_search_radius=float(arrangement_policy["row_search_radius"]), + ) + for steps in arrangement_groups.values() + ] + self.arrangements = {step.id: plan for plan in plans for step in plan.steps} + # Retain the singular attribute as a convenient introspection hook for + # the common one-arrangement case. + self.arrangement = plans[0] if len(plans) == 1 else None + placement_groups: dict[str, list[SemanticStep]] = {} + for step in program.semantic_steps: + if ( + step.operator == "place_relative" + and step.goal.get("relation") == "inside" + and isinstance(step.goal.get("reference_object"), str) + ): + placement_groups.setdefault( + str(step.goal["reference_object"]), + [], + ).append(step) + placement_plans = [ + LivePlacementPlan( + env, + steps, + clearance=float(runtime_policy.grounding["placement"]["clearance"]), + ) + for steps in placement_groups.values() + if len(steps) > 1 + ] + self.placements = { + step.id: plan for plan in placement_plans for step in plan.steps + } + self.adapter = AtomicActionAdapter( + env, + grasp_policy=runtime_policy.grasp, + planner_policy=runtime_policy.planner, + capability_registry=capability_registry, + scene_provider=scene_provider, + ) + self.grounder = ActionGrounder( + program, + env, + self.adapter.semantics, + self.arrangements, + self.placements, + runtime_policy=runtime_policy, + capability_registry=capability_registry, + ) + self._step_states: dict[tuple[str, str], ExecutionState] = {} + self._object_states: dict[tuple[str, str], ExecutionState] = {} + self._object_owners: dict[str, list[str | None]] = {} + self._arm_owners: dict[str, list[str | None]] = { + "left_arm": [None] * int(env.num_envs), + "right_arm": [None] * int(env.num_envs), + } + self._assignments: dict[str, list[str | None]] = {} + self._candidate_cache: dict[tuple[str, str], _Candidate] = {} + self._candidate_failures: dict[tuple[str, str], str] = {} + self._candidate_diagnostics: dict[str, tuple[str, ...]] = {} + self._candidate_blockers: dict[str, tuple[dict[str, Any], ...]] = {} + self._reported_candidates: set[str] = set() + self._pickup_retry_exclusions: dict[tuple[str, int], set[str]] = {} + self._targets: dict[str, torch.Tensor] = {} + self._target_poses: dict[str, torch.Tensor] = {} + self._orientation_references: dict[str, torch.Tensor] = {} + self._orientation_errors: dict[str, torch.Tensor] = {} + self._policies: dict[str, dict[str, Any]] = {} + self._payload_initial: dict[str, dict[str, torch.Tensor]] = {} + self._support_relations: dict[str, list[_SupportRelation | None]] = {} + self._placement_candidate_history: dict[tuple[str, str], set[int]] = {} + self._robot_lateral_axis_cache: torch.Tensor | None = None + self._transition_count = 0 + self._retry_counts = [0] * int(self.env.num_envs) + + def run( + self, + *, + run_id: str | None = None, + episode_index: int = 0, + ) -> ExecutionResult: + """Execute ready edges until the DAG completes or raises a structural error.""" + self._reset_runtime_state() + recorder = RuntimeRecorder( + self.program, + num_envs=int(self.env.num_envs), + run_id=run_id, + episode_index=episode_index, + output_root=self.record_root, + enabled=self.record_runtime, + runtime_policy=self.runtime_policy.as_mapping(), + runtime_policy_hash=runtime_policy_hash(self.runtime_policy), + ) + aggregate_failed = torch.zeros( + int(self.env.num_envs), + dtype=torch.bool, + device=self.env.device, + ) + edge_failures: dict[str, torch.Tensor] = {} + semantic_success: dict[str, torch.Tensor] = {} + failure_events: list[dict[str, Any]] = [] + completed: set[str] = set() + remaining = [edge.id for edge in self.program.edges] + executed_actions: list[torch.Tensor] = [] + error_message = None + try: + while remaining: + ready = [ + self.edges[edge_id] + for edge_id in remaining + if set(self.edges[edge_id].depends_on) <= completed + ] + if not ready: + raise RuntimeError( + "Execution program is deadlocked: no remaining edge is ready." + ) + ready_blocked = { + edge.id: self._dependency_failures(edge, edge_failures) + for edge in ready + } + batch = self._pack_ready_edges( + ready, + inactive=ready_blocked, + completed=completed, + ) + blocked = {edge.id: ready_blocked[edge.id] for edge in batch} + # A synchronized pair needs the same active rows. Execute a + # healthy independent branch separately when its peer is blocked. + if len(batch) == 2 and not torch.equal( + blocked[batch[0].id], blocked[batch[1].id] + ): + batch = (batch[0],) + self._consume_transitions(len(batch)) + + if len(batch) == 2: + posture_before = { + edge.id: self._object_not_fallen(self.step_by_edge[edge.id]) + for edge in batch + } + edge_results, _ = self._execute_parallel_pickups( + batch, + failed=blocked[batch[0].id], + ) + for edge in batch: + result = edge_results[edge.id] + step = self.step_by_edge[edge.id] + active = ~blocked[edge.id] + recorder.edge( + edge.id, + step, + assignments=self._assignments[step.id], + grounded=result.grounded, + active=active, + failed=result.failed, + action_steps=len(result.actions), + planner_traces=getattr(result, "planner_traces", ()), + diagnostics=self._edge_diagnostics( + step, + edge, + result.failed, + ), + ) + failure_events.extend( + self._failure_events( + edge, + step, + result.failed & ~blocked[edge.id], + postcondition=False, + executed=result.executed, + fallen_transition=self._fallen_transition( + step, + posture_before[edge.id], + result, + ), + planner_traces=result.planner_traces, + ) + ) + # Both edge records describe the same synchronized command + # stream. Store it once in the returned execution trace. + executed_actions.extend(edge_results[batch[0].id].actions) + for edge in batch: + edge_failures[edge.id] = edge_results[edge.id].failed.clone() + else: + edge = batch[0] + step = self.step_by_edge[edge.id] + branch_failed = blocked[edge.id] + self._ensure_assignment(step, branch_failed) + active = ~branch_failed + posture_before = self._object_not_fallen(step) + failure_policy = self._edge_failure_policy(edge) + try: + primary_result = self._execute_edge_with_retries( + edge, + step, + failed=branch_failed, + ) + except Exception as exc: + if failure_policy != "best_effort": + raise + primary_result = self._edge_exception_result( + edge, + step, + branch_failed, + exc, + ) + newly_failed = primary_result.failed & ~branch_failed + fallen_transition = self._fallen_transition( + step, + posture_before, + primary_result, + ) + recorder.edge( + edge.id, + step, + assignments=self._assignments[step.id], + grounded=primary_result.grounded, + active=active, + failed=primary_result.failed, + action_steps=len(primary_result.actions), + planner_traces=getattr(primary_result, "planner_traces", ()), + diagnostics=self._edge_diagnostics( + step, + edge, + primary_result.failed, + ), + phase="primary", + ) + edge_result = primary_result + if failure_policy == "task_required": + edge_result = self._recover_object_fallen( + edge, + step, + edge_result, + inherited_failed=branch_failed, + fallen_transition=fallen_transition, + recorder=recorder, + ) + if failure_policy == "best_effort": + # Best-effort parking is observable but cannot invalidate + # an already verified task or safety condition. + next_failed = branch_failed + else: + next_failed = edge_result.failed + executed_actions.extend(edge_result.actions) + edge_failures[edge.id] = next_failed + failure_events.extend( + self._failure_events( + edge, + step, + newly_failed & edge_result.failed, + postcondition=False, + executed=getattr(primary_result, "executed", None), + fallen_transition=fallen_transition, + planner_traces=getattr( + primary_result, "planner_traces", () + ), + ) + ) + + for edge in batch: + completed.add(edge.id) + remaining.remove(edge.id) + step = self.step_by_edge[edge.id] + if edge.id != step.edge_ids[-1]: + continue + prior_failed = edge_failures[edge.id] + verified_failed, step_success, observed = self._verify_step( + step, prior_failed + ) + postcondition_failed = verified_failed & ~prior_failed + recovery_covered = torch.zeros_like(verified_failed) + primary_step_recorded = False + if ( + self.placement_recovery_attempts + and bool(postcondition_failed.any()) + and step.operator == "place_relative" + and normalize_placement_relation( + step.goal.get("relation", "on") + ) + == "on" + ): + recorder.step( + step, + step_success, + observed=observed, + target=self._targets.get(step.id), + metadata=( + self._step_runtime_metadata(step) + if self.record_runtime + else None + ), + phase="primary", + ) + primary_step_recorded = True + recovery = self._recover_unstable_placement( + step, + postcondition_failed, + recorder=recorder, + ) + executed_actions.extend(recovery.actions) + failure_events.extend(recovery.failure_events) + recovery_covered = ( + torch.zeros_like(verified_failed) + if recovery.covered_failures is None + else recovery.covered_failures + ) + verified_failed = ( + verified_failed & ~postcondition_failed + ) | recovery.failed + step_success |= recovery.succeeded + observed = recovery.observed + failure_events.extend( + self._failure_events( + edge, + step, + verified_failed & ~prior_failed & ~recovery_covered, + postcondition=True, + executed=~prior_failed, + fallen_transition=None, + ) + ) + edge_failures[edge.id] = verified_failed + aggregate_failed |= ~step_success + semantic_success[step.id] = step_success + arrangement = self.arrangements.get(step.id) + if arrangement is not None: + arrangement.mark_completed(step.id, step_success) + if not primary_step_recorded: + recorder.step( + step, + step_success, + observed=observed, + target=self._targets.get(step.id), + metadata=( + self._step_runtime_metadata(step) + if self.record_runtime + else None + ), + ) + revalidation_failures = self._revalidate_support_relations() + for step_id, lost in revalidation_failures.items(): + step = self.steps[step_id] + edge = self.edges[step.edge_ids[-1]] + aggregate_failed |= lost + semantic_success[step_id] = semantic_success[step_id] & ~lost + edge_failures[edge.id] |= lost + failure_events.extend( + self._failure_events( + edge, + step, + lost, + postcondition=True, + executed=torch.ones_like(lost), + fallen_transition=None, + ) + ) + recorder.step( + step, + semantic_success[step_id], + observed=self._entity_pose(step.object_uid)[:, :3, 3], + target=self._targets.get(step.id), + metadata=( + self._step_runtime_metadata(step) + if self.record_runtime + else None + ), + phase="final_revalidation", + ) + record_dir = recorder.finalize(~aggregate_failed) + except BaseException as exc: + error_message = f"{type(exc).__name__}: {exc}" + recorder.finalize(~aggregate_failed, error=error_message) + raise + finally: + if error_message is not None: + log_warning(f"Action Engine execution aborted: {error_message}") + + return ExecutionResult( + actions=executed_actions, + success=~aggregate_failed, + semantic_success=semantic_success, + record_dir=record_dir, + retry_count=self.retry_count, + retry_counts=list(self._retry_counts), + recovery_count=( + 0 + if self.runtime_graph is None + else sum( + revision.kind == "insert_recovery" + for revision in self.runtime_graph.revisions + ) + ), + revision_count=( + 0 if self.runtime_graph is None else len(self.runtime_graph.revisions) + ), + failure_events=failure_events, + runtime_revisions=( + [] + if self.runtime_graph is None + else [ + { + "revision": revision.revision, + "kind": revision.kind, + "reason": revision.reason, + "failed_node_id": revision.failed_node_id, + "inserted_group_ids": list(revision.inserted_group_ids), + "replaced_group_ids": list(revision.replaced_group_ids), + "active_env_ids": list(revision.active_env_ids), + } + for revision in self.runtime_graph.revisions + ] + ), + ) + + def _failure_events( + self, + edge: ExecutionEdge, + step: SemanticStep, + failed: torch.Tensor, + *, + postcondition: bool, + executed: torch.Tensor | None, + fallen_transition: torch.Tensor | None, + planner_traces: Sequence[Mapping[str, Any]] = (), + ) -> list[dict[str, Any]]: + if not bool(failed.any()): + return [] + action = edge.actions[-1] + action_name = str(action["atomic_action_class"]) + capability = self.adapter.capabilities.get(action_name) + executed_mask = ( + torch.zeros_like(failed) + if executed is None + else torch.as_tensor( + executed, + dtype=torch.bool, + device=failed.device, + ).reshape(-1) + ) + if executed_mask.shape != failed.shape: + raise ValueError("Failure provenance mask must match failed rows.") + transitioned = ( + torch.zeros_like(failed) + if fallen_transition is None + else torch.as_tensor( + fallen_transition, + dtype=torch.bool, + device=failed.device, + ).reshape(-1) + ) + if transitioned.shape != failed.shape: + raise ValueError("Fallen-transition mask must match failed rows.") + failure_policy = ( + "task_required" if postcondition else self._edge_failure_policy(edge) + ) + fatal = failure_policy != "best_effort" + if postcondition: + classified = (("postcondition_failed", failed),) + else: + fallen = failed & executed_mask & transitioned + planning = failed & ~executed_mask + execution = failed & executed_mask & ~fallen + if capability.failure_classifier == "grasp": + execution_type = "grasp_missed" + elif capability.state_effect in {"preserve_hold", "transfer_hold"}: + execution_type = "object_dropped" + else: + execution_type = "plan_failed" + classified = ( + ("object_fallen", fallen), + ("search_exhausted", planning), + (execution_type, execution), + ) + result: list[dict[str, Any]] = [] + for failure_type, mask in classified: + env_ids = torch.nonzero(mask, as_tuple=False).flatten().tolist() + if not env_ids: + continue + if failure_type == "search_exhausted": + covered: set[int] = set() + for blocker in getattr(self, "_candidate_blockers", {}).get( + step.id, () + ): + env_id = int(blocker["env_id"]) + if env_id not in env_ids: + continue + assignment = self._assignments.get(step.id, [None] * len(failed))[ + env_id + ] + if assignment is not None and blocker.get("arm") != assignment: + continue + blocker_policy = str(blocker.get("failure_policy", failure_policy)) + result.append( + { + "node_id": blocker.get("node_id"), + "edge_id": edge.id, + "origin_edge_id": edge.id, + "blocking_edge_id": blocker["blocking_edge_id"], + "task_instance_id": step.id, + "atomic_action": blocker["atomic_action"], + "object_uid": step.object_uid, + "arm": blocker.get("arm"), + "failure_type": "search_exhausted", + "failure_policy": blocker_policy, + "fatal": blocker_policy != "best_effort", + "planning_stage": blocker["planning_stage"], + "search_strategy": blocker["search_strategy"], + "search_budget": deepcopy(blocker["search_budget"]), + "reason": ( + "Bounded candidate search exhausted without a " + "valid plan; this is not a geometric proof of " + "unreachability." + ), + "evidence": deepcopy(blocker["evidence"]), + "env_ids": [env_id], + } + ) + covered.add(env_id) + for env_id in (item for item in env_ids if item not in covered): + trace = next( + ( + item + for item in planner_traces + if str(item.get("arm", "")) + == str(self._assignments.get(step.id, [None])[env_id]) + ), + planner_traces[0] if planner_traces else {}, + ) + result.append( + { + "node_id": action.get("seed_node_id"), + "edge_id": edge.id, + "blocking_edge_id": edge.id, + "task_instance_id": step.id, + "atomic_action": action_name, + "arm": self._assignments.get(step.id, [None])[env_id], + "failure_type": "search_exhausted", + "failure_policy": failure_policy, + "fatal": fatal, + "planning_stage": "runtime_planning", + **self._planner_failure_details(trace, env_id), + "reason": ( + "Bounded runtime search exhausted without a valid " + "plan; this is not a geometric proof of " + "unreachability." + ), + "env_ids": [env_id], + } + ) + continue + result.append( + { + "node_id": action.get("seed_node_id"), + "edge_id": edge.id, + "blocking_edge_id": edge.id, + "task_instance_id": step.id, + "atomic_action": action_name, + "object_uid": step.object_uid, + "failure_type": failure_type, + "failure_policy": failure_policy, + "fatal": fatal, + "planning_stage": ( + "postcondition" if postcondition else "execution" + ), + "env_ids": env_ids, + } + ) + return result + + def _edge_exception_result( + self, + edge: ExecutionEdge, + step: SemanticStep, + inherited_failed: torch.Tensor, + exc: Exception, + ) -> _EdgeResult: + """Convert a planning exception into an auditable failed edge result.""" + action = edge.actions[0] + assignments = self._assignments.get(step.id, [None] * int(self.env.num_envs)) + arm = next((item for item in assignments if item is not None), None) + trace = { + "action_class": str(action.get("atomic_action_class")), + "arm": arm, + "primary_strategy": "planner_exception", + "primary_success": torch.zeros_like(inherited_failed), + "fallback_attempted": torch.zeros_like(inherited_failed), + "fallback_success": torch.zeros_like(inherited_failed), + "search_budget": self._planner_search_budget(), + "exception": f"{type(exc).__name__}: {exc}", + } + return _EdgeResult( + [], + torch.ones_like(inherited_failed), + [], + [trace], + torch.zeros_like(inherited_failed), + ) + + def _object_not_fallen(self, step: SemanticStep) -> torch.Tensor | None: + """Return the live posture predicate when the object supports it.""" + try: + return evaluate_predicate( + self.env, + {"type": "object_not_fallen", "object": step.object_uid}, + ) + except (TypeError, ValueError): + return None + + def _fallen_transition( + self, + step: SemanticStep, + before: torch.Tensor | None, + result: _EdgeResult, + ) -> torch.Tensor: + """Identify rows where an executed action changed upright to fallen.""" + result_executed = getattr(result, "executed", None) + if before is None or result_executed is None: + return torch.zeros_like(result.failed) + after = self._object_not_fallen(step) + if after is None: + return torch.zeros_like(result.failed) + executed = torch.as_tensor( + result_executed, + dtype=torch.bool, + device=result.failed.device, + ).reshape(-1) + return ( + executed & before.to(result.failed.device) & ~after.to(result.failed.device) + ) + + def _execute_edge_with_retries( + self, + edge: ExecutionEdge, + step: SemanticStep, + *, + failed: torch.Tensor, + ) -> _EdgeResult: + """Retry a complete AtomicAction with fresh Grounding on failed rows.""" + result = self._execute_edge(edge, step, failed=failed) + if self.runtime_graph is None or len(edge.actions) != 1: + return result + action = edge.actions[0] + node_id = action.get("seed_node_id") + if not isinstance(node_id, str): + return result + aggregate_actions = list(result.actions) + grounded = list(result.grounded) + planner_traces = list(getattr(result, "planner_traces", ())) + executed = ( + torch.zeros_like(result.failed) + if getattr(result, "executed", None) is None + else result.executed.clone() + ) + current_failed = result.failed.clone() + attempted_failure = current_failed & ~failed + while bool(attempted_failure.any()): + precondition = self._retry_precondition(node_id, attempted_failure) + decision = self.runtime_graph.record_failure( + node_id, + attempted_failure, + precondition_holds=precondition, + ) + if not bool(decision.retry.any()): + break + self.retry_count += int(decision.retry.sum()) + for env_id in ( + torch.nonzero(decision.retry, as_tuple=False).flatten().tolist() + ): + self._retry_counts[env_id] += 1 + self._consume_transitions(1) + for arm in ("left_arm", "right_arm"): + self._candidate_cache.pop((step.id, arm), None) + self._candidate_failures.pop((step.id, arm), None) + capability = self.adapter.capabilities.get( + str(action.get("atomic_action_class")) + ) + if capability.state_effect == "hold": + previous = list(self._assignments[step.id]) + for env_id in ( + torch.nonzero(decision.retry, as_tuple=False).flatten().tolist() + ): + arm = previous[env_id] + if step.actor.get("mode") == "auto" and arm in { + "left_arm", + "right_arm", + }: + self._pickup_retry_exclusions.setdefault( + (step.id, env_id), set() + ).add(str(arm)) + for arm in ("left_arm", "right_arm"): + self._step_states.pop((step.id, arm), None) + if step.actor.get("mode") == "auto": + self._assignments.pop(step.id, None) + self._ensure_assignment(step, ~decision.retry) + refreshed = self._assignments[step.id] + self._assignments[step.id] = [ + refreshed[index] if bool(decision.retry[index]) else assignment + for index, assignment in enumerate(previous) + ] + retry_result = self._execute_edge( + edge, + step, + failed=~decision.retry, + ) + aggregate_actions.extend(retry_result.actions) + grounded.extend(retry_result.grounded) + planner_traces.extend(getattr(retry_result, "planner_traces", ())) + if getattr(retry_result, "executed", None) is not None: + executed |= retry_result.executed + succeeded = decision.retry & ~retry_result.failed + current_failed &= ~succeeded + attempted_failure = decision.retry & retry_result.failed + return _EdgeResult( + aggregate_actions, + current_failed, + grounded, + planner_traces, + executed, + ) + + def _recover_object_fallen( + self, + edge: ExecutionEdge, + step: SemanticStep, + result: _EdgeResult, + *, + inherited_failed: torch.Tensor, + fallen_transition: torch.Tensor, + recorder: RuntimeRecorder, + ) -> _EdgeResult: + """Run the bounded E2 repair and replay only the failed vector rows.""" + if self.runtime_graph is None or len(edge.actions) != 1: + return result + node_id = edge.actions[0].get("seed_node_id") + if not isinstance(node_id, str) or not node_id: + return result + newly_failed = result.failed & ~inherited_failed + if not bool(newly_failed.any()): + return result + executed = ( + torch.zeros_like(result.failed) + if getattr(result, "executed", None) is None + else torch.as_tensor( + result.executed, + dtype=torch.bool, + device=result.failed.device, + ).reshape(-1) + ) + transition = torch.as_tensor( + fallen_transition, + dtype=torch.bool, + device=result.failed.device, + ).reshape(-1) + if ( + executed.shape != result.failed.shape + or transition.shape != result.failed.shape + ): + raise ValueError("Recovery provenance masks must match failed rows.") + fallen = newly_failed & executed & transition + if not bool(fallen.any()): + return result + + env_ids = torch.nonzero(fallen, as_tuple=False).flatten().tolist() + original_assignment = list( + self._assignments.get(step.id, [None] * int(self.env.num_envs)) + ) + try: + patched = self.runtime_graph.insert_default_recovery( + failed_node_id=node_id, + failure_type="object_fallen", + active_env_ids=env_ids, + resume_failed_group=True, + ) + revision = self.runtime_graph.revisions[-1] + recovery_group_id = revision.inserted_group_ids[0] + from .loader import load_execution_program + + recovery_program = load_execution_program( + patched, + registry=self.capability_registry, + require_executable=True, + ) + recovery_step = next( + item + for item in recovery_program.semantic_steps + if item.id == recovery_group_id + ) + recovery_spec = next( + item + for item in recovery_program.raw["semantic_steps"] + if str(item["id"]) == recovery_group_id + ) + recorder.register_step(recovery_step, recovery_spec) + recovery_edges = { + item.id: item + for item in recovery_program.edges + if item.id in set(recovery_step.edge_ids) + } + if set(recovery_edges) != set(recovery_step.edge_ids): + raise RuntimeError("Compiled recovery group is incomplete.") + except Exception as exc: + recorder.recovery( + failure_type="object_fallen", + failed_node_id=node_id, + active=fallen, + status="rejected", + error=f"{type(exc).__name__}: {exc}", + semantic_step_id=step.id, + ) + return result + + recorder.recovery( + failure_type="object_fallen", + failed_node_id=node_id, + active=fallen, + status="started", + recovery_group_id=recovery_group_id, + semantic_step_id=step.id, + ) + aggregate_actions = list(result.actions) + grounded = list(result.grounded) + planner_traces = list(getattr(result, "planner_traces", ())) + self._clear_recovery_rows(step, fallen) + + # Recovery edges are compiled from the revised graph but execute through + # this executor so live ownership, recorder, and simulator state remain + # continuous. They are removed from the scheduling maps afterwards. + installed_edge_ids: list[str] = [] + self.steps[recovery_step.id] = recovery_step + for recovery_edge in recovery_edges.values(): + self.edges[recovery_edge.id] = recovery_edge + self.step_by_edge[recovery_edge.id] = recovery_step + installed_edge_ids.append(recovery_edge.id) + recovery_failed = ~fallen + try: + self._assignments.pop(recovery_step.id, None) + self._ensure_assignment(recovery_step, recovery_failed) + for recovery_edge_id in recovery_step.edge_ids: + self._consume_transitions(1) + recovery_edge = recovery_edges[recovery_edge_id] + recovery_result = self._execute_edge_with_retries( + recovery_edge, + recovery_step, + failed=recovery_failed, + ) + recorder.edge( + recovery_edge.id, + recovery_step, + assignments=self._assignments[recovery_step.id], + grounded=recovery_result.grounded, + active=~recovery_failed, + failed=recovery_result.failed, + action_steps=len(recovery_result.actions), + planner_traces=recovery_result.planner_traces, + phase="recovery", + ) + aggregate_actions.extend(recovery_result.actions) + grounded.extend(recovery_result.grounded) + planner_traces.extend(recovery_result.planner_traces) + recovery_failed = recovery_result.failed + _, recovery_success, observed = self._verify_step( + recovery_step, + recovery_failed, + ) + recorder.step( + recovery_step, + recovery_success, + observed=observed, + target=self._targets.get(recovery_step.id), + metadata=( + self._step_runtime_metadata(recovery_step) + if self.record_runtime + else None + ), + phase="recovery", + ) + except Exception as exc: + recorder.recovery( + failure_type="object_fallen", + failed_node_id=node_id, + active=fallen, + status="failed", + recovery_group_id=recovery_group_id, + error=f"{type(exc).__name__}: {exc}", + semantic_step_id=step.id, + ) + return _EdgeResult( + aggregate_actions, + result.failed, + grounded, + planner_traces, + result.executed, + ) + finally: + for recovery_edge_id in installed_edge_ids: + self.edges.pop(recovery_edge_id, None) + self.step_by_edge.pop(recovery_edge_id, None) + self.steps.pop(recovery_step.id, None) + + recovered = fallen & recovery_success + if not bool(recovered.any()): + recorder.recovery( + failure_type="object_fallen", + failed_node_id=node_id, + active=fallen, + status="failed", + recovery_group_id=recovery_group_id, + semantic_step_id=step.id, + ) + return _EdgeResult( + aggregate_actions, + result.failed, + grounded, + planner_traces, + result.executed, + ) + + # Recompute this TaskGroup's assignment for recovered rows, retaining + # the untouched assignments of healthy vector rows. Replay the prefix + # through the failed edge; the ordinary main loop will then continue at + # the next edge and verify the TaskGroup exactly once. + try: + self._assignments.pop(step.id, None) + for arm in ("left_arm", "right_arm"): + self._candidate_cache.pop((step.id, arm), None) + self._candidate_failures.pop((step.id, arm), None) + self._ensure_assignment(step, ~recovered) + replay_assignment = self._assignments[step.id] + self._assignments[step.id] = [ + ( + replay_assignment[index] + if bool(recovered[index]) + else original_assignment[index] + ) + for index in range(int(self.env.num_envs)) + ] + replay_failed = ~recovered + for prefix_edge_id in step.edge_ids: + self._consume_transitions(1) + prefix_edge = self.edges[prefix_edge_id] + replay_active = ~replay_failed + prefix_result = self._execute_edge_with_retries( + prefix_edge, + step, + failed=replay_failed, + ) + recorder.edge( + prefix_edge.id, + step, + assignments=self._assignments[step.id], + grounded=prefix_result.grounded, + active=replay_active, + failed=prefix_result.failed, + action_steps=len(prefix_result.actions), + planner_traces=prefix_result.planner_traces, + diagnostics=self._edge_diagnostics( + step, + prefix_edge, + prefix_result.failed, + ), + phase="replay", + ) + aggregate_actions.extend(prefix_result.actions) + grounded.extend(prefix_result.grounded) + planner_traces.extend(prefix_result.planner_traces) + replay_failed = prefix_result.failed + if prefix_edge_id == edge.id: + break + except Exception as exc: + recorder.recovery( + failure_type="object_fallen", + failed_node_id=node_id, + active=fallen, + status="failed", + recovery_group_id=recovery_group_id, + error=f"{type(exc).__name__}: {exc}", + semantic_step_id=step.id, + ) + return _EdgeResult( + aggregate_actions, + result.failed, + grounded, + planner_traces, + result.executed, + ) + final_failed = result.failed.clone() + final_failed[fallen] = replay_failed[fallen] + recorder.recovery( + failure_type="object_fallen", + failed_node_id=node_id, + active=fallen, + status=("succeeded" if not bool(final_failed[fallen].any()) else "failed"), + recovery_group_id=recovery_group_id, + semantic_step_id=step.id, + ) + return _EdgeResult( + aggregate_actions, + final_failed, + grounded, + planner_traces, + result.executed, + ) + + def _clear_recovery_rows( + self, + step: SemanticStep, + mask: torch.Tensor, + ) -> None: + """Discard stale hold projections only for rows entering recovery.""" + owners = self._object_owners.setdefault( + step.object_uid, [None] * int(self.env.num_envs) + ) + for env_id in torch.nonzero(mask, as_tuple=False).flatten().tolist(): + owner = owners[env_id] + owners[env_id] = None + if owner in self._arm_owners and ( + self._arm_owners[str(owner)][env_id] == step.object_uid + ): + self._arm_owners[str(owner)][env_id] = None + for arm in ("left_arm", "right_arm"): + if self._arm_owners[arm][env_id] == step.object_uid: + self._arm_owners[arm][env_id] = None + + candidate_keys = [ + key for key in self._object_states if key[0] == step.object_uid + ] + step_keys = [key for key in self._step_states if key[0] == step.id] + for cache, keys in ( + (self._object_states, candidate_keys), + (self._step_states, step_keys), + ): + for key in keys: + state = cache[key] + delta = StateDelta( + held_object_updates={name: None for name in state.held_objects}, + ) + if delta.is_empty: + continue + cache[key] = ExecutionState.from_task_state( + delta.apply(state.to_task_state(), mask), + last_qpos=self.env.robot.get_qpos().clone(), + ) + + def _recover_unstable_placement( + self, + step: SemanticStep, + failed: torch.Tensor, + *, + recorder: RuntimeRecorder, + ) -> _PlacementRecoveryResult: + """Regrasp after release, then retry unused placement poses only.""" + pending = failed.clone() + recovered = torch.zeros_like(failed) + observed = self._entity_pose(step.object_uid)[:, :3, 3] + actions: list[torch.Tensor] = [] + blocking_failures: list[tuple[ExecutionEdge, _EdgeResult, torch.Tensor]] = [] + terminal_edge = self.edges[step.edge_ids[-1]] + failed_node_id = str( + terminal_edge.actions[-1].get("seed_node_id", terminal_edge.id) + ) + recorder.recovery( + failure_type="placement_unstable", + failed_node_id=failed_node_id, + active=failed, + status="started", + semantic_step_id=step.id, + ) + for _attempt in range(self.placement_recovery_attempts): + if not bool(pending.any()): + break + self._consume_transitions(len(step.edge_ids)) + attempt_active = pending.clone() + self._clear_recovery_rows(step, attempt_active) + self._assignments.pop(step.id, None) + for arm in ("left_arm", "right_arm"): + self._candidate_cache.pop((step.id, arm), None) + self._candidate_failures.pop((step.id, arm), None) + try: + self._ensure_assignment(step, ~attempt_active) + except Exception as exc: + blocking_edge = self.edges[step.edge_ids[0]] + blocking_result = self._edge_exception_result( + blocking_edge, + step, + ~attempt_active, + exc, + ) + blocking_failures.append( + (blocking_edge, blocking_result, attempt_active) + ) + recorder.edge( + blocking_edge.id, + step, + assignments=self._assignments.get( + step.id, + [None] * int(self.env.num_envs), + ), + grounded=(), + active=attempt_active, + failed=blocking_result.failed, + action_steps=0, + planner_traces=blocking_result.planner_traces, + diagnostics=self._edge_diagnostics( + step, + blocking_edge, + blocking_result.failed, + ), + phase="recovery", + ) + break + + replay_failed = ~attempt_active + for edge_id in step.edge_ids: + edge = self.edges[edge_id] + edge_active = ~replay_failed + try: + result = self._execute_edge_with_retries( + edge, + step, + failed=replay_failed, + ) + except Exception as exc: + result = self._edge_exception_result( + edge, + step, + replay_failed, + exc, + ) + actions.extend(result.actions) + recorder.edge( + edge.id, + step, + assignments=self._assignments[step.id], + grounded=result.grounded, + active=edge_active, + failed=result.failed, + action_steps=len(result.actions), + planner_traces=result.planner_traces, + diagnostics=self._edge_diagnostics(step, edge, result.failed), + phase="recovery", + ) + newly_failed = edge_active & result.failed + if bool(newly_failed.any()): + blocking_failures.append((edge, result, newly_failed)) + replay_failed = result.failed + if not bool((attempt_active & ~replay_failed).any()): + break + + execution_succeeded = attempt_active & ~replay_failed + if not bool(execution_succeeded.any()): + break + verified_failed, verified_success, observed = self._verify_step( + step, + ~execution_succeeded, + ) + del verified_failed + recovered_now = attempt_active & verified_success + recovered |= recovered_now + recorder.step( + step, + verified_success, + observed=observed, + target=self._targets.get(step.id), + metadata=( + self._step_runtime_metadata(step) if self.record_runtime else None + ), + phase="recovery", + ) + action_failed = attempt_active & replay_failed + pending &= ~recovered_now + if bool(action_failed.any()): + break + + final_failed = failed & ~recovered + recovery_events: list[dict[str, Any]] = [] + covered_failures = torch.zeros_like(failed) + for blocking_edge, blocking_result, blocking_rows in blocking_failures: + event_rows = final_failed & blocking_rows & ~covered_failures + if not bool(event_rows.any()): + continue + events = self._failure_events( + blocking_edge, + step, + event_rows, + postcondition=False, + executed=blocking_result.executed, + fallen_transition=None, + planner_traces=blocking_result.planner_traces, + ) + for event in events: + event["phase"] = "recovery" + event["origin_edge_id"] = terminal_edge.id + recovery_events.extend(events) + covered_failures |= event_rows + recorder.recovery( + failure_type="placement_unstable", + failed_node_id=failed_node_id, + active=failed, + status="failed" if bool(final_failed.any()) else "succeeded", + semantic_step_id=step.id, + ) + return _PlacementRecoveryResult( + failed=final_failed, + succeeded=recovered, + observed=observed, + actions=actions, + failure_events=recovery_events, + covered_failures=covered_failures, + ) + + def _retry_precondition( + self, + node_id: str, + failed: torch.Tensor, + ) -> torch.Tensor: + assert self.runtime_graph is not None + node = next( + item for item in self.runtime_graph.graph["nodes"] if item["id"] == node_id + ) + predicate = node.get("precondition", {}) + if not predicate: + return failed.clone() + try: + return failed & evaluate_predicate( + self.env, + predicate, + held_owners=self._object_owners, + held_states=self._object_states, + coordinated_state=self._step_states.get( + (str(node.get("task_instance_id", "")), "coordinated") + ), + ) + except (TypeError, ValueError): + return torch.zeros_like(failed) + + def _dependency_failures( + self, + edge: ExecutionEdge, + failures: Mapping[str, torch.Tensor], + ) -> torch.Tensor: + """Return success-required failures that can reach this edge.""" + result = torch.zeros( + int(self.env.num_envs), dtype=torch.bool, device=self.env.device + ) + for dependency in edge.depends_on: + if (dependency, edge.id) in self._completion_only_dependencies: + continue + result |= failures[dependency] + return result + + def _completion_only_dependency_edges(self) -> frozenset[tuple[str, str]]: + """Resolve linker-added resource ordering to executable edge pairs.""" + graph = self.program.seed_graph + if not isinstance(graph, Mapping): + return frozenset() + metadata = graph.get("metadata", {}) + if not isinstance(metadata, Mapping): + return frozenset() + + reasons_by_pair: dict[tuple[str, str], set[str]] = {} + sources = ( + ("action_contract_task_linker", "linked_dependencies"), + ("action_contract_linker", "group_dependencies"), + ) + for metadata_key, dependency_key in sources: + provenance = metadata.get(metadata_key, {}) + if not isinstance(provenance, Mapping): + continue + dependencies = provenance.get(dependency_key, ()) + if not isinstance(dependencies, Sequence) or isinstance( + dependencies, (str, bytes, bytearray) + ): + continue + for dependency in dependencies: + if not isinstance(dependency, Mapping): + continue + parent = dependency.get("from") + child = dependency.get("to") + reason = dependency.get("reason") + if not all( + isinstance(value, str) and value for value in (parent, child) + ): + continue + if reason not in {"causal", "resource"}: + continue + reasons_by_pair.setdefault((parent, child), set()).add(reason) + + completion_only_steps = { + pair for pair, reasons in reasons_by_pair.items() if reasons == {"resource"} + } + return frozenset( + (dependency, edge.id) + for edge in self.program.edges + for dependency in edge.depends_on + if ( + self.step_by_edge[dependency].id, + self.step_by_edge[edge.id].id, + ) + in completion_only_steps + ) + + def _reset_runtime_state(self) -> None: + self.retry_count = 0 + if self.program.seed_graph is not None: + execution = self.runtime_policy.execution + self.runtime_graph = RuntimeGraph( + self.program.seed_graph, + num_envs=int(self.env.num_envs), + max_retries=self.max_retries_per_action, + max_revisions=int(execution["max_graph_revisions"]), + max_recovery_actions=int(execution["max_recovery_actions"]), + registry=self.capability_registry, + ) + self._step_states.clear() + self._object_states.clear() + self._object_owners.clear() + for owners in self._arm_owners.values(): + owners[:] = [None] * int(self.env.num_envs) + self._assignments.clear() + self._candidate_cache.clear() + self._candidate_failures.clear() + self._candidate_diagnostics.clear() + self._candidate_blockers.clear() + self._reported_candidates.clear() + self._pickup_retry_exclusions.clear() + self._targets.clear() + self._target_poses.clear() + self._orientation_references.clear() + self._orientation_errors.clear() + self._policies.clear() + self._payload_initial.clear() + self._support_relations.clear() + self._placement_candidate_history.clear() + self._robot_lateral_axis_cache = None + self._transition_count = 0 + self._retry_counts = [0] * int(self.env.num_envs) + + def _consume_transitions(self, count: int) -> None: + """Charge ordinary, retry, and recovery edges to one runtime budget.""" + self._transition_count += int(count) + if self._transition_count > self.max_transitions: + raise RuntimeError("Execution exceeded max_transitions.") + + def _pack_ready_edges( + self, + ready: Sequence[ExecutionEdge], + *, + inactive: Mapping[str, torch.Tensor] | None = None, + completed: set[str] | None = None, + ) -> tuple[ExecutionEdge, ...]: + """Prefer progress on held payloads and pack only resource-safe pickups.""" + inactive = inactive or {} + completed = completed or set() + schedulable = [ + edge + for edge in ready + if not self._temporarily_resource_blocked( + edge, + inactive.get(edge.id), + ) + ] + started = [ + edge + for edge in schedulable + if any( + edge_id in completed for edge_id in self.step_by_edge[edge.id].edge_ids + ) + ] + candidates = started or schedulable or list(ready) + first = candidates[0] + if not self._parallel_pickup_candidate(first): + return (first,) + if not self._two_arms_available(inactive.get(first.id)): + return (first,) + first_step = self.step_by_edge[first.id] + for second in candidates[1:]: + if not self._parallel_pickup_candidate(second): + continue + second_step = self.step_by_edge[second.id] + if first_step.object_uid == second_step.object_uid: + continue + shared = set(first.resources) & set(second.resources) + same_group = self.group_by_step.get(first_step.id) is not None and ( + self.group_by_step.get(first_step.id) + is self.group_by_step.get(second_step.id) + ) + if same_group: + # A shared destination workspace constrains transport/place, + # not two independent pickups declared by this group. + conflicts = { + item + for item in shared + if item != "arm:auto" and not item.startswith("workspace:") + } + else: + conflicts = shared - {"arm:auto"} + if conflicts: + continue + required_opposite = ( + first_step.actor.get("mode") == "required" + and second_step.actor.get("mode") == "required" + and first_step.actor.get("arm") != second_step.actor.get("arm") + ) + if same_group or required_opposite: + return first, second + return (first,) + + def _temporarily_resource_blocked( + self, + edge: ExecutionEdge, + inactive: torch.Tensor | None, + ) -> bool: + """Defer a new pickup while its arm is carrying another payload.""" + if not self._parallel_pickup_candidate(edge): + return False + step = self.step_by_edge[edge.id] + mode = str(step.actor.get("mode", "auto")) + inactive_mask = ( + torch.zeros( + int(self.env.num_envs), dtype=torch.bool, device=self.env.device + ) + if inactive is None + else inactive + ) + for env_id in range(int(self.env.num_envs)): + if bool(inactive_mask[env_id]): + continue + if mode == "required": + arms = (str(step.actor["arm"]),) + else: + arms = ("left_arm", "right_arm") + if not any( + self._arm_owners[arm][env_id] in {None, step.object_uid} for arm in arms + ): + return True + return False + + def _two_arms_available(self, inactive: torch.Tensor | None) -> bool: + inactive_mask = ( + torch.zeros( + int(self.env.num_envs), dtype=torch.bool, device=self.env.device + ) + if inactive is None + else inactive + ) + for env_id in range(int(self.env.num_envs)): + if bool(inactive_mask[env_id]): + continue + free = sum( + self._arm_owners[arm][env_id] is None + for arm in ("left_arm", "right_arm") + ) + if free < 2: + return False + return True + + def _parallel_pickup_candidate(self, edge: ExecutionEdge) -> bool: + if len(edge.actions) != 1: + return False + step = self.step_by_edge[edge.id] + capability = self.adapter.capabilities.get( + str(edge.actions[0].get("atomic_action_class")) + ) + return ( + capability.state_effect == "hold" + and capability.resource_mode == "single_arm_object" + and step.actor.get("mode") in {"auto", "required"} + and step.operator != "orient_object" + ) + + def _preferred_in_place_arm( + self, + step: SemanticStep, + env_id: int, + ) -> str | None: + """Map a clearly sided in-place object to the robot-view arm slot.""" + if step.operator != "orient_object": + return None + initial = getattr(self.env, "agent_initial_object_poses", {}).get( + step.object_uid + ) + if initial is None: + entity = self.env.sim.get_rigid_object(step.object_uid) + if entity is None: + return None + initial = entity.get_local_pose(to_matrix=True) + pose = torch.as_tensor(initial, device=self.env.device) + if pose.ndim == 2: + pose = pose.unsqueeze(0) + center, _, lateral_axis = self._arm_selection_workspace(step) + index = min(env_id, pose.shape[0] - 1) + lateral = float( + torch.sum((pose[index, :2, 3] - center[index]) * lateral_axis[index]) + ) + if ( + abs(lateral) + <= self.runtime_policy.arm_selection.orient_object_preferred_arm_deadband + ): + return None + return "left_arm" if lateral > 0.0 else "right_arm" + + def _preferred_live_pickup_arm( + self, + step: SemanticStep, + env_id: int, + ) -> str | None: + """Return the object's same-side arm outside the central deadband.""" + pose = self._entity_pose(step.object_uid) + center, half_width, lateral_axis = self._arm_selection_workspace(step) + index = min(env_id, pose.shape[0] - 1) + lateral = float( + torch.sum((pose[index, :2, 3] - center[index]) * lateral_axis[index]) + ) + deadband = float(half_width[index]) * float( + self.runtime_policy.arm_selection.crossing_deadband_ratio + ) + if abs(lateral) <= deadband: + return None + return "left_arm" if lateral > 0.0 else "right_arm" + + def _auto_arm_is_allowed( + self, + step: SemanticStep, + arm: str, + env_id: int, + ) -> bool: + """Apply the same-side constraint to automatic arm allocation.""" + if step.actor.get("mode") != "auto": + return True + preferred = self._preferred_live_pickup_arm(step, env_id) + if preferred is None or arm == preferred: + return True + excluded = self._pickup_retry_exclusions.get((step.id, env_id), set()) + return bool( + self.runtime_policy.arm_selection.allow_cross_side_fallback + and preferred in excluded + ) + + def _ensure_assignment( + self, + step: SemanticStep, + failed: torch.Tensor, + *, + allow_rematch: bool = True, + ) -> None: + self._capture_orientation_reference(step) + if step.id in self._assignments: + return + if bool(failed.all()): + self._assignments[step.id] = [None] * int(self.env.num_envs) + return + mode = str(step.actor.get("mode", "auto")) + group = self.group_by_step.get(step.id) + if mode == "auto" and group is not None: + self._ensure_serial_group_assignments(group, failed) + if step.id in self._assignments: + return + if mode == "coordinated": + self._assignments[step.id] = [ + ( + None + if bool(failed[index]) + or self._arm_owners["left_arm"][index] is not None + or self._arm_owners["right_arm"][index] is not None + else "coordinated" + ) + for index in range(len(failed)) + ] + return + if mode == "required": + arm = str(step.actor["arm"]) + first_action = self.edges[step.edge_ids[0]].actions[0] + first_capability = self.adapter.capabilities.get( + str(first_action.get("atomic_action_class")) + ) + if step.operator == "handover" and first_capability.state_effect != "hold": + # A coordinated handover has an internal, multi-arm planner. + # Do not let a speculative single-arm suffix plan veto the + # real execution (or create a misleading downstream pickup + # error) once a predecessor already established the transfer + # hold. A standalone E4 starts with PickUp and still needs its + # cached candidate plan for that first action. + source_state = self._state_for(step, arm) + has_hold = ( + source_state.get_held_object(arm_control_part(self.env, arm)) + is not None + ) + self._assignments[step.id] = [ + arm if has_hold and not bool(failed[index]) else None + for index in range(len(failed)) + ] + return + candidate = self._candidate(step, arm, failed) + conflicts = self._resource_conflicts(step, arm) + self._assignments[step.id] = [ + ( + arm + if not bool(failed[index]) and not bool(conflicts[index]) + else None + ) + for index in range(len(failed)) + ] + self._report_candidates(step, (candidate,)) + return + + left = self._candidate(step, "left_arm", failed) + right = self._candidate(step, "right_arm", failed) + candidates = {"left_arm": left, "right_arm": right} + conflicts = { + arm: self._resource_conflicts(step, arm) + for arm in ("left_arm", "right_arm") + } + owners = self._object_owners.get(step.object_uid, [None] * len(failed)) + assignments: list[str | None] = [] + selection_failed = torch.zeros_like(failed) + for env_id in range(len(failed)): + if bool(failed[env_id]): + assignments.append(None) + continue + if owners[env_id] is not None: + owner = str(owners[env_id]) + excluded = self._pickup_retry_exclusions.get((step.id, env_id), set()) + if owner not in excluded and not bool(conflicts[owner][env_id]): + assignments.append(owner) + else: + assignments.append(None) + selection_failed[env_id] = True + continue + excluded = self._pickup_retry_exclusions.get((step.id, env_id), set()) + available = [ + arm + for arm in ("left_arm", "right_arm") + if arm not in excluded + and not bool(conflicts[arm][env_id]) + and self._auto_arm_is_allowed(step, arm, env_id) + ] + if not available: + assignments.append(None) + selection_failed[env_id] = True + continue + preferred = self._preferred_in_place_arm(step, env_id) + feasible = [ + arm for arm in available if bool(candidates[arm].feasible[env_id]) + ] + if preferred in feasible: + assignments.append(preferred) + elif feasible: + assignments.append( + min(feasible, key=lambda arm: float(candidates[arm].cost[env_id])) + ) + else: + live_preferred = self._preferred_live_pickup_arm(step, env_id) + assignments.append( + live_preferred if live_preferred in available else available[0] + ) + + if ( + allow_rematch + and bool(selection_failed.any()) + and step.id in self.arrangements + and step.goal.get("slot_constraint") == "free_reassignable" + and bool(self._rematch_arrangement(step, selection_failed, failed).any()) + ): + self._assignments.pop(step.id, None) + self._ensure_assignment(step, failed, allow_rematch=False) + return + self._assignments[step.id] = assignments + self._report_candidates(step, (left, right)) + + def _ensure_serial_group_assignments( + self, + group: Mapping[str, Any], + failed: torch.Tensor, + ) -> None: + """Bind a distinct-arm pair even when its operators execute serially.""" + step_ids = [str(value) for value in group.get("semantic_step_ids", ())] + if len(step_ids) != 2 or any( + step_id in self._assignments for step_id in step_ids + ): + return + steps = [self.steps[step_id] for step_id in step_ids] + for candidate_step in steps: + self._capture_orientation_reference(candidate_step) + candidates = { + (candidate_step.id, arm): self._candidate(candidate_step, arm, failed) + for candidate_step in steps + for arm in ("left_arm", "right_arm") + } + for candidate_step in steps: + self._report_candidates( + candidate_step, + ( + candidates[(candidate_step.id, "left_arm")], + candidates[(candidate_step.id, "right_arm")], + ), + ) + assignments = { + candidate_step.id: [None] * len(failed) for candidate_step in steps + } + permutations = (("left_arm", "right_arm"), ("right_arm", "left_arm")) + for env_id in range(len(failed)): + if bool(failed[env_id]): + continue + ranked: list[tuple[bool, bool, float, float, str, str]] = [] + for first_arm, second_arm in permutations: + first = candidates[(steps[0].id, first_arm)] + second = candidates[(steps[1].id, second_arm)] + feasible = bool(first.feasible[env_id] and second.feasible[env_id]) + required_match = all( + candidate_step.actor.get("mode") != "required" + or str(candidate_step.actor.get("arm")) == candidate_arm + for candidate_step, candidate_arm in ( + (steps[0], first_arm), + (steps[1], second_arm), + ) + ) + available = ( + required_match + and not bool( + self._resource_conflicts(steps[0], first_arm)[env_id] + or self._resource_conflicts(steps[1], second_arm)[env_id] + ) + and all( + self._auto_arm_is_allowed(candidate_step, candidate_arm, env_id) + for candidate_step, candidate_arm in ( + (steps[0], first_arm), + (steps[1], second_arm), + ) + ) + ) + preferred = ( + self._preferred_in_place_arm(steps[0], env_id), + self._preferred_in_place_arm(steps[1], env_id), + ) + side_penalty = float(first_arm != preferred[0]) if preferred[0] else 0.0 + side_penalty += ( + float(second_arm != preferred[1]) if preferred[1] else 0.0 + ) + ranked.append( + ( + not available, + not feasible, + side_penalty, + float(first.cost[env_id] + second.cost[env_id]), + first_arm, + second_arm, + ) + ) + ranked.sort() + unavailable, _, _, _, first_arm, second_arm = ranked[0] + if unavailable: + continue + assignments[steps[0].id][env_id] = first_arm + assignments[steps[1].id][env_id] = second_arm + self._assignments.update(assignments) + + def _candidate( + self, + step: SemanticStep, + arm: str, + failed: torch.Tensor, + ) -> _Candidate: + """Plan the complete semantic suffix before fixing an arm.""" + if step.actor.get("mode") == "required" and str(step.actor.get("arm")) != arm: + return _Candidate( + feasible=torch.zeros_like(failed), + cost=torch.full( + failed.shape, + torch.inf, + dtype=torch.float32, + device=self.env.device, + ), + plans={}, + ) + cached = self._candidate_cache.get((step.id, arm)) + if cached is not None: + return _Candidate( + feasible=cached.feasible & ~failed, + cost=cached.cost, + plans=cached.plans, + score_components=cached.score_components, + warnings=cached.warnings, + blockers=cached.blockers, + ) + feasible = ~failed.clone() & ~self._resource_conflicts(step, arm) + motion_cost = torch.zeros( + int(self.env.num_envs), + dtype=torch.float32, + device=self.env.device, + ) + source_pose = self._entity_pose(step.object_uid) + target_pose = None + state = self._state_for(step, arm) + reference_eef_pose = None + plans: dict[str, tuple[GroundedAction, ActionOutcome]] = {} + warnings: list[str] = [] + blockers: list[dict[str, Any]] = [] + try: + with _capture_speculative_warnings() as captured: + for edge_id in step.edge_ids: + edge = self.edges[edge_id] + if len(edge.actions) != 1: + raise ValueError( + "Auto/required arm candidates require one action per edge." + ) + action = edge.actions[0] + capability = self.adapter.capabilities.get( + str(action.get("atomic_action_class")) + ) + if ( + step.operator == "handover" + and capability.resource_mode == "coordinated_object" + ): + # A standalone E4 needs a speculative PickUp/staging + # prefix to choose and cache its transfer arm. The + # actual HandOver is coordinated, however, and must + # only be planned from the live post-staging state. + break + failure_policy = self._edge_failure_policy(edge) + try: + if capability.state_effect == "hold": + grounded = self.grounder.ground( + action, + step, + arm=arm, + state=state, + reference_eef_pose=reference_eef_pose, + orientation_reference_pose=self._orientation_references.get( + step.id + ), + ) + grounded = self._with_downstream_targets( + step, edge_id, arm, state, grounded + ) + outcome = self.adapter.plan(grounded, state) + else: + grounded, outcome = self._ground_and_plan_candidates( + action, + step, + arm=arm, + state=state, + active=feasible & ~failed, + reference_eef_pose=reference_eef_pose, + orientation_reference_pose=self._orientation_references.get( + step.id + ), + ) + except Exception as exc: + if failure_policy != "best_effort": + blockers.extend( + self._candidate_exception_blockers( + step, + edge, + arm, + failed, + exc, + ) + ) + raise + warnings.append( + f"{arm} best-effort action could not be planned at " + f"{edge_id} ({capability.name}): " + f"{type(exc).__name__}: {exc}" + ) + continue + plans[edge_id] = (grounded, outcome) + if failure_policy != "best_effort": + feasible &= outcome.success + motion_cost += outcome.cost + blockers.extend( + self._candidate_outcome_blockers( + step, + edge, + arm, + failed, + outcome, + ) + ) + elif not bool(outcome.success.all()): + warnings.append( + f"{arm} best-effort action degraded at {edge_id} " + f"({capability.name}); required suffix remains feasible." + ) + state = outcome.next_state + target = outcome.grounded.target_object_pose + if isinstance(target, torch.Tensor): + reference_eef_pose = self._eef_target(outcome) + binding = edge.actions[0].get("target_binding", {}) + if ( + binding.get("kind") + in { + "semantic_goal", + "coordinated_goal", + } + and binding.get("phase", "final") != "staging" + ): + target_pose = target + if not bool((feasible & ~failed).any()): + target = getattr(grounded.target, "xpos", None) + target_detail = "" + if isinstance(target, torch.Tensor) and target.shape[-2:] == ( + 4, + 4, + ): + target_z = target[..., 2, 3] + target_detail = ( + f" target_z=[{float(target_z.min()):.3f}, " + f"{float(target_z.max()):.3f}]" + ) + warnings.append( + f"{arm} candidate became infeasible at {edge_id} " + f"({capability.name}).{target_detail}" + ) + break + warnings.extend(captured) + except Exception as exc: + self._candidate_failures[(step.id, arm)] = f"{type(exc).__name__}: {exc}" + feasible = torch.zeros_like(failed) + motion_cost[:] = torch.inf + center_xy, half_width, lateral_axis = self._arm_selection_workspace(step) + score_components = _score_arm_candidate( + arm=arm, + motion_cost=motion_cost, + source_pose=source_pose, + target_pose=target_pose, + workspace_center_xy=center_xy, + workspace_half_width=half_width, + robot_lateral_axis=lateral_axis, + policy=self.runtime_policy.arm_selection, + ) + cost = score_components["total_cost"] + candidate = _Candidate( + feasible=feasible, + cost=cost, + plans=plans, + score_components=score_components, + warnings=tuple(warnings), + blockers=tuple(blockers), + ) + self._candidate_cache[(step.id, arm)] = candidate + return _Candidate( + feasible=feasible & ~failed, + cost=cost, + plans=plans, + score_components=score_components, + warnings=tuple(warnings), + blockers=tuple(blockers), + ) + + def _ground_and_plan_candidates( + self, + action: Mapping[str, Any], + step: SemanticStep, + *, + arm: str, + state: ExecutionState, + active: torch.Tensor, + reference_eef_pose: torch.Tensor | None = None, + orientation_reference_pose: torch.Tensor | None = None, + ) -> tuple[GroundedAction, ActionOutcome]: + """Plan live grounding candidates and retain the best bounded attempt.""" + groundings = self.grounder.ground_candidates( + action, + step, + arm=arm, + state=state, + reference_eef_pose=reference_eef_pose, + orientation_reference_pose=orientation_reference_pose, + ) + used = self._placement_candidate_history.get((step.id, arm), set()) + selected: tuple[GroundedAction, ActionOutcome] | None = None + selected_rank: tuple[int, float, int] | None = None + attempts: list[dict[str, Any]] = [] + last_error: Exception | None = None + for ordinal, grounded in enumerate(groundings): + candidate_index = int( + grounded.motion_policy.get("placement_candidate_index", ordinal) + ) + is_placement = "placement_candidate_index" in grounded.motion_policy + if is_placement and candidate_index in used: + attempts.append( + { + "candidate_index": candidate_index, + "status": "previously_released", + } + ) + continue + try: + outcome = self.adapter.plan(grounded, state) + except Exception as exc: + last_error = exc + attempts.append( + { + "candidate_index": candidate_index, + "status": "planning_error", + "error": f"{type(exc).__name__}: {exc}", + } + ) + continue + failed_count = int((active & ~outcome.success).sum()) + active_cost = ( + float(outcome.cost[active].sum()) if bool(active.any()) else 0.0 + ) + rank = (failed_count, active_cost, candidate_index) + attempts.append( + { + "candidate_index": candidate_index, + "status": "planned", + "failed_rows": failed_count, + "cost": active_cost, + } + ) + if selected is None or rank < selected_rank: + selected = (grounded, outcome) + selected_rank = rank + if failed_count == 0: + break + if selected is None: + if last_error is not None: + raise RuntimeError( + "All grounding candidates raised during planning." + ) from last_error + raise RuntimeError("No unused grounding candidate remains.") + grounded, outcome = selected + outcome = replace( + outcome, + planner_trace={ + **outcome.planner_trace, + "grounding_candidates": attempts, + "selected_grounding_candidate": int( + grounded.motion_policy.get("placement_candidate_index", 0) + ), + }, + ) + return grounded, outcome + + def _arm_selection_workspace( + self, + step: SemanticStep, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return workspace geometry along the robot's live lateral axis.""" + lateral_axis = self._robot_view_lateral_axis() + arrangement = self.arrangements.get(step.id) + if arrangement is not None: + minimum = arrangement.table_bounds[:, 0, :2] + maximum = arrangement.table_bounds[:, 1, :2] + center = (minimum + maximum) * 0.5 + half_extents = (maximum - minimum) * 0.5 + half_width = torch.sum(torch.abs(lateral_axis) * half_extents, dim=1) + return center, half_width, lateral_axis + count = int(self.env.num_envs) + centers = torch.zeros((count, 2), dtype=torch.float32, device=self.env.device) + half_widths = torch.full( + (count,), + float(self.runtime_policy.arm_selection.fallback_workspace_half_width), + dtype=torch.float32, + device=self.env.device, + ) + table = self.env.sim.get_rigid_object("table") + if table is None or not hasattr(table, "get_vertices"): + return centers, half_widths, lateral_axis + table_pose = self._entity_pose("table") + for env_id in range(count): + value = table.get_vertices(env_ids=[env_id], scale=True) + if isinstance(value, (list, tuple)): + value = value[0] + vertices = torch.as_tensor( + value, + dtype=torch.float32, + device=self.env.device, + ) + if vertices.ndim == 3 and vertices.shape[0] == 1: + vertices = vertices[0] + if vertices.ndim != 2 or vertices.shape[-1] != 3: + continue + world = ( + vertices @ table_pose[env_id, :3, :3].transpose(0, 1) + + table_pose[env_id, :3, 3] + ) + minimum = world[:, :2].min(dim=0).values + maximum = world[:, :2].max(dim=0).values + center = (minimum + maximum) * 0.5 + lateral = torch.sum((world[:, :2] - center) * lateral_axis[env_id], dim=1) + half_width = torch.max(torch.abs(lateral)) + if float(half_width) > 1.0e-6: + centers[env_id] = center + half_widths[env_id] = half_width + return centers, half_widths, lateral_axis + + def _robot_view_lateral_axis(self) -> torch.Tensor: + """Return the normalized world-space axis pointing right-arm to left-arm.""" + if self._robot_lateral_axis_cache is not None: + return self._robot_lateral_axis_cache + _, self._robot_lateral_axis_cache = robot_frame_axes(self.env) + return self._robot_lateral_axis_cache + + def _report_candidates( + self, + step: SemanticStep, + candidates: Sequence[_Candidate], + ) -> None: + if step.id in self._reported_candidates: + return + warning_count = sum(len(item.warnings) for item in candidates) + failures = [ + message + for (step_id, _), message in self._candidate_failures.items() + if step_id == step.id + ] + diagnostics = tuple( + dict.fromkeys(message for item in candidates for message in item.warnings) + ) + tuple(dict.fromkeys(failures)) + diagnostics = tuple(dict.fromkeys(diagnostics)) + if diagnostics: + self._candidate_diagnostics[step.id] = diagnostics + blockers = tuple( + deepcopy(item) + for candidate in candidates + for item in getattr(candidate, "blockers", ()) + ) + if blockers: + self._candidate_blockers[step.id] = blockers + if warning_count or failures: + feasible = ", ".join( + f"{int(item.feasible.sum())}/{len(item.feasible)}" + for item in candidates + ) + log_info( + f"Speculative arm candidates for {step.id}: feasible=[{feasible}], " + f"suppressed_warnings={warning_count}, exceptions={len(failures)}." + ) + edge_failures = tuple( + message + for message in diagnostics + if "candidate became infeasible" in message + ) + prioritized = tuple( + dict.fromkeys((*failures, *edge_failures, *diagnostics)) + ) + for message in prioritized[:3]: + log_warning(f"Candidate planning for {step.id}: {message}") + self._reported_candidates.add(step.id) + + def _candidate_outcome_blockers( + self, + step: SemanticStep, + edge: ExecutionEdge, + arm: str, + inherited_failed: torch.Tensor, + outcome: ActionOutcome, + ) -> list[dict[str, Any]]: + """Capture the real suffix edge that exhausted bounded planning.""" + failed = ~outcome.success & ~inherited_failed + action = edge.actions[0] + return [ + { + "env_id": int(env_id), + "node_id": action.get("seed_node_id"), + "blocking_edge_id": edge.id, + "atomic_action": str(action.get("atomic_action_class")), + "arm": arm, + "failure_policy": self._edge_failure_policy(edge), + "planning_stage": "candidate_suffix", + **self._planner_failure_details(outcome.planner_trace, env_id), + } + for env_id in torch.nonzero(failed, as_tuple=False).flatten().tolist() + ] + + def _candidate_exception_blockers( + self, + step: SemanticStep, + edge: ExecutionEdge, + arm: str, + inherited_failed: torch.Tensor, + exc: Exception, + ) -> list[dict[str, Any]]: + """Record a bounded candidate-planning exception without claiming proof.""" + del step + action = edge.actions[0] + budget = self._planner_search_budget() + return [ + { + "env_id": int(env_id), + "node_id": action.get("seed_node_id"), + "blocking_edge_id": edge.id, + "atomic_action": str(action.get("atomic_action_class")), + "arm": arm, + "failure_policy": self._edge_failure_policy(edge), + "planning_stage": "candidate_suffix", + "search_strategy": "planner_exception", + "search_budget": budget, + "evidence": {"exception": f"{type(exc).__name__}: {exc}"}, + } + for env_id in torch.nonzero(~inherited_failed, as_tuple=False) + .flatten() + .tolist() + ] + + def _planner_search_budget(self) -> dict[str, Any]: + """Return the configured finite search budget used by motion planning.""" + runtime_policy = getattr(self, "runtime_policy", None) + planner = getattr(runtime_policy, "planner", {}) + curobo = planner.get("curobo", {}) if isinstance(planner, Mapping) else {} + return { + "primary_max_attempts": int(curobo.get("max_attempts", 1)), + "fallback_enabled": bool(planner.get("allow_fallback", False)), + } + + def _planner_failure_details( + self, + trace: Mapping[str, Any], + env_id: int, + ) -> dict[str, Any]: + """Extract compact row-local evidence from one planner trace.""" + reachability = trace.get("reachability_search") + reachability = reachability if isinstance(reachability, Mapping) else {} + strategy = str( + reachability.get("strategy") or trace.get("primary_strategy") or "unknown" + ) + budget = deepcopy( + dict(trace.get("search_budget", self._planner_search_budget())) + ) + attempts = reachability.get("attempts", ()) + evidence: dict[str, Any] = { + "primary_success": bool( + self._row_trace_value(trace.get("primary_success", False), env_id) + ), + "fallback_attempted": bool( + self._row_trace_value(trace.get("fallback_attempted", False), env_id) + ), + "fallback_success": bool( + self._row_trace_value(trace.get("fallback_success", False), env_id) + ), + } + if trace.get("exception") is not None: + evidence["exception"] = str(trace["exception"]) + if isinstance(attempts, Sequence) and not isinstance( + attempts, (str, bytes, bytearray) + ): + evidence["reachability_attempts"] = [ + { + "candidate": str(item.get("candidate", "")), + "target_z": self._row_trace_value(item.get("target_z"), env_id), + "success": bool( + self._row_trace_value(item.get("success", False), env_id) + ), + } + for item in attempts + if isinstance(item, Mapping) + ] + budget["reachability_candidate_count"] = len( + evidence["reachability_attempts"] + ) + return { + "search_strategy": strategy, + "search_budget": budget, + "evidence": evidence, + } + + @staticmethod + def _row_trace_value(value: Any, env_id: int) -> Any: + """Detach one environment row from JSON-like or tensor trace data.""" + if isinstance(value, torch.Tensor): + detached = value.detach().cpu() + if detached.ndim == 0: + return detached.item() + row = detached[min(env_id, detached.shape[0] - 1)] + return row.item() if row.ndim == 0 else row.tolist() + if isinstance(value, Sequence) and not isinstance( + value, (str, bytes, bytearray) + ): + if not value: + return None + return deepcopy(value[min(env_id, len(value) - 1)]) + return deepcopy(value) + + def _edge_diagnostics( + self, + step: SemanticStep, + edge: ExecutionEdge, + failed: torch.Tensor, + ) -> tuple[str, ...]: + if edge.id != step.edge_ids[0] or not bool(failed.any()): + return () + return self._candidate_diagnostics.get(step.id, ()) + + def _with_downstream_targets( + self, + step: SemanticStep, + pickup_edge_id: str, + arm: str, + state: ExecutionState, + grounded: GroundedAction, + ) -> GroundedAction: + """Screen grasp poses against every later held-object target. + + A handover is split across semantic steps: its staging ``MoveHeldObject`` + edge is not part of the pickup step's local edge suffix. Include that + first exchange pose here so ``PickUp`` can reject a grasp whose + ``object_to_eef`` transform makes the later transfer arm unreachable. + This keeps the screening speculative and bounded; no simulator steps + are sent while a candidate is being built. + """ + targets: list[torch.Tensor] = [] + start = step.edge_ids.index(pickup_edge_id) + 1 + for edge_id in step.edge_ids[start:]: + edge = self.edges[edge_id] + if len(edge.actions) != 1: + continue + action = edge.actions[0] + if ( + self.adapter.capabilities.get( + str(action.get("atomic_action_class")) + ).target_materializer + != "semantic_held_object" + ): + continue + future = self.grounder.ground( + action, + step, + arm=arm, + state=state, + orientation_reference_pose=self._orientation_references.get(step.id), + ) + if future.target_object_pose is not None: + targets.append(future.target_object_pose) + targets.extend(self._handover_successor_targets(step, arm, state)) + if not targets: + return grounded + existing = tuple(grounded.cfg.get("downstream_object_target_poses", ())) + return replace( + grounded, + cfg={ + **grounded.cfg, + "downstream_object_target_poses": existing + tuple(targets), + }, + ) + + def _handover_successor_targets( + self, + step: SemanticStep, + arm: str, + state: ExecutionState, + ) -> list[torch.Tensor]: + """Return staging poses for handovers downstream of a pickup. + + ``SemanticStep.depends_on`` contains semantic IDs rather than edge IDs, + so walk the small dependency graph instead of assuming the handover is + an immediate child. Only a handover that transfers this object from + the selected pickup arm is relevant to the grasp screen. + """ + reachable = {step.id} + changed = True + while changed: + changed = False + for candidate in self.steps.values(): + if candidate.id in reachable: + continue + if any(dependency in reachable for dependency in candidate.depends_on): + reachable.add(candidate.id) + changed = True + + targets: list[torch.Tensor] = [] + for successor in self.steps.values(): + if ( + successor.id not in reachable + or successor.id == step.id + or successor.operator != "handover" + or successor.object_uid != step.object_uid + ): + continue + for edge_id in successor.edge_ids: + edge = self.edges[edge_id] + if len(edge.actions) != 1: + continue + action = edge.actions[0] + binding = action.get("target_binding", {}) + if not isinstance(binding, Mapping): + continue + if binding.get("kind") != "handover_staging": + continue + transfer_arm = str( + binding.get( + "transfer_arm", + successor.goal.get("transfer_arm", ""), + ) + ) + if transfer_arm != arm: + break + try: + grounded = self.grounder.ground( + action, + successor, + arm=arm, + state=state, + orientation_reference_pose=self._orientation_references.get( + successor.id, + self._orientation_references.get(step.id), + ), + ) + except (AttributeError, KeyError, ValueError): + # A malformed/incomplete successor must not make an + # otherwise valid pickup candidate disappear. The normal + # successor execution will report that grounding error. + break + if grounded.target_object_pose is not None: + targets.append(grounded.target_object_pose) + break + return targets + + def _eef_target(self, outcome: ActionOutcome) -> torch.Tensor | None: + state = outcome.next_state + held_object = state.get_held_object( + arm_control_part(self.env, outcome.grounded.arm) + ) + object_target = outcome.grounded.target_object_pose + if object_target is not None and held_object is not None: + object_to_eef = held_object.object_to_eef.to( + device=object_target.device, + dtype=object_target.dtype, + ) + return torch.bmm(object_target, object_to_eef) + if held_object is not None: + return held_object.grasp_xpos + target = outcome.grounded.target + return getattr(target, "xpos", None) + + def _state_for(self, step: SemanticStep, arm: str) -> ExecutionState: + """Refresh qpos while retaining holds across TaskGroup boundaries.""" + cached = self._step_states.get((step.id, arm)) + if cached is None: + cached = self._object_states.get((step.object_uid, arm)) + live_qpos = self.env.robot.get_qpos().clone() + if cached is None: + return ExecutionState(last_qpos=live_qpos) + return cached.with_updates(last_qpos=live_qpos) + + def _resource_conflicts( + self, + step: SemanticStep, + arm: str, + ) -> torch.Tensor: + object_owners = self._object_owners.get( + step.object_uid, [None] * int(self.env.num_envs) + ) + arm_owners = self._arm_owners[arm] + return torch.tensor( + [ + (object_owner not in {None, arm}) + or (arm_owner not in {None, step.object_uid}) + for object_owner, arm_owner in zip(object_owners, arm_owners) + ], + dtype=torch.bool, + device=self.env.device, + ) + + def _update_ownership( + self, + step: SemanticStep, + arm: str, + action_class: str, + state: ExecutionState, + successful: torch.Tensor, + ) -> None: + capability = self.adapter.capabilities.get(action_class) + owners = self._object_owners.setdefault( + step.object_uid, [None] * int(self.env.num_envs) + ) + if capability.state_effect == "release": + for env_id in torch.nonzero(successful, as_tuple=False).flatten().tolist(): + if owners[env_id] == arm: + owners[env_id] = None + if self._arm_owners[arm][env_id] == step.object_uid: + self._arm_owners[arm][env_id] = None + if arm not in owners: + self._object_states.pop((step.object_uid, arm), None) + return + held_object = state.get_held_object(arm_control_part(self.env, arm)) + if held_object is None or not bool(successful.any()): + return + self._object_states[(step.object_uid, arm)] = state + if capability.state_effect == "hold": + self._clear_support_relation(step.object_uid, successful) + for env_id in torch.nonzero(successful, as_tuple=False).flatten().tolist(): + owners[env_id] = arm + self._arm_owners[arm][env_id] = step.object_uid + + def _rematch_arrangement( + self, + trigger_step: SemanticStep, + trigger: torch.Tensor, + failed: torch.Tensor, + ) -> torch.Tensor: + """Globally rematch unfinished objects to feasible free slots.""" + from scipy.optimize import linear_sum_assignment + + arrangement = self.arrangements[trigger_step.id] + changed = torch.zeros_like(trigger) + for env_id in torch.nonzero(trigger, as_tuple=False).flatten().tolist(): + step_ids = arrangement.remaining(env_id) + slots = arrangement.available_slots(env_id) + if len(step_ids) != len(slots): + continue + original = { + step_id: int(arrangement.assignments[step_id][env_id]) + for step_id in step_ids + } + costs = np.full((len(step_ids), len(slots)), np.inf, dtype=np.float64) + isolate = torch.ones_like(failed) + isolate[env_id] = failed[env_id] + for row, step_id in enumerate(step_ids): + for column, slot_id in enumerate(slots): + arrangement.assignments[step_id][env_id] = slot_id + self._candidate_cache.pop((step_id, "left_arm"), None) + self._candidate_cache.pop((step_id, "right_arm"), None) + arm_costs = [] + for arm in ("left_arm", "right_arm"): + candidate = self._candidate(self.steps[step_id], arm, isolate) + if bool(candidate.feasible[env_id]): + arm_costs.append(float(candidate.cost[env_id])) + if arm_costs: + costs[row, column] = min(arm_costs) + arrangement.assignments[step_id][env_id] = original[step_id] + self._candidate_cache.pop((step_id, "left_arm"), None) + self._candidate_cache.pop((step_id, "right_arm"), None) + if not np.isfinite(costs).any(axis=1).all(): + continue + rows, columns = linear_sum_assignment( + np.where(np.isfinite(costs), costs, 1.0e12) + ) + if not np.isfinite(costs[rows, columns]).all(): + continue + arrangement.assign( + env_id, + { + step_ids[int(row)]: slots[int(column)] + for row, column in zip(rows, columns) + }, + ) + changed[env_id] = True + return changed + + def _execute_edge( + self, + edge: ExecutionEdge, + step: SemanticStep, + *, + failed: torch.Tensor, + ) -> _EdgeResult: + if step.goal.get("payloads"): + # Capture before the first physical action, including an ordinary + # single-arm pickup. Verification then measures whether every + # direct payload stayed fixed relative to its carrier. + self._capture_payloads(step) + if ( + len(edge.actions) == 1 + and self.adapter.capabilities.get( + str(edge.actions[0].get("atomic_action_class")) + ).resource_mode + == "coordinated_object" + ): + return self._execute_coordinated(edge, step, failed) + if len(edge.actions) == 2: + return self._execute_explicit_dual(edge, step, failed) + if len(edge.actions) != 1: + raise ValueError( + f"Edge {edge.id!r} must contain one action or an explicit dual pair." + ) + assignments = self._assignments[step.id] + outcomes: dict[str, ActionOutcome | None] = { + "left_arm": None, + "right_arm": None, + } + masks = { + arm: torch.tensor( + [assignment == arm for assignment in assignments], + dtype=torch.bool, + device=self.env.device, + ) + & ~failed + for arm in outcomes + } + grounded_items: list[GroundedAction] = [] + planner_traces: list[dict[str, Any]] = [] + planning_failed = torch.zeros_like(failed) + action_class = str(edge.actions[0]["atomic_action_class"]) + capability = self.adapter.capabilities.get(action_class) + for arm in outcomes: + if not bool(masks[arm].any()): + continue + state = self._state_for(step, arm) + if capability.state_effect == "hold": + try: + grounded, outcome = self._plan_live_hold(edge, step, arm) + except Exception as exc: + planning_failed |= masks[arm] + planner_traces.append( + self._live_hold_failure_trace(edge, step, arm, exc) + ) + continue + else: + # Re-ground transport and placement from live simulator state. + grounded, outcome = self._ground_and_plan_candidates( + edge.actions[0], + step, + arm=arm, + state=state, + active=masks[arm], + orientation_reference_pose=self._orientation_references.get( + step.id + ), + ) + grounded = outcome.grounded + outcomes[arm] = outcome + grounded_items.append(grounded) + planner_traces.append(outcome.planner_trace) + self._remember_target(step, grounded) + placement_index = grounded.motion_policy.get("placement_candidate_index") + if placement_index is not None and bool( + (masks[arm] & outcome.success).any() + ): + self._placement_candidate_history.setdefault((step.id, arm), set()).add( + int(placement_index) + ) + assigned = masks["left_arm"] | masks["right_arm"] + if not grounded_items: + return _EdgeResult( + [], + failed | (~failed & ~assigned) | planning_failed, + [], + planner_traces, + executed=torch.zeros_like(failed), + ) + trajectory, action_success = self.adapter.combine(outcomes, masks) + active = assigned & action_success & ~failed & ~planning_failed + actions = self.adapter.execute_trajectory(trajectory, active=active) + physical_failed = torch.zeros_like(failed) + for arm, outcome in outcomes.items(): + if outcome is not None: + successful = masks[arm] & outcome.success & active + if capability.state_effect == "hold": + physical = self._physical_pickup( + step.object_uid, arm, outcome.next_state, successful + ) + physical_failed |= successful & ~physical + successful = physical + elif capability.state_effect == "preserve_hold": + physical = self._physical_hold( + step.object_uid, arm, outcome.next_state, successful + ) + lost = successful & ~physical + physical_failed |= lost + self._release_ownership(step.object_uid, arm, lost) + successful = physical + if capability.verifier_hook is not None: + verified = torch.as_tensor( + capability.verifier_hook( + executor=self, + step=step, + arm=arm, + outcome=outcome, + attempted=successful, + ), + dtype=torch.bool, + device=self.env.device, + ).reshape(-1) + if verified.numel() != int(self.env.num_envs): + raise ValueError( + f"AtomicAction {action_class!r} verifier returned " + "an invalid vectorized mask." + ) + physical_failed |= successful & ~verified + successful &= verified + committed_state = outcome.state_after(successful) + if capability.state_effect in {"hold", "preserve_hold"}: + committed_state = self._rebase_held_state( + step.object_uid, + arm, + committed_state, + successful, + from_planned_qpos=capability.state_effect == "preserve_hold", + ) + self._step_states[(step.id, arm)] = committed_state + self._update_ownership( + step, + arm, + action_class, + committed_state, + successful, + ) + if capability.verifier == "pressed": + semantic_states = getattr( + self.env, + "action_engine_semantic_states", + None, + ) + if semantic_states is None: + semantic_states = {} + self.env.action_engine_semantic_states = semantic_states + semantic_states[(step.object_uid, "pressed")] = successful.clone() + edge_failed = ( + failed + | (~failed & ~assigned) + | (assigned & ~action_success) + | planning_failed + | physical_failed + ) + return _EdgeResult( + actions, + edge_failed, + grounded_items, + planner_traces, + active, + ) + + def _plan_live_hold( + self, + edge: ExecutionEdge, + step: SemanticStep, + arm: str, + ) -> tuple[GroundedAction, ActionOutcome]: + """Replace a speculative hold plan with one grounded at execution time.""" + candidate = self._candidate_cache.get((step.id, arm)) + cached_plan_available = candidate is not None and edge.id in candidate.plans + update_obj_info = getattr(self.env, "update_obj_info", None) + if callable(update_obj_info): + update_obj_info() + object_pose = self._entity_pose(step.object_uid).detach().clone() + state = self._state_for(step, arm) + grounded = self.grounder.ground( + edge.actions[0], + step, + arm=arm, + state=state, + orientation_reference_pose=self._orientation_references.get(step.id), + ) + grounded = self._with_downstream_targets(step, edge.id, arm, state, grounded) + outcome = self.adapter.plan(grounded, state) + return grounded, replace( + outcome, + planner_trace={ + **outcome.planner_trace, + "execution_replanned_from_live_state": True, + "speculative_candidate_available": cached_plan_available, + "speculative_candidate_replaced": cached_plan_available, + "execution_object_pose": object_pose, + }, + ) + + def _live_hold_failure_trace( + self, + edge: ExecutionEdge, + step: SemanticStep, + arm: str, + exc: Exception, + ) -> dict[str, Any]: + """Describe a live PickUp planning exception without aborting the task.""" + candidate = self._candidate_cache.get((step.id, arm)) + return { + "action_class": str(edge.actions[0].get("atomic_action_class")), + "arm": arm, + "primary_strategy": "live_pickup_replan", + "primary_success": torch.zeros( + int(self.env.num_envs), + dtype=torch.bool, + device=self.env.device, + ), + "execution_replanned_from_live_state": True, + "speculative_candidate_available": ( + candidate is not None and edge.id in candidate.plans + ), + "speculative_candidate_replaced": False, + "execution_object_pose": self._entity_pose(step.object_uid) + .detach() + .clone(), + "exception": f"{type(exc).__name__}: {exc}", + } + + def _physical_pickup( + self, + uid: str, + arm: str, + state: ExecutionState, + attempted: torch.Tensor, + ) -> torch.Tensor: + owners = list(self._object_owners.get(uid, [None] * int(self.env.num_envs))) + for env_id in torch.nonzero(attempted, as_tuple=False).flatten().tolist(): + owners[env_id] = arm + states = dict(self._object_states) + states[(uid, arm)] = state + physical = attempted & evaluate_predicate( + self.env, + { + "type": "object_held", + "object": uid, + "position_tolerance": self.runtime_policy.predicate_fallbacks[ + "held_position_tolerance" + ], + }, + held_owners={**self._object_owners, uid: owners}, + held_states=states, + ) + return physical + + def _physical_hold( + self, + uid: str, + arm: str, + state: ExecutionState, + attempted: torch.Tensor, + *, + owners: Mapping[str, Sequence[str | None]] | None = None, + states: Mapping[tuple[str, str], ExecutionState] | None = None, + position_tolerance: float | None = None, + ) -> torch.Tensor: + candidate_states = dict(self._object_states if states is None else states) + candidate_states[(uid, arm)] = state + return attempted & evaluate_predicate( + self.env, + { + "type": "object_held", + "object": uid, + "position_tolerance": ( + self.runtime_policy.predicate_fallbacks["held_position_tolerance"] + if position_tolerance is None + else float(position_tolerance) + ), + "arm": arm, + }, + held_owners=self._object_owners if owners is None else owners, + held_states=candidate_states, + ) + + def _release_ownership( + self, + uid: str, + arm: str, + lost: torch.Tensor, + ) -> None: + owners = self._object_owners.get(uid) + if owners is None: + return + for env_id in torch.nonzero(lost, as_tuple=False).flatten().tolist(): + if owners[env_id] == arm: + owners[env_id] = None + if self._arm_owners[arm][env_id] == uid: + self._arm_owners[arm][env_id] = None + if arm not in owners: + self._object_states.pop((uid, arm), None) + + def _execute_coordinated( + self, + edge: ExecutionEdge, + step: SemanticStep, + failed: torch.Tensor, + ) -> _EdgeResult: + action = edge.actions[0] + action_name = str(action["atomic_action_class"]) + capability = self.adapter.capabilities.get(action_name) + binding = action.get("target_binding", {}) + transfer_arm = str(binding.get("transfer_arm", "left_arm")) + accepted_assignments = ( + {"coordinated", transfer_arm} + if capability.state_effect == "transfer_hold" + else {"coordinated"} + ) + assigned = torch.tensor( + [item in accepted_assignments for item in self._assignments[step.id]], + dtype=torch.bool, + device=self.env.device, + ) + receiver_arm = str(binding.get("receive_arm", "right_arm")) + receiver_conflict = torch.tensor( + [ + owner not in {None, step.object_uid} + for owner in self._arm_owners[receiver_arm] + ], + dtype=torch.bool, + device=self.env.device, + ) + active = assigned & ~failed & ~receiver_conflict + if not bool(active.any()): + return _EdgeResult( + [], + failed | (~failed & ~assigned) | receiver_conflict, + [], + executed=torch.zeros_like(failed), + ) + state_key = ( + transfer_arm + if capability.state_effect == "transfer_hold" + else "coordinated" + ) + state = self._state_for(step, state_key) + if capability.state_effect == "coordinated_release": + held_objects = dict(state.held_objects) + for arm in ("left_arm", "right_arm"): + arm_state = self._step_states.get((step.id, arm)) + if arm_state is None: + continue + control_part = arm_control_part(self.env, arm) + held_object = arm_state.get_held_object(control_part) + if held_object is not None: + held_objects[control_part] = held_object + state = state.with_updates(held_objects=held_objects) + update_obj_info = getattr(self.env, "update_obj_info", None) + if callable(update_obj_info): + update_obj_info() + groundings = self.grounder.ground_candidates( + action, + step, + arm="coordinated", + state=state, + orientation_reference_pose=self._orientation_references.get(step.id), + ) + selected: tuple[GroundedAction, ActionOutcome] | None = None + selected_warnings: tuple[str, ...] = () + best_failure_count = int(active.sum()) + 1 + rejected_warning_count = 0 + for candidate in groundings: + with _capture_speculative_warnings() as captured: + candidate_outcome = self.adapter.plan(candidate, state) + failure_count = int((active & ~candidate_outcome.success).sum()) + if selected is None or failure_count < best_failure_count: + selected = (candidate, candidate_outcome) + selected_warnings = tuple(captured) + best_failure_count = failure_count + if failure_count == 0: + if rejected_warning_count: + log_info( + "Selected a feasible coordinated grounding after " + f"suppressing {rejected_warning_count} warnings from " + "rejected candidates." + ) + break + rejected_warning_count += len(captured) + if selected is None: + raise RuntimeError("Coordinated action grounding produced no candidates.") + if best_failure_count: + for message in dict.fromkeys(selected_warnings): + log_warning(message) + grounded, outcome = selected + if capability.state_effect == "transfer_hold": + outcome = replace( + outcome, + planner_trace={ + **outcome.planner_trace, + "execution_replanned_from_live_state": True, + "execution_object_pose": self._entity_pose(step.object_uid) + .detach() + .clone(), + }, + ) + self._remember_target(step, grounded) + successful = active & outcome.success + actions = self.adapter.execute_trajectory( + outcome.trajectory, + active=successful, + ) + physical_failed = torch.zeros_like(failed) + committed_state = outcome.state_after(successful) + if capability.state_effect == "coordinated_hold": + self._clear_support_relation(step.object_uid, successful) + if capability.state_effect == "transfer_hold": + if bool(successful.any()): + current_owners = list( + self._object_owners.get( + step.object_uid, + [None] * int(self.env.num_envs), + ) + ) + tentative_owners = list(current_owners) + for env_id in ( + torch.nonzero(successful, as_tuple=False).flatten().tolist() + ): + tentative_owners[env_id] = receiver_arm + tentative_states = dict(self._object_states) + tentative_states[(step.object_uid, receiver_arm)] = outcome.next_state + physical = self._physical_hold( + step.object_uid, + receiver_arm, + outcome.next_state, + successful, + owners={ + **self._object_owners, + step.object_uid: tentative_owners, + }, + states=tentative_states, + position_tolerance=min( + float( + self.runtime_policy.predicate_fallbacks[ + "held_position_tolerance" + ] + ), + float( + grounded.motion_policy.get( + "held_position_tolerance", + self.runtime_policy.predicate_fallbacks[ + "held_position_tolerance" + ], + ) + ), + ), + ) + lost = successful & ~physical + physical_failed |= lost + successful = physical + committed_state = outcome.state_after(successful) + committed_state = self._rebase_held_state( + step.object_uid, + receiver_arm, + committed_state, + successful, + from_planned_qpos=True, + ) + committed_owners = list(current_owners) + for env_id in ( + torch.nonzero( + active & outcome.success, + as_tuple=False, + ) + .flatten() + .tolist() + ): + if bool(physical[env_id]): + committed_owners[env_id] = receiver_arm + self._arm_owners[receiver_arm][env_id] = step.object_uid + else: + committed_owners[env_id] = None + self._arm_owners[transfer_arm][env_id] = None + self._object_owners[step.object_uid] = committed_owners + if any(owner == receiver_arm for owner in committed_owners): + self._step_states[(step.id, receiver_arm)] = committed_state + self._object_states[(step.object_uid, receiver_arm)] = ( + committed_state + ) + else: + self._object_states.pop((step.object_uid, receiver_arm), None) + if not any(owner == transfer_arm for owner in committed_owners): + self._object_states.pop((step.object_uid, transfer_arm), None) + self._step_states[(step.id, "coordinated")] = committed_state + return _EdgeResult( + actions, + failed + | (~failed & ~assigned) + | (active & ~outcome.success) + | physical_failed, + [grounded], + [outcome.planner_trace], + active & outcome.success, + ) + + def _rebase_held_state( + self, + uid: str, + arm: str, + state: ExecutionState, + mask: torch.Tensor, + *, + from_planned_qpos: bool = True, + ) -> ExecutionState: + """Refresh a held object's object-to-EEF transform after execution.""" + control_part = arm_control_part(self.env, arm) + held = state.get_held_object(control_part) + entity = self.env.sim.get_rigid_object(uid) + if held is None or entity is None or not bool(mask.any()): + return state + if from_planned_qpos: + # Preserve-hold planning must stay in its terminal qpos/FK frame; + # get_current_xpos_agent() may still expose the previous command. + joint_ids = self.env.robot.get_joint_ids(name=control_part) + eef_pose = self.env.robot.compute_fk( + state.last_qpos[:, joint_ids], + name=control_part, + to_matrix=True, + ) + else: + eef_poses = self.env.get_current_xpos_agent() + eef_pose = eef_poses[0 if arm == "left_arm" else 1] + eef_pose = torch.as_tensor( + eef_pose, + dtype=held.object_to_eef.dtype, + device=held.object_to_eef.device, + ) + object_pose = torch.as_tensor( + entity.get_local_pose(to_matrix=True), + dtype=eef_pose.dtype, + device=eef_pose.device, + ) + if eef_pose.ndim == 2: + eef_pose = eef_pose.unsqueeze(0).repeat(int(self.env.num_envs), 1, 1) + if object_pose.ndim == 2: + object_pose = object_pose.unsqueeze(0).repeat(int(self.env.num_envs), 1, 1) + selector = mask[:, None, None] + live_object_to_eef = torch.bmm(torch.linalg.inv(object_pose), eef_pose) + rebased = HeldObjectState( + semantics=held.semantics, + object_to_eef=torch.where( + selector, + live_object_to_eef, + held.object_to_eef, + ), + grasp_xpos=torch.where(selector, eef_pose, held.grasp_xpos), + env_mask=held.env_mask, + ) + held_objects = dict(state.held_objects) + held_objects[control_part] = rebased + return state.with_updates(held_objects=held_objects) + + def _execute_explicit_dual( + self, + edge: ExecutionEdge, + step: SemanticStep, + failed: torch.Tensor, + ) -> _EdgeResult: + assigned = torch.tensor( + [item == "coordinated" for item in self._assignments[step.id]], + dtype=torch.bool, + device=self.env.device, + ) + if not bool((assigned & ~failed).any()): + return _EdgeResult( + [], + failed | (~failed & ~assigned), + [], + executed=torch.zeros_like(failed), + ) + outcomes: dict[str, ActionOutcome | None] = { + "left_arm": None, + "right_arm": None, + } + masks = {arm: assigned & ~failed for arm in outcomes} + grounded_items = [] + coordinated_state = self._state_for(step, "coordinated") + for action in edge.actions: + actor = action.get("actor", {}) + arm = str(actor.get("arm", "")) + if arm not in outcomes or outcomes[arm] is not None: + raise ValueError( + f"Explicit dual edge {edge.id!r} must bind each arm once." + ) + state = self._step_states.get((step.id, arm)) + if state is None: + state = coordinated_state + else: + state = self._state_for(step, arm) + grounded = self.grounder.ground( + action, + step, + arm=arm, + state=state, + orientation_reference_pose=self._orientation_references.get(step.id), + ) + outcome = self.adapter.plan(grounded, state) + outcomes[arm] = outcome + grounded_items.append(grounded) + trajectory, action_success = self.adapter.combine(outcomes, masks) + active = assigned & ~failed & action_success + actions = self.adapter.execute_trajectory(trajectory, active=active) + is_coordinated_release = { + str( + action.get("target_binding", {}).get( + "coordinated_release_role", + "", + ) + ) + for action in edge.actions + } == {"participant", "commit"} and all( + action.get("control") == "hand" + and action.get("target_binding", {}).get("kind") == "joint_state" + and action.get("target_binding", {}).get("source") == "gripper_open" + for action in edge.actions + ) + physical_failed = torch.zeros_like(failed) + if is_coordinated_release: + opened = evaluate_predicate(self.env, {"type": "both_grippers_open"}) + released = active & opened + physical_failed = active & ~opened + control_parts = ( + arm_control_part(self.env, "left_arm"), + arm_control_part(self.env, "right_arm"), + ) + released_task = StateDelta( + held_object_updates={name: None for name in control_parts} + ).apply(coordinated_state.to_task_state(), released) + released_state = ExecutionState.from_task_state( + released_task, + last_qpos=self.env.robot.get_qpos().clone(), + ) + for key in ("coordinated", "left_arm", "right_arm"): + self._step_states[(step.id, key)] = released_state + else: + for arm, outcome in outcomes.items(): + if outcome is not None: + self._step_states[(step.id, arm)] = outcome.state_after( + active & outcome.success + ) + return _EdgeResult( + actions, + failed + | (~failed & ~assigned) + | (assigned & ~action_success) + | physical_failed, + grounded_items, + [ + outcome.planner_trace + for outcome in outcomes.values() + if outcome is not None + ], + active, + ) + + def _execute_parallel_pickups( + self, + edges: Sequence[ExecutionEdge], + *, + failed: torch.Tensor, + ) -> tuple[dict[str, _EdgeResult], torch.Tensor]: + steps = [self.step_by_edge[edge.id] for edge in edges] + for step in steps: + self._capture_orientation_reference(step) + candidates = { + (step.id, arm): self._candidate(step, arm, failed) + for step in steps + for arm in ("left_arm", "right_arm") + } + for step in steps: + self._report_candidates( + step, + ( + candidates[(step.id, "left_arm")], + candidates[(step.id, "right_arm")], + ), + ) + assignments = {step.id: [None] * len(failed) for step in steps} + selection_failed = torch.zeros_like(failed) + permutations = ( + ("left_arm", "right_arm"), + ("right_arm", "left_arm"), + ) + for env_id in range(len(failed)): + if bool(failed[env_id]): + continue + ranked: list[tuple[bool, bool, float, float, str, str]] = [] + for first_arm, second_arm in permutations: + first = candidates[(steps[0].id, first_arm)] + second = candidates[(steps[1].id, second_arm)] + feasible = bool(first.feasible[env_id] and second.feasible[env_id]) + required_match = all( + candidate_step.actor.get("mode") != "required" + or str(candidate_step.actor.get("arm")) == candidate_arm + for candidate_step, candidate_arm in ( + (steps[0], first_arm), + (steps[1], second_arm), + ) + ) + available = ( + required_match + and not bool( + self._resource_conflicts(steps[0], first_arm)[env_id] + or self._resource_conflicts(steps[1], second_arm)[env_id] + ) + and all( + self._auto_arm_is_allowed(candidate_step, candidate_arm, env_id) + for candidate_step, candidate_arm in ( + (steps[0], first_arm), + (steps[1], second_arm), + ) + ) + ) + first_preferred = self._preferred_in_place_arm(steps[0], env_id) + second_preferred = self._preferred_in_place_arm(steps[1], env_id) + side_penalty = ( + float(first_arm != first_preferred) if first_preferred else 0.0 + ) + side_penalty += ( + float(second_arm != second_preferred) if second_preferred else 0.0 + ) + cost = float(first.cost[env_id] + second.cost[env_id]) + ranked.append( + ( + not available, + not feasible, + side_penalty, + cost, + first_arm, + second_arm, + ) + ) + ranked.sort() + unavailable, _, _, _, first_arm, second_arm = ranked[0] + if unavailable: + selection_failed[env_id] = True + continue + assignments[steps[0].id][env_id] = first_arm + assignments[steps[1].id][env_id] = second_arm + self._assignments.update(assignments) + + base_failed = failed | selection_failed + results = { + edge.id: _EdgeResult( + [], + base_failed.clone(), + [], + executed=torch.zeros_like(failed), + ) + for edge in edges + } + for first_arm, second_arm in permutations: + partition = torch.tensor( + [ + assignments[steps[0].id][env_id] == first_arm + and assignments[steps[1].id][env_id] == second_arm + for env_id in range(len(failed)) + ], + dtype=torch.bool, + device=self.env.device, + ) + if not bool(partition.any()): + continue + outcomes: dict[str, ActionOutcome | None] = { + "left_arm": None, + "right_arm": None, + } + masks = { + "left_arm": partition, + "right_arm": partition, + } + edge_by_arm = {first_arm: edges[0], second_arm: edges[1]} + parallel_planning_failed = False + for arm, edge in edge_by_arm.items(): + step = self.step_by_edge[edge.id] + try: + grounded, outcome = self._plan_live_hold(edge, step, arm) + except Exception as exc: + parallel_planning_failed = True + results[edge.id].planner_traces.append( + self._live_hold_failure_trace(edge, step, arm, exc) + ) + continue + outcomes[arm] = outcome + results[edge.id].grounded.append(outcome.grounded) + results[edge.id].planner_traces.append(outcome.planner_trace) + if bool((partition & ~outcome.success).any()): + parallel_planning_failed = True + if parallel_planning_failed: + serial_actions: list[torch.Tensor] = [] + for edge in edges: + step = self.step_by_edge[edge.id] + serial = self._execute_edge_with_retries( + edge, + step, + failed=~partition, + ) + serial_actions.extend(serial.actions) + results[edge.id].grounded.extend(serial.grounded) + results[edge.id].planner_traces.extend(serial.planner_traces) + results[edge.id].failed = torch.where( + partition, + serial.failed, + results[edge.id].failed, + ) + assert results[edge.id].executed is not None + if serial.executed is not None: + results[edge.id].executed |= serial.executed + for edge in edges: + results[edge.id].actions.extend(serial_actions) + continue + trajectory, action_success = self.adapter.combine(outcomes, masks) + active = partition & ~base_failed & action_success + commands = self.adapter.execute_trajectory(trajectory, active=active) + for edge in edges: + assert results[edge.id].executed is not None + results[edge.id].executed |= active + for arm, edge in edge_by_arm.items(): + step = self.step_by_edge[edge.id] + outcome = outcomes[arm] + assert outcome is not None + attempted = active & outcome.success + physical = self._physical_pickup( + step.object_uid, arm, outcome.next_state, attempted + ) + committed_state = outcome.state_after(physical) + committed_state = self._rebase_held_state( + step.object_uid, + arm, + committed_state, + physical, + from_planned_qpos=False, + ) + self._step_states[(step.id, arm)] = committed_state + results[edge.id].failed |= partition & ~physical + self._update_ownership( + step, + arm, + str(edge.actions[0]["atomic_action_class"]), + committed_state, + physical, + ) + for edge in edges: + # Both edge records refer to this one synchronized stream. The + # run loop adds only the first copy to its returned trace. + results[edge.id].actions.extend(commands) + aggregate_failed = torch.zeros_like(failed) + for result in results.values(): + aggregate_failed |= result.failed + return results, aggregate_failed + + def _remember_target( + self, + step: SemanticStep, + grounded: GroundedAction, + ) -> None: + target = grounded.target_object_pose + if target is not None: + self._targets[step.id] = target[:, :3, 3].clone() + self._target_poses[step.id] = target.clone() + self._policies[step.id] = grounded.motion_policy + + def _capture_orientation_reference(self, step: SemanticStep) -> None: + """Freeze preserve orientation before speculative pickup can disturb it.""" + if ( + compile_orientation_constraint(step.goal).requires_reference + and step.id not in self._orientation_references + ): + predecessor_references = [ + self._orientation_references[predecessor.id] + for dependency in step.depends_on + if (predecessor := self.steps.get(dependency)) is not None + and predecessor.object_uid == step.object_uid + and predecessor.id in self._orientation_references + ] + if predecessor_references: + self._orientation_references[step.id] = predecessor_references[ + 0 + ].clone() + return + self._orientation_references[step.id] = self._entity_pose( + step.object_uid + ).clone() + + def _step_runtime_metadata(self, step: SemanticStep) -> list[dict[str, Any]]: + """Expose the live grounding and allocation decisions for diagnosis.""" + observed_pose = self._entity_pose(step.object_uid) + assignments = self._assignments.get( + step.id, + [None] * int(self.env.num_envs), + ) + target_pose = self._target_poses.get(step.id) + orientation_reference = self._orientation_references.get(step.id) + orientation_error = self._orientation_errors.get(step.id) + arrangement = self.arrangements.get(step.id) + result = [] + for env_id, assignment in enumerate(assignments): + same_side_arm = self._preferred_live_pickup_arm(step, env_id) + physical_part = assignment + if assignment in {"left_arm", "right_arm"}: + physical_part = arm_control_part(self.env, assignment) + candidate_scores = {} + for arm in ("left_arm", "right_arm"): + candidate = self._candidate_cache.get((step.id, arm)) + if candidate is None: + candidate_scores[arm] = None + continue + scores = { + name: float(values[env_id]) + for name, values in candidate.score_components.items() + } + candidate_scores[arm] = { + "feasible": bool(candidate.feasible[env_id]), + **scores, + "failure": self._candidate_failures.get((step.id, arm)), + } + item: dict[str, Any] = { + "assigned_arm": assignment, + "physical_control_part": physical_part, + "same_side_arm": same_side_arm, + "inside_arm_deadband": same_side_arm is None, + "cross_side_fallback_allowed": bool( + self.runtime_policy.arm_selection.allow_cross_side_fallback + ), + "cross_side_fallback_used": bool( + same_side_arm is not None + and assignment in {"left_arm", "right_arm"} + and assignment != same_side_arm + ), + "observed_object_pose": observed_pose[env_id], + "final_target_pose": ( + None if target_pose is None else target_pose[env_id] + ), + "orientation_reference_pose": ( + None + if orientation_reference is None + else orientation_reference[env_id] + ), + "orientation_error": ( + None + if orientation_error is None + else float(orientation_error[env_id]) + ), + "candidate_scores": candidate_scores, + } + if arrangement is not None: + item["arrangement"] = arrangement.metadata(step, env_id) + result.append(item) + return result + + def _verify_step( + self, + step: SemanticStep, + failed: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if self.settle_steps < 0: + raise ValueError("settle_steps must be non-negative.") + if self.settle_steps and bool((~failed).any()): + self.env.sim.update(step=self.settle_steps) + entity = self.env.sim.get_rigid_object(step.object_uid) + if entity is None: + raise ValueError(f"Unknown semantic object {step.object_uid!r}.") + observed_pose = torch.as_tensor( + entity.get_local_pose(to_matrix=True), + dtype=torch.float32, + device=self.env.device, + ) + observed = observed_pose[:, :3, 3] + active = ~failed + if not bool(active.any()): + success = torch.zeros_like(failed) + log_info(f"Skipped verification for {step.id}: no active environments.") + return failed, success, observed + relation = ( + normalize_placement_relation(step.goal.get("relation", "on")) + if step.operator == "place_relative" + else str(step.goal.get("relation", "")) + ) + reference = self._support_reference_uid(step) + postcondition_type = step.postcondition.get("type") + if postcondition_type in {"object_held", "handover_complete"}: + # A planned hover target is not evidence that the object remains + # grasped. Verify live TCP/object geometry and gripper closure. + satisfied = evaluate_predicate( + self.env, + step.postcondition, + held_owners=self._object_owners, + held_states=self._object_states, + ) + satisfied &= self._placement_orientation_satisfied(step, observed_pose) + elif postcondition_type in { + "held_by_both_grippers", + "object_held_by_both_grippers", + }: + satisfied = evaluate_predicate( + self.env, + step.postcondition, + coordinated_state=self._step_states.get((step.id, "coordinated")), + ) + target = self._targets.get(step.id) + if target is not None: + policy = self._policies.get(step.id, {}) + tolerance = float( + policy.get( + "postcondition_tolerance", + self.runtime_policy.predicate_fallbacks["position_tolerance"], + ) + ) + target = target.to(device=observed.device, dtype=observed.dtype) + satisfied &= ( + torch.linalg.vector_norm(observed - target, dim=1) <= tolerance + ) + elif postcondition_type == "pressed": + satisfied = evaluate_predicate(self.env, step.postcondition) + elif relation == "inside" and isinstance(reference, str): + satisfied = evaluate_predicate( + self.env, + { + "type": "object_in_container", + "object": step.object_uid, + "container": reference, + }, + ) + elif relation in {"on", "on_top", "on_top_of"} and isinstance(reference, str): + satisfied = self._support_stable_for(step, reference, active) + satisfied &= self._support_cycle_free( + step.object_uid, + reference, + active, + ) + elif step.operator == "orient_object": + position_anchor = str(step.goal.get("position_anchor", "initial_xy")) + anchor_pose = None + if position_anchor == "initial_xy": + anchor_pose = getattr(self.env, "agent_initial_object_poses", {}).get( + step.object_uid + ) + if anchor_pose is None: + anchor_pose = self._targets.get(step.id) + if anchor_pose is None: + raise ValueError( + f"orient_object step {step.id!r} has no {position_anchor} anchor." + ) + anchor_pose = torch.as_tensor( + anchor_pose, + dtype=observed.dtype, + device=observed.device, + ) + if anchor_pose.ndim == 2 and anchor_pose.shape == (4, 4): + anchor_pose = anchor_pose.unsqueeze(0).repeat( + int(self.env.num_envs), 1, 1 + ) + target_xy = ( + anchor_pose[:, :2, 3] if anchor_pose.ndim == 3 else anchor_pose[:, :2] + ) + policy = self._policies.get(step.id, {}) + fallbacks = self.runtime_policy.predicate_fallbacks + upright = evaluate_predicate( + self.env, + { + "type": "object_upright", + "object": step.object_uid, + "local_axis": policy.get("upright_local_axis", "long_axis"), + "max_tilt": float( + policy.get("upright_max_tilt", fallbacks["upright_max_tilt"]) + ), + }, + ) + xy_near_initial = evaluate_predicate( + self.env, + { + "type": "object_xy_near", + "object": step.object_uid, + "target_xy": target_xy, + "tolerance": float( + policy.get("upright_xy_tolerance", fallbacks["xy_tolerance"]) + ), + }, + ) + satisfied = upright & xy_near_initial + elif step.id in self._targets: + target = self._targets[step.id].to( + device=observed.device, + dtype=observed.dtype, + ) + policy = self._policies.get(step.id, {}) + fallbacks = self.runtime_policy.predicate_fallbacks + tolerance = float( + policy.get("postcondition_tolerance", fallbacks["position_tolerance"]) + ) + arrangement = self.arrangements.get(step.id) + if arrangement is not None: + # Line membership is a planar relation. Height changes after + # release (for example a can settling onto another stable face) + # must not invalidate an otherwise correct row placement. + delta = torch.abs(observed - target) + axis_tolerance = float( + policy.get( + "line_axis_tolerance", + fallbacks["line_axis_tolerance"], + ) + ) + perpendicular_tolerance = float( + policy.get( + "line_perpendicular_tolerance", + fallbacks["line_perpendicular_tolerance"], + ) + ) + satisfied = (delta[:, arrangement.axis_index] <= axis_tolerance) & ( + delta[:, arrangement.perpendicular_index] <= perpendicular_tolerance + ) + elif relation in DIRECTIONAL_RELATIONS: + # Left/right/front/behind constrain the support plane. The + # grounded release height is a transport target and may differ + # from the stable height after the object settles. + satisfied = ( + torch.linalg.vector_norm(observed[:, :2] - target[:, :2], dim=-1) + <= tolerance + ) + else: + satisfied = ( + torch.linalg.vector_norm(observed - target, dim=-1) <= tolerance + ) + else: + satisfied = evaluate_predicate(self.env, step.postcondition) + if relation in DIRECTIONAL_RELATIONS and isinstance(reference, str): + policy = self._policies.get(step.id, {}) + satisfied &= evaluate_predicate( + self.env, + { + "type": "object_relative_position", + "object": step.object_uid, + "reference_object": reference, + "relation": relation, + "relation_frame": step.goal.get("relation_frame", "world"), + "minimum_distance": float(policy.get("relation_clearance", 0.01)), + }, + ) + verifies_placement_orientation = bool( + compile_orientation_constraint(step.goal).terms + ) and ( + postcondition_type == "semantic_goal" + or self.arrangements.get(step.id) is not None + ) + if verifies_placement_orientation: + satisfied &= self._placement_orientation_satisfied(step, observed_pose) + if step.goal.get("payloads"): + satisfied &= self._verify_payloads(step) + success = active & satisfied + failed = failed | (active & ~satisfied) + if relation in {"on", "on_top", "on_top_of"} and isinstance(reference, str): + self._commit_support_relation(step, reference, success) + log_info( + f"Verified {step.id}: {int(success.sum())}/{len(success)} envs succeeded." + ) + return failed, success, observed + + def _capture_payloads(self, step: SemanticStep) -> None: + if step.id in self._payload_initial: + return + carrier = self._entity_pose(step.object_uid) + record = {"carrier_rotation": carrier[:, :3, :3].clone()} + for payload in step.goal.get("payloads", []): + uid = str(payload["object"]) + record[uid] = torch.bmm(torch.linalg.inv(carrier), self._entity_pose(uid)) + self._payload_initial[step.id] = record + + def _verify_payloads(self, step: SemanticStep) -> torch.Tensor: + record = self._payload_initial.get(step.id) + if record is None: + return torch.zeros( + int(self.env.num_envs), + dtype=torch.bool, + device=self.env.device, + ) + carrier = self._entity_pose(step.object_uid) + initial_up = record["carrier_rotation"][:, :3, 2] + live_up = carrier[:, :3, 2] + fallbacks = self.runtime_policy.predicate_fallbacks + tilt_ok = torch.sum(initial_up * live_up, dim=-1) >= float( + fallbacks["payload_minimum_upright_cosine"] + ) + result = tilt_ok + carrier_entity = self.env.sim.get_rigid_object(step.object_uid) + for payload in step.goal["payloads"]: + uid = str(payload["object"]) + expected = torch.bmm(carrier, record[uid]) + observed = self._entity_pose(uid) + drift_ok = torch.linalg.vector_norm( + observed[:, :3, 3] - expected[:, :3, 3], + dim=-1, + ) <= float(fallbacks["payload_position_tolerance"]) + support_ok = torch.ones_like(drift_ok) + for env_id in range(int(self.env.num_envs)): + vertices = carrier_entity.get_vertices( + env_ids=[env_id], + scale=True, + ) + if isinstance(vertices, (list, tuple)): + vertices = vertices[0] + vertices = torch.as_tensor( + vertices, + dtype=carrier.dtype, + device=carrier.device, + ) + if vertices.ndim == 3: + vertices = vertices[0] + world = ( + vertices @ carrier[env_id, :3, :3].transpose(0, 1) + + carrier[env_id, :3, 3] + ) + position = observed[env_id, :2, 3] + margin = float(fallbacks["payload_support_margin"]) + lower = world[:, :2].min(dim=0).values - margin + upper = world[:, :2].max(dim=0).values + margin + support_ok[env_id] = bool( + torch.all(position >= lower) and torch.all(position <= upper) + ) + result &= drift_ok & support_ok + return result + + def _entity_pose(self, uid: str) -> torch.Tensor: + entity = self.env.sim.get_rigid_object(uid) + if entity is None: + raise ValueError(f"Unknown rigid object {uid!r}.") + pose = torch.as_tensor( + entity.get_local_pose(to_matrix=True), + dtype=torch.float32, + device=self.env.device, + ) + if pose.ndim == 2: + pose = pose.unsqueeze(0).repeat(int(self.env.num_envs), 1, 1) + return pose + + def _entity_motion_stable(self, uid: str) -> torch.Tensor: + entity = self.env.sim.get_rigid_object(uid) + if entity is None: + raise ValueError(f"Unknown rigid object {uid!r}.") + + def velocity(value: Any, name: str) -> torch.Tensor | None: + if callable(value): + value = value() + if value is None: + return None + tensor = torch.as_tensor( + value, + dtype=torch.float32, + device=self.env.device, + ) + if tensor.ndim == 1: + tensor = tensor.unsqueeze(0).repeat(int(self.env.num_envs), 1) + if tensor.shape != (int(self.env.num_envs), 3): + raise ValueError( + f"Rigid object {uid!r} {name} must have shape " + f"({int(self.env.num_envs)}, 3)." + ) + return tensor + + linear = velocity(getattr(entity, "lin_vel", None), "lin_vel") + angular = velocity(getattr(entity, "ang_vel", None), "ang_vel") + if linear is None or angular is None: + body_state = getattr(entity, "body_state", None) + if callable(body_state): + body_state = body_state() + if body_state is not None: + state = torch.as_tensor( + body_state, + dtype=torch.float32, + device=self.env.device, + ) + if state.ndim == 1: + state = state.unsqueeze(0).repeat(int(self.env.num_envs), 1) + if state.shape == (int(self.env.num_envs), 13): + linear = state[:, 7:10] + angular = state[:, 10:13] + if linear is None or angular is None: + body_data = getattr(entity, "body_data", None) + if body_data is not None: + if linear is None: + linear = velocity(getattr(body_data, "lin_vel", None), "lin_vel") + if angular is None: + angular = velocity(getattr(body_data, "ang_vel", None), "ang_vel") + if linear is None or angular is None: + return torch.zeros( + int(self.env.num_envs), + dtype=torch.bool, + device=self.env.device, + ) + return ( + torch.linalg.vector_norm(linear, dim=1) + <= self.support_linear_velocity_tolerance + ) & ( + torch.linalg.vector_norm(angular, dim=1) + <= self.support_angular_velocity_tolerance + ) + + def _support_stable_for( + self, + step: SemanticStep, + support_uid: str, + active: torch.Tensor, + ) -> torch.Tensor: + """Require the support relation and low motion across a time window.""" + stable = active.clone() + for sample_index in range(self.support_stability_samples): + supported = evaluate_predicate( + self.env, + { + "type": "object_supported_by", + "object": step.object_uid, + "support": support_uid, + }, + ) + stable &= ( + supported + & self._entity_motion_stable(step.object_uid) + & self._entity_motion_stable(support_uid) + ) + if ( + sample_index + 1 < self.support_stability_samples + and self.support_stability_interval_steps + and bool(active.any()) + ): + self.env.sim.update(step=self.support_stability_interval_steps) + return stable + + def _clear_support_relation(self, object_uid: str, mask: torch.Tensor) -> None: + relations = self._support_relations.get(object_uid) + if relations is None: + return + for env_id in torch.nonzero(mask, as_tuple=False).flatten().tolist(): + relations[env_id] = None + if not any(relation is not None for relation in relations): + self._support_relations.pop(object_uid, None) + + def _support_cycle_free( + self, + object_uid: str, + support_uid: str, + active: torch.Tensor, + ) -> torch.Tensor: + result = active.clone() + for env_id in torch.nonzero(active, as_tuple=False).flatten().tolist(): + current = support_uid + visited: set[str] = set() + while current and current not in visited: + if current == object_uid: + result[env_id] = False + break + visited.add(current) + relations = self._support_relations.get(current) + relation = None if relations is None else relations[env_id] + current = "" if relation is None else relation.support_uid + return result + + def _commit_support_relation( + self, + step: SemanticStep, + support_uid: str, + successful: torch.Tensor, + ) -> None: + relations = self._support_relations.setdefault( + step.object_uid, + [None] * int(self.env.num_envs), + ) + relation = _SupportRelation( + support_uid=support_uid, + semantic_step_id=step.id, + ) + for env_id in torch.nonzero(successful, as_tuple=False).flatten().tolist(): + relations[env_id] = relation + + def _placement_orientation_satisfied( + self, + step: SemanticStep, + observed_pose: torch.Tensor, + ) -> torch.Tensor: + constraint = compile_orientation_constraint(step.goal) + satisfied = torch.ones( + int(self.env.num_envs), + dtype=torch.bool, + device=self.env.device, + ) + if not constraint.terms or ( + step.goal.get("orientation_goal") == "preserve" + and step.goal.get("relation") == "inside" + ): + return satisfied + policy = self._policies.get(step.id, {}) + fallbacks = self.runtime_policy.predicate_fallbacks + errors = [] + for term in constraint.terms: + if isinstance(term, AlignAxisConstraint): + if term.target_axis != "world_up": + raise ValueError( + f"Unsupported orientation target axis {term.target_axis!r}." + ) + satisfied &= evaluate_predicate( + self.env, + { + "type": "object_upright", + "object": step.object_uid, + "local_axis": term.local_axis, + "directed": term.directed, + "max_tilt": float( + term.tolerance + if term.tolerance is not None + else policy.get( + "upright_max_tilt", fallbacks["upright_max_tilt"] + ) + ), + }, + ) + continue + if not isinstance(term, MatchRotationConstraint): + raise TypeError(f"Unsupported orientation term {type(term)!r}.") + reference_pose = ( + self._orientation_references.get(step.id) + if term.reference == "step_start" + else self._target_poses.get(step.id) + ) + if reference_pose is None: + satisfied &= False + continue + reference_rotation = reference_pose[:, :3, :3].to( + device=observed_pose.device, + dtype=observed_pose.dtype, + ) + relative = torch.bmm( + reference_rotation.transpose(1, 2), + observed_pose[:, :3, :3], + ) + cosine = (relative.diagonal(dim1=-2, dim2=-1).sum(dim=-1) - 1.0) * 0.5 + error = torch.acos(cosine.clamp(-1.0, 1.0)) + errors.append(error) + satisfied &= error <= float( + term.tolerance + if term.tolerance is not None + else policy.get( + "preserve_orientation_tolerance", + fallbacks["preserve_orientation_tolerance"], + ) + ) + if errors: + self._orientation_errors[step.id] = torch.stack(errors).amax(dim=0) + return satisfied + + def _revalidate_support_relations(self) -> dict[str, torch.Tensor]: + active_by_step: dict[str, torch.Tensor] = {} + for relations in self._support_relations.values(): + for env_id, relation in enumerate(relations): + if relation is None: + continue + active = active_by_step.setdefault( + relation.semantic_step_id, + torch.zeros( + int(self.env.num_envs), + dtype=torch.bool, + device=self.env.device, + ), + ) + active[env_id] = True + failures: dict[str, torch.Tensor] = {} + for step_id, active in active_by_step.items(): + step = self.steps[step_id] + support_uid = self._support_reference_uid(step) + if support_uid is None: + failures[step_id] = active + continue + observed_pose = self._entity_pose(step.object_uid) + valid = self._support_stable_for(step, support_uid, active) + valid &= self._placement_orientation_satisfied(step, observed_pose) + lost = active & ~valid + if bool(lost.any()): + failures[step_id] = lost + return failures + + @staticmethod + def _support_reference_uid(step: SemanticStep) -> str | None: + value = step.goal.get("reference_object", step.goal.get("support_object")) + if isinstance(value, str) and value: + return value + if ( + step.postcondition.get("type") == "stack_layer_supported" + and int(step.goal.get("layer_index", -1)) == 0 + ): + return "table" + return None + + @staticmethod + def _edge_failure_policy(edge: ExecutionEdge) -> str: + """Return the persisted node policy for one synchronized edge.""" + policies = { + str(action.get("failure_policy", "task_required")) + for action in edge.actions + } + if not policies <= {"task_required", "safety_required", "best_effort"}: + raise ValueError( + f"Edge {edge.id!r} contains unknown failure policies {policies}." + ) + if len(policies) != 1: + raise ValueError( + f"Edge {edge.id!r} mixes incompatible failure policies {policies}." + ) + return next(iter(policies)) diff --git a/embodichain/gen_sim/action_engine/runtime/recording.py b/embodichain/gen_sim/action_engine/runtime/recording.py new file mode 100644 index 000000000..4c38a64a2 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/recording.py @@ -0,0 +1,391 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Append compact per-environment execution events and a final summary.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +from datetime import datetime, timezone +import json +import os +from pathlib import Path +import re +import tempfile +from typing import Any + +import torch + +from embodichain.gen_sim.action_engine.domain import ( + execution_program_hash, + seed_graph_hash, +) +from embodichain.utils.logger import log_warning + +from .models import ExecutionProgram, GroundedAction, SemanticStep + +__all__ = ["RuntimeRecorder"] + +_SAFE_NAME = re.compile(r"[^0-9A-Za-z._-]+") + + +def _safe_name(value: str) -> str: + result = _SAFE_NAME.sub("_", value).strip("._") + if not result: + raise ValueError("Runtime record path component must not be empty.") + return result + + +def _default_output_root() -> Path: + for parent in Path(__file__).resolve().parents: + if (parent / "setup.py").is_file() and (parent / "embodichain").is_dir(): + return parent / "outputs" / "action_engine" + return Path.cwd() / "outputs" / "action_engine" + + +def _jsonable(value: Any, env_id: int | None = None) -> Any: + if isinstance(value, torch.Tensor): + item = value + if env_id is not None and item.ndim > 0 and item.shape[0] > env_id: + item = item[env_id] + return item.detach().cpu().tolist() + if isinstance(value, dict): + return {str(key): _jsonable(item, env_id) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_jsonable(item, env_id) for item in value] + if isinstance(value, Path): + return value.as_posix() + return value + + +class RuntimeRecorder: + """Record execution decisions without copying the whole program per step.""" + + def __init__( + self, + program: ExecutionProgram, + *, + num_envs: int, + run_id: str | None = None, + episode_index: int = 0, + output_root: str | Path | None = None, + enabled: bool = True, + runtime_policy: Mapping[str, Any] | None = None, + runtime_policy_hash: str | None = None, + ) -> None: + self.enabled = enabled + self.num_envs = int(num_envs) + self.run_id = _safe_name( + run_id or datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ") + ) + root = ( + Path(output_root).expanduser().resolve() + if output_root is not None + else _default_output_root() + ) + self.output_dir = ( + root + / _safe_name(program.task) + / self.run_id + / f"episode_{int(episode_index):04d}" + ) + # The validated source graph remains untouched. Runtime documents are + # built from a detached copy and extend it only with a runtime envelope. + self.seed_topology = deepcopy(program.seed_graph or program.raw) + self.program_hash = ( + seed_graph_hash(program.seed_graph) + if program.seed_graph is not None + else execution_program_hash(program.raw) + ) + self.step_specs = { + str(step["id"]): deepcopy(step) for step in program.raw["semantic_steps"] + } + self.step_ordinals = { + step.id: index for index, step in enumerate(program.semantic_steps, start=1) + } + self.events: list[list[dict[str, Any]]] = [[] for _ in range(self.num_envs)] + self.program_metadata = { + "schema_version": "action_engine_runtime_record_v2", + "task": program.task, + "run_id": self.run_id, + "episode_index": int(episode_index), + "program_schema_version": self.seed_topology.get("schema_version"), + "seed_graph_hash": self.program_hash, + } + if runtime_policy is not None: + if not isinstance(runtime_policy_hash, str) or not runtime_policy_hash: + raise ValueError("Recorded runtime policy requires a non-empty hash.") + self.program_metadata["runtime_policy"] = deepcopy(dict(runtime_policy)) + self.program_metadata["runtime_policy_hash"] = runtime_policy_hash + + def register_step( + self, + step: SemanticStep, + spec: Mapping[str, Any], + ) -> None: + """Register a semantic step inserted by a runtime graph revision.""" + if not self.enabled: + return + raw = deepcopy(dict(spec)) + if str(raw.get("id")) != step.id: + raise ValueError("Runtime step spec ID must match the semantic step ID.") + existing = self.step_specs.get(step.id) + if existing is not None: + if existing != raw: + raise ValueError( + f"Runtime step {step.id!r} was registered with a different spec." + ) + return + self.step_specs[step.id] = raw + self.step_ordinals[step.id] = max(self.step_ordinals.values(), default=0) + 1 + + def edge( + self, + edge_id: str, + step: SemanticStep, + *, + assignments: list[str | None], + grounded: list[GroundedAction], + active: torch.Tensor, + failed: torch.Tensor, + action_steps: int, + planner_traces: Sequence[Mapping[str, Any]] = (), + diagnostics: Sequence[str] = (), + phase: str = "primary", + ) -> None: + if not self.enabled: + return + if phase not in {"primary", "recovery", "replay", "final_revalidation"}: + raise ValueError(f"Unknown execution phase {phase!r}.") + for env_id in range(self.num_envs): + event = { + "event": "edge", + "phase": phase, + "edge_id": edge_id, + "semantic_step_id": step.id, + "operator": step.operator, + "object": step.object_uid, + "arm": assignments[env_id], + "status": ( + "skipped" + if not bool(active[env_id]) + else ("failed" if bool(failed[env_id]) else "executed") + ), + "actions": [ + { + "class": item.action_class, + "control": item.control, + "target_object_pose": _jsonable( + item.target_object_pose, env_id + ), + "motion_policy": _jsonable(item.motion_policy), + } + for item in grounded + ], + "trajectory_steps": (int(action_steps) if bool(active[env_id]) else 0), + "time_utc": datetime.now(timezone.utc).isoformat(), + } + if diagnostics: + event["diagnostics"] = [str(item) for item in diagnostics] + if planner_traces: + event["planner_attempts"] = _jsonable(planner_traces, env_id) + self.events[env_id].append(event) + + def step( + self, + step: SemanticStep, + success: torch.Tensor, + *, + observed: torch.Tensor | None, + target: torch.Tensor | None, + metadata: Sequence[Mapping[str, Any]] | None = None, + phase: str = "primary", + ) -> None: + if not self.enabled: + return + if phase not in {"primary", "recovery", "replay", "final_revalidation"}: + raise ValueError(f"Unknown execution phase {phase!r}.") + if metadata is not None and len(metadata) != self.num_envs: + raise ValueError("Runtime step metadata must match num_envs.") + for env_id in range(self.num_envs): + event = { + "event": "semantic_step", + "phase": phase, + "semantic_step_id": step.id, + "status": "success" if bool(success[env_id]) else "failed", + "observed_position": _jsonable(observed, env_id), + "target_position": _jsonable(target, env_id), + "time_utc": datetime.now(timezone.utc).isoformat(), + } + if metadata is not None: + event.update(_jsonable(dict(metadata[env_id]))) + self.events[env_id].append(event) + self._write_step_checkpoint(env_id, step, event) + + def recovery( + self, + *, + failure_type: str, + failed_node_id: str, + active: torch.Tensor, + status: str, + recovery_group_id: str | None = None, + error: str | None = None, + semantic_step_id: str | None = None, + ) -> None: + """Record one bounded local-recovery phase for the selected rows.""" + if not self.enabled: + return + if status not in {"started", "succeeded", "failed", "rejected"}: + raise ValueError(f"Unknown recovery status {status!r}.") + for env_id in range(self.num_envs): + if not bool(active[env_id]): + continue + event = { + "event": "local_recovery", + "phase": "recovery", + "failure_type": str(failure_type), + "failed_node_id": str(failed_node_id), + "recovery_group_id": recovery_group_id, + "status": status, + "error": error, + "time_utc": datetime.now(timezone.utc).isoformat(), + } + if semantic_step_id is not None: + event["semantic_step_id"] = str(semantic_step_id) + self.events[env_id].append(event) + + def _env_dir(self, env_id: int) -> Path: + return self.output_dir / f"env_{env_id:04d}" + + def _write_step_checkpoint( + self, + env_id: int, + step: SemanticStep, + event: dict[str, Any], + ) -> None: + """Atomically publish one closed-loop semantic-step checkpoint.""" + related_events = [ + deepcopy(item) + for item in self.events[env_id] + if item.get("semantic_step_id") == step.id + ] + checkpoint = { + "schema_version": "action_engine_semantic_checkpoint_v2", + "seed_graph_hash": self.program_hash, + "task": self.program_metadata["task"], + "run_id": self.run_id, + "episode_index": self.program_metadata["episode_index"], + "env_id": env_id, + "semantic_step": deepcopy(self.step_specs[step.id]), + "status": event["status"], + "events": related_events, + "checkpointed_at_utc": event["time_utc"], + } + ordinal = self.step_ordinals[step.id] + filename = f"step_{ordinal:04d}_{_safe_name(step.id)}.json" + _write_json_atomic( + self._env_dir(env_id) / "checkpoints" / filename, + checkpoint, + ) + + def finalize( + self, + success: torch.Tensor, + *, + error: str | None = None, + ) -> str | None: + if not self.enabled: + return None + from embodichain.gen_sim.action_engine.graph_visualization import ( + render_task_graph_png, + ) + + finished_at = datetime.now(timezone.utc).isoformat() + for env_id in range(self.num_envs): + runtime = { + **self.program_metadata, + "env_id": env_id, + "status": ( + "aborted" + if error is not None + else ("success" if bool(success[env_id]) else "failed") + ), + "error": error, + "events": self.events[env_id], + "finished_at_utc": finished_at, + } + document = deepcopy(self.seed_topology) + document["runtime"] = runtime + env_dir = self._env_dir(env_id) + _write_json_atomic( + env_dir / "task_graph.json", + document, + ) + try: + png = render_task_graph_png(document) + if not isinstance(png, bytes): + raise TypeError("render_task_graph_png must return bytes.") + _write_bytes_atomic(env_dir / "task_graph.png", png) + except Exception as exc: + runtime["visualization_error"] = f"{type(exc).__name__}: {exc}" + document["runtime"] = runtime + _write_json_atomic(env_dir / "task_graph.json", document) + log_warning( + "Unable to render Action Engine runtime graph for " + f"env {env_id}: {exc}" + ) + return self.output_dir.as_posix() + + +def _write_json_atomic(path: Path, value: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + descriptor, temp_name = tempfile.mkstemp( + dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + text=True, + ) + temp_path = Path(temp_name) + try: + with os.fdopen(descriptor, "w", encoding="utf-8") as stream: + json.dump(value, stream, ensure_ascii=False, indent=2) + stream.write("\n") + stream.flush() + os.fsync(stream.fileno()) + os.replace(temp_path, path) + finally: + temp_path.unlink(missing_ok=True) + + +def _write_bytes_atomic(path: Path, value: bytes) -> None: + """Write one binary artifact without exposing a partial destination.""" + path.parent.mkdir(parents=True, exist_ok=True) + descriptor, temp_name = tempfile.mkstemp( + dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + ) + temp_path = Path(temp_name) + try: + with os.fdopen(descriptor, "wb") as stream: + stream.write(value) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temp_path, path) + finally: + temp_path.unlink(missing_ok=True) diff --git a/embodichain/gen_sim/action_engine/runtime/recovery.py b/embodichain/gen_sim/action_engine/runtime/recovery.py new file mode 100644 index 000000000..cbe3caa7a --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/recovery.py @@ -0,0 +1,649 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Bounded retry decisions and auditable RuntimeGraph revisions.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass +from typing import Any + +import torch + +from embodichain.gen_sim.action_engine.capabilities import ( + AtomicCapabilityRegistry, + build_atomic_capability_registry, + capability_precondition, +) +from embodichain.gen_sim.action_engine.domain import motion_policy, validate_seed_graph + +__all__ = [ + "FAILURE_TYPES", + "GraphRevision", + "RetryDecision", + "RuntimeGraph", + "build_upright_recovery", + "classify_failure", +] + +FAILURE_TYPES = frozenset( + { + "plan_failed", + "search_exhausted", + "grasp_missed", + "object_fallen", + "object_dropped", + "postcondition_failed", + } +) + + +@dataclass(frozen=True) +class RetryDecision: + """Per-environment result of one failed full-AtomicAction attempt.""" + + retry: torch.Tensor + recover: torch.Tensor + exhausted: torch.Tensor + attempts: tuple[int, ...] + + +@dataclass(frozen=True) +class GraphRevision: + """One immutable patch record over the original SeedGraph.""" + + revision: int + kind: str + reason: str + failed_node_id: str | None + inserted_group_ids: tuple[str, ...] + replaced_group_ids: tuple[str, ...] + active_env_ids: tuple[int, ...] = () + + +class RuntimeGraph: + """Keep SeedGraph immutable while applying bounded, validated revisions.""" + + def __init__( + self, + seed_graph: Mapping[str, Any], + *, + num_envs: int, + max_retries: int = 2, + max_revisions: int = 8, + max_recovery_actions: int = 12, + registry: AtomicCapabilityRegistry | None = None, + ) -> None: + if num_envs < 1: + raise ValueError("RuntimeGraph num_envs must be positive.") + self.registry = registry or build_atomic_capability_registry() + self.seed_graph = validate_seed_graph( + seed_graph, + known_actions=self.registry.names(), + ) + from embodichain.gen_sim.action_engine.planning.linker import ( + validate_persisted_contracts, + ) + + validate_persisted_contracts(self.seed_graph, self.registry) + self._graph = deepcopy(self.seed_graph) + self.num_envs = int(num_envs) + self.max_retries = int(max_retries) + self.max_revisions = int(max_revisions) + self.max_recovery_actions = int(max_recovery_actions) + if min(self.max_retries, self.max_revisions, self.max_recovery_actions) < 0: + raise ValueError("RuntimeGraph budgets must be non-negative.") + self._attempts: dict[str, list[int]] = {} + self._recovery_action_count = 0 + self.revisions: list[GraphRevision] = [] + + @property + def graph(self) -> dict[str, Any]: + """Return the current detached RuntimeGraph snapshot.""" + return deepcopy(self._graph) + + def record_failure( + self, + node_id: str, + failed: torch.Tensor, + *, + precondition_holds: torch.Tensor, + ) -> RetryDecision: + """Consume attempt budgets and distinguish retry from recovery.""" + failed = _mask(failed, self.num_envs) + precondition_holds = _mask(precondition_holds, self.num_envs) + attempts = self._attempts.setdefault(node_id, [1] * self.num_envs) + retry = torch.zeros_like(failed) + recover = torch.zeros_like(failed) + exhausted = torch.zeros_like(failed) + node = _node(self._graph, node_id) + capability = self.registry.get(str(node["atomic_action"])) + for env_id in torch.nonzero(failed, as_tuple=False).flatten().tolist(): + attempts[env_id] += 1 + can_retry = ( + capability.retry_mode != "non_retryable" + and bool(precondition_holds[env_id]) + and attempts[env_id] <= self.max_retries + 1 + ) + if can_retry: + retry[env_id] = True + elif capability.retry_mode != "non_retryable": + recover[env_id] = True + else: + exhausted[env_id] = True + return RetryDecision(retry, recover, exhausted, tuple(attempts)) + + def insert_recovery_subgraph( + self, + *, + failed_node_id: str, + recovery_nodes: Sequence[Mapping[str, Any]], + recovery_group: Mapping[str, Any], + failure_type: str, + active_env_ids: Sequence[int] | None = None, + preserve_failed_group_suffix: bool = False, + ) -> dict[str, Any]: + """Insert a complete recovery TaskGroup and rewire the unfinished suffix.""" + if failure_type not in FAILURE_TYPES: + raise ValueError(f"Unknown failure type {failure_type!r}.") + if len(self.revisions) >= self.max_revisions: + raise RuntimeError("RuntimeGraph revision budget exhausted.") + if ( + self._recovery_action_count + len(recovery_nodes) + > self.max_recovery_actions + ): + raise RuntimeError("RuntimeGraph recovery-action budget exhausted.") + env_ids = tuple( + sorted( + set( + range(self.num_envs) + if active_env_ids is None + else (int(env_id) for env_id in active_env_ids) + ) + ) + ) + if not env_ids or env_ids[0] < 0 or env_ids[-1] >= self.num_envs: + raise ValueError( + "Recovery active_env_ids are outside the environment range." + ) + failed_node = _node(self._graph, failed_node_id) + failed_group_id = str(failed_node["task_instance_id"]) + group = deepcopy(dict(recovery_group)) + if group.get("role") != "recovery": + raise ValueError("Inserted recovery TaskGroup must use role='recovery'.") + group_id = str(group.get("id", "")) + if not group_id: + raise ValueError("Inserted recovery TaskGroup requires an ID.") + if any(item["id"] == group_id for item in self._graph["task_groups"]): + raise ValueError(f"RuntimeGraph already contains TaskGroup {group_id!r}.") + nodes = [deepcopy(dict(node)) for node in recovery_nodes] + if not nodes: + raise ValueError("Recovery subgraph must contain at least one node.") + for node in nodes: + node.pop("contract", None) + node.pop("resources", None) + node["role"] = "cleanup" if node.get("role") == "cleanup" else "recovery" + node["task_instance_id"] = group_id + node["task_type"] = group["task_type"] + recovery_ids = {str(node["id"]) for node in nodes} + if len(recovery_ids) != len(nodes): + raise ValueError("Recovery node IDs must be unique.") + node_by_id = {str(node["id"]): node for node in self._graph["nodes"]} + children_by_id = {node_id: [] for node_id in node_by_id} + for node_id, node in node_by_id.items(): + for dependency in node["depends_on"]: + children_by_id[str(dependency)].append(node_id) + descendants: set[str] = set() + pending = list(children_by_id[failed_node_id]) + while pending: + node_id = pending.pop() + if node_id in descendants: + continue + descendants.add(node_id) + pending.extend(children_by_id[node_id]) + same_group_descendants = { + node_id + for node_id in descendants + if str(node_by_id[node_id]["task_instance_id"]) == failed_group_id + } + cleanup_suffix_ids: set[str] = set() + if ( + str(failed_node["atomic_action"]) == "HandOver" + and not preserve_failed_group_suffix + ): + # A failed handover leaves ownership indeterminate. Its + # transfer-arm retreat/home tail must not execute from a stale + # handover pose; recovery owns the cleanup before replanning. + cleanup_suffix_ids = { + node_id + for node_id in same_group_descendants + if node_by_id[node_id]["role"] == "cleanup" + } + non_cleanup_dependents = [ + node_id + for node_id in same_group_descendants - cleanup_suffix_ids + if any( + dependency in cleanup_suffix_ids + for dependency in node_by_id[node_id]["depends_on"] + ) + ] + if non_cleanup_dependents: + raise ValueError( + "Cannot remove the HandOver cleanup suffix because it feeds " + f"same-group non-cleanup nodes: {sorted(non_cleanup_dependents)}." + ) + first = [ + node + for node in nodes + if not any(dep in recovery_ids for dep in node.get("depends_on", [])) + ] + if not first: + raise ValueError("Recovery subgraph has no entry node.") + for node in first: + node["depends_on"] = list( + dict.fromkeys([*node.get("depends_on", []), failed_node_id]) + ) + terminal_ids = _terminal_ids(nodes) + + patched = deepcopy(self._graph) + for node in patched["nodes"]: + if node["id"] in recovery_ids: + raise ValueError(f"RuntimeGraph already contains node {node['id']!r}.") + node_id = str(node["id"]) + if ( + preserve_failed_group_suffix + or node_id not in descendants + or node_id in same_group_descendants + ): + continue + node["depends_on"] = list( + dict.fromkeys( + [ + dependency + for dependency in node["depends_on"] + if dependency != failed_node_id + and dependency not in cleanup_suffix_ids + ] + + terminal_ids + ) + ) + if cleanup_suffix_ids: + patched["nodes"] = [ + node + for node in patched["nodes"] + if str(node["id"]) not in cleanup_suffix_ids + ] + failed_group = next( + item + for item in patched["task_groups"] + if str(item["id"]) == failed_group_id + ) + failed_group["node_ids"] = [ + node_id + for node_id in failed_group["node_ids"] + if node_id not in cleanup_suffix_ids + ] + group["depends_on"] = list( + dict.fromkeys([failed_group_id, *group.get("depends_on", [])]) + ) + group.pop("contract", None) + group["node_ids"] = [str(node["id"]) for node in nodes] + for downstream in patched["task_groups"]: + if ( + not preserve_failed_group_suffix + and failed_group_id in downstream["depends_on"] + ): + downstream["depends_on"] = [ + dependency + for dependency in downstream["depends_on"] + if dependency != failed_group_id + ] + [group_id] + patched["nodes"].extend(nodes) + patched["task_groups"].append(group) + patched["metadata"] = { + **patched.get("metadata", {}), + "runtime_revision": len(self.revisions) + 1, + } + from embodichain.gen_sim.action_engine.planning.linker import link_seed_graph + + self._graph = link_seed_graph( + patched, + registry=self.registry, + ) + self._recovery_action_count += len(nodes) + self.revisions.append( + GraphRevision( + revision=len(self.revisions) + 1, + kind="insert_recovery", + reason=failure_type, + failed_node_id=failed_node_id, + inserted_group_ids=(group_id,), + replaced_group_ids=(), + active_env_ids=env_ids, + ) + ) + return self.graph + + def insert_default_recovery( + self, + *, + failed_node_id: str, + failure_type: str, + active_env_ids: Sequence[int] | None = None, + resume_failed_group: bool = False, + ) -> dict[str, Any]: + """Insert one of the deliberately small built-in recovery strategies.""" + if failure_type != "object_fallen": + raise ValueError( + f"No bounded default recovery is registered for {failure_type!r}." + ) + nodes, group = build_upright_recovery( + self._graph, + failed_node_id=failed_node_id, + revision=len(self.revisions) + 1, + resume_failed_group=resume_failed_group, + ) + return self.insert_recovery_subgraph( + failed_node_id=failed_node_id, + recovery_nodes=nodes, + recovery_group=group, + failure_type=failure_type, + active_env_ids=active_env_ids, + preserve_failed_group_suffix=resume_failed_group, + ) + + def replace_unfinished_suffix( + self, + replacement: Mapping[str, Any], + *, + completed_group_ids: Sequence[str], + reason: str, + ) -> dict[str, Any]: + """Install a fully replanned suffix while preserving completed groups.""" + if len(self.revisions) >= self.max_revisions: + raise RuntimeError("RuntimeGraph revision budget exhausted.") + candidate = validate_seed_graph( + replacement, + known_actions=self.registry.names(), + ) + from embodichain.gen_sim.action_engine.planning.linker import ( + validate_persisted_contracts, + ) + + validate_persisted_contracts(candidate, self.registry) + if candidate["task_id"] != self.seed_graph["task_id"]: + raise ValueError("Suffix replanning cannot change the task_id.") + if candidate["capability_catalog_hash"] != self.registry.catalog_hash(): + raise ValueError( + "Replanned suffix capability catalog does not match runtime." + ) + current_groups = {group["id"]: group for group in self._graph["task_groups"]} + replacement_groups = {group["id"]: group for group in candidate["task_groups"]} + current_nodes = {node["id"]: node for node in self._graph["nodes"]} + replacement_nodes = {node["id"]: node for node in candidate["nodes"]} + completed = set(completed_group_ids) + for group_id in completed: + current_group = current_groups.get(group_id) + replacement_group = replacement_groups.get(group_id) + if current_group is None or replacement_group != current_group: + raise ValueError( + f"Replanning changed completed TaskGroup {group_id!r}." + ) + if any( + replacement_nodes.get(node_id) != current_nodes[node_id] + for node_id in current_group["node_ids"] + ): + raise ValueError( + f"Replanning changed nodes of completed TaskGroup {group_id!r}." + ) + replaced = tuple(sorted(set(current_groups) - completed)) + self._graph = candidate + self.revisions.append( + GraphRevision( + revision=len(self.revisions) + 1, + kind="replan_suffix", + reason=str(reason), + failed_node_id=None, + inserted_group_ids=tuple( + sorted(set(replacement_groups) - set(current_groups)) + ), + replaced_group_ids=replaced, + active_env_ids=tuple(range(self.num_envs)), + ) + ) + return self.graph + + +def classify_failure( + action_name: str, + *, + planning_succeeded: bool, + postcondition_succeeded: bool | None = None, + object_fallen: bool = False, + held_before: bool = False, + held_after: bool = False, + registry: AtomicCapabilityRegistry | None = None, +) -> str: + """Classify only the bounded common recovery cases supported by v2.""" + capability = (registry or build_atomic_capability_registry()).get(action_name) + if capability.failure_classifier_hook is not None: + result = capability.failure_classifier_hook( + action_name=action_name, + planning_succeeded=planning_succeeded, + postcondition_succeeded=postcondition_succeeded, + object_fallen=object_fallen, + held_before=held_before, + held_after=held_after, + ) + if result not in FAILURE_TYPES: + raise ValueError( + f"AtomicAction {action_name!r} failure classifier returned {result!r}." + ) + return result + if not planning_succeeded: + return "search_exhausted" + if object_fallen: + return "object_fallen" + if held_before and not held_after: + return "object_dropped" + if capability.failure_classifier == "grasp" and not held_after: + return "grasp_missed" + if postcondition_succeeded is False: + return "postcondition_failed" + return "postcondition_failed" + + +def build_upright_recovery( + graph: Mapping[str, Any], + *, + failed_node_id: str, + revision: int, + resume_failed_group: bool = False, +) -> tuple[list[dict[str, Any]], dict[str, Any]]: + """Build a coordinate-free E2 recovery group for a fallen rigid object.""" + failed = _node(graph, failed_node_id) + object_uid = str(failed["object_uid"]) + group_id = f"recovery_e2_{int(revision):02d}_{failed_node_id}" + actor = _recovery_actor(graph, failed) + held_consumer_arm = None + if not resume_failed_group: + held_consumer_arm = _downstream_held_consumer_arm(graph, failed, object_uid) + if held_consumer_arm is not None and not ( + actor.get("mode") == "required" and actor.get("arm") == held_consumer_arm + ): + raise ValueError( + "Recovery cannot satisfy the downstream held-object contract without " + "changing the failed TaskGroup actor; resume and replay the failed " + "TaskGroup instead." + ) + hold_for_downstream = ( + held_consumer_arm is not None + and actor.get("mode") == "required" + and actor.get("arm") == held_consumer_arm + ) + upright = motion_policy(("orientation", "upright")) + full_specs = ( + ("PickUp", {"kind": "object", "object": object_uid}, upright), + ( + "MoveHeldObject", + {"kind": "semantic_goal", "semantic_step": group_id, "phase": "final"}, + upright, + ), + ("Place", {"kind": "current_held_pose"}, upright), + ( + "MoveEndEffector", + {"kind": "policy_pose", "source": "release", "operation": "retreat"}, + upright, + ), + ( + "MoveJoints", + {"kind": "joint_state", "source": "initial"}, + motion_policy(), + ), + ) + specs = full_specs[:2] if hold_for_downstream else full_specs + nodes = [] + registry = build_atomic_capability_registry() + dependencies: list[str] = [] + for index, (action, binding, policy_spec) in enumerate(specs, start=1): + node_id = f"{group_id}__a{index:02d}" + node = { + "id": node_id, + "atomic_action": action, + "object_uid": object_uid, + "actor": actor, + "control": "arm", + "target_binding": binding, + "depends_on": dependencies, + "task_instance_id": group_id, + "task_type": "E2", + "role": "recovery" if index <= 3 else "cleanup", + "precondition": {}, + "postcondition": {}, + "motion_policy": deepcopy(dict(policy_spec)), + } + node["precondition"] = capability_precondition( + registry.get(action), + object_uid=object_uid, + actor=actor, + target_binding=binding, + ) + nodes.append(node) + dependencies = [node_id] + group = { + "id": group_id, + "task_type": "E2", + "role": "recovery", + "operator": "orient_object", + "object_uid": object_uid, + "actor": actor, + "goal": { + "relation": "none", + "reference_state": "live", + "orientation_goal": "upright", + "orientation_axis": "none", + "position_anchor": "live_xy", + "support_object": "table", + "upright_local_axis": "long_axis", + "terminal_behavior": "hold" if hold_for_downstream else "place", + }, + "depends_on": [], + "parent_task_instance_id": str(failed["task_instance_id"]), + "node_ids": [node["id"] for node in nodes], + "success": {"type": "object_upright", "object": object_uid}, + } + return nodes, group + + +def _recovery_actor( + graph: Mapping[str, Any], + failed: Mapping[str, Any], +) -> dict[str, Any]: + """Preserve the failed TaskGroup's arm-selection contract.""" + group_id = str(failed["task_instance_id"]) + group = next( + (item for item in graph["task_groups"] if str(item["id"]) == group_id), + None, + ) + source = (group or failed).get("actor", {"mode": "auto"}) + if not isinstance(source, Mapping): + raise ValueError(f"Failed TaskGroup {group_id!r} has an invalid actor.") + actor = deepcopy(dict(source)) + mode = str(actor.get("mode", "auto")) + if mode == "required": + if actor.get("arm") not in {"left_arm", "right_arm"}: + raise ValueError( + f"Failed TaskGroup {group_id!r} has an invalid required arm." + ) + elif mode == "auto": + actor = {"mode": "auto"} + elif mode == "coordinated": + raise ValueError( + "The single-arm upright recovery cannot inherit a coordinated actor." + ) + else: + raise ValueError(f"The upright recovery cannot inherit actor mode {mode!r}.") + return actor + + +def _downstream_held_consumer_arm( + graph: Mapping[str, Any], + failed: Mapping[str, Any], + object_uid: str, +) -> str | None: + failed_group_id = str(failed["task_instance_id"]) + nodes = {str(node["id"]): node for node in graph["nodes"]} + for group in graph["task_groups"]: + if failed_group_id not in {str(item) for item in group.get("depends_on", ())}: + continue + if str(group.get("object_uid")) != object_uid: + continue + node_ids = {str(item) for item in group["node_ids"]} + for node_id in group["node_ids"]: + node = nodes[str(node_id)] + if any(str(parent) in node_ids for parent in node["depends_on"]): + continue + for requirement in node.get("contract", {}).get("requires", ()): + if ( + requirement.get("predicate") == "object_held" + and requirement.get("object_uid") == object_uid + and requirement.get("arm") in {"left_arm", "right_arm"} + ): + return str(requirement["arm"]) + return None + + +def _node(graph: Mapping[str, Any], node_id: str) -> Mapping[str, Any]: + try: + return next(node for node in graph["nodes"] if node["id"] == node_id) + except StopIteration as error: + raise ValueError(f"RuntimeGraph contains no node {node_id!r}.") from error + + +def _terminal_ids(nodes: Sequence[Mapping[str, Any]]) -> list[str]: + depended = { + dependency for node in nodes for dependency in node.get("depends_on", []) + } + return [str(node["id"]) for node in nodes if node["id"] not in depended] + + +def _mask(value: torch.Tensor, count: int) -> torch.Tensor: + result = torch.as_tensor(value, dtype=torch.bool).reshape(-1) + if result.numel() != count: + raise ValueError(f"Expected a mask with {count} values.") + return result diff --git a/embodichain/gen_sim/action_engine/runtime/reporting.py b/embodichain/gen_sim/action_engine/runtime/reporting.py new file mode 100644 index 000000000..ea17ef151 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/reporting.py @@ -0,0 +1,348 @@ +# ---------------------------------------------------------------------------- +# 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 validation and atomic publication for execution reports.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +import json +import os +from pathlib import Path +import platform +import subprocess +import tempfile +from typing import Any + +from embodichain import __version__ as embodichain_version + +from .models import ExecutionReport + +__all__ = [ + "EXECUTION_REPORT_FILENAME", + "EXECUTION_REPORT_SCHEMA", + "build_execution_provenance", + "validate_execution_report", + "write_execution_report", +] + +EXECUTION_REPORT_SCHEMA = "action_engine_execution_report_v2" +EXECUTION_REPORT_FILENAME = "execution_report.json" + + +def build_execution_provenance( + *, + episode_seed: int | None = None, + runtime_arguments: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + """Capture the minimum code and runtime context needed to reproduce a run.""" + git_commit, git_dirty = _git_code_state() + provenance = { + "episode_seed": episode_seed, + "embodichain_version": str(embodichain_version), + "python_version": platform.python_version(), + "git_commit": git_commit, + "git_dirty": git_dirty, + "runtime_arguments": deepcopy(dict(runtime_arguments or {})), + } + return _validate_execution_provenance(provenance) + + +def validate_execution_report(value: Mapping[str, Any]) -> dict[str, Any]: + """Validate the tensor-free, strict-JSON Action Agent result protocol.""" + result = _mapping(value, "ExecutionReport") + keys = { + "schema_version", + "task_id", + "plan_hash", + "action_graph_hash", + "status", + "run_id", + "episode_id", + "provenance", + "environments", + "action_count", + "retry_count", + "recovery_count", + "revision_count", + "failure_events", + "graph_revisions", + "record_dir", + "error", + } + _keys(result, keys, "ExecutionReport") + if result.get("schema_version") != EXECUTION_REPORT_SCHEMA: + raise ValueError( + "ExecutionReport.schema_version must be " f"{EXECUTION_REPORT_SCHEMA!r}." + ) + for key in ("task_id", "run_id", "episode_id"): + result[key] = _nonempty(result.get(key), f"ExecutionReport.{key}") + result["provenance"] = _validate_execution_provenance(result.get("provenance")) + for key in ("plan_hash", "action_graph_hash"): + result[key] = _digest(result.get(key), f"ExecutionReport.{key}") + result["status"] = _enum( + result.get("status"), + {"succeeded", "failed", "rejected", "aborted"}, + "ExecutionReport.status", + ) + + env_keys = { + "env_id", + "success", + "semantic_success", + "action_count", + "retry_count", + "recovery_count", + "revision_count", + "failures", + } + environments = [] + for index, raw in enumerate( + _sequence(result.get("environments"), "ExecutionReport.environments") + ): + context = f"ExecutionReport.environments[{index}]" + environment = _mapping(raw, context) + _keys(environment, env_keys, context) + environment["env_id"] = _string(environment.get("env_id"), f"{context}.env_id") + if not isinstance(environment.get("success"), bool): + raise ValueError(f"{context}.success must be a boolean.") + semantic_success = _mapping( + environment.get("semantic_success"), f"{context}.semantic_success" + ) + if any(not isinstance(item, bool) for item in semantic_success.values()): + raise ValueError(f"{context}.semantic_success values must be booleans.") + environment["semantic_success"] = semantic_success + for key in ("action_count", "retry_count", "recovery_count", "revision_count"): + environment[key] = _integer( + environment.get(key), f"{context}.{key}", minimum=0 + ) + environment["failures"] = _mapping_sequence( + environment.get("failures"), f"{context}.failures" + ) + environments.append(environment) + result["environments"] = environments + + for key in ("action_count", "retry_count", "recovery_count", "revision_count"): + result[key] = _integer(result.get(key), f"ExecutionReport.{key}", minimum=0) + result["failure_events"] = _mapping_sequence( + result.get("failure_events"), "ExecutionReport.failure_events" + ) + result["graph_revisions"] = _mapping_sequence( + result.get("graph_revisions"), "ExecutionReport.graph_revisions" + ) + for key in ("record_dir", "error"): + if result.get(key) is not None: + result[key] = _string(result.get(key), f"ExecutionReport.{key}") + + if result["status"] == "rejected" and result["action_count"] != 0: + raise ValueError("A rejected ExecutionReport must have action_count=0.") + successes = [environment["success"] for environment in environments] + if result["status"] == "succeeded" and ( + not successes or not all(successes) or result.get("error") is not None + ): + raise ValueError( + "A succeeded ExecutionReport requires successful environments and no error." + ) + if result["status"] == "failed" and ( + not successes or all(successes) or result.get("error") is not None + ): + raise ValueError( + "A failed ExecutionReport requires at least one failed environment and no error." + ) + if result["status"] in {"rejected", "aborted"} and not result.get("error"): + raise ValueError( + f"A {result['status']} ExecutionReport requires a non-empty error." + ) + _json_safe(result, "ExecutionReport") + return result + + +def write_execution_report(output_dir: str | Path, value: Any) -> Path: + """Atomically write a validated execution report into a record directory.""" + payload = value.as_mapping() if isinstance(value, ExecutionReport) else value + validated = validate_execution_report(payload) + root = Path(output_dir).expanduser().resolve() + root.mkdir(parents=True, exist_ok=True) + path = root / EXECUTION_REPORT_FILENAME + encoded = ( + json.dumps(validated, ensure_ascii=False, indent=2, allow_nan=False) + "\n" + ).encode("utf-8") + temporary: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="wb", + dir=root, + prefix=f".{path.name}.", + suffix=".tmp", + delete=False, + ) as stream: + temporary = Path(stream.name) + stream.write(encoded) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path) + temporary = None + finally: + if temporary is not None: + temporary.unlink(missing_ok=True) + return path + + +def _validate_execution_provenance(value: Any) -> dict[str, Any]: + context = "ExecutionReport.provenance" + result = _mapping(value, context) + _keys( + result, + { + "episode_seed", + "embodichain_version", + "python_version", + "git_commit", + "git_dirty", + "runtime_arguments", + }, + context, + ) + seed = result.get("episode_seed") + if seed is not None and (not isinstance(seed, int) or isinstance(seed, bool)): + raise ValueError(f"{context}.episode_seed must be an integer or null.") + result["embodichain_version"] = _nonempty( + result.get("embodichain_version"), f"{context}.embodichain_version" + ) + result["python_version"] = _nonempty( + result.get("python_version"), f"{context}.python_version" + ) + commit = result.get("git_commit") + if commit is not None: + commit = _string(commit, f"{context}.git_commit") + if len(commit) not in {40, 64} or any( + character not in "0123456789abcdef" for character in commit + ): + raise ValueError( + f"{context}.git_commit must be a lowercase Git object ID or null." + ) + result["git_commit"] = commit + dirty = result.get("git_dirty") + if dirty is not None and not isinstance(dirty, bool): + raise ValueError(f"{context}.git_dirty must be a boolean or null.") + arguments = _mapping( + result.get("runtime_arguments"), f"{context}.runtime_arguments" + ) + if any(not isinstance(key, str) or not key for key in arguments): + raise ValueError(f"{context}.runtime_arguments keys must be non-empty strings.") + _json_safe(arguments, f"{context}.runtime_arguments") + result["runtime_arguments"] = arguments + return result + + +def _git_code_state() -> tuple[str | None, bool | None]: + repository = Path(__file__).resolve().parents[4] + try: + commit = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=repository, + check=False, + capture_output=True, + text=True, + timeout=2.0, + ) + except (OSError, subprocess.SubprocessError): + return None, None + commit_id = commit.stdout.strip().lower() + if commit.returncode != 0 or len(commit_id) not in {40, 64}: + return None, None + try: + status = subprocess.run( + ["git", "status", "--porcelain"], + cwd=repository, + check=False, + capture_output=True, + text=True, + timeout=2.0, + ) + except (OSError, subprocess.SubprocessError): + return commit_id, None + dirty = bool(status.stdout.strip()) if status.returncode == 0 else None + return commit_id, dirty + + +def _mapping(value: Any, context: str) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise ValueError(f"{context} must be a mapping.") + return deepcopy(dict(value)) + + +def _sequence(value: Any, context: str) -> list[Any]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise ValueError(f"{context} must be a list.") + return list(value) + + +def _keys(value: Mapping[str, Any], expected: set[str], context: str) -> None: + if set(value) != expected: + raise ValueError( + f"{context} requires exactly fields {sorted(expected)}; " + f"received {sorted(value)}." + ) + + +def _string(value: Any, context: str) -> str: + if not isinstance(value, str): + raise ValueError(f"{context} must be a string.") + return value.strip() + + +def _nonempty(value: Any, context: str) -> str: + result = _string(value, context) + if not result: + raise ValueError(f"{context} must not be empty.") + return result + + +def _enum(value: Any, choices: set[str], context: str) -> str: + result = _string(value, context) + if result not in choices: + raise ValueError(f"{context} must be one of {sorted(choices)}.") + return result + + +def _integer(value: Any, context: str, *, minimum: int) -> int: + if not isinstance(value, int) or isinstance(value, bool) or value < minimum: + raise ValueError(f"{context} must be an integer in the allowed range.") + return value + + +def _digest(value: Any, context: str) -> str: + result = _string(value, context) + if len(result) != 64 or any( + character not in "0123456789abcdef" for character in result + ): + raise ValueError(f"{context} must be a lowercase SHA-256 digest.") + return result + + +def _mapping_sequence(value: Any, context: str) -> list[dict[str, Any]]: + result = [_mapping(item, context) for item in _sequence(value, context)] + _json_safe(result, context) + return result + + +def _json_safe(value: Any, context: str) -> None: + try: + json.dumps(value, allow_nan=False) + except (TypeError, ValueError) as error: + raise ValueError(f"{context} must be finite and JSON serializable.") from error diff --git a/embodichain/gen_sim/action_engine/unbound.py b/embodichain/gen_sim/action_engine/unbound.py new file mode 100644 index 000000000..1b0a4af77 --- /dev/null +++ b/embodichain/gen_sim/action_engine/unbound.py @@ -0,0 +1,205 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Scene-independent Action Engine draft produced before final UID binding.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +import json +from typing import Any, Final, TypeAlias + +from embodichain.gen_sim.action_engine.domain.task_contracts import TASK_CONTRACTS + +__all__ = [ + "ActionCapabilityError", + "UNBOUND_ACTION_PLAN_SCHEMA", + "UnboundActionPlan", + "build_unbound_action_plan", + "validate_unbound_action_plan", +] + +UNBOUND_ACTION_PLAN_SCHEMA: Final = "embodichain.unbound-action-plan/v1" +UnboundActionPlan: TypeAlias = dict[str, Any] + + +class ActionCapabilityError(ValueError): + """A required AtomicAction is missing or not executable.""" + + +_PLAN_KEYS = frozenset( + { + "schema_version", + "task_id", + "candidate_id", + "instruction", + "steps", + "required_actions", + } +) +_STEP_KEYS = frozenset( + {"step_id", "task_type", "object", "target", "depends_on", "actions"} +) + + +def build_unbound_action_plan(candidate: Mapping[str, Any]) -> UnboundActionPlan: + """Lower a TaskCandidate into an Action-owned plan without scene UIDs. + + Args: + candidate: Validated Task Engine candidate or an equivalent mapping. + + Returns: + A strict JSON plan whose selectors remain logical references. + + Raises: + TypeError: If the candidate or draft is not a mapping. + ValueError: If the draft references an unsupported task type. + """ + value = _mapping(candidate, "candidate") + draft = _mapping(value.get("draft"), "candidate.draft") + task_id = _nonempty(draft.get("task_id"), "candidate.draft.task_id") + instruction = _nonempty(draft.get("instruction"), "candidate.draft.instruction") + candidate_id = _nonempty(value.get("candidate_id"), "candidate.candidate_id") + steps = [] + required_actions: set[str] = set() + for index, raw in enumerate(_sequence(draft.get("steps"), "candidate.draft.steps")): + step = _mapping(raw, f"candidate.draft.steps[{index}]") + task_type = _nonempty( + step.get("task_type"), f"candidate.draft.steps[{index}].task_type" + ) + contract = TASK_CONTRACTS.get(task_type) + if contract is None: + raise ValueError(f"Action Engine does not support task type {task_type!r}.") + actions = [str(name) for name in contract.core_actions] + required_actions.update(actions) + steps.append( + { + "step_id": _nonempty( + step.get("id"), f"candidate.draft.steps[{index}].id" + ), + "task_type": task_type, + "object": deepcopy(step.get("object")), + "target": deepcopy(step.get("target")), + "depends_on": deepcopy(step.get("depends_on", [])), + "actions": actions, + } + ) + return validate_unbound_action_plan( + { + "schema_version": UNBOUND_ACTION_PLAN_SCHEMA, + "task_id": task_id, + "candidate_id": candidate_id, + "instruction": instruction, + "steps": steps, + "required_actions": sorted(required_actions), + } + ) + + +def validate_unbound_action_plan( + value: Mapping[str, Any], +) -> UnboundActionPlan: + """Validate and detach one scene-independent Action plan. + + Args: + value: Candidate plan mapping. + + Returns: + A strict JSON-safe detached plan. + + Raises: + TypeError: If a mapping or sequence field has the wrong type. + ValueError: If the schema, dependency graph, or actions are invalid. + """ + result = _mapping(value, "UnboundActionPlan") + if set(result) != _PLAN_KEYS: + raise ValueError("UnboundActionPlan fields are invalid.") + if result.get("schema_version") != UNBOUND_ACTION_PLAN_SCHEMA: + raise ValueError("UnboundActionPlan.schema_version is invalid.") + for key in ("task_id", "candidate_id", "instruction"): + result[key] = _nonempty(result.get(key), f"UnboundActionPlan.{key}") + + steps = [] + seen: set[str] = set() + actions_used: set[str] = set() + for index, raw in enumerate(_sequence(result.get("steps"), "steps")): + context = f"UnboundActionPlan.steps[{index}]" + step = _mapping(raw, context) + if set(step) != _STEP_KEYS: + raise ValueError(f"{context} fields are invalid.") + step_id = _nonempty(step.get("step_id"), f"{context}.step_id") + if step_id in seen: + raise ValueError("UnboundActionPlan step IDs must be unique.") + task_type = _nonempty(step.get("task_type"), f"{context}.task_type") + contract = TASK_CONTRACTS.get(task_type) + if contract is None: + raise ValueError(f"{context}.task_type is unsupported.") + dependencies = _strings(step.get("depends_on"), f"{context}.depends_on") + if any(dependency not in seen for dependency in dependencies): + raise ValueError( + f"{context}.depends_on must reference preceding unbound steps." + ) + actions = _strings(step.get("actions"), f"{context}.actions") + if actions != [str(name) for name in contract.core_actions]: + raise ValueError(f"{context}.actions do not match the task contract.") + for selector_name in ("object", "target"): + if not isinstance(step.get(selector_name), Mapping): + raise TypeError(f"{context}.{selector_name} must be a mapping.") + step[selector_name] = deepcopy(dict(step[selector_name])) + step["step_id"] = step_id + step["task_type"] = task_type + step["depends_on"] = dependencies + step["actions"] = actions + steps.append(step) + seen.add(step_id) + actions_used.update(actions) + if not steps: + raise ValueError("UnboundActionPlan.steps must not be empty.") + required = _strings(result.get("required_actions"), "required_actions") + if required != sorted(actions_used): + raise ValueError("UnboundActionPlan.required_actions is not canonical.") + result["steps"] = steps + result["required_actions"] = required + json.dumps(result, ensure_ascii=False, allow_nan=False) + return result + + +def _mapping(value: Any, context: str) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise TypeError(f"{context} must be a mapping.") + return deepcopy(dict(value)) + + +def _sequence(value: Any, context: str) -> list[Any]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise TypeError(f"{context} must be a sequence.") + return list(value) + + +def _strings(value: Any, context: str) -> list[str]: + result = _sequence(value, context) + if any(not isinstance(item, str) or not item for item in result): + raise ValueError(f"{context} must contain non-empty strings.") + if len(set(result)) != len(result): + raise ValueError(f"{context} must not contain duplicates.") + return list(result) + + +def _nonempty(value: Any, context: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{context} must be a non-empty string.") + return value.strip() diff --git a/tests/gen_sim/action_engine/cli/test_run_agent.py b/tests/gen_sim/action_engine/cli/test_run_agent.py new file mode 100644 index 000000000..53b2dc99b --- /dev/null +++ b/tests/gen_sim/action_engine/cli/test_run_agent.py @@ -0,0 +1,219 @@ +# ---------------------------------------------------------------------------- +# 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 json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from embodichain.gen_sim.action_engine.cli.run_agent import ( + _ABWorkerConfig, + _SerializedABBranch, + _capture_ab_initial_frame, + _prepare_ab_branches, + _publish_task_engine_report, + _task_engine_exit_code, +) +from embodichain.gen_sim.action_engine.runtime import ( + ExecutionReport, + build_execution_provenance, +) + + +class record_camera_data: + def __init__(self) -> None: + self.calls = [] + + def __call__(self, *args, **kwargs) -> None: + self.calls.append((args, kwargs)) + + +class _FakeEnv: + def __init__(self, recorder=None) -> None: + self.unwrapped = self + self.event_manager = SimpleNamespace( + _mode_functor_cfgs={ + "interval": ( + [ + SimpleNamespace( + func=recorder, + params={"name": "record_cam_audience_view"}, + ) + ] + if recorder is not None + else [] + ) + } + ) + + +def test_capture_ab_initial_frame_invokes_only_audience_recorder() -> None: + recorder = record_camera_data() + env = _FakeEnv(recorder) + + _capture_ab_initial_frame(env) + + assert len(recorder.calls) == 1 + args, kwargs = recorder.calls[0] + assert args == (env, None) + assert kwargs == {"name": "record_cam_audience_view"} + + +def test_capture_ab_initial_frame_requires_audience_recorder() -> None: + with pytest.raises(RuntimeError, match="audience recorder"): + _capture_ab_initial_frame(_FakeEnv()) + + +def _worker_config(route: str) -> _ABWorkerConfig: + return _ABWorkerConfig( + route=route, + gym_config={}, + env_options={}, + gym_id="ActionEngine-v1", + agent_config={}, + agent_config_path="agent_config.json", + task_name="smoke", + runtime_backend="independent", + seed=7, + camera_uids=("vlm_front",), + staging_dir=f"/tmp/ab/{route}/video", + ) + + +class _MemoryAwareFakeWorker: + instances = [] + + def __init__(self, config: _ABWorkerConfig) -> None: + self.config = config + self.closed = False + self.startup_snapshot = { + "robot_qpos": [0.0, 1.0], + "object_poses": {"object": [0.0, 0.0, 0.0]}, + } + self.startup_observation = {"route": config.route} + self.events = [] + self.instances.append(self) + if ( + config.route == "online" + and Path(config.staging_dir).parent.name == config.route + ): + raise RuntimeError("CUDA out of memory") + + def preflight(self, graph): + self.events.append(("preflight", graph)) + return True + + def run(self, graph, **kwargs): + self.events.append(("run", graph, kwargs)) + return SimpleNamespace(success=True) + + def finalize(self, branch_dir: Path, *, episode_index: int): + self.events.append(("finalize", branch_dir, episode_index)) + return [(branch_dir / "video.mp4").as_posix()] + + def close(self): + self.closed = True + + +def test_ab_serializes_workers_after_startup_oom() -> None: + _MemoryAwareFakeWorker.instances = [] + branches, snapshots = _prepare_ab_branches( + {"offline": _worker_config("offline"), "online": _worker_config("online")}, + worker_factory=_MemoryAwareFakeWorker, + prefer_serial=False, + ) + + assert set(branches) == {"offline", "online"} + assert all(isinstance(branch, _SerializedABBranch) for branch in branches.values()) + assert snapshots["offline"] == snapshots["online"] + for route, branch in branches.items(): + assert branch.preflight({"route": route}) is True + branch.run( + {"route": route}, + run_id=f"run-{route}", + episode_index=0, + record_root=Path("/tmp/ab/runtime"), + ) + assert branch.finalize(Path(f"/tmp/ab/{route}"), episode_index=0) == [ + f"/tmp/ab/{route}/video.mp4" + ] + branch.close() + + phases = [ + Path(worker.config.staging_dir).parent.name + for worker in _MemoryAwareFakeWorker.instances + if worker.config.route == "offline" + ] + assert phases == ["offline", "probe", "preflight", "execute"] + + +@pytest.mark.parametrize( + ("status", "success"), + [("succeeded", True), ("failed", False)], +) +def test_task_engine_report_is_mirrored_into_bundle_only_when_enabled( + tmp_path: Path, + status: str, + success: bool, +) -> None: + bundle = tmp_path / "bundle" + bundle.mkdir() + agent_config = bundle / "agent_config.json" + report = ExecutionReport( + task_id="task", + plan_hash="0" * 64, + action_graph_hash="1" * 64, + status=status, + run_id="run", + episode_id="0", + provenance=build_execution_provenance(episode_seed=7), + environments=( + { + "env_id": "0", + "success": success, + "semantic_success": {"task_01": success}, + "action_count": 3, + "retry_count": 0, + "recovery_count": 0, + "revision_count": 0, + "failures": [], + }, + ), + action_count=3, + record_dir=(tmp_path / "runtime-records").as_posix(), + ) + + assert _publish_task_engine_report(agent_config, report, enabled=False) is None + assert not (bundle / "execution_report.json").exists() + + path = _publish_task_engine_report(agent_config, report, enabled=True) + + assert path == bundle / "execution_report.json" + payload = json.loads(path.read_text(encoding="utf-8")) + assert payload["status"] == status + assert payload["record_dir"] == report.record_dir + + +def test_task_engine_exit_code_uses_report_status() -> None: + success = SimpleNamespace(status="succeeded") + failure = SimpleNamespace(status="failed") + + assert _task_engine_exit_code(False, [success]) == 0 + assert _task_engine_exit_code(False, [success, failure]) == 1 + assert _task_engine_exit_code(True, []) == 1 diff --git a/tests/gen_sim/action_engine/evaluation/__init__.py b/tests/gen_sim/action_engine/evaluation/__init__.py new file mode 100644 index 000000000..adfe7f7b4 --- /dev/null +++ b/tests/gen_sim/action_engine/evaluation/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# 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 + +"""Action Engine evaluation tests.""" diff --git a/tests/gen_sim/action_engine/evaluation/test_ab.py b/tests/gen_sim/action_engine/evaluation/test_ab.py new file mode 100644 index 000000000..72baa5e47 --- /dev/null +++ b/tests/gen_sim/action_engine/evaluation/test_ab.py @@ -0,0 +1,445 @@ +# ---------------------------------------------------------------------------- +# 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 copy import deepcopy +import json +from types import SimpleNamespace + +import pytest +import torch + +from embodichain.gen_sim.action_engine.evaluation import run_strict_ab, state_digest +from embodichain.gen_sim.action_engine.evaluation.ab import _graph_difference +from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph + +from ..task_fixtures import make_task_level, make_task_spec + + +class _Env: + def __init__(self, route: str, seed: int, config: dict) -> None: + self.route = route + self.seed = seed + self.config = config + self.closed = False + + def reset(self, *, seed: int) -> None: + self.seed = seed + + def close(self) -> None: + self.closed = True + + +class _Executor: + def __init__(self, graph: dict, env: _Env) -> None: + self.graph = graph + self.env = env + + def run(self, **_kwargs): + return SimpleNamespace( + success=torch.tensor([True]), + actions=[torch.tensor([[0.0]]), torch.tensor([[1.0]])], + retry_count=0, + recovery_count=0, + revision_count=0, + runtime_revisions=[], + record_dir=f"records/{self.env.route}", + ) + + +def _inputs(): + task, requirements = make_task_spec("E1") + bindings = { + item["role_id"]: f"uid_{item['role_id']}" for item in requirements["objects"] + } + offline = instantiate_seed_graph(task, bindings) + online = deepcopy(offline) + online["planner_route"] = "online" + return task, offline, online + + +def test_state_digest_is_mapping_order_stable() -> None: + assert state_digest( + {"qpos": torch.tensor([1.0]), "objects": {"a": [2.0]}} + ) == state_digest({"objects": {"a": [2.0]}, "qpos": torch.tensor([1.0])}) + + +def test_strict_ab_writes_isolated_branches_and_comparison(tmp_path) -> None: + task, offline, online = _inputs() + created = [] + + def env_factory(**kwargs): + env = _Env(**kwargs) + created.append(env) + return env + + result = run_strict_ab( + task, + offline, + online, + env_factory=env_factory, + executor_factory=lambda graph, env: _Executor(graph, env), + snapshot_reader=lambda env: { + "robot_qpos": torch.tensor([float(env.seed)]), + "object_poses": {"object": torch.eye(4)}, + }, + output_dir=tmp_path, + seed=123, + shared_config={"robot": "same"}, + planning_metrics={ + "offline": {"planning_seconds": 0.1, "vlm_call_count": 0}, + "online": {"planning_seconds": 0.2, "vlm_call_count": 2}, + }, + ) + + assert result.comparison_path.is_file() + assert (result.offline_dir / "seed_task_graph.json").is_file() + assert (result.online_dir / "seed_task_graph.json").is_file() + assert result.comparison["initial_state_digest"] == result.initial_state_digest + assert result.comparison["branches"]["offline"]["planning_seconds"] == 0.1 + assert result.comparison["branches"]["online"]["vlm_call_count"] == 2 + assert all(env.closed for env in created) + + +def test_graph_difference_reports_changed_task_group_fields() -> None: + _task_spec, offline, online = _inputs() + online["task_groups"][0]["goal"] = deepcopy(online["task_groups"][0]["goal"]) + online["task_groups"][0]["goal"]["relation"] = "right_of" + + difference = _graph_difference(offline, online) + + assert difference["changed_task_groups"] == [ + {"id": offline["task_groups"][0]["id"], "changed_fields": ["goal"]} + ] + assert ( + difference["task_group_difference"]["changed_groups"] + == difference["changed_task_groups"] + ) + + +def test_strict_ab_finalizes_two_branch_videos_and_revision_files(tmp_path) -> None: + task, offline, online = _inputs() + finalized = [] + + def finalizer(**kwargs): + route = kwargs["route"] + path = kwargs["branch_dir"] / "video.mp4" + path.write_bytes(route.encode("ascii")) + finalized.append(route) + return [path.as_posix()] + + result = run_strict_ab( + task, + offline, + online, + env_factory=lambda **kwargs: _Env(**kwargs), + executor_factory=lambda graph, env: _Executor(graph, env), + snapshot_reader=lambda env: { + "robot_qpos": torch.tensor([float(env.seed)]), + "object_poses": {"object": torch.eye(4)}, + }, + branch_finalizer=finalizer, + output_dir=tmp_path, + seed=11, + ) + + assert finalized == ["offline", "online"] + for route, branch_dir in ( + ("offline", result.offline_dir), + ("online", result.online_dir), + ): + assert (branch_dir / "video.mp4").read_bytes() == route.encode("ascii") + assert (branch_dir / "runtime_revisions.json").is_file() + assert result.comparison["branches"][route]["video_paths"] == [ + (branch_dir / "video.mp4").as_posix() + ] + + +def test_strict_ab_aborts_before_execution_on_state_mismatch(tmp_path) -> None: + task, offline, online = _inputs() + executions = [] + + with pytest.raises(RuntimeError, match="initial state mismatch"): + run_strict_ab( + task, + offline, + online, + env_factory=lambda **kwargs: _Env(**kwargs), + executor_factory=lambda graph, env: executions.append((graph, env)), + snapshot_reader=lambda env: { + "robot_qpos": torch.tensor([float(env.seed)]), + "object_poses": {env.route: torch.eye(4)}, + }, + output_dir=tmp_path, + seed=5, + ) + assert executions == [] + + +def test_strict_ab_rejects_incomplete_state_snapshot(tmp_path) -> None: + task, offline, online = _inputs() + + with pytest.raises(ValueError, match="missing required state"): + run_strict_ab( + task, + offline, + online, + env_factory=lambda **kwargs: _Env(**kwargs), + executor_factory=lambda graph, env: _Executor(graph, env), + snapshot_reader=lambda env: {"robot_qpos": torch.tensor([0.0])}, + output_dir=tmp_path, + seed=5, + ) + + +def test_strict_ab_requires_articulation_and_camera_digest_components(tmp_path) -> None: + task, offline, online = _inputs() + + with pytest.raises(ValueError, match="articulation_state.*camera_calibration"): + run_strict_ab( + task, + offline, + online, + env_factory=lambda **kwargs: _Env(**kwargs), + executor_factory=lambda graph, env: _Executor(graph, env), + snapshot_reader=lambda env: { + "robot_qpos": torch.tensor([float(env.seed)]), + "object_poses": {"object": torch.eye(4)}, + }, + output_dir=tmp_path, + seed=5, + strict_state_digest=True, + ) + + +def test_strict_ab_reuses_prepared_identical_resets(tmp_path) -> None: + task, offline, online = _inputs() + environments = { + route: _Env(route=route, seed=17, config={}) for route in ("offline", "online") + } + snapshots = { + route: { + "robot_qpos": torch.tensor([17.0]), + "object_poses": {"object": torch.eye(4)}, + } + for route in environments + } + + result = run_strict_ab( + task, + offline, + online, + executor_factory=lambda graph, env: _Executor(graph, env), + snapshot_reader=lambda _env: pytest.fail("prepared snapshots must be reused"), + output_dir=tmp_path, + seed=17, + prepared_environments=environments, + prepared_snapshots=snapshots, + ) + + assert result.initial_state_digest == state_digest(snapshots["offline"]) + assert all(env.closed for env in environments.values()) + + +def test_strict_ab_stops_both_branches_when_global_preflight_fails(tmp_path) -> None: + task, offline, online = _inputs() + runs: list[str] = [] + + class PreflightExecutor(_Executor): + def preflight(self) -> bool: + if self.env.route == "online": + raise ValueError("online capability unavailable") + return True + + def run(self, **kwargs): + runs.append(self.env.route) + return super().run(**kwargs) + + with pytest.raises(RuntimeError, match="no branch was allowed to move"): + run_strict_ab( + task, + offline, + online, + env_factory=lambda **kwargs: _Env(**kwargs), + executor_factory=lambda graph, env: PreflightExecutor(graph, env), + snapshot_reader=lambda env: { + "robot_qpos": torch.tensor([float(env.seed)]), + "object_poses": {"object": torch.eye(4)}, + }, + output_dir=tmp_path, + seed=17, + ) + assert runs == [] + + +def test_strict_ab_keeps_other_branch_running_after_execution_failure(tmp_path) -> None: + task, offline, online = _inputs() + runs: list[str] = [] + + class IsolatedExecutor(_Executor): + def preflight(self) -> bool: + return True + + def run(self, **kwargs): + runs.append(self.env.route) + if self.env.route == "offline": + raise RuntimeError("offline execution failed") + return super().run(**kwargs) + + result = run_strict_ab( + task, + offline, + online, + env_factory=lambda **kwargs: _Env(**kwargs), + executor_factory=lambda graph, env: IsolatedExecutor(graph, env), + snapshot_reader=lambda env: { + "robot_qpos": torch.tensor([float(env.seed)]), + "object_poses": {"object": torch.eye(4)}, + }, + output_dir=tmp_path, + seed=17, + ) + + assert runs == ["offline", "online"] + assert result.comparison["branches"]["offline"]["success_rate"] == 0.0 + assert result.comparison["branches"]["online"]["success_rate"] == 1.0 + + +def test_strict_ab_surfaces_video_finalizer_failure_and_closes(tmp_path) -> None: + task, offline, online = _inputs() + environments = [] + + def factory(**kwargs): + environment = _Env(**kwargs) + environments.append(environment) + return environment + + def finalizer(**kwargs): + if kwargs["route"] == "offline": + raise OSError("recorder did not produce a file") + return [] + + with pytest.raises(RuntimeError, match="branch video finalization failed"): + run_strict_ab( + task, + offline, + online, + env_factory=factory, + executor_factory=lambda graph, env: _Executor(graph, env), + snapshot_reader=lambda env: { + "robot_qpos": torch.tensor([float(env.seed)]), + "object_poses": {"object": torch.eye(4)}, + }, + branch_finalizer=finalizer, + output_dir=tmp_path, + seed=19, + ) + + assert len(environments) == 2 + assert all(environment.closed for environment in environments) + comparison = json.loads((tmp_path / "comparison.json").read_text()) + assert set(comparison["video_finalization_errors"]) == {"offline"} + + +def test_strict_ab_require_branch_videos_checks_normalized_artifact(tmp_path) -> None: + task, offline, online = _inputs() + + with pytest.raises(RuntimeError, match="video.mp4"): + run_strict_ab( + task, + offline, + online, + env_factory=lambda **kwargs: _Env(**kwargs), + executor_factory=lambda graph, env: _Executor(graph, env), + snapshot_reader=lambda env: { + "robot_qpos": torch.tensor([float(env.seed)]), + "object_poses": {"object": torch.eye(4)}, + }, + branch_finalizer=lambda **_kwargs: [], + output_dir=tmp_path, + seed=19, + require_branch_videos=True, + ) + + +def test_strict_ab_closes_prepared_environments_on_graph_validation_error( + tmp_path, +) -> None: + task, offline, online = _inputs() + environments = { + route: _Env(route=route, seed=23, config={}) for route in ("offline", "online") + } + invalid_online = deepcopy(online) + invalid_online["planner_route"] = "offline" + + with pytest.raises(ValueError, match="explicit offline and online"): + run_strict_ab( + task, + offline, + invalid_online, + executor_factory=lambda graph, env: _Executor(graph, env), + snapshot_reader=lambda _env: pytest.fail("validation must happen first"), + output_dir=tmp_path, + seed=23, + prepared_environments=environments, + ) + + assert all(environment.closed for environment in environments.values()) + + +def test_strict_l4_ab_requires_and_records_private_oracle(tmp_path) -> None: + task, requirements = make_task_level("L4") + bindings = { + item["role_id"]: f"uid_{item['role_id']}" for item in requirements["objects"] + } + offline = instantiate_seed_graph(task, bindings) + online = deepcopy(offline) + online["planner_route"] = "online" + + with pytest.raises(ValueError, match="private-oracle"): + run_strict_ab( + task, + offline, + online, + env_factory=lambda **kwargs: _Env(**kwargs), + executor_factory=lambda graph, env: _Executor(graph, env), + snapshot_reader=lambda env: { + "robot_qpos": torch.tensor([float(env.seed)]), + "object_poses": {"object": torch.eye(4)}, + }, + output_dir=tmp_path, + seed=7, + ) + + result = run_strict_ab( + task, + offline, + online, + env_factory=lambda **kwargs: _Env(**kwargs), + executor_factory=lambda graph, env: _Executor(graph, env), + snapshot_reader=lambda env: { + "robot_qpos": torch.tensor([float(env.seed)]), + "object_poses": {"object": torch.eye(4)}, + }, + success_evaluator=lambda **_kwargs: torch.tensor([True]), + output_dir=tmp_path, + seed=7, + ) + assert all( + branch["success_source"] == "private_oracle" + for branch in result.comparison["branches"].values() + ) diff --git a/tests/gen_sim/action_engine/runtime/test_recovery_v2.py b/tests/gen_sim/action_engine/runtime/test_recovery_v2.py new file mode 100644 index 000000000..5544aea45 --- /dev/null +++ b/tests/gen_sim/action_engine/runtime/test_recovery_v2.py @@ -0,0 +1,720 @@ +# ---------------------------------------------------------------------------- +# 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 copy import deepcopy +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +import embodichain.gen_sim.action_engine.runtime.executor as executor_module +from embodichain.gen_sim.action_engine.capabilities import ( + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.runtime import ( + DynamicRecoveryController, + ProgramExecutor, + RuntimeGraph, + build_upright_recovery, + classify_failure, + load_execution_program, +) +from embodichain.gen_sim.action_engine.runtime.executor import _EdgeResult +from embodichain.gen_sim.action_engine.runtime.grounding import ActionGrounder +from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph + +from ..task_fixtures import make_task_spec + + +def _graph(task_type: str) -> dict: + task, requirements = make_task_spec(task_type) + bindings = { + item["role_id"]: f"uid_{item['role_id']}" for item in requirements["objects"] + } + return instantiate_seed_graph(task, bindings) + + +def _handover_then_place_graph() -> dict: + task = { + "schema_version": "action_engine_task_spec_v2", + "task_id": "handover_then_place_recovery", + "level": "L3", + "instruction": "Hand the yellow can from the left arm to the right arm.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "task_01", + "task_type": "E4", + "params": { + "object_role": "yellow_can", + "transfer_arm": "left_arm", + "receive_arm": "right_arm", + }, + "depends_on": [], + "role": "primary", + }, + { + "id": "task_02", + "task_type": "E1", + "params": { + "object_role": "yellow_can", + "target_role": "purple_can", + "relation": "right_of", + }, + "depends_on": ["task_01"], + "role": "primary", + }, + ], + "success": { + "op": "all", + "terms": [ + {"type": "handover_complete", "task_instance_id": "task_01"}, + {"type": "semantic_goal", "task_instance_id": "task_02"}, + ], + }, + "oracle": {"task_order": ["task_01", "task_02"]}, + "metadata": {}, + } + return instantiate_seed_graph( + task, + { + "yellow_can": "interact_yellow_can", + "purple_can": "interact_purple_can", + }, + ) + + +class _RecoveryRecorder: + def __init__(self) -> None: + self.recovery_events: list[dict[str, Any]] = [] + self.edge_events: list[dict[str, Any]] = [] + + def recovery(self, **event: Any) -> None: + self.recovery_events.append(event) + + def edge(self, edge_id: str, step: Any, **event: Any) -> None: + self.edge_events.append({"edge_id": edge_id, "step_id": step.id, **event}) + + def step(self, *_args: Any, **_kwargs: Any) -> None: + return None + + def register_step(self, *_args: Any, **_kwargs: Any) -> None: + return None + + +def _local_recovery_harness( + graph: dict[str, Any], + *, + num_envs: int, + max_transitions: int = 100, + max_revisions: int = 8, +) -> tuple[ProgramExecutor, Any, Any, list[tuple[str, str, list[bool]]]]: + program = load_execution_program(graph, require_executable=True) + step = next(item for item in program.semantic_steps if item.id == "task_01") + failed_node = next( + node + for node in graph["nodes"] + if node["task_instance_id"] == step.id and node["atomic_action"] == "HandOver" + ) + edge = next( + item + for item in program.edges + if item.actions[0].get("seed_node_id") == failed_node["id"] + ) + executor = object.__new__(ProgramExecutor) + executor.runtime_graph = RuntimeGraph( + graph, + num_envs=num_envs, + max_revisions=max_revisions, + ) + executor.env = SimpleNamespace( + num_envs=num_envs, + device=torch.device("cpu"), + robot=SimpleNamespace(get_qpos=lambda: torch.zeros((num_envs, 4))), + ) + executor.capability_registry = None + executor.steps = {item.id: item for item in program.semantic_steps} + executor.edges = {item.id: item for item in program.edges} + executor.step_by_edge = { + edge_id: item for item in program.semantic_steps for edge_id in item.edge_ids + } + executor._assignments = {step.id: ["left_arm"] * num_envs} + executor._candidate_cache = {} + executor._candidate_failures = {} + executor._candidate_diagnostics = {} + executor._object_states = {} + executor._step_states = {} + executor._object_owners = {} + executor._arm_owners = { + "left_arm": [None] * num_envs, + "right_arm": [None] * num_envs, + } + executor._targets = {} + executor.record_runtime = False + executor.max_transitions = max_transitions + executor._transition_count = 0 + executor.retry_count = 0 + call_log: list[tuple[str, str, list[bool]]] = [] + + def execute_edge(current_edge: Any, current_step: Any, *, failed: torch.Tensor): + call_log.append((current_step.id, current_edge.id, failed.tolist())) + return _EdgeResult([], failed.clone(), []) + + def ensure_assignment(current_step: Any, failed: torch.Tensor) -> None: + actor = current_step.actor + assignment = ( + str(actor["arm"]) if actor.get("mode") == "required" else "right_arm" + ) + executor._assignments[current_step.id] = [ + None if bool(failed[index]) else assignment for index in range(num_envs) + ] + + executor._execute_edge_with_retries = execute_edge + executor._ensure_assignment = ensure_assignment + executor._clear_recovery_rows = lambda *_args, **_kwargs: None + executor._verify_step = lambda _step, failed: ( + failed.clone(), + ~failed, + torch.zeros((num_envs, 3)), + ) + return executor, step, edge, call_log + + +def test_runtime_graph_retries_twice_then_requests_recovery() -> None: + graph = _graph("E4") + runtime = RuntimeGraph(graph, num_envs=2, max_retries=2) + handover = next( + node for node in graph["nodes"] if node["atomic_action"] == "HandOver" + ) + failed = torch.tensor([True, False]) + holds = torch.tensor([True, True]) + + first = runtime.record_failure(handover["id"], failed, precondition_holds=holds) + second = runtime.record_failure(handover["id"], failed, precondition_holds=holds) + third = runtime.record_failure(handover["id"], failed, precondition_holds=holds) + + assert first.retry.tolist() == [True, False] + assert second.retry.tolist() == [True, False] + assert third.recover.tolist() == [True, False] + assert runtime.seed_graph == graph + + +def test_recovery_insertion_revises_runtime_graph_not_seed_graph() -> None: + graph = _graph("E4") + runtime = RuntimeGraph(graph, num_envs=1) + failed_node = next( + node for node in graph["nodes"] if node["atomic_action"] == "HandOver" + ) + recovery_source = _graph("E2") + source_group = recovery_source["task_groups"][0] + recovery_group_id = "recovery_upright_01" + recovery_nodes = [] + id_map = { + node["id"]: f"recovery_{index:02d}" + for index, node in enumerate(recovery_source["nodes"], start=1) + } + for node in recovery_source["nodes"]: + item = deepcopy(node) + item["id"] = id_map[node["id"]] + item["object_uid"] = failed_node["object_uid"] + item["target_binding"] = deepcopy(item["target_binding"]) + if item["target_binding"].get("kind") == "object": + item["target_binding"]["object"] = failed_node["object_uid"] + item["depends_on"] = [id_map.get(dep, dep) for dep in node["depends_on"]] + recovery_nodes.append(item) + recovery_group = deepcopy(source_group) + recovery_group.update( + { + "id": recovery_group_id, + "role": "recovery", + "object_uid": failed_node["object_uid"], + "node_ids": [node["id"] for node in recovery_nodes], + "depends_on": [], + "parent_task_instance_id": failed_node["task_instance_id"], + } + ) + recovery_group["success"] = { + "type": "object_upright", + "object": failed_node["object_uid"], + } + + patched = runtime.insert_recovery_subgraph( + failed_node_id=failed_node["id"], + recovery_nodes=recovery_nodes, + recovery_group=recovery_group, + failure_type="object_fallen", + ) + + assert graph == runtime.seed_graph + assert any(group["id"] == recovery_group_id for group in patched["task_groups"]) + assert not any( + node["task_instance_id"] == failed_node["task_instance_id"] + and node["target_binding"].get("source") == "handover" + for node in patched["nodes"] + ) + assert runtime.revisions[0].kind == "insert_recovery" + assert ( + classify_failure("PickUp", planning_succeeded=True, held_after=False) + == "grasp_missed" + ) + + +def test_recovery_rejects_downstream_contract_that_requires_actor_switch() -> None: + graph = _handover_then_place_graph() + runtime = RuntimeGraph(graph, num_envs=1) + handover = next( + node for node in graph["nodes"] if node["atomic_action"] == "HandOver" + ) + cleanup_ids = { + node["id"] + for node in graph["nodes"] + if node["task_instance_id"] == handover["task_instance_id"] + and node["role"] == "cleanup" + } + assert cleanup_ids + + with pytest.raises(ValueError, match="without changing.*actor"): + runtime.insert_default_recovery( + failed_node_id=handover["id"], + failure_type="object_fallen", + ) + + assert runtime.graph == graph + assert runtime.revisions == [] + + +def test_failed_group_resume_recovery_places_before_prefix_replay() -> None: + graph = _handover_then_place_graph() + runtime = RuntimeGraph(graph, num_envs=1) + handover = next( + node for node in graph["nodes"] if node["atomic_action"] == "HandOver" + ) + + patched = runtime.insert_default_recovery( + failed_node_id=handover["id"], + failure_type="object_fallen", + resume_failed_group=True, + ) + + group_id = runtime.revisions[-1].inserted_group_ids[0] + group = next(item for item in patched["task_groups"] if item["id"] == group_id) + nodes = [node for node in patched["nodes"] if node["id"] in group["node_ids"]] + original_cleanup = { + node["id"] + for node in graph["nodes"] + if node["task_instance_id"] == handover["task_instance_id"] + and node["role"] == "cleanup" + } + assert group["goal"]["terminal_behavior"] == "place" + assert [node["atomic_action"] for node in nodes] == [ + "PickUp", + "MoveHeldObject", + "Place", + "MoveEndEffector", + "MoveJoints", + ] + assert original_cleanup <= {node["id"] for node in patched["nodes"]} + + +@pytest.mark.parametrize( + "actor", + ( + {"mode": "required", "arm": "left_arm"}, + {"mode": "required", "arm": "right_arm"}, + {"mode": "auto"}, + ), +) +def test_upright_recovery_inherits_failed_group_actor(actor: dict[str, Any]) -> None: + graph = _graph("E2") + failed_group = graph["task_groups"][0] + failed_group["actor"] = deepcopy(actor) + for node in graph["nodes"]: + if node["task_instance_id"] == failed_group["id"]: + node["actor"] = deepcopy(actor) + + nodes, recovery_group = build_upright_recovery( + graph, + failed_node_id=failed_group["node_ids"][0], + revision=1, + resume_failed_group=True, + ) + + assert recovery_group["actor"] == actor + assert all(node["actor"] == actor for node in nodes) + + +def test_local_recovery_replays_failed_group_prefix_and_preserves_seed_graph( + monkeypatch: pytest.MonkeyPatch, +) -> None: + graph = _handover_then_place_graph() + original = deepcopy(graph) + executor, step, edge, calls = _local_recovery_harness(graph, num_envs=1) + recorder = _RecoveryRecorder() + monkeypatch.setattr( + executor_module, + "evaluate_predicate", + lambda *_args, **_kwargs: torch.tensor([False]), + ) + + result = executor._recover_object_fallen( + edge, + step, + _EdgeResult( + [], + torch.tensor([True]), + [], + executed=torch.tensor([True]), + ), + inherited_failed=torch.tensor([False]), + fallen_transition=torch.tensor([True]), + recorder=recorder, + ) + + replayed = [edge_id for step_id, edge_id, _failed in calls if step_id == step.id] + expected_prefix = list(step.edge_ids[: step.edge_ids.index(edge.id) + 1]) + assert result.failed.tolist() == [False] + assert replayed == expected_prefix + assert executor.runtime_graph.seed_graph == original + assert graph == original + assert [event["status"] for event in recorder.recovery_events] == [ + "started", + "succeeded", + ] + recovery_edges = [ + event + for event in recorder.edge_events + if event["step_id"].startswith("recovery_e2_") + ] + replay_edges = [ + event for event in recorder.edge_events if event["step_id"] == step.id + ] + assert recovery_edges + assert all(event["phase"] == "recovery" for event in recovery_edges) + assert replay_edges + assert all(event["phase"] == "replay" for event in replay_edges) + + +def test_local_recovery_only_executes_and_rebinds_failed_vector_row( + monkeypatch: pytest.MonkeyPatch, +) -> None: + graph = _handover_then_place_graph() + executor, step, edge, calls = _local_recovery_harness(graph, num_envs=2) + recorder = _RecoveryRecorder() + monkeypatch.setattr( + executor_module, + "evaluate_predicate", + lambda *_args, **_kwargs: torch.tensor([True, False]), + ) + + result = executor._recover_object_fallen( + edge, + step, + _EdgeResult( + [], + torch.tensor([False, True]), + [], + executed=torch.tensor([False, True]), + ), + inherited_failed=torch.tensor([False, False]), + fallen_transition=torch.tensor([False, True]), + recorder=recorder, + ) + + assert result.failed.tolist() == [False, False] + assert all(failed == [True, False] for _step_id, _edge_id, failed in calls) + assert executor._assignments[step.id] == ["left_arm", "left_arm"] + assert executor.runtime_graph.revisions[-1].active_env_ids == (1,) + assert all( + event["active"].tolist() == [False, True] for event in recorder.recovery_events + ) + + +def test_local_recovery_failure_does_not_replay_prefix( + monkeypatch: pytest.MonkeyPatch, +) -> None: + graph = _handover_then_place_graph() + executor, step, edge, calls = _local_recovery_harness(graph, num_envs=1) + executor._verify_step = lambda _step, failed: ( + torch.ones_like(failed), + torch.zeros_like(failed), + torch.zeros((1, 3)), + ) + recorder = _RecoveryRecorder() + monkeypatch.setattr( + executor_module, + "evaluate_predicate", + lambda *_args, **_kwargs: torch.tensor([False]), + ) + + result = executor._recover_object_fallen( + edge, + step, + _EdgeResult( + [], + torch.tensor([True]), + [], + executed=torch.tensor([True]), + ), + inherited_failed=torch.tensor([False]), + fallen_transition=torch.tensor([True]), + recorder=recorder, + ) + + assert result.failed.tolist() == [True] + assert not any(step_id == step.id for step_id, _edge_id, _failed in calls) + assert recorder.recovery_events[-1]["status"] == "failed" + + +def test_local_recovery_budget_exhaustion_terminates_with_original_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + graph = _handover_then_place_graph() + executor, step, edge, calls = _local_recovery_harness( + graph, + num_envs=1, + max_transitions=0, + ) + recorder = _RecoveryRecorder() + monkeypatch.setattr( + executor_module, + "evaluate_predicate", + lambda *_args, **_kwargs: torch.tensor([False]), + ) + + result = executor._recover_object_fallen( + edge, + step, + _EdgeResult( + [], + torch.tensor([True]), + [], + executed=torch.tensor([True]), + ), + inherited_failed=torch.tensor([False]), + fallen_transition=torch.tensor([True]), + recorder=recorder, + ) + + assert result.failed.tolist() == [True] + assert calls == [] + assert recorder.recovery_events[-1]["status"] == "failed" + assert "max_transitions" in recorder.recovery_events[-1]["error"] + + +def test_non_fallen_failure_does_not_create_recovery_revision( + monkeypatch: pytest.MonkeyPatch, +) -> None: + graph = _handover_then_place_graph() + executor, step, edge, calls = _local_recovery_harness(graph, num_envs=1) + recorder = _RecoveryRecorder() + monkeypatch.setattr( + executor_module, + "evaluate_predicate", + lambda *_args, **_kwargs: torch.tensor([True]), + ) + + result = executor._recover_object_fallen( + edge, + step, + _EdgeResult( + [], + torch.tensor([True]), + [], + executed=torch.tensor([True]), + ), + inherited_failed=torch.tensor([False]), + fallen_transition=torch.tensor([False]), + recorder=recorder, + ) + + assert result.failed.tolist() == [True] + assert calls == [] + assert executor.runtime_graph.revisions == [] + assert recorder.recovery_events == [] + + +def test_initially_fallen_planning_failure_does_not_trigger_recovery() -> None: + graph = _handover_then_place_graph() + executor, step, edge, calls = _local_recovery_harness(graph, num_envs=1) + recorder = _RecoveryRecorder() + + result = executor._recover_object_fallen( + edge, + step, + _EdgeResult( + [], + torch.tensor([True]), + [], + executed=torch.tensor([False]), + ), + inherited_failed=torch.tensor([False]), + fallen_transition=torch.tensor([False]), + recorder=recorder, + ) + + assert result.failed.tolist() == [True] + assert calls == [] + assert executor.runtime_graph.revisions == [] + assert recorder.recovery_events == [] + + +def test_failure_provenance_distinguishes_planning_from_execution_caused_fall() -> None: + graph = _handover_then_place_graph() + executor, step, edge, _ = _local_recovery_harness(graph, num_envs=1) + executor.adapter = SimpleNamespace(capabilities=build_atomic_capability_registry()) + failed = torch.tensor([True]) + + planning = executor._failure_events( + edge, + step, + failed, + postcondition=False, + executed=torch.tensor([False]), + fallen_transition=torch.tensor([False]), + ) + execution = executor._failure_events( + edge, + step, + failed, + postcondition=False, + executed=torch.tensor([True]), + fallen_transition=torch.tensor([True]), + ) + + assert [event["failure_type"] for event in planning] == ["search_exhausted"] + assert [event["failure_type"] for event in execution] == ["object_fallen"] + + +def test_offline_and_online_dynamic_replanners_are_route_isolated() -> None: + for mode in ("offline_dynamic", "online_dynamic"): + graph = _graph("E4") + runtime = RuntimeGraph(graph, num_envs=1) + calls = [] + + def replanner(**kwargs): + calls.append((mode, kwargs["failure_type"])) + return kwargs["graph"] + + controller = DynamicRecoveryController( + runtime, + mode=mode, + offline_replanner=replanner if mode == "offline_dynamic" else None, + online_replanner=replanner if mode == "online_dynamic" else None, + ) + failed_node = next( + node for node in graph["nodes"] if node["atomic_action"] == "HandOver" + ) + directive = controller.handle_failure( + failed_node_id=failed_node["id"], + failure_type="postcondition_failed", + ) + completed = [group["id"] for group in graph["task_groups"]] + controller.replan( + directive, + completed_group_ids=completed, + recovery_succeeded=False, + ) + + assert calls == [(mode, "postcondition_failed")] + assert runtime.revisions[-1].kind == "replan_suffix" + + +def test_dynamic_recovery_consumes_per_environment_failure_events() -> None: + graph = _graph("E4") + runtime = RuntimeGraph(graph, num_envs=2) + controller = DynamicRecoveryController( + runtime, + mode="offline_dynamic", + offline_replanner=lambda **kwargs: kwargs["graph"], + ) + failed_node = next( + node for node in graph["nodes"] if node["atomic_action"] == "HandOver" + ) + result = SimpleNamespace( + failure_events=[ + { + "node_id": failed_node["id"], + "failure_type": "object_fallen", + "env_ids": [1], + } + ] + ) + + directive = controller.handle_execution_result(result) + + assert directive.active_env_ids == (1,) + assert runtime.revisions[-1].active_env_ids == (1,) + + +def test_runtime_graph_stops_at_revision_and_recovery_budgets() -> None: + graph = _graph("E4") + failed_node = next( + node for node in graph["nodes"] if node["atomic_action"] == "HandOver" + ) + + with pytest.raises(RuntimeError, match="revision budget"): + RuntimeGraph(graph, num_envs=1, max_revisions=0).insert_default_recovery( + failed_node_id=failed_node["id"], + failure_type="object_fallen", + ) + with pytest.raises(RuntimeError, match="recovery-action budget"): + RuntimeGraph(graph, num_envs=1, max_recovery_actions=0).insert_default_recovery( + failed_node_id=failed_node["id"], + failure_type="object_fallen", + ) + + +def test_visual_constraint_grounding_reads_fresh_camera_depth() -> None: + class Sensor: + def __init__(self) -> None: + self.depth = torch.ones((1, 4, 4, 1)) + + def get_data(self): + return {"depth": self.depth} + + def get_intrinsics(self): + return torch.tensor([[[2.0, 0.0, 1.5], [0.0, 2.0, 1.5], [0.0, 0.0, 1.0]]]) + + def get_arena_pose(self, *, to_matrix: bool): + assert to_matrix + return torch.eye(4).unsqueeze(0) + + sensor = Sensor() + env = SimpleNamespace( + num_envs=1, + device=torch.device("cpu"), + sim=SimpleNamespace(get_sensor=lambda uid: sensor if uid == "front" else None), + get_current_xpos_agent=lambda: ( + torch.eye(4).unsqueeze(0), + torch.eye(4).unsqueeze(0), + ), + ) + grounder = object.__new__(ActionGrounder) + grounder.env = env + binding = {"camera_uid": "front", "normalized_keypoint": [0.5, 0.5]} + + first = grounder._visual_target(binding, "left_arm") + sensor.depth.fill_(2.0) + second = grounder._visual_target( + {"camera_uid": "front", "normalized_bbox": [0.4, 0.4, 0.6, 0.6]}, + "left_arm", + ) + + assert first[0, 2, 3].item() == 1.0 + assert second[0, 2, 3].item() == 2.0 diff --git a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py new file mode 100644 index 000000000..61ab63fcd --- /dev/null +++ b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py @@ -0,0 +1,5871 @@ +# ---------------------------------------------------------------------------- +# 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 copy import deepcopy +from dataclasses import replace +import hashlib +import json +from pathlib import Path +import sys +from types import ModuleType, SimpleNamespace +from typing import Any + +import numpy as np +import pytest +import torch + +from embodichain.gen_sim.action_engine.config import ( + RuntimePolicyCfg, + default_runtime_policy, + resolve_agent_runtime_policy, + runtime_policy_hash, +) +from embodichain.gen_sim.action_engine.compiler import ( + compile_task_agent, + compile_task_agent_v2, +) +from embodichain.gen_sim.action_engine.cli.run_agent import ( + build_parser as build_run_parser, +) +from embodichain.gen_sim.action_engine.domain import ( + TASK_AGENT_SCHEMA, + execution_program_hash, + motion_policy, + validate_execution_program, +) +from embodichain.gen_sim.action_engine.environment import agent_env as env_module +from embodichain.gen_sim.action_engine.runtime.actions import AtomicActionAdapter +from embodichain.gen_sim.action_engine.runtime.executor import ( + ProgramExecutor, + _EdgeResult, + _score_arm_candidate, +) +from embodichain.gen_sim.action_engine.runtime.frames import ( + relation_offset, + robot_frame_axes, +) +from embodichain.gen_sim.action_engine.runtime.grounding import ActionGrounder +from embodichain.gen_sim.action_engine.runtime.loader import ( + load_agent_execution_program, + load_execution_program as _load_execution_program, +) +from embodichain.gen_sim.action_engine.runtime.models import ( + ActionOutcome, + ExecutionEdge, + ExecutionProgram, + ExecutionResult, + GroundedAction, + SemanticStep, +) +from embodichain.gen_sim.action_engine.runtime.motion_policy import ( + resolve_motion_policy, +) +from embodichain.gen_sim.action_engine.runtime.predicates import evaluate_predicate +from embodichain.gen_sim.action_engine.runtime.recording import RuntimeRecorder +from embodichain.gen_sim.action_engine.runtime.state import ExecutionState +from embodichain.gen_sim.action_engine.runtime import solver_compat +from embodichain.gen_sim.action_engine.protocol import ( + SEED_GRAPH_SCHEMA, + TASK_SPEC_SCHEMA, +) +from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph +from embodichain.lab.gym.envs import EmbodiedEnv +from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + Affordance, + AntipodalAffordance, + CoordinatedPickGoal, + CoordinatedPlacementGoal, + CoordinatedPlacementOptions, + HandOverOptions, + HeldObjectPoseGoal, + HeldObjectState, + ObjectSemantics, + PickUpOptions, + PressAffordance, + PressGoal, + PressOptions, +) + +from ..task_fixtures import make_task_spec +from embodichain.lab.sim.solvers import URSolverCfg + + +def _task_agent(*steps: dict[str, Any]) -> dict[str, Any]: + return { + "schema_version": TASK_AGENT_SCHEMA, + "task": "runtime_contract", + "goal": "Exercise the deterministic runtime contract.", + "semantic_steps": list(steps), + } + + +def load_execution_program(source: Any, **kwargs: Any) -> ExecutionProgram: + """Adapt legacy compiler fixtures without weakening the production loader.""" + if isinstance(source, dict) and source.get("schema_version") != SEED_GRAPH_SCHEMA: + return ExecutionProgram.from_mapping(validate_execution_program(source)) + return _load_execution_program(source, **kwargs) + + +def _hold_step(step_id: str, object_uid: str, arm: str) -> dict[str, Any]: + return { + "id": step_id, + "operator": "hold_hover", + "object": object_uid, + "actor": {"mode": "required", "arm": arm}, + "goal": {}, + "depends_on": [], + } + + +class _FakeEntity: + def __init__( + self, + uid: str, + pose: torch.Tensor, + vertices: torch.Tensor, + ) -> None: + self.uid = uid + self._pose = pose + self._vertices = vertices + self._triangles = torch.tensor( + [[0, 1, 2], [0, 2, 3]], + dtype=torch.int64, + ) + self.lin_vel = torch.zeros(pose.shape[0], 3) + self.ang_vel = torch.zeros(pose.shape[0], 3) + + def get_local_pose(self, *, to_matrix: bool) -> torch.Tensor: + assert to_matrix + return self._pose.clone() + + def get_vertices( + self, + *, + env_ids: list[int], + scale: bool, + ) -> torch.Tensor: + del env_ids, scale + return self._vertices.clone() + + def get_triangles(self, *, env_ids: list[int]) -> torch.Tensor: + del env_ids + return self._triangles.clone() + + +class _FakeSim: + def __init__(self, entities: dict[str, _FakeEntity]) -> None: + self.entities = entities + + def get_rigid_object(self, uid: str) -> _FakeEntity | None: + return self.entities.get(uid) + + def get_rigid_object_uid_list(self) -> list[str]: + return list(self.entities) + + def update(self, *, step: int) -> None: + del step + + +class _FakeRobot: + def __init__(self, num_envs: int = 1) -> None: + self.uid = "fake_robot" + self.dof = 8 + self._qpos = torch.zeros(num_envs, self.dof) + self.control_parts = { + "physical_left_arm": ["l0", "l1"], + "physical_left_eef": ["lh0", "lh1"], + "physical_right_arm": ["r0", "r1"], + "physical_right_eef": ["rh0", "rh1"], + "dual_arm": ["l0", "l1", "r0", "r1"], + } + self._ids = { + "physical_left_arm": [0, 1], + "physical_left_eef": [2, 3], + "physical_right_arm": [4, 5], + "physical_right_eef": [6, 7], + "dual_arm": [0, 1, 4, 5], + } + + def get_qpos(self) -> torch.Tensor: + return self._qpos.clone() + + def get_joint_ids(self, *, name: str) -> list[int]: + return list(self._ids[name]) + + def get_control_part_base_pose(self, *, name: str, to_matrix: bool) -> torch.Tensor: + del name + assert to_matrix + return torch.eye(4).repeat(self._qpos.shape[0], 1, 1) + + def get_solver(self, *, name: str) -> SimpleNamespace: + return SimpleNamespace(root_link_name=name.replace("_arm", "_base")) + + def get_link_pose(self, *, link_name: str, to_matrix: bool) -> torch.Tensor: + assert to_matrix + pose = torch.eye(4).repeat(self._qpos.shape[0], 1, 1) + pose[:, 1, 3] = -0.3 if link_name == "physical_left_base" else 0.3 + return pose + + def compute_fk( + self, + qpos: torch.Tensor, + *, + name: str, + to_matrix: bool, + ) -> torch.Tensor: + del name + assert to_matrix + return torch.eye(4).repeat(qpos.shape[0], 1, 1) + + +class _FakeEnv: + def __init__(self, entities: dict[str, _FakeEntity] | None = None) -> None: + self.num_envs = 1 + self.device = torch.device("cpu") + self.robot = _FakeRobot(self.num_envs) + self.sim = _FakeSim(entities or {}) + self.left_arm_joints = [0, 1] + self.left_eef_joints = [2, 3] + self.right_arm_joints = [4, 5] + self.right_eef_joints = [6, 7] + self.open_state = torch.tensor([0.0, 0.0]) + self.close_state = torch.tensor([0.7, -0.7]) + + def get_agent_arm_control_part(self, is_left: bool) -> str: + return "physical_left_arm" if is_left else "physical_right_arm" + + def get_agent_eef_control_part(self, is_left: bool) -> str: + return "physical_left_eef" if is_left else "physical_right_eef" + + def get_current_xpos_agent(self) -> tuple[torch.Tensor, torch.Tensor]: + left = torch.eye(4).repeat(self.num_envs, 1, 1) + right = left.clone() + left[:, 1, 3] = -0.2 + right[:, 1, 3] = 0.2 + return left, right + + def get_current_qpos_agent(self) -> tuple[torch.Tensor, torch.Tensor]: + qpos = self.robot.get_qpos() + return qpos[:, self.left_arm_joints], qpos[:, self.right_arm_joints] + + def get_current_gripper_state_agent(self) -> tuple[torch.Tensor, torch.Tensor]: + qpos = self.robot.get_qpos() + return qpos[:, self.left_eef_joints], qpos[:, self.right_eef_joints] + + +def _box_vertices(half_extent: float) -> torch.Tensor: + h = float(half_extent) + return torch.tensor( + [ + [-h, -h, -h], + [h, -h, -h], + [h, h, h], + [-h, h, h], + ], + dtype=torch.float32, + ) + + +def _rect_vertices(x: float, y: float, z: float) -> torch.Tensor: + return torch.tensor( + [ + [sx * x, sy * y, sz * z] + for sx in (-1.0, 1.0) + for sy in (-1.0, 1.0) + for sz in (-1.0, 1.0) + ], + dtype=torch.float32, + ) + + +def _pose(x: float, y: float, z: float) -> torch.Tensor: + result = torch.eye(4).unsqueeze(0) + result[:, :3, 3] = torch.tensor([x, y, z]) + return result + + +def test_press_grounding_adapts_top_surface_and_depth_to_mainline_contract() -> None: + entity = _FakeEntity("button", _pose(0.1, -0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"button": entity}) + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "press", + "operator": "press", + "object": "button", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": {}, + "depends_on": [], + } + ) + ) + ) + semantics = ObjectSemantics( + affordance=Affordance(), + geometry={}, + label="button", + entity=entity, + ) + step = program.semantic_steps[0] + action = program.edges[0].actions[0] + + grounded = ActionGrounder(program, env, lambda _uid: semantics).ground( + action, + step, + arm="left_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + + assert isinstance(grounded.target, PressGoal) + assert isinstance(grounded.target.semantics.affordance, PressAffordance) + contact = grounded.target.semantics.affordance.get_press_pose( + grounded.target.target_pose + ) + assert torch.allclose(contact[0, :3, 3], torch.tensor([0.1, -0.2, 0.78])) + assert torch.allclose(contact[0, :3, 2], torch.tensor([0.0, 0.0, -1.0])) + + options = AtomicActionAdapter(env)._build_config(grounded, PressOptions) + + assert options.press_distance == pytest.approx(0.004) + + +def test_loader_regenerates_in_memory_without_execution_artifact( + tmp_path: Path, +) -> None: + task = _task_agent(_hold_step("hold", "can", "left_arm")) + graph = compile_task_agent_v2(task) + task_spec = { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "runtime_contract", + "level": "L1", + "instruction": "Hold the can.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "hold", + "task_type": "E1", + "params": {"object_role": "can"}, + "depends_on": [], + "role": "primary", + } + ], + "success": {"type": "object_held", "object": "can"}, + "oracle": {"reference_seed_graph": graph}, + "metadata": {"role_bindings": {"can": "can"}}, + } + task_path = tmp_path / "task_spec.json" + task_path.write_text(json.dumps(task_spec), encoding="utf-8") + agent_config = { + "schema_version": "action_engine_config_v2", + "task_spec": task_path.name, + "seed_task_graph": "not_written.json", + } + config_path = tmp_path / "agent_config.json" + config_path.write_text(json.dumps(agent_config), encoding="utf-8") + + program = load_agent_execution_program( + agent_config, + agent_config_path=config_path, + regenerate=True, + ) + + assert program.task == "runtime_contract" + assert program.semantic_steps[0].operator == "hold_hover" + assert not (tmp_path / "not_written.json").exists() + + +def test_production_loader_rejects_legacy_mapping() -> None: + legacy = compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + + with pytest.raises(ValueError, match="regenerate"): + _load_execution_program(legacy) + + +def test_documented_run_command_arguments_remain_compatible() -> None: + args = build_run_parser().parse_args( + [ + "--task_name", + "task4_2", + "--gym_config", + "/tmp/fast_gym_config.json", + "--agent_config", + "/tmp/agent_config.json", + "--regenerate", + "--headless", + "--seed", + "17", + ] + ) + + assert args.task_name == "task4_2" + assert args.regenerate is True + assert args.headless is True + assert args.seed == 17 + assert args.runtime_backend == "independent" + + +def test_dual_ur5_policy_uses_short_reach_upright_lifts() -> None: + upright = motion_policy(("orientation", "upright")) + ur5_pickup = resolve_motion_policy("dual_ur5", "PickUp", upright) + ur5_transport = resolve_motion_policy("dual_ur5", "MoveHeldObject", upright) + ur10_pickup = resolve_motion_policy("dual_ur10", "PickUp", upright) + ur10_transport = resolve_motion_policy("dual_ur10", "MoveHeldObject", upright) + + assert ur5_pickup["lift_height"] == pytest.approx(0.12) + assert ur5_transport["staging_lift_height"] == pytest.approx(0.12) + assert ur10_pickup["lift_height"] == pytest.approx(0.30) + assert ur10_transport["staging_lift_height"] == pytest.approx(0.25) + + +def test_joint_state_binding_selects_hand_timing_without_a_named_policy() -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + env.agent_initial_object_poses = {"can": entity.get_local_pose(to_matrix=True)} + program = load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ) + step = program.semantic_steps[0] + action = next( + action + for edge in program.edges + for action in edge.actions + if action["target_binding"].get("source") == "gripper_closed" + ) + grounder = ActionGrounder(program, env, lambda _uid: None) + + grounded = grounder.ground( + action, + step, + arm="left_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + + assert grounded.cfg["sample_interval"] == 10 + + +def test_runtime_policy_discards_legacy_support_z_fallbacks() -> None: + snapshot = default_runtime_policy("dual_franka").as_mapping() + snapshot["predicate_fallbacks"].update( + { + "support_min_z_offset": 0.02, + "support_max_z_offset": 0.35, + } + ) + + policy = RuntimePolicyCfg.from_mapping(snapshot) + + assert "support_min_z_offset" not in policy.predicate_fallbacks + assert "support_max_z_offset" not in policy.predicate_fallbacks + + +def test_runtime_policy_v4_migrates_grasp_direction_count() -> None: + snapshot = default_runtime_policy("dual_franka").as_mapping() + snapshot["schema_version"] = "action_engine_runtime_policy_v4" + snapshot["grasp"].pop("n_deviated_approach_directions") + snapshot_hash = hashlib.sha256( + json.dumps( + snapshot, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ).encode("utf-8") + ).hexdigest() + + policy = resolve_agent_runtime_policy( + { + "robot_profile": "dual_franka", + "runtime_policy": snapshot, + "runtime_policy_hash": snapshot_hash, + } + ) + + assert policy.schema_version == "action_engine_runtime_policy_v6" + assert policy.grasp["n_deviated_approach_directions"] == 4 + + +def test_runtime_policy_v5_migrates_support_geometry_thresholds() -> None: + snapshot = default_runtime_policy("dual_franka").as_mapping() + snapshot["schema_version"] = "action_engine_runtime_policy_v5" + snapshot["grounding"]["placement"]["clearance"] = 0.019 + for key in ( + "candidate_count", + "candidate_offset_fraction", + "support_margin", + "recovery_attempts", + ): + snapshot["grounding"]["placement"].pop(key) + for key in ( + "support_stability_samples", + "support_stability_interval_steps", + "support_linear_velocity_tolerance", + "support_angular_velocity_tolerance", + ): + snapshot["execution"].pop(key) + for key in ( + "support_com_margin", + "support_max_vertical_gap", + "support_max_penetration", + "support_min_overlap_ratio", + ): + snapshot["predicate_fallbacks"].pop(key) + snapshot_hash = hashlib.sha256( + json.dumps( + snapshot, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ).encode("utf-8") + ).hexdigest() + + policy = resolve_agent_runtime_policy( + { + "robot_profile": "dual_franka", + "runtime_policy": snapshot, + "runtime_policy_hash": snapshot_hash, + } + ) + + assert policy.schema_version == "action_engine_runtime_policy_v6" + assert policy.predicate_fallbacks["support_min_overlap_ratio"] == 0.25 + assert policy.grounding["placement"]["clearance"] == 0.019 + assert policy.grounding["placement"]["candidate_count"] == 5 + + +def test_runtime_recorder_writes_checkpoints_and_rendered_env_graphs( + tmp_path: Path, + monkeypatch: Any, +) -> None: + program = load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ) + original_seed = deepcopy(program.raw) + runtime_policy = default_runtime_policy("dual_ur10") + recorder = RuntimeRecorder( + program, + num_envs=2, + run_id="run-1", + episode_index=3, + output_root=tmp_path, + runtime_policy=runtime_policy.as_mapping(), + runtime_policy_hash=runtime_policy_hash(runtime_policy), + ) + step = program.semantic_steps[0] + recorder.edge( + program.edges[0].id, + step, + assignments=["left_arm", None], + grounded=[ + GroundedAction( + action_class="PickUp", + arm="left_arm", + control="arm", + target=None, + cfg={}, + motion_policy={"obj_upright_direction": torch.tensor([0.0, 0.0, 1.0])}, + ) + ], + active=torch.tensor([True, False]), + failed=torch.tensor([False, True]), + action_steps=4, + planner_traces=[ + { + "primary_strategy": "motion_gen", + "primary_success": torch.tensor([True, False]), + "fallback_used": torch.tensor([False, True]), + "planned_trajectory": torch.arange(24, dtype=torch.float32).reshape( + 2, 3, 4 + ), + } + ], + ) + recorder.step( + step, + torch.tensor([True, False]), + observed=torch.tensor([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]), + target=torch.tensor([[0.1, 0.2, 0.3], [0.0, 0.0, 0.0]]), + metadata=[ + { + "assigned_arm": "left_arm", + "physical_control_part": "physical_right_arm", + }, + {"assigned_arm": None, "physical_control_part": None}, + ], + ) + + episode_dir = tmp_path / "runtime_contract" / "run-1" / "episode_0003" + checkpoint_paths = sorted(episode_dir.glob("env_*/checkpoints/*.json")) + assert len(checkpoint_paths) == 2 + checkpoint = json.loads(checkpoint_paths[0].read_text(encoding="utf-8")) + assert checkpoint["semantic_step"]["id"] == "hold" + assert checkpoint["status"] == "success" + assert [item["event"] for item in checkpoint["events"]] == [ + "edge", + "semantic_step", + ] + assert checkpoint["events"][0]["actions"][0]["motion_policy"][ + "obj_upright_direction" + ] == [0.0, 0.0, 1.0] + assert checkpoint["events"][0]["planner_attempts"] == [ + { + "primary_strategy": "motion_gen", + "primary_success": True, + "fallback_used": False, + "planned_trajectory": [ + [0.0, 1.0, 2.0, 3.0], + [4.0, 5.0, 6.0, 7.0], + [8.0, 9.0, 10.0, 11.0], + ], + } + ] + assert checkpoint["events"][1]["assigned_arm"] == "left_arm" + assert checkpoint["events"][1]["physical_control_part"] == "physical_right_arm" + + rendered_documents: list[dict[str, Any]] = [] + visualization = ModuleType("embodichain.gen_sim.action_engine.graph_visualization") + + def render_task_graph_png(document: dict[str, Any]) -> bytes: + rendered_documents.append(deepcopy(document)) + return b"\x89PNG\r\n\x1a\nruntime-graph" + + visualization.render_task_graph_png = render_task_graph_png + monkeypatch.setitem(sys.modules, visualization.__name__, visualization) + output_dir = recorder.finalize(torch.tensor([True, False])) + + assert output_dir == episode_dir.as_posix() + assert program.raw == original_seed + expected_hash = execution_program_hash(original_seed) + for env_id, expected_status in enumerate(("success", "failed")): + env_dir = episode_dir / f"env_{env_id:04d}" + document = json.loads((env_dir / "task_graph.json").read_text(encoding="utf-8")) + assert document["schema_version"] == original_seed["schema_version"] + assert document["nodes"] == original_seed["nodes"] + assert document["edges"] == original_seed["edges"] + assert document["runtime"]["status"] == expected_status + assert document["runtime"]["seed_graph_hash"] == expected_hash + assert document["runtime"]["runtime_policy"] == runtime_policy.as_mapping() + assert document["runtime"]["runtime_policy_hash"] == runtime_policy_hash( + runtime_policy + ) + assert (env_dir / "task_graph.png").read_bytes().startswith(b"\x89PNG") + assert len(rendered_documents) == 2 + assert not list(episode_dir.rglob("*.tmp")) + + +def test_runtime_recorder_separates_dynamic_recovery_and_replay_phases( + tmp_path: Path, +) -> None: + program = load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ) + recorder = RuntimeRecorder( + program, + num_envs=1, + run_id="phased-recovery", + output_root=tmp_path, + ) + primary = program.semantic_steps[0] + recovery = replace( + primary, + id="recovery_e2_hold", + parent_step_id=primary.id, + ) + recovery_spec = deepcopy(program.raw["semantic_steps"][0]) + recovery_spec.update( + { + "id": recovery.id, + "parent_step_id": primary.id, + "role": "recovery", + } + ) + recorder.register_step(recovery, recovery_spec) + active = torch.tensor([True]) + recorder.edge( + "edge_recovery", + recovery, + assignments=["left_arm"], + grounded=[], + active=active, + failed=torch.tensor([False]), + action_steps=4, + phase="recovery", + ) + recorder.step( + recovery, + torch.tensor([True]), + observed=torch.zeros((1, 3)), + target=None, + phase="recovery", + ) + recorder.edge( + program.edges[0].id, + primary, + assignments=["left_arm"], + grounded=[], + active=active, + failed=torch.tensor([False]), + action_steps=3, + phase="replay", + ) + recorder.step( + primary, + torch.tensor([True]), + observed=torch.zeros((1, 3)), + target=None, + ) + + checkpoints = sorted( + (recorder.output_dir / "env_0000" / "checkpoints").glob("*.json") + ) + assert len(checkpoints) == 2 + recovery_checkpoint = next( + json.loads(path.read_text(encoding="utf-8")) + for path in checkpoints + if "recovery_e2_hold" in path.name + ) + primary_checkpoint = next( + json.loads(path.read_text(encoding="utf-8")) + for path in checkpoints + if path.name.endswith("_hold.json") and "recovery_e2" not in path.name + ) + assert {event["phase"] for event in recovery_checkpoint["events"]} == {"recovery"} + assert primary_checkpoint["events"][0]["phase"] == "replay" + assert primary_checkpoint["events"][-1]["phase"] == "primary" + + +def test_runtime_recorder_does_not_mask_execution_when_png_rendering_fails( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + program = load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ) + recorder = RuntimeRecorder( + program, + num_envs=1, + run_id="render_failure", + output_root=tmp_path, + ) + visualization = ModuleType("embodichain.gen_sim.action_engine.graph_visualization") + + def fail_render(_document: dict[str, Any]) -> bytes: + raise ValueError("broken renderer") + + visualization.render_task_graph_png = fail_render + monkeypatch.setitem(sys.modules, visualization.__name__, visualization) + + output_dir = recorder.finalize(torch.tensor([False])) + + record = json.loads( + (Path(output_dir) / "env_0000" / "task_graph.json").read_text(encoding="utf-8") + ) + assert record["runtime"]["status"] == "failed" + assert record["runtime"]["visualization_error"] == ("ValueError: broken renderer") + + +def test_ready_scheduler_packs_only_declared_opposite_arm_pickups() -> None: + compiled = compile_task_agent( + _task_agent( + _hold_step("left", "can_a", "left_arm"), + _hold_step("right", "can_b", "right_arm"), + ) + ) + executor = ProgramExecutor( + load_execution_program(compiled), + _FakeEnv(), + record_runtime=False, + ) + ready = [edge for edge in executor.program.edges if not edge.depends_on] + + packed = executor._pack_ready_edges(ready) + + assert len(packed) == 2 + assert {executor.step_by_edge[edge.id].object_uid for edge in packed} == { + "can_a", + "can_b", + } + + +def test_ready_scheduler_serializes_contact_sensitive_orient_pickups() -> None: + steps = [ + { + "id": step_id, + "operator": "orient_object", + "object": object_uid, + "actor": {"mode": "auto"}, + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + }, + "depends_on": [], + } + for step_id, object_uid in (("left", "can_a"), ("right", "can_b")) + ] + task_agent = _task_agent(*steps) + task_agent["allocation_groups"] = [ + { + "id": "dual_arms_1", + "semantic_step_ids": ["left", "right"], + "arm_constraint": "distinct_arms", + } + ] + executor = ProgramExecutor( + load_execution_program(compile_task_agent(task_agent)), + _FakeEnv(), + record_runtime=False, + ) + ready = [edge for edge in executor.program.edges if not edge.depends_on] + + assert len(executor._pack_ready_edges(ready)) == 1 + + +def test_ready_scheduler_defers_pickups_until_a_carried_payload_is_released() -> None: + entities = { + "table": _FakeEntity( + "table", + _pose(0.0, 0.0, 0.70), + _rect_vertices(0.60, 0.40, 0.02), + ), + **{ + uid: _FakeEntity( + uid, + _pose(x, 0.0, 0.76), + _rect_vertices(0.03, 0.03, 0.06), + ) + for uid, x in (("can_a", -0.2), ("can_b", 0.0), ("can_c", 0.2)) + }, + } + task = _task_agent( + { + "id": "line", + "operator": "arrange_line", + "objects": ["can_a", "can_b", "can_c"], + "actor": {"mode": "auto"}, + "goal": {"axis": "world_x", "order_constraint": "free"}, + "depends_on": [], + } + ) + executor = ProgramExecutor( + load_execution_program(compile_task_agent(task)), + _FakeEnv(entities), + record_runtime=False, + ) + pickup_edges = [ + edge + for edge in executor.program.edges + if executor._parallel_pickup_candidate(edge) + ] + completed = {pickup_edges[0].id, pickup_edges[1].id} + ready = [ + edge + for edge in executor.program.edges + if edge.id not in completed and set(edge.depends_on) <= completed + ] + executor._arm_owners["left_arm"][0] = "can_a" + executor._arm_owners["right_arm"][0] = "can_b" + + packed = executor._pack_ready_edges(ready, completed=completed) + + assert not executor._parallel_pickup_candidate(packed[0]) + + executor._arm_owners["right_arm"][0] = None + packed = executor._pack_ready_edges(ready, completed=completed) + assert len(packed) == 1 + assert not executor._parallel_pickup_candidate(packed[0]) + + +def test_parallel_pickups_plan_each_arm_at_execution_time( + monkeypatch: pytest.MonkeyPatch, +) -> None: + first = _hold_step("first", "can_a", "left_arm") + second = _hold_step("second", "can_b", "right_arm") + first["actor"] = {"mode": "auto"} + second["actor"] = {"mode": "auto"} + executor = ProgramExecutor( + load_execution_program(compile_task_agent(_task_agent(first, second))), + _FakeEnv( + { + "can_a": _FakeEntity( + "can_a", _pose(0.0, 0.2, 0.75), _box_vertices(0.03) + ), + "can_b": _FakeEntity( + "can_b", _pose(0.0, -0.2, 0.75), _box_vertices(0.03) + ), + } + ), + record_runtime=False, + ) + edges = tuple( + next( + edge + for edge in executor.program.edges + if edge.id in step.edge_ids + and edge.actions[0]["atomic_action_class"] == "PickUp" + ) + for step in executor.program.semantic_steps + ) + estimate = SimpleNamespace( + feasible=torch.tensor([True]), + cost=torch.tensor([0.0]), + ) + monkeypatch.setattr(executor, "_candidate", lambda *_args, **_kwargs: estimate) + monkeypatch.setattr(executor, "_report_candidates", lambda *_args: None) + live_calls: list[tuple[str, str]] = [] + + def plan_live(edge, step, arm): + live_calls.append((step.id, arm)) + grounded = GroundedAction( + action_class="PickUp", + arm=arm, + control="arm", + target=SimpleNamespace(), + cfg={}, + ) + return grounded, ActionOutcome( + trajectory=torch.zeros(1, 1, executor.env.robot.dof), + success=torch.tensor([True]), + next_state=ExecutionState(last_qpos=executor.env.robot.get_qpos()), + grounded=grounded, + ) + + monkeypatch.setattr(executor, "_plan_live_hold", plan_live) + monkeypatch.setattr(executor.adapter, "execute_trajectory", lambda *_a, **_k: []) + monkeypatch.setattr( + executor, "_physical_pickup", lambda _u, _a, _s, attempted: attempted + ) + monkeypatch.setattr( + executor, "_rebase_held_state", lambda _u, _a, state, *_args, **_kwargs: state + ) + monkeypatch.setattr(executor, "_update_ownership", lambda *_args, **_kwargs: None) + + _, failed = executor._execute_parallel_pickups( + edges, + failed=torch.tensor([False]), + ) + + assert set(live_calls) == { + ("first", "right_arm"), + ("second", "left_arm"), + } + assert not bool(failed[0]) + + +def test_required_arm_rejects_wrong_candidate_without_planning() -> None: + compiled = compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + executor = ProgramExecutor( + load_execution_program(compiled), + _FakeEnv(), + record_runtime=False, + ) + failed = torch.zeros(1, dtype=torch.bool) + + candidate = executor._candidate( + executor.program.semantic_steps[0], + "right_arm", + failed, + ) + + assert not bool(candidate.feasible.any()) + assert bool(torch.isinf(candidate.cost).all()) + + +def test_required_arm_speculative_failure_still_reaches_live_planning( + monkeypatch: pytest.MonkeyPatch, +) -> None: + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ), + _FakeEnv( + {"can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03))} + ), + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + estimate = SimpleNamespace( + feasible=torch.tensor([False]), + cost=torch.tensor([torch.inf]), + ) + monkeypatch.setattr(executor, "_candidate", lambda *_args, **_kwargs: estimate) + monkeypatch.setattr(executor, "_report_candidates", lambda *_args: None) + + executor._ensure_assignment(step, torch.tensor([False])) + + assert executor._assignments[step.id] == ["left_arm"] + + +def test_auto_pickup_outside_deadband_requires_same_side_arm( + monkeypatch: pytest.MonkeyPatch, +) -> None: + step_mapping = _hold_step("hold", "can", "left_arm") + step_mapping["actor"] = {"mode": "auto"} + executor = ProgramExecutor( + load_execution_program(compile_task_agent(_task_agent(step_mapping))), + _FakeEnv( + {"can": _FakeEntity("can", _pose(0.0, -0.2, 0.75), _box_vertices(0.03))} + ), + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + + def candidate(_step, arm, _failed): + cost = 10.0 if arm == "left_arm" else 1.0 + return SimpleNamespace( + feasible=torch.tensor([True]), + cost=torch.tensor([cost]), + ) + + monkeypatch.setattr(executor, "_candidate", candidate) + monkeypatch.setattr(executor, "_report_candidates", lambda *_args: None) + + executor._ensure_assignment(step, torch.tensor([False])) + + assert executor._preferred_live_pickup_arm(step, 0) == "left_arm" + assert executor._assignments[step.id] == ["left_arm"] + + +def test_auto_pickup_inside_deadband_selects_lower_cost_arm( + monkeypatch: pytest.MonkeyPatch, +) -> None: + step_mapping = _hold_step("hold", "can", "left_arm") + step_mapping["actor"] = {"mode": "auto"} + executor = ProgramExecutor( + load_execution_program(compile_task_agent(_task_agent(step_mapping))), + _FakeEnv( + {"can": _FakeEntity("can", _pose(0.0, 0.01, 0.75), _box_vertices(0.03))} + ), + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + + def candidate(_step, arm, _failed): + cost = 10.0 if arm == "left_arm" else 1.0 + return SimpleNamespace( + feasible=torch.tensor([True]), + cost=torch.tensor([cost]), + ) + + monkeypatch.setattr(executor, "_candidate", candidate) + monkeypatch.setattr(executor, "_report_candidates", lambda *_args: None) + + executor._ensure_assignment(step, torch.tensor([False])) + + assert executor._preferred_live_pickup_arm(step, 0) is None + assert executor._assignments[step.id] == ["right_arm"] + + +def test_auto_pickup_retry_does_not_cross_sides_by_default( + monkeypatch: pytest.MonkeyPatch, +) -> None: + step_mapping = _hold_step("hold", "can", "left_arm") + step_mapping["actor"] = {"mode": "auto"} + executor = ProgramExecutor( + load_execution_program(compile_task_agent(_task_agent(step_mapping))), + _FakeEnv( + {"can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03))} + ), + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + estimate = SimpleNamespace( + feasible=torch.tensor([False]), + cost=torch.tensor([torch.inf]), + ) + monkeypatch.setattr(executor, "_candidate", lambda *_args, **_kwargs: estimate) + monkeypatch.setattr(executor, "_report_candidates", lambda *_args: None) + monkeypatch.setattr( + executor, + "_preferred_live_pickup_arm", + lambda *_args: "left_arm", + ) + + executor._ensure_assignment(step, torch.tensor([False])) + assert executor._assignments[step.id] == ["left_arm"] + + executor._pickup_retry_exclusions[(step.id, 0)] = {"left_arm"} + executor._assignments.pop(step.id) + executor._ensure_assignment(step, torch.tensor([False])) + + assert executor._assignments[step.id] == [None] + + +def test_auto_pickup_retry_can_explicitly_cross_sides( + monkeypatch: pytest.MonkeyPatch, +) -> None: + step_mapping = _hold_step("hold", "can", "left_arm") + step_mapping["actor"] = {"mode": "auto"} + executor = ProgramExecutor( + load_execution_program(compile_task_agent(_task_agent(step_mapping))), + _FakeEnv( + {"can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03))} + ), + record_runtime=False, + ) + executor.runtime_policy.arm_selection.allow_cross_side_fallback = True + step = executor.program.semantic_steps[0] + estimate = SimpleNamespace( + feasible=torch.tensor([False]), + cost=torch.tensor([torch.inf]), + ) + monkeypatch.setattr(executor, "_candidate", lambda *_args, **_kwargs: estimate) + monkeypatch.setattr(executor, "_report_candidates", lambda *_args: None) + monkeypatch.setattr( + executor, + "_preferred_live_pickup_arm", + lambda *_args: "left_arm", + ) + + executor._pickup_retry_exclusions[(step.id, 0)] = {"left_arm"} + executor._ensure_assignment(step, torch.tensor([False])) + + assert executor._assignments[step.id] == ["right_arm"] + + +def test_auto_pickup_runtime_retry_uses_the_other_arm( + monkeypatch: pytest.MonkeyPatch, +) -> None: + step_mapping = _hold_step("hold", "can", "left_arm") + step_mapping["actor"] = {"mode": "auto"} + executor = ProgramExecutor( + load_execution_program(compile_task_agent(_task_agent(step_mapping))), + _FakeEnv( + {"can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03))} + ), + record_runtime=False, + ) + executor.runtime_policy.arm_selection.allow_cross_side_fallback = True + step = executor.program.semantic_steps[0] + original_edge = executor.edges[step.edge_ids[0]] + action = {**original_edge.actions[0], "seed_node_id": "pickup_node"} + edge = replace(original_edge, actions=(action,)) + estimate = SimpleNamespace( + feasible=torch.tensor([False]), + cost=torch.tensor([torch.inf]), + ) + monkeypatch.setattr(executor, "_candidate", lambda *_args, **_kwargs: estimate) + monkeypatch.setattr(executor, "_report_candidates", lambda *_args: None) + monkeypatch.setattr( + executor, + "_preferred_live_pickup_arm", + lambda *_args: "left_arm", + ) + executor._ensure_assignment(step, torch.tensor([False])) + attempts: list[str | None] = [] + + def execute(_edge, _step, *, failed): + arm = executor._assignments[step.id][0] + attempts.append(arm) + return _EdgeResult( + actions=[], + failed=torch.tensor([arm == "left_arm"]) | failed, + grounded=[], + planner_traces=[], + executed=torch.tensor([False]), + ) + + decisions = 0 + + def record_failure(*_args, **_kwargs): + nonlocal decisions + decisions += 1 + return SimpleNamespace(retry=torch.tensor([decisions == 1])) + + executor.runtime_graph = SimpleNamespace( + graph={"nodes": [{"id": "pickup_node", "precondition": {}}]}, + record_failure=record_failure, + ) + monkeypatch.setattr(executor, "_execute_edge", execute) + + result = executor._execute_edge_with_retries( + edge, + step, + failed=torch.tensor([False]), + ) + + assert attempts == ["left_arm", "right_arm"] + assert executor.retry_count == 1 + assert not bool(result.failed[0]) + + +def _held_state( + env: _FakeEnv, + entity: _FakeEntity, + *, + arm: str = "left_arm", +) -> ExecutionState: + semantics = ObjectSemantics( + affordance=Affordance(), + geometry={}, + label=entity.uid, + entity=entity, + ) + left_eef, right_eef = env.get_current_xpos_agent() + eef = left_eef if arm == "left_arm" else right_eef + object_pose = entity.get_local_pose(to_matrix=True) + return ExecutionState( + last_qpos=env.robot.get_qpos(), + held_objects={ + f"physical_{arm}": HeldObjectState( + semantics=semantics, + object_to_eef=torch.bmm(torch.linalg.inv(object_pose), eef), + grasp_xpos=eef, + ) + }, + ) + + +def _coordinated_held_state( + env: _FakeEnv, + entity: _FakeEntity, +) -> ExecutionState: + semantics = ObjectSemantics( + affordance=Affordance(), + geometry={}, + label=entity.uid, + entity=entity, + ) + left_eef, right_eef = env.get_current_xpos_agent() + object_pose = entity.get_local_pose(to_matrix=True) + return ExecutionState( + last_qpos=env.robot.get_qpos(), + held_objects={ + "physical_left_arm": HeldObjectState( + semantics=semantics, + object_to_eef=torch.bmm(torch.linalg.inv(object_pose), left_eef), + grasp_xpos=left_eef, + env_mask=torch.ones(env.num_envs, dtype=torch.bool), + ), + "physical_right_arm": HeldObjectState( + semantics=semantics, + object_to_eef=torch.bmm(torch.linalg.inv(object_pose), right_eef), + grasp_xpos=right_eef, + env_mask=torch.ones(env.num_envs, dtype=torch.bool), + ), + }, + ) + + +@pytest.mark.parametrize( + ("opens", "expected_failed", "expect_held"), + ((True, False, False), (False, True, True)), +) +def test_explicit_dual_gripper_release_commits_only_after_both_hands_open( + opens: bool, + expected_failed: bool, + expect_held: bool, +) -> None: + entity = _FakeEntity("tray", _pose(0.0, 0.0, 0.75), _box_vertices(0.2)) + env = _FakeEnv({"tray": entity}) + env.robot._qpos[:, env.left_eef_joints] = env.close_state + env.robot._qpos[:, env.right_eef_joints] = env.close_state + state = _coordinated_held_state(env, entity) + executor = object.__new__(ProgramExecutor) + executor.env = env + executor._assignments = {"task_01": ["coordinated"]} + executor._step_states = {("task_01", "coordinated"): state} + executor._object_states = {} + executor._orientation_references = {} + + def ground( + action: dict[str, Any], + _step: Any, + *, + arm: str, + **_kwargs: Any, + ) -> GroundedAction: + return GroundedAction( + action_class=str(action["atomic_action_class"]), + arm=arm, + control="hand", + target=None, + cfg={}, + ) + + def plan(grounded: GroundedAction, current: ExecutionState) -> ActionOutcome: + return ActionOutcome( + trajectory=torch.zeros(1, 2, env.robot.dof), + success=torch.ones(1, dtype=torch.bool), + next_state=current, + grounded=grounded, + ) + + def execute_trajectory( + _trajectory: torch.Tensor, + *, + active: torch.Tensor, + ) -> list[torch.Tensor]: + if opens and bool(active.any()): + env.robot._qpos[:, env.left_eef_joints] = env.open_state + env.robot._qpos[:, env.right_eef_joints] = env.open_state + elif bool(active.any()): + env.robot._qpos[:, env.left_eef_joints] = env.open_state + return [] + + executor.grounder = SimpleNamespace(ground=ground) + executor.adapter = SimpleNamespace( + plan=plan, + combine=lambda _outcomes, _masks: ( + torch.zeros(1, 2, env.robot.dof), + torch.ones(1, dtype=torch.bool), + ), + execute_trajectory=execute_trajectory, + ) + actions = [ + { + "atomic_action_class": "MoveJoints", + "actor": {"arm": arm}, + "control": "hand", + "target_binding": { + "kind": "joint_state", + "source": "gripper_open", + "coordinated_release_role": role, + }, + } + for arm, role in ( + ("left_arm", "participant"), + ("right_arm", "commit"), + ) + ] + + result = executor._execute_explicit_dual( + SimpleNamespace(id="release", actions=actions), + SimpleNamespace(id="task_01"), + torch.zeros(1, dtype=torch.bool), + ) + + released_state = executor._step_states[("task_01", "coordinated")] + left_held = released_state.get_held_object("physical_left_arm") + right_held = released_state.get_held_object("physical_right_arm") + assert result.failed.tolist() == [expected_failed] + assert (left_held is not None and right_held is not None) is expect_held + + +def _handover_held_state( + env: _FakeEnv, + entity: _FakeEntity, + *, + arm: str = "left_arm", +) -> ExecutionState: + """Build a fixture grasp on the side assigned to the transfer arm.""" + state = _held_state(env, entity, arm=arm) + held = state.get_held_object(f"physical_{arm}") + assert held is not None + _, lateral = robot_frame_axes(env) + role_axis = lateral if arm == "left_arm" else -lateral + offset = torch.cat((role_axis, role_axis.new_zeros((int(env.num_envs), 1))), dim=1) + object_to_eef = held.object_to_eef.clone() + object_to_eef[:, :3, 3] = offset * 0.02 + replacement = HeldObjectState( + semantics=held.semantics, + object_to_eef=object_to_eef, + grasp_xpos=torch.bmm(entity.get_local_pose(to_matrix=True), object_to_eef), + env_mask=held.env_mask, + ) + held_objects = dict(state.held_objects) + held_objects[f"physical_{arm}"] = replacement + return state.with_updates(held_objects=held_objects) + + +def test_handover_grounding_uses_bottom_region_and_diagonal_receive() -> None: + entities = { + "can": _FakeEntity( + "can", + _pose(0.0, 0.2, 1.2), + _rect_vertices(0.03, 0.03, 0.10), + ), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + task = { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "handover", + "level": "L1", + "instruction": "Hand over the can.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "task_01", + "task_type": "E4", + "params": { + "object_role": "can", + "transfer_arm": "left_arm", + "receive_arm": "right_arm", + }, + "depends_on": [], + "role": "primary", + } + ], + "success": {"type": "handover_complete"}, + "oracle": {}, + "metadata": {}, + } + program = load_execution_program(instantiate_seed_graph(task, {"can": "can"})) + step = program.semantic_steps[0] + edge = next( + edge + for edge in program.edges + if edge.actions[0]["atomic_action_class"] == "HandOver" + ) + state = _handover_held_state(env, entities["can"]) + held = state.get_held_object("physical_left_arm") + assert held is not None + grounder = ActionGrounder(program, env, lambda _uid: held.semantics) + + grounded = grounder.ground( + edge.actions[0], + step, + arm="coordinated", + state=state, + ) + middle = grounded.cfg["middle_object_pose"] + final = grounded.cfg["final_object_pose"] + cfg = AtomicActionAdapter(env)._build_config(grounded, HandOverOptions) + + staging_edge = next( + edge + for edge in program.edges + if edge.actions[0]["target_binding"].get("kind") == "handover_staging" + ) + staging = grounder.ground( + staging_edge.actions[0], + step, + arm="left_arm", + state=state, + ) + + assert middle[0, 1, 3] == pytest.approx(0.0) + torch.testing.assert_close(final, middle) + torch.testing.assert_close(cfg.middle_object_pose, cfg.final_object_pose) + assert cfg.receive_pick_object_part == "bottom" + assert cfg.receive_approach_direction[1] < 0.0 + assert cfg.receive_approach_direction[2] < 0.0 + assert staging.motion_policy["upright_yaw_samples"] == 8 + + +def test_handover_rejects_receiver_motion_during_internal_final_phase() -> None: + adapter = AtomicActionAdapter(_FakeEnv()) + grounded = GroundedAction( + action_class="HandOver", + arm="coordinated", + control="coordinated", + target=SimpleNamespace(), + cfg={"transfer_arm": "left_arm"}, + ) + options = HandOverOptions(retreat_steps=4) + trajectory = torch.zeros(1, 12, adapter.env.robot.dof) + + assert bool( + adapter._handover_receiver_hold_mask( + trajectory, + grounded, + options, + tolerance=1.0e-3, + )[0] + ) + + trajectory[0, -1, 4] = 0.02 + assert not bool( + adapter._handover_receiver_hold_mask( + trajectory, + grounded, + options, + tolerance=1.0e-3, + )[0] + ) + + +def _handover_then_place_task() -> dict[str, Any]: + return { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "handover_then_place", + "level": "L3", + "instruction": "Hand over the can and place it beside the target.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "handover", + "task_type": "E4", + "params": { + "object_role": "can", + "transfer_arm": "left_arm", + "receive_arm": "right_arm", + "orientation_goal": "preserve", + }, + "depends_on": [], + "role": "primary", + }, + { + "id": "place", + "task_type": "E1", + "params": { + "object_role": "can", + "target_role": "target", + "relation": "right_of", + }, + "depends_on": ["handover"], + "role": "primary", + }, + ], + "success": { + "op": "all", + "terms": [ + {"type": "handover_complete", "task_instance_id": "handover"}, + {"type": "semantic_goal", "task_instance_id": "place"}, + ], + }, + "oracle": {}, + "metadata": {}, + } + + +def test_handover_continuation_uses_stable_upright_policies() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.2, 0.0, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + task = deepcopy(_handover_then_place_task()) + task["task_instances"][1]["params"]["orientation_goal"] = "upright" + program = load_execution_program( + instantiate_seed_graph( + task, + {"can": "can", "target": "target"}, + ) + ) + step = next( + candidate + for candidate in program.semantic_steps + if candidate.operator == "place_relative" + ) + state = _held_state(env, entities["can"], arm="right_arm") + held = state.get_held_object("physical_right_arm") + assert held is not None + grounder = ActionGrounder(program, env, lambda _uid: held.semantics) + edges = [edge for edge in program.edges if edge.id in step.edge_ids] + staging = next( + edge + for edge in edges + if edge.actions[0]["target_binding"].get("phase") == "staging" + ) + final = next( + edge + for edge in edges + if edge.actions[0]["target_binding"].get("phase") == "final" + ) + release = next( + edge for edge in edges if edge.actions[0]["atomic_action_class"] == "Place" + ) + retreat = next( + edge + for edge in edges + if edge.actions[0]["atomic_action_class"] == "MoveEndEffector" + ) + home = next( + edge for edge in edges if edge.actions[0]["atomic_action_class"] == "MoveJoints" + ) + + grounded_staging = grounder.ground( + staging.actions[0], step, arm="right_arm", state=state + ) + grounded_final = grounder.ground( + final.actions[0], step, arm="right_arm", state=state + ) + supported_reference = _pose(0.0, 0.0, 0.90) + grounded_final_with_reference = grounder.ground( + final.actions[0], + step, + arm="right_arm", + state=state, + orientation_reference_pose=supported_reference, + ) + grounded_release = grounder.ground( + release.actions[0], step, arm="right_arm", state=state + ) + grounded_retreat = grounder.ground( + retreat.actions[0], step, arm="right_arm", state=state + ) + grounded_home = grounder.ground(home.actions[0], step, arm="right_arm", state=state) + upright = motion_policy(("orientation", "upright")) + release_defaults = resolve_motion_policy("dual_ur10", "Place", upright) + retreat_defaults = resolve_motion_policy("dual_ur10", "MoveEndEffector", upright) + + assert grounded_staging.cfg["upright_yaw_samples"] == 8 + assert grounded_final.cfg["upright_yaw_samples"] == 8 + assert grounded_final_with_reference.target_object_pose is not None + assert grounded_final_with_reference.target_object_pose[0, 2, 3] == pytest.approx( + 0.90 + ) + assert ( + grounded_release.cfg["sample_interval"] == release_defaults["sample_interval"] + ) + assert ( + grounded_release.cfg["post_hold_steps"] == release_defaults["post_hold_steps"] + ) + assert ( + grounded_retreat.cfg["sample_interval"] == retreat_defaults["sample_interval"] + ) + assert grounded_retreat.cfg["retreat_height"] == pytest.approx( + retreat_defaults["retreat_height"] + ) + assert grounded_retreat.cfg["retreat_height"] == pytest.approx(0.30) + assert grounded_retreat.motion_policy["clearance_object_uid"] == "can" + assert grounded_retreat.motion_policy["collision_exclusion_uids"] == [ + "can", + "target", + ] + assert grounded_retreat.motion_policy["collision_safety"] == "required" + assert grounded_home.motion_policy["collision_safety"] == "required" + + +def test_preserve_handover_continuation_does_not_enable_yaw_search() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.2, 0.0, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + task = deepcopy(_handover_then_place_task()) + task["task_instances"][1]["params"]["orientation_goal"] = "preserve" + program = load_execution_program( + instantiate_seed_graph(task, {"can": "can", "target": "target"}) + ) + step = next( + candidate + for candidate in program.semantic_steps + if candidate.operator == "place_relative" + ) + state = _held_state(env, entities["can"], arm="right_arm") + held = state.get_held_object("physical_right_arm") + assert held is not None + grounder = ActionGrounder(program, env, lambda _uid: held.semantics) + final = next( + edge + for edge in program.edges + if edge.id in step.edge_ids + and edge.actions[0]["target_binding"].get("phase") == "final" + ) + reference = _pose(0.0, 0.0, 0.90) + + grounded = grounder.ground( + final.actions[0], + step, + arm="right_arm", + state=state, + orientation_reference_pose=reference, + ) + + assert "upright_yaw_samples" not in grounded.cfg + assert grounded.target_object_pose is not None + torch.testing.assert_close( + grounded.target_object_pose[:, :3, :3], + reference[:, :3, :3], + ) + + +def test_dual_franka_handover_uses_explicit_exchange_clearance() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 1.20), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.2, 0.0, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + env.agent_robot_profile = "dual_franka" + program = load_execution_program( + instantiate_seed_graph( + _handover_then_place_task(), + {"can": "can", "target": "target"}, + ) + ) + step = next( + candidate for candidate in program.semantic_steps if candidate.id == "handover" + ) + state = _handover_held_state(env, entities["can"], arm="left_arm") + held = state.get_held_object("physical_left_arm") + assert held is not None + grounder = ActionGrounder(program, env, lambda _uid: held.semantics) + staging_edge = next( + edge + for edge in program.edges + if edge.id in step.edge_ids + and edge.actions[0]["target_binding"].get("kind") == "handover_staging" + ) + handover_edge = next( + edge + for edge in program.edges + if edge.id in step.edge_ids + and edge.actions[0]["target_binding"].get("kind") == "handover_goal" + ) + staging = grounder.ground( + staging_edge.actions[0], step, arm="left_arm", state=state + ) + handover = grounder.ground( + handover_edge.actions[0], step, arm="coordinated", state=state + ) + + assert staging.motion_policy["exchange_clearance"] > 0.0 + assert handover.motion_policy["exchange_clearance"] > 0.0 + assert handover.motion_policy["lift_height"] > 0.0 + assert staging.target_object_pose is not None + live_object_pose = entities["can"].get_local_pose(to_matrix=True) + assert handover.motion_policy["middle_object_pose"][0, 2, 3] == pytest.approx( + live_object_pose[0, 2, 3] + ) + + +def test_handover_candidates_avoid_occupied_table_center_and_lift_payload() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 1.03), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.35, 0.0, 1.03), _box_vertices(0.03)), + "notebook": _FakeEntity("notebook", _pose(0.0, 0.0, 1.04), _box_vertices(0.05)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + env.agent_robot_profile = "dual_franka" + program = load_execution_program( + instantiate_seed_graph( + _handover_then_place_task(), + {"can": "can", "target": "target"}, + ) + ) + step = next( + candidate for candidate in program.semantic_steps if candidate.id == "handover" + ) + state = _handover_held_state(env, entities["can"], arm="left_arm") + held = state.get_held_object("physical_left_arm") + assert held is not None + grounder = ActionGrounder(program, env, lambda _uid: held.semantics) + staging_edge = next( + edge + for edge in program.edges + if edge.id in step.edge_ids + and edge.actions[0]["target_binding"].get("kind") == "handover_staging" + ) + handover_edge = next( + edge + for edge in program.edges + if edge.id in step.edge_ids + and edge.actions[0]["target_binding"].get("kind") == "handover_goal" + ) + + staging = grounder.ground( + staging_edge.actions[0], step, arm="left_arm", state=state + ) + candidates = grounder.ground_candidates( + handover_edge.actions[0], step, arm="coordinated", state=state + ) + + assert staging.target_object_pose is not None + assert torch.linalg.vector_norm(staging.target_object_pose[0, :2, 3]) > 0.10 + assert float(staging.target_object_pose[0, 2, 3]) >= 1.15 + assert len(candidates) == 4 + assert all( + torch.linalg.vector_norm(candidate.cfg["middle_object_pose"][0, :2, 3]) > 0.10 + for candidate in candidates[:2] + ) + + +def test_on_placement_grounding_samples_bounded_live_support_poses() -> None: + entities = { + "payload": _FakeEntity( + "payload", + _pose(0.0, -0.20, 0.79), + _rect_vertices(0.03, 0.03, 0.03), + ), + "support": _FakeEntity( + "support", + _pose(0.0, 0.0, 0.75), + _rect_vertices(0.20, 0.15, 0.01), + ), + } + env = _FakeEnv(entities) + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "payload", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "support", + "relation": "on", + "orientation_goal": "none", + }, + "depends_on": [], + } + ) + ) + ) + step = program.semantic_steps[0] + edge = next( + item + for item in program.edges + if item.id in step.edge_ids + and item.actions[0]["target_binding"].get("kind") == "semantic_goal" + and item.actions[0]["target_binding"].get("phase", "final") != "staging" + ) + grounder = ActionGrounder(program, env, lambda _uid: None) + + candidates = grounder.ground_candidates( + edge.actions[0], + step, + arm="left_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + + assert len(candidates) == 5 + assert [item.motion_policy["placement_candidate_index"] for item in candidates] == [ + 0, + 1, + 2, + 3, + 4, + ] + offsets = [item.motion_policy["placement_xy_offset"][0] for item in candidates] + assert len({tuple(float(value) for value in offset) for offset in offsets}) == 5 + support_lower = torch.tensor([-0.20, -0.15]) + support_upper = torch.tensor([0.20, 0.15]) + for item in candidates: + center = item.target_object_pose[0, :2, 3] + assert torch.all(center >= support_lower) + assert torch.all(center <= support_upper) + + +def test_on_placement_candidates_respect_support_geometry_origin() -> None: + support_vertices = _rect_vertices(0.10, 0.08, 0.01) + support_vertices[:, 0] += 0.25 + entities = { + "payload": _FakeEntity( + "payload", + _pose(0.0, -0.20, 0.79), + _rect_vertices(0.03, 0.03, 0.03), + ), + "support": _FakeEntity( + "support", + _pose(0.10, 0.0, 0.75), + support_vertices, + ), + } + env = _FakeEnv(entities) + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "payload", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "support", + "relation": "on", + "orientation_goal": "none", + }, + "depends_on": [], + } + ) + ) + ) + step = program.semantic_steps[0] + edge = next( + item + for item in program.edges + if item.id in step.edge_ids + and item.actions[0]["target_binding"].get("kind") == "semantic_goal" + and item.actions[0]["target_binding"].get("phase", "final") != "staging" + ) + candidates = ActionGrounder(program, env, lambda _uid: None).ground_candidates( + edge.actions[0], + step, + arm="left_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + + support_world = support_vertices[:, :2] + torch.tensor([0.10, 0.0]) + lower = support_world.min(dim=0).values + 0.002 + upper = support_world.max(dim=0).values - 0.002 + payload_local = entities["payload"]._vertices[:, :2] + for candidate in candidates: + origin = candidate.target_object_pose[0, :2, 3] + assert torch.all(origin + payload_local.min(dim=0).values >= lower) + assert torch.all(origin + payload_local.max(dim=0).values <= upper) + + +def test_build_stack_root_compiles_to_generic_table_support() -> None: + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "stack", + "operator": "build_stack", + "objects": ["base", "nested"], + "actor": {"mode": "auto"}, + "goal": { + "anchor": "table_center", + "stack_mode": "nested", + "orientation_goal": "none", + }, + "depends_on": [], + } + ) + ) + ) + + root, child = program.semantic_steps + assert root.goal["relation"] == "on" + assert root.goal["reference_object"] == "table" + assert root.postcondition["reference_object"] == "table" + assert child.goal["relation"] == "inside" + assert child.goal["reference_object"] == "base" + + +def test_executor_tries_next_placement_pose_after_planning_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entities = { + "payload": _FakeEntity( + "payload", + _pose(0.0, -0.20, 0.79), + _rect_vertices(0.03, 0.03, 0.03), + ), + "support": _FakeEntity( + "support", + _pose(0.0, 0.0, 0.75), + _rect_vertices(0.20, 0.15, 0.01), + ), + } + env = _FakeEnv(entities) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "payload", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "support", + "relation": "on", + "orientation_goal": "none", + }, + "depends_on": [], + } + ) + ) + ), + env, + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + edge = next( + item + for item in executor.program.edges + if item.id in step.edge_ids + and item.actions[0]["target_binding"].get("kind") == "semantic_goal" + and item.actions[0]["target_binding"].get("phase", "final") != "staging" + ) + state = ExecutionState(last_qpos=env.robot.get_qpos()) + + def plan(grounded: GroundedAction, _state: ExecutionState) -> ActionOutcome: + index = int(grounded.motion_policy["placement_candidate_index"]) + return ActionOutcome( + trajectory=torch.zeros(1, 1, env.robot.dof), + success=torch.tensor([index == 1]), + next_state=state, + grounded=grounded, + ) + + monkeypatch.setattr(executor.adapter, "plan", plan) + grounded, outcome = executor._ground_and_plan_candidates( + edge.actions[0], + step, + arm="left_arm", + state=state, + active=torch.tensor([True]), + ) + + assert bool(outcome.success[0]) + assert grounded.motion_policy["placement_candidate_index"] == 1 + assert outcome.planner_trace["selected_grounding_candidate"] == 1 + assert len(outcome.planner_trace["grounding_candidates"]) == 2 + + +def test_post_release_candidate_search_skips_the_released_pose( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entities = { + "payload": _FakeEntity( + "payload", + _pose(0.0, -0.20, 0.79), + _rect_vertices(0.03, 0.03, 0.03), + ), + "support": _FakeEntity( + "support", + _pose(0.0, 0.0, 0.75), + _rect_vertices(0.20, 0.15, 0.01), + ), + } + env = _FakeEnv(entities) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "payload", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "support", + "relation": "on", + "orientation_goal": "none", + }, + "depends_on": [], + } + ) + ) + ), + env, + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + edge = next( + item + for item in executor.program.edges + if item.id in step.edge_ids + and item.actions[0]["target_binding"].get("kind") == "semantic_goal" + and item.actions[0]["target_binding"].get("phase", "final") != "staging" + ) + state = ExecutionState(last_qpos=env.robot.get_qpos()) + executor._placement_candidate_history[(step.id, "left_arm")] = {0} + + def plan(grounded: GroundedAction, _state: ExecutionState) -> ActionOutcome: + return ActionOutcome( + trajectory=torch.zeros(1, 1, env.robot.dof), + success=torch.tensor([True]), + next_state=state, + grounded=grounded, + ) + + monkeypatch.setattr(executor.adapter, "plan", plan) + grounded, outcome = executor._ground_and_plan_candidates( + edge.actions[0], + step, + arm="left_arm", + state=state, + active=torch.tensor([True]), + ) + + assert grounded.motion_policy["placement_candidate_index"] == 1 + assert outcome.planner_trace["grounding_candidates"][0] == { + "candidate_index": 0, + "status": "previously_released", + } + + +def test_unstable_placement_recovery_replays_pick_before_another_release( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entities = { + "payload": _FakeEntity( + "payload", + _pose(0.0, 0.0, 0.79), + _rect_vertices(0.03, 0.03, 0.03), + ), + "support": _FakeEntity( + "support", + _pose(0.0, 0.0, 0.75), + _rect_vertices(0.10, 0.08, 0.01), + ), + } + env = _FakeEnv(entities) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "payload", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "support", + "relation": "on", + "orientation_goal": "none", + }, + "depends_on": [], + } + ) + ) + ), + env, + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + replayed_actions: list[str] = [] + verification_count = 0 + + def ensure_assignment(_step: SemanticStep, failed: torch.Tensor) -> None: + executor._assignments[_step.id] = [ + None if bool(failed[env_id]) else "left_arm" + for env_id in range(len(failed)) + ] + + def execute( + edge: ExecutionEdge, + _step: SemanticStep, + *, + failed: torch.Tensor, + ) -> _EdgeResult: + replayed_actions.append(str(edge.actions[0]["atomic_action_class"])) + return _EdgeResult([], failed.clone(), [], executed=~failed) + + def verify( + _step: SemanticStep, + failed: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + nonlocal verification_count + verification_count += 1 + success = torch.tensor([verification_count == 2]) & ~failed + return failed | ~success, success, executor._entity_pose("payload")[:, :3, 3] + + monkeypatch.setattr(executor, "_ensure_assignment", ensure_assignment) + monkeypatch.setattr(executor, "_execute_edge_with_retries", execute) + monkeypatch.setattr(executor, "_verify_step", verify) + recorder = RuntimeRecorder( + executor.program, + num_envs=1, + enabled=False, + ) + + recovery = executor._recover_unstable_placement( + step, + torch.tensor([True]), + recorder=recorder, + ) + + first_action = str( + executor.edges[step.edge_ids[0]].actions[0]["atomic_action_class"] + ) + assert replayed_actions.count(first_action) == 2 + assert verification_count == 2 + assert bool(recovery.succeeded[0]) + assert not bool(recovery.failed[0]) + assert recovery.failure_events == [] + + +def test_unstable_placement_recovery_reports_its_own_planning_blocker( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entities = { + "payload": _FakeEntity( + "payload", + _pose(0.0, 0.0, 0.79), + _rect_vertices(0.03, 0.03, 0.03), + ), + "support": _FakeEntity( + "support", + _pose(0.0, 0.0, 0.75), + _rect_vertices(0.10, 0.08, 0.01), + ), + } + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "payload", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "support", + "relation": "on", + "orientation_goal": "none", + }, + "depends_on": [], + } + ) + ) + ), + _FakeEnv(entities), + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + + def fail_assignment(*_args: Any, **_kwargs: Any) -> None: + raise RuntimeError("no plan") + + monkeypatch.setattr(executor, "_ensure_assignment", fail_assignment) + + recovery = executor._recover_unstable_placement( + step, + torch.tensor([True]), + recorder=RuntimeRecorder(executor.program, num_envs=1, enabled=False), + ) + + assert bool(recovery.failed[0]) + assert bool(recovery.covered_failures[0]) + assert len(recovery.failure_events) == 1 + event = recovery.failure_events[0] + assert event["failure_type"] == "search_exhausted" + assert event["phase"] == "recovery" + assert event["origin_edge_id"] == step.edge_ids[-1] + assert event["blocking_edge_id"] == step.edge_ids[0] + + +def test_handover_height_accounts_for_obstacle_and_tool_envelope() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 1.03), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.35, 0.0, 1.03), _box_vertices(0.03)), + "shelf": _FakeEntity( + "shelf", + _pose(0.0, 0.0, 1.05), + _rect_vertices(0.40, 0.35, 0.05), + ), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + env.agent_robot_profile = "dual_franka" + program = load_execution_program( + instantiate_seed_graph( + _handover_then_place_task(), + {"can": "can", "target": "target"}, + ) + ) + step = next( + candidate for candidate in program.semantic_steps if candidate.id == "handover" + ) + state = _handover_held_state(env, entities["can"]) + held = state.get_held_object("physical_left_arm") + assert held is not None + grounder = ActionGrounder(program, env, lambda _uid: held.semantics) + staging = next( + edge.actions[0] + for edge in program.edges + if edge.id in step.edge_ids + and edge.actions[0]["target_binding"].get("kind") == "handover_staging" + ) + + grounded = grounder.ground(staging, step, arm="left_arm", state=state) + + assert grounded.target_object_pose is not None + obstacle_top = 1.10 + object_bottom = -0.03 + object_clearance = 0.06 + tool_vertical_envelope = 0.025 + 0.04 + expected_height = ( + obstacle_top + object_clearance + tool_vertical_envelope - object_bottom + ) + assert grounded.target_object_pose[0, 2, 3] == pytest.approx(expected_height) + + +def test_handover_workspace_rejects_points_outside_shared_reach() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 1.20), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.2, 0.0, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + program = load_execution_program( + instantiate_seed_graph( + _handover_then_place_task(), + {"can": "can", "target": "target"}, + ) + ) + step = next( + candidate for candidate in program.semantic_steps if candidate.id == "handover" + ) + state = _handover_held_state(env, entities["can"]) + held = state.get_held_object("physical_left_arm") + assert held is not None + grounder = ActionGrounder(program, env, lambda _uid: held.semantics) + staging = next( + edge.actions[0] + for edge in program.edges + if edge.id in step.edge_ids + and edge.actions[0]["target_binding"].get("kind") == "handover_staging" + ) + staging = { + **staging, + "motion_policy_config": {"exchange_maximum_reach": 0.20}, + } + + with pytest.raises(ValueError, match="reachable intersection"): + grounder.ground(staging, step, arm="left_arm", state=state) + + +def test_handover_grounding_preserves_the_original_object_affordance() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 1.20), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.2, 0.0, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + program = load_execution_program( + instantiate_seed_graph( + _handover_then_place_task(), + {"can": "can", "target": "target"}, + ) + ) + step = next( + candidate for candidate in program.semantic_steps if candidate.id == "handover" + ) + state = _handover_held_state(env, entities["can"]) + held = state.get_held_object("physical_left_arm") + assert held is not None + held.object_to_eef[:, 1, 3] *= -1.0 + grounder = ActionGrounder(program, env, lambda _uid: held.semantics) + action = next( + edge.actions[0] + for edge in program.edges + if edge.id in step.edge_ids + and edge.actions[0]["target_binding"].get("kind") == "handover_goal" + ) + + grounded = grounder.ground(action, step, arm="coordinated", state=state) + + assert grounded.target.semantics.affordance is held.semantics.affordance + + +def test_robot_relative_left_uses_live_right_to_left_arm_axis() -> None: + env = _FakeEnv() + + forward, lateral = robot_frame_axes(env) + offset = relation_offset( + env, + "left_of", + frame="robot", + forward_distance=0.10, + lateral_distance=0.12, + dtype=torch.float32, + device=env.device, + ) + + torch.testing.assert_close(forward, torch.tensor([[-1.0, 0.0]])) + torch.testing.assert_close(lateral, torch.tensor([[0.0, -1.0]])) + assert offset is not None + torch.testing.assert_close(offset, torch.tensor([[0.0, -0.12, 0.0]])) + + +def test_directional_verification_rejects_grounded_target_on_wrong_side() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.12, 0.75), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.0, 0.0, 0.75), _box_vertices(0.03)), + } + env = _FakeEnv(entities) + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "target", + "relation": "left_of", + }, + "depends_on": [], + } + ) + ) + ) + executor = ProgramExecutor(program, env, settle_steps=0, record_runtime=False) + step = program.semantic_steps[0] + step.goal["relation_frame"] = "robot" + executor._targets[step.id] = entities["can"].get_local_pose(to_matrix=True)[ + :, :3, 3 + ] + executor._policies[step.id] = { + "postcondition_tolerance": 0.08, + "relation_clearance": 0.01, + } + + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + + assert bool(failed[0]) + assert not bool(success[0]) + + +def test_directional_verification_accepts_support_height_settling() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, -0.12, 0.75), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.0, 0.0, 0.75), _box_vertices(0.03)), + } + env = _FakeEnv(entities) + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "target", + "relation": "left_of", + }, + "depends_on": [], + } + ) + ) + ) + executor = ProgramExecutor(program, env, settle_steps=0, record_runtime=False) + step = program.semantic_steps[0] + step.goal["relation_frame"] = "robot" + executor._targets[step.id] = torch.tensor([[0.0, -0.12, 0.90]]) + executor._policies[step.id] = { + "postcondition_tolerance": 0.08, + "relation_clearance": 0.01, + } + + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + + assert not bool(failed[0]) + assert bool(success[0]) + + +def test_legacy_released_above_relation_verifies_as_physical_support() -> None: + entities = { + "payload": _FakeEntity( + "payload", + _pose(0.0, 0.0, 0.79), + _rect_vertices(0.03, 0.03, 0.03), + ), + "support": _FakeEntity( + "support", + _pose(0.0, 0.0, 0.75), + _rect_vertices(0.10, 0.08, 0.01), + ), + } + env = _FakeEnv(entities) + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "payload", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "support", + "relation": "on", + }, + "depends_on": [], + } + ) + ) + ) + executor = ProgramExecutor(program, env, settle_steps=0, record_runtime=False) + step = program.semantic_steps[0] + step.goal["relation"] = "above" + executor._targets[step.id] = torch.tensor([[0.0, 0.0, 1.0]]) + + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + + assert not bool(failed[0]) + assert bool(success[0]) + + +def test_handover_retreat_clears_exchange_toward_transfer_workspace() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 1.106), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.2, 0.0, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + env.agent_robot_profile = "dual_franka" + high_left = _pose(0.0, 0.2, 1.106) + right = _pose(0.0, -0.2, 0.8) + env.get_current_xpos_agent = lambda: (high_left, right) + program = load_execution_program( + instantiate_seed_graph( + _handover_then_place_task(), + {"can": "can", "target": "target"}, + ) + ) + step = next( + candidate for candidate in program.semantic_steps if candidate.id == "handover" + ) + state = _held_state(env, entities["can"], arm="left_arm") + held = state.get_held_object("physical_left_arm") + assert held is not None + grounder = ActionGrounder(program, env, lambda _uid: held.semantics) + retreat_edge = next( + edge + for edge in program.edges + if edge.id in step.edge_ids + and edge.actions[0]["target_binding"].get("source") == "handover" + ) + + grounded = grounder.ground( + retreat_edge.actions[0], step, arm="left_arm", state=state + ) + + torch.testing.assert_close( + grounded.target.xpos[0, :2, 3], + torch.tensor([0.0, 0.10]), + ) + assert grounded.target.xpos[0, 2, 3] == pytest.approx(1.206) + assert grounded.cfg["retreat_distance"] == pytest.approx(0.10) + assert grounded.cfg["maximum_eef_height"] == pytest.approx(1.50) + + +def test_handover_retreat_and_home_block_receiver_continuation() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.2, 0.0, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + program = load_execution_program( + instantiate_seed_graph( + _handover_then_place_task(), + {"can": "can", "target": "target"}, + ) + ) + executor = ProgramExecutor(program, _FakeEnv(entities), record_runtime=False) + handover = next( + step for step in program.semantic_steps if step.operator == "handover" + ) + handover_edges = [edge for edge in program.edges if edge.id in handover.edge_ids] + retreat = next( + edge + for edge in handover_edges + if edge.actions[0]["target_binding"].get("source") == "handover" + ) + home = next( + edge + for edge in handover_edges + if edge.actions[0]["target_binding"].get("operation") == "handover_home" + ) + + assert executor._edge_failure_policy(retreat) == "safety_required" + assert executor._edge_failure_policy(home) == "best_effort" + + +def test_release_retreat_is_required_and_exact_home_is_best_effort() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.2, 0.0, 0.75), _box_vertices(0.03)), + } + program = load_execution_program( + compile_task_agent_v2( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "target", + "relation": "left_of", + }, + "depends_on": [], + } + ) + ) + ) + executor = ProgramExecutor(program, _FakeEnv(entities), record_runtime=False) + step = program.semantic_steps[0] + edges = [edge for edge in program.edges if edge.id in step.edge_ids] + retreat = next( + edge + for edge in edges + if edge.actions[0]["atomic_action_class"] == "MoveEndEffector" + ) + home = next( + edge for edge in edges if edge.actions[0]["atomic_action_class"] == "MoveJoints" + ) + + assert executor._edge_failure_policy(retreat) == "safety_required" + assert executor._edge_failure_policy(home) == "best_effort" + + +def test_best_effort_home_does_not_veto_required_arm_candidate( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, -0.2, 0.75), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.0, 0.0, 0.75), _box_vertices(0.03)), + } + graph = compile_task_agent_v2( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": {"reference_object": "target", "relation": "left_of"}, + "depends_on": [], + } + ) + ) + executor = ProgramExecutor( + load_execution_program(graph), + _FakeEnv(entities), + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + + def ground(action: dict[str, Any], *_args: Any, **_kwargs: Any) -> GroundedAction: + return GroundedAction( + action_class=str(action["atomic_action_class"]), + arm="left_arm", + control=str(action["control"]), + target=SimpleNamespace(xpos=None), + cfg={}, + ) + + def plan(grounded: GroundedAction, state: ExecutionState) -> ActionOutcome: + return ActionOutcome( + trajectory=torch.zeros(1, 1, executor.env.robot.dof), + success=torch.tensor([grounded.action_class != "MoveJoints"]), + next_state=state, + grounded=grounded, + planner_trace={"primary_strategy": "motion_gen"}, + ) + + monkeypatch.setattr(executor.grounder, "ground", ground) + monkeypatch.setattr(executor, "_with_downstream_targets", lambda *args: args[-1]) + monkeypatch.setattr(executor.adapter, "plan", plan) + + candidate = executor._candidate(step, "left_arm", torch.tensor([False])) + + assert bool(candidate.feasible[0]) + assert any("best-effort action degraded" in item for item in candidate.warnings) + + +def test_best_effort_home_exception_does_not_fail_semantic_step( + monkeypatch: pytest.MonkeyPatch, +) -> None: + graph = compile_task_agent_v2( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": {"reference_object": "target", "relation": "left_of"}, + "depends_on": [], + } + ) + ) + executor = ProgramExecutor( + load_execution_program(graph), + _FakeEnv(), + settle_steps=0, + record_runtime=False, + ) + monkeypatch.setattr( + executor, + "_ensure_assignment", + lambda step, _failed: executor._assignments.setdefault(step.id, ["left_arm"]), + ) + + def execute(edge: ExecutionEdge, _step: SemanticStep, *, failed: torch.Tensor): + if executor._edge_failure_policy(edge) == "best_effort": + raise RuntimeError("home search failed") + return SimpleNamespace( + actions=[], + failed=failed.clone(), + grounded=[], + planner_traces=[], + executed=~failed, + ) + + monkeypatch.setattr(executor, "_execute_edge_with_retries", execute) + monkeypatch.setattr( + executor, + "_verify_step", + lambda _step, failed: (failed, ~failed, torch.zeros(1, 3)), + ) + + result = executor.run() + + assert bool(result.success[0]) + assert len(result.failure_events) == 1 + assert result.failure_events[0]["failure_type"] == "search_exhausted" + assert result.failure_events[0]["failure_policy"] == "best_effort" + assert result.failure_events[0]["fatal"] is False + assert result.failure_events[0]["evidence"]["exception"].endswith( + "home search failed" + ) + + +def test_candidate_failure_reports_real_blocking_safety_edge( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, -0.2, 0.75), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.0, 0.0, 0.75), _box_vertices(0.03)), + } + executor = ProgramExecutor( + load_execution_program( + compile_task_agent_v2( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "target", + "relation": "left_of", + }, + "depends_on": [], + } + ) + ) + ), + _FakeEnv(entities), + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + + def ground(action: dict[str, Any], *_args: Any, **_kwargs: Any) -> GroundedAction: + return GroundedAction( + action_class=str(action["atomic_action_class"]), + arm="left_arm", + control=str(action["control"]), + target=SimpleNamespace(xpos=None), + cfg={}, + ) + + def plan(grounded: GroundedAction, state: ExecutionState) -> ActionOutcome: + return ActionOutcome( + trajectory=torch.zeros(1, 1, executor.env.robot.dof), + success=torch.tensor([grounded.action_class != "MoveEndEffector"]), + next_state=state, + grounded=grounded, + planner_trace={"primary_strategy": "motion_gen"}, + ) + + monkeypatch.setattr(executor.grounder, "ground", ground) + monkeypatch.setattr(executor, "_with_downstream_targets", lambda *args: args[-1]) + monkeypatch.setattr(executor.adapter, "plan", plan) + candidate = executor._candidate(step, "left_arm", torch.tensor([False])) + executor._assignments[step.id] = [None] + executor._report_candidates(step, (candidate,)) + first_edge = executor.edges[step.edge_ids[0]] + + events = executor._failure_events( + first_edge, + step, + torch.tensor([True]), + postcondition=False, + executed=torch.tensor([False]), + fallen_transition=torch.tensor([False]), + ) + + assert len(events) == 1 + event = events[0] + assert event["failure_type"] == "search_exhausted" + assert event["failure_policy"] == "safety_required" + assert event["atomic_action"] == "MoveEndEffector" + assert event["blocking_edge_id"] != first_edge.id + assert event["planning_stage"] == "candidate_suffix" + assert "not a geometric proof" in event["reason"] + + +def test_on_relation_rejects_preserve_orientation_drift() -> None: + rotated = _pose(0.0, 0.0, 0.82) + rotated[:, :3, :3] = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + entities = { + "can": _FakeEntity("can", rotated, _rect_vertices(0.03, 0.03, 0.06)), + "notebook": _FakeEntity( + "notebook", + _pose(0.0, 0.0, 0.75), + _rect_vertices(0.10, 0.08, 0.01), + ), + } + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "notebook", + "relation": "on", + "orientation_goal": "preserve", + }, + "depends_on": [], + } + ) + ) + ), + _FakeEnv(entities), + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + executor._orientation_references[step.id] = _pose(0.0, 0.0, 0.82) + executor._policies[step.id] = { + "preserve_orientation_tolerance": torch.pi / 12, + } + + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + + assert bool(failed[0]) + assert not bool(success[0]) + assert executor._orientation_errors[step.id][0] > torch.pi / 12 + + +def test_inside_relation_accepts_settling_orientation_drift() -> None: + rotated = _pose(0.02, -0.02, 0.72) + rotated[:, :3, :3] = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + entities = { + "can": _FakeEntity("can", rotated, _rect_vertices(0.03, 0.03, 0.06)), + "basket": _FakeEntity( + "basket", + _pose(0.0, 0.0, 0.70), + _rect_vertices(0.10, 0.10, 0.08), + ), + } + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "basket", + "relation": "inside", + "orientation_goal": "preserve", + }, + "depends_on": [], + } + ) + ) + ), + _FakeEnv(entities), + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + executor._orientation_references[step.id] = _pose(0.02, -0.02, 0.72) + + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + + assert not bool(failed[0]) + assert bool(success[0]) + assert step.id not in executor._orientation_errors + + +@pytest.mark.parametrize( + ("orientation_goal", "expected_success"), + (("upright", False), ("none", True)), +) +def test_on_relation_applies_only_the_requested_orientation_goal( + orientation_goal: str, + expected_success: bool, +) -> None: + fallen = _pose(0.0, 0.0, 0.79) + fallen[:, :3, :3] = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + entities = { + "can": _FakeEntity("can", fallen, _rect_vertices(0.03, 0.03, 0.06)), + "notebook": _FakeEntity( + "notebook", + _pose(0.0, 0.0, 0.75), + _rect_vertices(0.10, 0.08, 0.01), + ), + } + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "notebook", + "relation": "on", + "orientation_goal": orientation_goal, + }, + "depends_on": [], + } + ) + ) + ), + _FakeEnv(entities), + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + + assert bool(success[0]) is expected_success + assert bool(failed[0]) is not expected_success + + +def test_support_stability_window_rejects_motion_after_initial_contact() -> None: + payload = _FakeEntity( + "payload", + _pose(0.0, 0.0, 0.79), + _rect_vertices(0.03, 0.03, 0.03), + ) + support = _FakeEntity( + "support", + _pose(0.0, 0.0, 0.75), + _rect_vertices(0.10, 0.08, 0.01), + ) + env = _FakeEnv({"payload": payload, "support": support}) + update_count = 0 + + def update(*, step: int) -> None: + nonlocal update_count + del step + update_count += 1 + payload.lin_vel[:, 0] = 0.10 + + env.sim.update = update + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "payload", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "support", + "relation": "on", + "orientation_goal": "none", + }, + "depends_on": [], + } + ) + ) + ), + env, + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + + assert update_count == executor.support_stability_samples - 1 + assert bool(failed[0]) + assert not bool(success[0]) + + +def test_support_stability_reads_real_rigid_object_body_state() -> None: + payload = _FakeEntity( + "payload", + _pose(0.0, 0.0, 0.79), + _rect_vertices(0.03, 0.03, 0.03), + ) + del payload.lin_vel + del payload.ang_vel + payload.body_state = torch.zeros(1, 13) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "payload", "left_arm"))) + ), + _FakeEnv({"payload": payload}), + settle_steps=0, + record_runtime=False, + ) + + assert bool(executor._entity_motion_stable("payload")[0]) + payload.body_state[:, 7] = 0.10 + assert not bool(executor._entity_motion_stable("payload")[0]) + + +def test_final_support_revalidation_detects_later_chain_damage() -> None: + payload = _FakeEntity( + "payload", + _pose(0.0, 0.0, 0.79), + _rect_vertices(0.03, 0.03, 0.03), + ) + support = _FakeEntity( + "support", + _pose(0.0, 0.0, 0.75), + _rect_vertices(0.10, 0.08, 0.01), + ) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "payload", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "support", + "relation": "on", + "orientation_goal": "none", + }, + "depends_on": [], + } + ) + ) + ), + _FakeEnv({"payload": payload, "support": support}), + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + assert not bool(failed[0]) + assert bool(success[0]) + + payload._pose[:, 2, 3] += 0.20 + failures = executor._revalidate_support_relations() + + assert bool(failures[step.id][0]) + + +def test_support_relation_state_rejects_cycles() -> None: + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place_b", + "operator": "place_relative", + "object": "b", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "a", + "relation": "on", + "orientation_goal": "none", + }, + "depends_on": [], + } + ) + ) + ), + _FakeEnv( + { + "a": _FakeEntity("a", _pose(0.0, 0.0, 0.75), _box_vertices(0.03)), + "b": _FakeEntity("b", _pose(0.0, 0.0, 0.81), _box_vertices(0.03)), + } + ), + settle_steps=0, + record_runtime=False, + ) + step_b = executor.program.semantic_steps[0] + step_a = replace(step_b, id="prior", object_uid="a") + executor._commit_support_relation(step_a, "b", torch.tensor([True])) + + cycle_free = executor._support_cycle_free("b", "a", torch.tensor([True])) + + assert not bool(cycle_free[0]) + + +def test_standalone_handover_assigns_its_pickup_candidate( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + task = _handover_then_place_task() + task["level"] = "L1" + task["task_instances"] = [task["task_instances"][0]] + task["success"] = {"type": "handover_complete"} + program = load_execution_program(instantiate_seed_graph(task, {"can": "can"})) + executor = ProgramExecutor(program, _FakeEnv(entities), record_runtime=False) + step = program.semantic_steps[0] + calls: list[str] = [] + + def candidate(_step: SemanticStep, arm: str, _failed: torch.Tensor) -> Any: + calls.append(arm) + return SimpleNamespace(feasible=torch.tensor([True])) + + monkeypatch.setattr(executor, "_candidate", candidate) + monkeypatch.setattr(executor, "_report_candidates", lambda *_args: None) + + executor._ensure_assignment(step, torch.zeros(1, dtype=torch.bool)) + + assert calls == ["left_arm"] + assert executor._assignments[step.id] == ["left_arm"] + + +def test_standalone_handover_candidate_stops_before_coordinated_transfer( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + task = _handover_then_place_task() + task["level"] = "L1" + task["task_instances"] = [task["task_instances"][0]] + task["success"] = {"type": "handover_complete"} + env = _FakeEnv(entities) + program = load_execution_program(instantiate_seed_graph(task, {"can": "can"})) + executor = ProgramExecutor(program, env, record_runtime=False) + step = program.semantic_steps[0] + planned_actions: list[str] = [] + + def ground( + action: dict[str, Any], + _step: SemanticStep, + *, + arm: str, + state: ExecutionState, + reference_eef_pose: torch.Tensor | None = None, + orientation_reference_pose: torch.Tensor | None = None, + ) -> GroundedAction: + del state, reference_eef_pose, orientation_reference_pose + return GroundedAction( + action_class=str(action["atomic_action_class"]), + arm=arm, + control=str(action["control"]), + target=SimpleNamespace(), + cfg={}, + ) + + def plan(grounded: GroundedAction, state: ExecutionState) -> ActionOutcome: + planned_actions.append(grounded.action_class) + return ActionOutcome( + trajectory=torch.zeros(1, 1, env.robot.dof), + success=torch.tensor([True]), + next_state=state, + grounded=grounded, + ) + + monkeypatch.setattr(executor.grounder, "ground", ground) + monkeypatch.setattr(executor, "_with_downstream_targets", lambda *args: args[-1]) + monkeypatch.setattr(executor.adapter, "plan", plan) + + candidate = executor._candidate( + step, + "left_arm", + torch.zeros(1, dtype=torch.bool), + ) + + assert bool(candidate.feasible[0]) + assert planned_actions == ["PickUp", "MoveHeldObject"] + assert set(candidate.plans) == set(step.edge_ids[:2]) + + +def test_failed_handover_keeps_transfer_ownership( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + task = _handover_then_place_task() + task["level"] = "L1" + task["task_instances"] = [task["task_instances"][0]] + task["success"] = {"type": "handover_complete"} + program = load_execution_program( + instantiate_seed_graph( + task, + {"can": "can"}, + ) + ) + step = program.semantic_steps[0] + edge = next( + candidate + for candidate in program.edges + if candidate.actions[0]["atomic_action_class"] == "HandOver" + ) + state = _held_state(env, entities["can"], arm="left_arm") + executor = ProgramExecutor(program, env, record_runtime=False) + executor._assignments[step.id] = ["coordinated"] + executor._object_owners["can"] = ["left_arm"] + executor._arm_owners["left_arm"] = ["can"] + executor._object_states[("can", "left_arm")] = state + grounded = GroundedAction( + action_class="HandOver", + arm="coordinated", + control="coordinated", + target=SimpleNamespace(), + cfg={}, + ) + failed_outcome = ActionOutcome( + trajectory=torch.zeros(1, 1, env.robot.dof), + success=torch.tensor([False]), + next_state=state, + grounded=grounded, + ) + monkeypatch.setattr( + executor.grounder, + "ground", + lambda *_args, **_kwargs: grounded, + ) + monkeypatch.setattr( + executor.adapter, "plan", lambda *_args, **_kwargs: failed_outcome + ) + monkeypatch.setattr( + executor.adapter, + "execute_trajectory", + lambda *_args, **_kwargs: [], + ) + + result = executor._execute_coordinated(edge, step, torch.tensor([False])) + + assert bool(result.failed[0]) + assert executor._object_owners["can"] == ["left_arm"] + assert executor._arm_owners["left_arm"] == ["can"] + assert executor._arm_owners["right_arm"] == [None] + assert ("can", "left_arm") in executor._object_states + assert ("can", "right_arm") not in executor._object_states + + +def test_handover_commits_receiver_ownership_only_after_physical_verification( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + task = _handover_then_place_task() + task["level"] = "L1" + task["task_instances"] = [task["task_instances"][0]] + task["success"] = {"type": "handover_complete"} + program = load_execution_program(instantiate_seed_graph(task, {"can": "can"})) + step = program.semantic_steps[0] + edge = next( + candidate + for candidate in program.edges + if candidate.actions[0]["atomic_action_class"] == "HandOver" + ) + transfer_state = _held_state(env, entities["can"], arm="left_arm") + receiver_state = _held_state(env, entities["can"], arm="right_arm") + executor = ProgramExecutor(program, env, record_runtime=False) + executor._assignments[step.id] = ["coordinated"] + executor._object_owners["can"] = ["left_arm"] + executor._arm_owners["left_arm"] = ["can"] + executor._object_states[("can", "left_arm")] = transfer_state + grounded = GroundedAction( + action_class="HandOver", + arm="coordinated", + control="coordinated", + target=SimpleNamespace(), + cfg={}, + motion_policy={"held_position_tolerance": 0.03}, + ) + successful_outcome = ActionOutcome( + trajectory=torch.zeros(1, 1, env.robot.dof), + success=torch.tensor([True]), + next_state=receiver_state, + grounded=grounded, + ) + observed_poses: list[torch.Tensor] = [] + + def ground_candidates(*_args, **_kwargs): + observed_poses.append(entities["can"].get_local_pose(to_matrix=True)) + return (grounded,) + + monkeypatch.setattr( + executor.grounder, + "ground_candidates", + ground_candidates, + ) + monkeypatch.setattr( + executor.adapter, "plan", lambda *_args, **_kwargs: successful_outcome + ) + monkeypatch.setattr( + executor.adapter, + "execute_trajectory", + lambda *_args, **_kwargs: [], + ) + + entities["can"]._pose[:, 0, 3] += 0.30 + result = executor._execute_coordinated(edge, step, torch.tensor([False])) + + assert observed_poses[0][0, 0, 3] == pytest.approx(0.30) + assert result.planner_traces[0]["execution_replanned_from_live_state"] is True + assert bool(result.failed[0]) + assert executor._object_owners["can"] == [None] + assert executor._arm_owners["left_arm"] == [None] + assert executor._arm_owners["right_arm"] == [None] + assert ("can", "right_arm") not in executor._object_states + + +def test_handover_defers_clearance_verification_to_retreat_action() -> None: + adapter = AtomicActionAdapter(_FakeEnv()) + + assert adapter.capabilities.get("HandOver").verifier_hook is None + assert adapter.capabilities.get("MoveEndEffector").verifier_hook is not None + + +def test_orient_then_handover_reacquires_with_a_separate_transfer_policy() -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv( + { + "can": entity, + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + ) + task = { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "orient_then_handover", + "level": "L3", + "instruction": "Orient the can, then hand it over.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "task_01", + "task_type": "E2", + "params": { + "object_role": "can", + "required_arm": "right_arm", + "orientation_goal": "upright", + "support_role": "table", + }, + "depends_on": [], + "role": "primary", + }, + { + "id": "task_02", + "task_type": "E4", + "params": { + "object_role": "can", + "transfer_arm": "right_arm", + "receive_arm": "left_arm", + }, + "depends_on": ["task_01"], + "role": "primary", + }, + ], + "success": {"type": "handover_complete", "task_instance_id": "task_02"}, + "oracle": {}, + "metadata": {}, + } + program = load_execution_program(instantiate_seed_graph(task, {"can": "can"})) + orient_step = next( + candidate for candidate in program.semantic_steps if candidate.id == "task_01" + ) + handover_step = next( + candidate for candidate in program.semantic_steps if candidate.id == "task_02" + ) + orient_edge = next( + candidate + for candidate in program.edges + if candidate.id in orient_step.edge_ids + and candidate.actions[0]["atomic_action_class"] == "PickUp" + ) + handover_edge = next( + candidate + for candidate in program.edges + if candidate.id in handover_step.edge_ids + if candidate.actions[0]["atomic_action_class"] == "PickUp" + ) + semantics = ObjectSemantics( + affordance=AntipodalAffordance(object_label="can"), + geometry={}, + label="can", + entity=entity, + ) + grounder = ActionGrounder(program, env, lambda _uid: semantics) + + orient_pickup = grounder.ground( + orient_edge.actions[0], + orient_step, + arm="right_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + handover_pickup = grounder.ground( + handover_edge.actions[0], + handover_step, + arm="right_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + + assert "approach_direction_mode" not in orient_pickup.cfg + assert "approach_direction_mode" not in handover_pickup.cfg + assert handover_pickup.cfg["pick_object_part"] == "top" + assert orient_pickup.target.grasp_xpos is None + assert orient_pickup.target.semantics.affordance is semantics.affordance + assert handover_pickup.target.grasp_xpos is None + assert isinstance( + handover_pickup.target.semantics.affordance, + AntipodalAffordance, + ) + assert handover_pickup.target.semantics.affordance is semantics.affordance + + +def test_handover_clearance_verifier_checks_distance_and_transfer_side() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, -0.2, 1.0), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.2, 0.0, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + program = load_execution_program( + instantiate_seed_graph( + _handover_then_place_task(), + {"can": "can", "target": "target"}, + ) + ) + executor = ProgramExecutor(program, env, record_runtime=False) + hook = executor.adapter.capabilities.get("MoveEndEffector").verifier_hook + assert hook is not None + _, lateral = robot_frame_axes(env) + grounded = GroundedAction( + action_class="MoveEndEffector", + arm="left_arm", + control="arm", + target=SimpleNamespace(), + cfg={}, + motion_policy={ + "clearance_object_uid": "can", + "transfer_arm": "left_arm", + "transfer_role_axis": torch.cat( + (lateral, lateral.new_zeros((1, 1))), dim=1 + ), + "minimum_transfer_clearance": 0.10, + "minimum_transfer_lateral_clearance": 0.06, + }, + ) + outcome = SimpleNamespace(grounded=grounded) + attempted = torch.tensor([True]) + + assert not bool( + hook( + executor=executor, + step=program.semantic_steps[0], + arm="left_arm", + outcome=outcome, + attempted=attempted, + )[0] + ) + + clear_left = _pose(0.0, -0.4, 1.0) + env.get_current_xpos_agent = lambda: (clear_left, _pose(0.0, 0.2, 1.0)) + assert bool( + hook( + executor=executor, + step=program.semantic_steps[0], + arm="left_arm", + outcome=outcome, + attempted=attempted, + )[0] + ) + + +@pytest.mark.parametrize("arm", ["left_arm", "right_arm"]) +def test_handover_source_policy_uses_pickup_default_top_down_approach( + arm: str, +) -> None: + action = GroundedAction( + action_class="PickUp", + arm=arm, + control="arm", + target=SimpleNamespace(), + cfg={"pick_object_part": "top"}, + ) + + cfg = AtomicActionAdapter(_FakeEnv())._build_config(action, PickUpOptions) + + assert cfg.pick_object_part == "top" + torch.testing.assert_close( + cfg.approach_direction, + torch.tensor([0.0, 0.0, -1.0]), + ) + + +@pytest.mark.parametrize( + ("transfer_arm", "expected_world_y"), + [("left_arm", -1.0), ("right_arm", 1.0)], +) +def test_handover_receiver_uses_the_mirrored_diagonal_approach( + transfer_arm: str, + expected_world_y: float, +) -> None: + env = _FakeEnv() + action = GroundedAction( + action_class="HandOver", + arm="coordinated", + control="coordinated", + target=SimpleNamespace(), + cfg={ + "transfer_arm": transfer_arm, + "middle_object_pose": torch.eye(4).unsqueeze(0), + "final_object_pose": torch.eye(4).unsqueeze(0), + }, + ) + + cfg = AtomicActionAdapter(env)._build_config(action, HandOverOptions) + + diagonal = 2.0**-0.5 + assert cfg.receive_approach_direction[0] == pytest.approx(0.0) + assert cfg.receive_approach_direction[1] == pytest.approx( + expected_world_y * diagonal + ) + assert cfg.receive_approach_direction[2] == pytest.approx(-diagonal) + _, lateral = robot_frame_axes(env) + receive_side = "right_arm" if transfer_arm == "left_arm" else "left_arm" + receiver_outward = lateral[0] if receive_side == "left_arm" else -lateral[0] + pre_grasp_offset = -cfg.receive_approach_direction[:2] * cfg.pre_grasp_distance + assert torch.dot(pre_grasp_offset, receiver_outward) > 0.0 + + +@pytest.mark.parametrize( + ("transfer_arm", "expected_x"), + [("left_arm", 1.0), ("right_arm", -1.0)], +) +def test_handover_receiver_approach_tracks_rotated_robot_lateral_axis( + monkeypatch: pytest.MonkeyPatch, + transfer_arm: str, + expected_x: float, +) -> None: + env = _FakeEnv() + + def get_link_pose(*, link_name: str, to_matrix: bool) -> torch.Tensor: + assert to_matrix + pose = torch.eye(4).unsqueeze(0) + pose[:, 0, 3] = 0.3 if link_name == "physical_left_base" else -0.3 + return pose + + monkeypatch.setattr(env.robot, "get_link_pose", get_link_pose) + action = GroundedAction( + action_class="HandOver", + arm="coordinated", + control="coordinated", + target=SimpleNamespace(), + cfg={ + "transfer_arm": transfer_arm, + "middle_object_pose": torch.eye(4).unsqueeze(0), + "final_object_pose": torch.eye(4).unsqueeze(0), + }, + ) + + cfg = AtomicActionAdapter(env)._build_config(action, HandOverOptions) + + diagonal = 2.0**-0.5 + assert cfg.receive_approach_direction[0] == pytest.approx(expected_x * diagonal) + assert cfg.receive_approach_direction[1] == pytest.approx(0.0) + assert cfg.receive_approach_direction[2] == pytest.approx(-diagonal) + + +@pytest.mark.parametrize( + ("arm", "outward_x"), + [("left_arm", 1.0), ("right_arm", -1.0)], +) +def test_legacy_handover_transfer_mode_tracks_a_rotated_live_base_line( + monkeypatch: pytest.MonkeyPatch, + arm: str, + outward_x: float, +) -> None: + env = _FakeEnv() + + def get_link_pose(*, link_name: str, to_matrix: bool) -> torch.Tensor: + assert to_matrix + pose = torch.eye(4).unsqueeze(0) + pose[:, 0, 3] = 0.3 if link_name == "physical_left_base" else -0.3 + return pose + + monkeypatch.setattr(env.robot, "get_link_pose", get_link_pose) + action = GroundedAction( + action_class="PickUp", + arm=arm, + control="arm", + target=SimpleNamespace(), + cfg={"approach_direction_mode": "handover_transfer"}, + ) + + cfg = AtomicActionAdapter(env)._build_config(action, PickUpOptions) + + outward = torch.tensor([outward_x, 0.0]) + assert torch.dot(cfg.approach_direction[:2], outward) < 0.0 + assert cfg.approach_direction[2] < 0.0 + + +def test_pickup_is_replanned_from_live_pose_and_screens_downstream_targets( + monkeypatch: Any, +) -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ), + env, + record_runtime=False, + ) + plan_calls: list[GroundedAction] = [] + + def ground( + action, + step, + *, + arm, + state, + reference_eef_pose=None, + orientation_reference_pose=None, + ): + del reference_eef_pose, orientation_reference_pose + action_class = action["atomic_action_class"] + target_pose = ( + _pose(0.0, 0.2, 0.85) if action_class == "MoveHeldObject" else None + ) + return GroundedAction( + action_class=action_class, + arm=arm, + control=str(action.get("control", "arm")), + target=SimpleNamespace(xpos=None), + cfg={"planned_object_pose": entity.get_local_pose(to_matrix=True)}, + target_object_pose=target_pose, + ) + + def plan(grounded: GroundedAction, state: ExecutionState) -> ActionOutcome: + plan_calls.append(grounded) + next_state = ( + _held_state(env, entity) if grounded.action_class == "PickUp" else state + ) + return ActionOutcome( + trajectory=torch.zeros(1, 1, env.robot.dof), + success=torch.tensor([True]), + next_state=next_state, + grounded=grounded, + ) + + monkeypatch.setattr(executor.grounder, "ground", ground) + monkeypatch.setattr(executor.adapter, "plan", plan) + monkeypatch.setattr(executor.adapter, "execute_trajectory", lambda *a, **k: []) + step = executor.program.semantic_steps[0] + failed = torch.tensor([False]) + + executor._ensure_assignment(step, failed) + planned_call_count = len(plan_calls) + executor._candidate_cache.clear() + entity._pose[:, 0, 3] += 0.25 + edge_result = executor._execute_edge( + executor.edges[step.edge_ids[0]], step, failed=failed + ) + + assert len(plan_calls) == planned_call_count + 1 + assert planned_call_count == len(step.edge_ids) + assert len(plan_calls[0].cfg["downstream_object_target_poses"]) == 1 + assert plan_calls[-1].cfg["planned_object_pose"][0, 0, 3] == pytest.approx(0.25) + assert plan_calls[-1].cfg["downstream_object_target_poses"] + assert edge_result.planner_traces[0]["execution_replanned_from_live_state"] + assert not edge_result.planner_traces[0]["speculative_candidate_available"] + assert bool(edge_result.failed[0]) + assert executor._object_owners["can"] == [None] + + +def test_live_pickup_planning_exception_is_a_retryable_edge_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ), + _FakeEnv({"can": entity}), + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + edge = executor.edges[step.edge_ids[0]] + executor._assignments[step.id] = ["left_arm"] + monkeypatch.setattr( + executor.grounder, + "ground", + lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("no IK")), + ) + + result = executor._execute_edge(edge, step, failed=torch.tensor([False])) + + assert bool(result.failed[0]) + assert result.actions == [] + assert result.planner_traces[0]["primary_strategy"] == "live_pickup_replan" + assert result.planner_traces[0]["exception"] == "RuntimeError: no IK" + + +def test_pickup_candidate_screens_handover_successor_target( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A later handover staging pose participates in pickup grasp screening.""" + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("pickup", "can", "left_arm"))) + ), + env, + record_runtime=False, + ) + pickup_step = executor.program.semantic_steps[0] + handover_step = SemanticStep( + id="handover", + parent_step_id="handover", + operator="handover", + object_uid="can", + actor={"mode": "required", "arm": "left_arm"}, + goal={"transfer_arm": "left_arm", "receive_arm": "right_arm"}, + depends_on=(pickup_step.id,), + postcondition={}, + edge_ids=("handover_staging",), + ) + handover_edge = ExecutionEdge( + id="handover_staging", + source="pickup_done", + target="handover_done", + actions=( + { + "atomic_action_class": "MoveHeldObject", + "actor": {"mode": "required", "arm": "left_arm"}, + "control": "arm", + "target_binding": { + "kind": "handover_staging", + "transfer_arm": "left_arm", + "receive_arm": "right_arm", + }, + "motion_policy": motion_policy(), + }, + ), + ) + executor.steps[handover_step.id] = handover_step + executor.edges[handover_edge.id] = handover_edge + + existing_target = _pose(0.0, 0.2, 0.85) + handover_target = _pose(0.0, 0.0, 1.15) + grounded = GroundedAction( + action_class="PickUp", + arm="left_arm", + control="arm", + target=SimpleNamespace(), + cfg={"downstream_object_target_poses": (existing_target,)}, + ) + + def ground( + _action: Any, + candidate: SemanticStep, + *, + arm: str, + state: ExecutionState, + reference_eef_pose: torch.Tensor | None = None, + orientation_reference_pose: torch.Tensor | None = None, + ) -> GroundedAction: + del arm, state, reference_eef_pose, orientation_reference_pose + target = handover_target if candidate.id == handover_step.id else None + return GroundedAction( + action_class="MoveHeldObject", + arm="left_arm", + control="arm", + target=SimpleNamespace(), + cfg={}, + target_object_pose=target, + ) + + monkeypatch.setattr(executor.grounder, "ground", ground) + result = executor._with_downstream_targets( + pickup_step, + pickup_step.edge_ids[0], + "left_arm", + ExecutionState(last_qpos=env.robot.get_qpos()), + grounded, + ) + + targets = result.cfg["downstream_object_target_poses"] + assert len(targets) == 2 + assert torch.equal(targets[0], existing_target) + assert torch.equal(targets[1], handover_target) + + +def test_object_held_predicate_checks_live_gripper_and_tcp_geometry() -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + env.robot._qpos[:, env.left_eef_joints] = env.close_state + state = _held_state(env, entity) + left_eef, _ = env.get_current_xpos_agent() + env.get_current_xpos_agent = lambda: (left_eef, None) + + held = evaluate_predicate( + env, + {"type": "object_held", "object": "can"}, + held_owners={"can": ["left_arm"]}, + held_states={("can", "left_arm"): state}, + ) + + assert bool(held[0]) + env.robot._qpos[:, env.left_eef_joints] = env.open_state + assert not bool( + evaluate_predicate( + env, + {"type": "object_held", "object": "can"}, + held_owners={"can": ["left_arm"]}, + held_states={("can", "left_arm"): state}, + )[0] + ) + env.robot._qpos[:, env.left_eef_joints] = (env.open_state + env.close_state) / 2 + assert bool( + evaluate_predicate( + env, + {"type": "object_held", "object": "can"}, + held_owners={"can": ["left_arm"]}, + held_states={("can", "left_arm"): state}, + )[0] + ) + env.close_state = torch.tensor([0.0, 0.0]) + env.open_state = torch.tensor([0.04, 0.04]) + env.robot._qpos[:, env.left_eef_joints] = env.open_state + assert not bool( + evaluate_predicate( + env, + {"type": "object_held", "object": "can"}, + held_owners={"can": ["left_arm"]}, + held_states={("can", "left_arm"): state}, + )[0] + ) + + +def test_coordinated_held_predicate_uses_per_arm_held_relations() -> None: + entity = _FakeEntity("tray", _pose(0.0, 0.0, 0.75), _box_vertices(0.2)) + env = _FakeEnv({"tray": entity}) + env.robot._qpos[:, env.left_eef_joints] = env.close_state + env.robot._qpos[:, env.right_eef_joints] = env.close_state + state = _coordinated_held_state(env, entity) + predicate = {"type": "object_held_by_both_grippers", "object": "tray"} + + assert bool(evaluate_predicate(env, predicate, coordinated_state=state)[0]) + + env.robot._qpos[:, env.right_eef_joints] = env.open_state + assert not bool(evaluate_predicate(env, predicate, coordinated_state=state)[0]) + + +def test_object_supported_by_requires_overlap_and_vertical_contact() -> None: + support_z = 0.75 + payload_z = support_z + 0.05 + 0.02 + 0.005 + payload = _FakeEntity( + "payload", + _pose(0.002, -0.002, payload_z), + _box_vertices(0.02), + ) + support = _FakeEntity( + "support", + _pose(0.0, 0.0, support_z), + _box_vertices(0.05), + ) + env = _FakeEnv({"payload": payload, "support": support}) + predicate = { + "type": "object_supported_by", + "object": "payload", + "support": "support", + } + + assert bool(evaluate_predicate(env, predicate)[0]) + + payload._pose = _pose(0.002, -0.002, payload_z - 1.0) + assert not bool(evaluate_predicate(env, predicate)[0]) + + payload._pose = _pose(0.002, -0.002, payload_z + 1.0) + assert not bool(evaluate_predicate(env, predicate)[0]) + + payload._pose = _pose(0.081, 0.0, payload_z) + assert not bool(evaluate_predicate(env, predicate)[0]) + + +def test_object_supported_by_uses_local_not_mesh_wide_support_height() -> None: + support_vertices = torch.tensor( + [ + [-0.10, -0.10, -0.05], + [-0.02, -0.02, 0.05], + [0.02, -0.02, 0.05], + [0.02, 0.02, 0.05], + [-0.02, 0.02, 0.05], + [0.40, 0.00, 0.40], + ], + dtype=torch.float32, + ) + payload = _FakeEntity( + "payload", + _pose(0.0, 0.0, 0.075), + _box_vertices(0.02), + ) + support = _FakeEntity("support", _pose(0.0, 0.0, 0.0), support_vertices) + env = _FakeEnv({"payload": payload, "support": support}) + + supported = evaluate_predicate( + env, + { + "type": "object_supported_by", + "object": "payload", + "support": "support", + }, + ) + + assert bool(supported[0]) + + +def test_object_supported_by_uses_live_center_of_mass_projection() -> None: + payload = _FakeEntity( + "payload", + _pose(0.04, 0.0, 0.125), + _rect_vertices(0.08, 0.02, 0.02), + ) + payload.body_data = SimpleNamespace( + com_pose=torch.tensor([[0.04, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]]) + ) + support = _FakeEntity( + "support", + _pose(0.0, 0.0, 0.05), + _rect_vertices(0.05, 0.05, 0.05), + ) + env = _FakeEnv({"payload": payload, "support": support}) + + supported = evaluate_predicate( + env, + { + "type": "object_supported_by", + "object": "payload", + "support": "support", + }, + ) + + assert not bool(supported[0]) + + +def test_physical_pickup_rebases_a_compliant_grasp_from_live_pose() -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + env.robot._qpos[:, env.left_eef_joints] = env.close_state + state = _held_state(env, entity) + entity._pose[:, 0, 3] += 0.055 + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ), + env, + record_runtime=False, + ) + + physical = executor._physical_pickup( + "can", + "left_arm", + state, + torch.tensor([True]), + ) + state = executor._rebase_held_state( + "can", + "left_arm", + state, + physical, + from_planned_qpos=False, + ) + + assert bool(physical[0]) + left_eef, _ = env.get_current_xpos_agent() + rebased_eef = torch.bmm( + entity.get_local_pose(to_matrix=True), + state.get_held_object("physical_left_arm").object_to_eef, + ) + assert torch.allclose(rebased_eef, left_eef) + + +def test_physical_hold_accepts_configured_held_position_tolerance() -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + env.robot._qpos[:, env.left_eef_joints] = env.close_state + state = _held_state(env, entity) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ), + env, + record_runtime=False, + ) + executor._object_owners["can"] = ["left_arm"] + entity._pose[:, 0, 3] += 0.055 + + held = executor._physical_hold( + "can", + "left_arm", + state, + torch.tensor([True]), + ) + + assert bool(held[0]) + + +def test_physical_pickup_rejects_large_grasp_slip() -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + env.robot._qpos[:, env.left_eef_joints] = env.close_state + state = _held_state(env, entity) + entity._pose[:, 0, 3] += 0.08 + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ), + env, + record_runtime=False, + ) + + physical = executor._physical_pickup( + "can", + "left_arm", + state, + torch.tensor([True]), + ) + + assert not bool(physical[0]) + + +def test_physical_pickup_rejects_offset_even_when_object_was_lifted() -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + env.robot._qpos[:, env.left_eef_joints] = env.close_state + state = _held_state(env, entity) + entity._pose[:, 0, 3] += 0.08 + entity._pose[:, 2, 3] += 0.08 + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ), + env, + record_runtime=False, + ) + + physical = executor._physical_pickup( + "can", + "left_arm", + state, + torch.tensor([True]), + ) + + assert not bool(physical[0]) + + +def test_physical_hold_detects_loss_and_releases_runtime_ownership() -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + env.robot._qpos[:, env.left_eef_joints] = env.close_state + state = _held_state(env, entity) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ), + env, + record_runtime=False, + ) + executor._object_owners["can"] = ["left_arm"] + executor._arm_owners["left_arm"] = ["can"] + executor._object_states[("can", "left_arm")] = state + entity._pose[:, 0, 3] += 0.08 + + held = executor._physical_hold( + "can", + "left_arm", + state, + torch.tensor([True]), + ) + executor._release_ownership("can", "left_arm", ~held) + + assert not bool(held[0]) + assert executor._object_owners["can"] == [None] + assert executor._arm_owners["left_arm"] == [None] + assert ("can", "left_arm") not in executor._object_states + + +def test_rebase_held_state_uses_fk_qpos_not_stale_eef_cache( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + state = _held_state(env, entity, arm="left_arm") + expected_eef = _pose(0.31, -0.17, 1.06) + + def compute_fk( + qpos: torch.Tensor, + *, + name: str, + to_matrix: bool, + ) -> torch.Tensor: + del name + assert to_matrix + return expected_eef.repeat(qpos.shape[0], 1, 1) + + monkeypatch.setattr(env.robot, "compute_fk", compute_fk) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ), + env, + record_runtime=False, + ) + + state = executor._rebase_held_state( + "can", + "left_arm", + state, + torch.tensor([True]), + ) + held = state.get_held_object("physical_left_arm") + assert held is not None + expected_relation = torch.bmm( + torch.linalg.inv(entity.get_local_pose(to_matrix=True)), + expected_eef, + ) + assert torch.allclose(held.object_to_eef, expected_relation) + + +def test_upright_transport_state_tracks_selected_target_pose( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + state = _held_state(env, entity, arm="left_arm") + expected_eef = _pose(0.27, -0.11, 1.04) + target_pose = _pose(0.05, 0.01, 0.92) + + def compute_fk( + qpos: torch.Tensor, + *, + name: str, + to_matrix: bool, + ) -> torch.Tensor: + del name + assert to_matrix + return expected_eef.repeat(qpos.shape[0], 1, 1) + + monkeypatch.setattr(env.robot, "compute_fk", compute_fk) + entity._pose = target_pose.clone() + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ), + env, + record_runtime=False, + ) + synchronized = executor._rebase_held_state( + "can", + "left_arm", + state, + torch.tensor([True]), + ) + + held = synchronized.get_held_object("physical_left_arm") + assert held is not None + assert torch.allclose( + held.object_to_eef, + torch.bmm(torch.linalg.inv(target_pose), expected_eef), + ) + assert torch.allclose(held.grasp_xpos, expected_eef) + + +def test_existing_object_owner_reserves_same_arm(monkeypatch: Any) -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ), + env, + record_runtime=False, + ) + original = executor.program.semantic_steps[0] + continuation = replace( + original, + id="continuation", + actor={"mode": "auto"}, + ) + executor._object_owners["can"] = ["left_arm"] + executor._arm_owners["left_arm"] = ["can"] + executor._object_states[("can", "left_arm")] = _held_state(env, entity) + + monkeypatch.setattr( + executor, + "_candidate", + lambda step, arm, failed: SimpleNamespace( + feasible=torch.tensor([True]), + cost=torch.tensor([0.0 if arm == "right_arm" else 10.0]), + warnings=(), + ), + ) + executor._ensure_assignment(continuation, torch.tensor([False])) + + assert executor._assignments["continuation"] == ["left_arm"] + assert bool(executor._resource_conflicts(continuation, "right_arm")[0]) + + +def test_new_task_group_hydrates_predecessor_held_state() -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "right_arm"))) + ), + env, + record_runtime=False, + ) + held_state = _held_state(env, entity, arm="right_arm") + executor._object_states[("can", "right_arm")] = held_state + continuation = replace( + executor.program.semantic_steps[0], + id="place_after_handover", + actor={"mode": "required", "arm": "right_arm"}, + ) + + hydrated = executor._state_for(continuation, "right_arm") + + assert hydrated.get_held_object("physical_right_arm") is not None + assert torch.equal(hydrated.last_qpos, env.robot.get_qpos()) + + +def test_place_uses_preceding_or_live_eef_pose_not_original_grasp() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)), + "basket": _FakeEntity("basket", _pose(0.0, 0.0, 0.70), _box_vertices(0.10)), + } + env = _FakeEnv(entities) + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "basket", + "relation": "inside", + }, + "depends_on": [], + } + ) + ) + ) + state = _held_state(env, entities["can"]) + held_object = state.get_held_object("physical_left_arm") + assert held_object is not None + replacement = HeldObjectState( + semantics=held_object.semantics, + object_to_eef=held_object.object_to_eef, + grasp_xpos=_pose(0.0, -0.3, 0.75), + env_mask=held_object.env_mask, + ) + held_objects = dict(state.held_objects) + held_objects["physical_left_arm"] = replacement + state = state.with_updates(held_objects=held_objects) + held_object = replacement + edge = next( + edge + for edge in program.edges + if edge.actions[0]["atomic_action_class"] == "Place" + ) + grounder = ActionGrounder( + program, + env, + lambda uid: held_object.semantics, + ) + reference = _pose(0.0, 0.4, 0.85) + + planned = grounder.ground( + edge.actions[0], + program.semantic_steps[0], + arm="left_arm", + state=state, + reference_eef_pose=reference, + ) + live = grounder.ground( + edge.actions[0], + program.semantic_steps[0], + arm="left_arm", + state=state, + ) + + assert torch.equal(planned.target.xpos, reference) + assert torch.equal(live.target.xpos, env.get_current_xpos_agent()[0]) + assert not torch.equal(live.target.xpos, held_object.grasp_xpos) + + +def test_inside_target_preserves_pre_pick_supported_height() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 1.40), _box_vertices(0.03)), + "basket": _FakeEntity("basket", _pose(0.0, 0.0, 0.70), _box_vertices(0.10)), + } + env = _FakeEnv(entities) + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "basket", + "relation": "inside", + }, + "depends_on": [], + } + ) + ) + ) + step = program.semantic_steps[0] + final = next( + edge + for edge in program.edges + if edge.actions[0]["atomic_action_class"] == "MoveHeldObject" + ) + state = _held_state(env, entities["can"]) + held = state.get_held_object("physical_left_arm") + assert held is not None + grounder = ActionGrounder(program, env, lambda _uid: held.semantics) + supported_pose = _pose(0.0, 0.2, 0.75) + + grounded = grounder.ground( + final.actions[0], + step, + arm="left_arm", + state=state, + orientation_reference_pose=supported_pose, + ) + + assert grounded.target_object_pose is not None + assert grounded.target_object_pose[0, 2, 3] == pytest.approx(0.75) + + +def test_coordinated_step_rejects_an_arm_reserved_by_terminal_hold() -> None: + entity = _FakeEntity("shared_box", _pose(0.0, 0.0, 0.75), _box_vertices(0.05)) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "transport", + "operator": "coordinated_transport", + "object": "shared_box", + "actor": { + "mode": "coordinated", + "arms": ["left_arm", "right_arm"], + }, + "goal": { + "direction": "front", + "terminal_behavior": "hold", + }, + "depends_on": [], + } + ) + ) + ), + _FakeEnv({"shared_box": entity}), + record_runtime=False, + ) + executor._arm_owners["left_arm"] = ["held_can"] + step = executor.program.semantic_steps[0] + + executor._ensure_assignment(step, torch.tensor([False])) + + assert executor._assignments[step.id] == [None] + + +def test_failure_propagates_only_to_dependent_branch(monkeypatch: Any) -> None: + program = load_execution_program( + compile_task_agent( + _task_agent( + _hold_step("left", "can_a", "left_arm"), + _hold_step("right", "can_b", "right_arm"), + ) + ) + ) + executor = ProgramExecutor( + program, _FakeEnv(), settle_steps=0, record_runtime=False + ) + monkeypatch.setattr( + executor, + "_pack_ready_edges", + lambda ready, **_kwargs: (ready[0],), + ) + monkeypatch.setattr( + executor, + "_ensure_assignment", + lambda step, failed: executor._assignments.setdefault( + step.id, [step.actor["arm"]] + ), + ) + active_by_step: dict[str, list[bool]] = {"left": [], "right": []} + + def execute(edge, step, *, failed): + active_by_step[step.id].append(not bool(failed[0])) + action_failed = failed.clone() + if step.id == "left" and edge.id == step.edge_ids[0]: + action_failed[:] = True + return SimpleNamespace(actions=[], failed=action_failed, grounded=[]) + + monkeypatch.setattr(executor, "_execute_edge", execute) + monkeypatch.setattr( + executor, + "_verify_step", + lambda step, failed: ( + failed, + ~failed, + torch.zeros(1, 3), + ), + ) + + result = executor.run() + + assert any(active_by_step["right"]) + assert bool(result.semantic_success["right"][0]) + assert not bool(result.success[0]) + + +def test_resource_ordering_waits_without_propagating_semantic_failure() -> None: + task = { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "resource_ordering", + "level": "L3", + "instruction": "Stand both cans, then hand over the second can.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "task_01", + "task_type": "E2", + "params": { + "object_role": "first", + "required_arm": "right_arm", + }, + "depends_on": [], + "role": "primary", + }, + { + "id": "task_02", + "task_type": "E2", + "params": { + "object_role": "second", + "required_arm": "left_arm", + }, + "depends_on": [], + "role": "primary", + }, + { + "id": "task_03", + "task_type": "E4", + "params": { + "object_role": "second", + "transfer_arm": "left_arm", + "receive_arm": "right_arm", + }, + "depends_on": ["task_02"], + "role": "primary", + }, + ], + "success": {"type": "all_complete"}, + "oracle": {}, + "metadata": {}, + } + program = load_execution_program( + instantiate_seed_graph( + task, + {"first": "first_can", "second": "second_can"}, + ) + ) + executor = ProgramExecutor(program, _FakeEnv(), record_runtime=False) + handover_entry = next( + edge + for edge in program.edges + if executor.step_by_edge[edge.id].id == "task_03" + and all( + executor.step_by_edge[dependency].id != "task_03" + for dependency in edge.depends_on + ) + ) + dependencies = { + executor.step_by_edge[dependency].id: dependency + for dependency in handover_entry.depends_on + } + failures = { + dependency: torch.tensor([step_id == "task_01"]) + for step_id, dependency in dependencies.items() + } + + assert not bool(executor._dependency_failures(handover_entry, failures)[0]) + + failures[dependencies["task_02"]][:] = True + assert bool(executor._dependency_failures(handover_entry, failures)[0]) + + +def test_v2_executor_retries_one_complete_atomic_action_twice( + monkeypatch: Any, +) -> None: + task, requirements = make_task_spec("E9") + bindings = { + item["role_id"]: f"uid_{item['role_id']}" for item in requirements["objects"] + } + program = load_execution_program(instantiate_seed_graph(task, bindings)) + executor = ProgramExecutor( + program, + _FakeEnv(), + settle_steps=0, + record_runtime=False, + ) + monkeypatch.setattr( + executor, + "_ensure_assignment", + lambda step, failed: executor._assignments.setdefault(step.id, ["left_arm"]), + ) + attempts = 0 + + def execute(_edge, _step, *, failed): + nonlocal attempts + attempts += 1 + action_failed = failed.clone() + if attempts < 3: + action_failed[:] = True + return SimpleNamespace(actions=[], failed=action_failed, grounded=[]) + + monkeypatch.setattr(executor, "_execute_edge", execute) + monkeypatch.setattr( + executor, + "_verify_step", + lambda step, failed: (failed, ~failed, torch.zeros(1, 3)), + ) + + result = executor.run() + + assert attempts == 3 + assert result.retry_count == 2 + assert bool(result.success[0]) + assert result.failure_events == [] + + +def test_v2_executor_stops_at_transition_budget() -> None: + task, requirements = make_task_spec("E1") + bindings = { + item["role_id"]: f"uid_{item['role_id']}" for item in requirements["objects"] + } + program = load_execution_program(instantiate_seed_graph(task, bindings)) + executor = ProgramExecutor( + program, + _FakeEnv(), + max_transitions=0, + settle_steps=0, + record_runtime=False, + ) + + with pytest.raises(RuntimeError, match="max_transitions"): + executor.run() + + +def test_failed_arrangement_records_candidate_diagnostics_without_marker_error( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + entities = { + "table": _FakeEntity( + "table", + _pose(0.0, 0.0, 0.70), + _rect_vertices(0.60, 0.40, 0.02), + ), + "can_a": _FakeEntity( + "can_a", + _pose(-0.15, 0.0, 0.76), + _rect_vertices(0.03, 0.03, 0.06), + ), + "can_b": _FakeEntity( + "can_b", + _pose(0.15, 0.0, 0.76), + _rect_vertices(0.03, 0.03, 0.06), + ), + } + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "line", + "operator": "arrange_line", + "objects": ["can_a", "can_b"], + "actor": {"mode": "auto"}, + "goal": { + "axis": "world_x", + "order_constraint": "free", + }, + "depends_on": [], + } + ) + ) + ) + executor = ProgramExecutor( + program, + _FakeEnv(entities), + settle_steps=0, + record_root=tmp_path, + ) + infeasible = SimpleNamespace( + feasible=torch.tensor([False]), + cost=torch.tensor([torch.inf]), + plans={}, + warnings=("No IK solutions found for downstream target poses.",), + ) + monkeypatch.setattr(executor, "_candidate", lambda *_args, **_kwargs: infeasible) + + result = executor.run(run_id="no_candidate") + + assert not bool(result.success[0]) + record_path = Path(result.record_dir) / "env_0000" / "task_graph.json" + record = json.loads(record_path.read_text(encoding="utf-8")) + assert record["runtime"]["status"] == "failed" + first_event = record["runtime"]["events"][0] + assert first_event["status"] == "failed" + assert first_event["diagnostics"] == [ + "No IK solutions found for downstream target poses." + ] + + +def test_arrange_line_builds_live_slots_for_compiler_operator_name() -> None: + entities = { + "table": _FakeEntity( + "table", + _pose(0.0, 0.0, 0.70), + _rect_vertices(0.60, 0.40, 0.02), + ), + "can_a": _FakeEntity( + "can_a", + _pose(-0.15, 0.0, 0.76), + _rect_vertices(0.03, 0.03, 0.06), + ), + "can_b": _FakeEntity( + "can_b", + _pose(0.15, 0.0, 0.76), + _rect_vertices(0.03, 0.03, 0.06), + ), + } + compiled = compile_task_agent( + _task_agent( + { + "id": "line", + "operator": "arrange_line", + "objects": ["can_a", "can_b"], + "actor": {"mode": "auto"}, + "goal": { + "axis": "world_x", + "order_constraint": "free", + }, + "depends_on": [], + } + ) + ) + + executor = ProgramExecutor( + load_execution_program(compiled), + _FakeEnv(entities), + record_runtime=False, + ) + + assert executor.arrangement is not None + assert executor.arrangement.positions.shape == (1, 2, 3) + assert {step.operator for step in executor.program.semantic_steps} == { + "arrange_line" + } + + +def test_free_arrangement_matches_live_object_order_without_crossing() -> None: + entities = { + "table": _FakeEntity( + "table", + _pose(0.0, 0.0, 0.70), + _rect_vertices(0.60, 0.40, 0.02), + ), + "can_a": _FakeEntity( + "can_a", + _pose(0.0, 0.20, 0.78), + _rect_vertices(0.03, 0.03, 0.06), + ), + "can_b": _FakeEntity( + "can_b", + _pose(0.0, 0.00, 0.78), + _rect_vertices(0.03, 0.03, 0.06), + ), + "can_c": _FakeEntity( + "can_c", + _pose(0.0, -0.20, 0.78), + _rect_vertices(0.03, 0.03, 0.06), + ), + } + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "line", + "operator": "arrange_line", + "objects": ["can_a", "can_b", "can_c"], + "actor": {"mode": "auto"}, + "goal": { + "axis": "world_y", + "order_constraint": "free", + }, + "depends_on": [], + } + ) + ) + ), + _FakeEnv(entities), + record_runtime=False, + ) + arrangement = executor.arrangement + + assert arrangement is not None + assert int(arrangement.assignments["line__01"][0]) == 2 + assert int(arrangement.assignments["line__02"][0]) == 1 + assert int(arrangement.assignments["line__03"][0]) == 0 + assert arrangement.spacing[0] == pytest.approx(0.1648528) + + +def test_arm_candidate_score_softly_penalizes_cross_zone_motion() -> None: + source = _pose(0.0, -0.30, 0.78) + target = _pose(0.0, -0.20, 0.78) + kwargs = { + "motion_cost": torch.tensor([torch.pi]), + "source_pose": source, + "target_pose": target, + "workspace_center_xy": torch.tensor([[0.0, 0.0]]), + "workspace_half_width": torch.tensor([0.40]), + "robot_lateral_axis": torch.tensor([[0.0, -1.0]]), + "policy": default_runtime_policy("dual_ur10").arm_selection, + } + + left = _score_arm_candidate(arm="left_arm", **kwargs) + right = _score_arm_candidate(arm="right_arm", **kwargs) + + assert left["normalized_motion_cost"][0] == pytest.approx(1.0) + assert left["pickup_crossing_penalty"][0] == pytest.approx(0.0) + assert left["placement_crossing_penalty"][0] == pytest.approx(0.0) + assert right["pickup_crossing_penalty"][0] > 0.0 + assert right["placement_crossing_penalty"][0] > 0.0 + assert right["total_cost"][0] > left["total_cost"][0] + + +def test_preserve_grounding_uses_pre_pickup_orientation_reference() -> None: + entities = { + "table": _FakeEntity( + "table", + _pose(0.0, 0.0, 0.70), + _rect_vertices(0.60, 0.40, 0.02), + ), + "can_a": _FakeEntity( + "can_a", + _pose(0.0, -0.10, 0.78), + _rect_vertices(0.03, 0.03, 0.06), + ), + "can_b": _FakeEntity( + "can_b", + _pose(0.0, 0.10, 0.78), + _rect_vertices(0.03, 0.03, 0.06), + ), + } + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "line", + "operator": "arrange_line", + "objects": ["can_a", "can_b"], + "actor": {"mode": "auto"}, + "goal": { + "axis": "world_y", + "order_constraint": "free", + "orientation_goal": "preserve", + }, + "depends_on": [], + } + ) + ) + ), + _FakeEnv(entities), + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + final_edge = next( + edge + for edge in executor.program.edges + if edge.id in step.edge_ids + and edge.actions[0]["atomic_action_class"] == "MoveHeldObject" + and edge.actions[0]["target_binding"].get("phase") == "final" + ) + reference = entities[step.object_uid].get_local_pose(to_matrix=True) + disturbed = reference.clone() + disturbed[:, :3, :3] = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + entities[step.object_uid]._pose = disturbed + + grounded = executor.grounder.ground( + final_edge.actions[0], + step, + arm="left_arm", + state=ExecutionState(last_qpos=executor.env.robot.get_qpos()), + orientation_reference_pose=reference, + ) + + assert grounded.target_object_pose is not None + assert torch.allclose( + grounded.target_object_pose[:, :3, :3], + reference[:, :3, :3], + ) + + +def test_arrange_line_verifies_planar_slot_without_height_coupling() -> None: + entities = { + "table": _FakeEntity( + "table", + _pose(0.0, 0.0, 0.70), + _rect_vertices(0.60, 0.40, 0.02), + ), + "can_a": _FakeEntity( + "can_a", + _pose(-0.055, -0.190, 0.755), + _rect_vertices(0.03, 0.03, 0.06), + ), + "can_b": _FakeEntity( + "can_b", + _pose(0.15, 0.0, 0.76), + _rect_vertices(0.03, 0.03, 0.06), + ), + } + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "line", + "operator": "arrange_line", + "objects": ["can_a", "can_b"], + "actor": {"mode": "auto"}, + "goal": { + "axis": "world_y", + "order_constraint": "free", + }, + "depends_on": [], + } + ) + ) + ), + _FakeEnv(entities), + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + executor._targets[step.id] = torch.tensor([[0.0, -0.216, 0.842]]) + executor._policies[step.id] = { + "line_axis_tolerance": 0.06, + "line_perpendicular_tolerance": 0.06, + } + + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + + assert not bool(failed[0]) + assert bool(success[0]) + + +def test_arrange_line_rejects_preserve_orientation_drift() -> None: + rotated = _pose(0.0, -0.190, 0.755) + rotated[:, :3, :3] = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + entities = { + "table": _FakeEntity( + "table", + _pose(0.0, 0.0, 0.70), + _rect_vertices(0.60, 0.40, 0.02), + ), + "can_a": _FakeEntity( + "can_a", + rotated, + _rect_vertices(0.03, 0.03, 0.06), + ), + "can_b": _FakeEntity( + "can_b", + _pose(0.15, 0.0, 0.76), + _rect_vertices(0.03, 0.03, 0.06), + ), + } + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "line", + "operator": "arrange_line", + "objects": ["can_a", "can_b"], + "actor": {"mode": "auto"}, + "goal": { + "axis": "world_y", + "order_constraint": "free", + "orientation_goal": "preserve", + }, + "depends_on": [], + } + ) + ) + ), + _FakeEnv(entities), + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + executor._targets[step.id] = rotated[:, :3, 3].clone() + executor._orientation_references[step.id] = _pose(0.0, -0.190, 0.755) + executor._policies[step.id] = { + "line_axis_tolerance": 0.06, + "line_perpendicular_tolerance": 0.06, + "preserve_orientation_tolerance": torch.pi / 12, + } + + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + + assert bool(failed[0]) + assert not bool(success[0]) + + +def test_shared_container_placements_receive_non_overlapping_live_slots() -> None: + entities = { + "basket": _FakeEntity( + "basket", + _pose(0.0, 0.0, 0.72), + _rect_vertices(0.25, 0.18, 0.08), + ), + "cube": _FakeEntity( + "cube", + _pose(-0.15, 0.0, 0.76), + _rect_vertices(0.03, 0.03, 0.03), + ), + "cup": _FakeEntity( + "cup", + _pose(0.15, 0.0, 0.76), + _rect_vertices(0.035, 0.035, 0.06), + ), + } + steps = [ + { + "id": f"place_{uid}", + "operator": "place_relative", + "object": uid, + "actor": {"mode": "auto"}, + "goal": {"relation": "inside", "reference_object": "basket"}, + "depends_on": [], + } + for uid in ("cube", "cup") + ] + task = { + "schema_version": TASK_AGENT_SCHEMA, + "task": "two_in_basket", + "goal": "Place both objects in the basket.", + "semantic_steps": steps, + } + executor = ProgramExecutor( + load_execution_program(compile_task_agent(task)), + _FakeEnv(entities), + record_runtime=False, + ) + targets = [ + executor.placements[step.id].positions[step.id][0, :2] + for step in executor.program.semantic_steps + ] + + assert set(executor.placements) == {"place_cube", "place_cup"} + assert torch.linalg.vector_norm(targets[0] - targets[1]) > 0.05 + + +@pytest.mark.parametrize( + ("direction", "expected_position"), + ( + ("front_left", (0.16, 0.16, 0.85)), + ("up", (0.0, 0.0, 0.91)), + ), +) +def test_coordinated_transport_direction_is_grounded_from_live_pose( + direction: str, + expected_position: tuple[float, float, float], +) -> None: + entities = { + "shared_box": _FakeEntity( + "shared_box", + _pose(0.0, 0.0, 0.75), + _rect_vertices(0.10, 0.06, 0.03), + ) + } + env = _FakeEnv(entities) + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "transport", + "operator": "coordinated_transport", + "object": "shared_box", + "actor": { + "mode": "coordinated", + "arms": ["left_arm", "right_arm"], + }, + "goal": { + "direction": direction, + "terminal_behavior": "hold", + }, + "depends_on": [], + } + ) + ) + ) + + def semantics(uid: str) -> ObjectSemantics: + entity = entities[uid] + return ObjectSemantics( + affordance=Affordance(), + geometry={ + "mesh_vertices": entity.get_vertices(env_ids=[0], scale=True), + "mesh_triangles": entity.get_triangles(env_ids=[0]), + }, + label=uid, + entity=entity, + ) + + step = program.semantic_steps[0] + edge = program.edges[0] + grounded = ActionGrounder(program, env, semantics).ground( + edge.actions[0], + step, + arm="coordinated", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + + assert isinstance(grounded.target, CoordinatedPickGoal) + assert torch.allclose( + grounded.target.object_target_pose[0, :3, 3], + torch.tensor(expected_position), + ) + + +def test_coordinated_payload_monitor_rejects_drift_and_carrier_tilt() -> None: + class _BatchedVerticesEntity(_FakeEntity): + def get_vertices( + self, + *, + env_ids: list[int], + scale: bool, + ) -> torch.Tensor: + return super().get_vertices(env_ids=env_ids, scale=scale).unsqueeze(0) + + entities = { + "tray": _BatchedVerticesEntity( + "tray", + _pose(0.0, 0.0, 0.75), + _rect_vertices(0.20, 0.14, 0.02), + ), + "bottle": _FakeEntity( + "bottle", + _pose(0.0, 0.0, 0.80), + _rect_vertices(0.03, 0.03, 0.08), + ), + } + program = compile_task_agent( + _task_agent( + { + "id": "carry", + "operator": "coordinated_transport", + "object": "tray", + "actor": { + "mode": "coordinated", + "arms": ["left_arm", "right_arm"], + }, + "goal": { + "terminal_behavior": "place", + "payloads": [{"object": "bottle", "slot": "center"}], + }, + "depends_on": [], + } + ) + ) + executor = ProgramExecutor( + load_execution_program(program), + _FakeEnv(entities), + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + executor._capture_payloads(step) + + assert bool(executor._verify_payloads(step)[0]) + entities["bottle"]._pose[:, 0, 3] += 0.20 + assert not bool(executor._verify_payloads(step)[0]) + entities["bottle"]._pose = _pose(0.0, 0.0, 0.80) + entities["tray"]._pose[:, :3, :3] = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + assert not bool(executor._verify_payloads(step)[0]) + + +def test_lay_flat_surface_height_uses_rotated_live_mesh() -> None: + entities = { + "table": _FakeEntity( + "table", + _pose(0.0, 0.0, 0.70), + _rect_vertices(0.60, 0.40, 0.02), + ), + "rod": _FakeEntity( + "rod", + _pose(0.2, 0.0, 0.80), + _rect_vertices(0.02, 0.03, 0.10), + ), + } + env = _FakeEnv(entities) + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "rod", + "actor": {"mode": "auto"}, + "goal": { + "reference_object": "table", + "relation": "on", + "orientation_goal": "lay_flat", + "orientation_axis": "none", + }, + "depends_on": [], + } + ) + ) + ) + step = program.semantic_steps[0] + edge = next( + edge + for edge in program.edges + if edge.actions[0]["atomic_action_class"] == "MoveHeldObject" + ) + grounder = ActionGrounder( + program, + env, + lambda uid: ObjectSemantics( + affordance=Affordance(), + geometry={}, + label=uid, + entity=entities[uid], + ), + ) + grounded = grounder.ground( + edge.actions[0], + step, + arm="left_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + + assert isinstance(grounded.target, HeldObjectPoseGoal) + table_top = 0.70 + 0.02 + rotated_rod_half_height = 0.02 + surface_clearance = 0.005 + expected_surface_z = table_top + rotated_rod_half_height + surface_clearance + assert grounded.target.object_target_pose[0, 2, 3] == pytest.approx( + expected_surface_z, + abs=1.0e-5, + ) + + +def test_orient_object_anchors_final_pose_to_support_not_live_lift_height() -> None: + entities = { + "table": _FakeEntity( + "table", + _pose(0.0, 0.0, 0.70), + _rect_vertices(0.60, 0.40, 0.02), + ), + "bottle": _FakeEntity( + "bottle", + _pose(0.10, 0.20, 1.30), + _rect_vertices(0.02, 0.03, 0.10), + ), + } + env = _FakeEnv(entities) + env.agent_initial_object_poses = {"bottle": _pose(0.10, 0.20, 0.78)} + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "orient", + "operator": "orient_object", + "object": "bottle", + "actor": {"mode": "auto"}, + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + }, + "depends_on": [], + } + ) + ) + ) + step = program.semantic_steps[0] + edges = { + edge.actions[0]["target_binding"].get("phase"): edge + for edge in program.edges + if edge.actions[0]["atomic_action_class"] == "MoveHeldObject" + } + grounder = ActionGrounder( + program, + env, + lambda uid: ObjectSemantics( + affordance=Affordance(), + geometry={}, + label=uid, + entity=entities[uid], + ), + ) + + staging = grounder.ground( + edges["staging"].actions[0], + step, + arm="left_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + final = grounder.ground( + edges["final"].actions[0], + step, + arm="left_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + expected_final_z = 0.70 + 0.02 + 0.10 + 0.05 + assert final.target_object_pose[0, 2, 3] == pytest.approx(expected_final_z) + assert final.target_object_pose[0, :2, 3].tolist() == pytest.approx([0.10, 0.20]) + assert staging.target_object_pose[0, 2, 3] > final.target_object_pose[0, 2, 3] + assert staging.target_object_pose[0, :2, 3].tolist() == pytest.approx([0.10, 0.20]) + + +def test_orient_grounding_uses_mature_robot_profile_policy() -> None: + entities = { + "table": _FakeEntity( + "table", + _pose(0.0, 0.0, 0.70), + _rect_vertices(0.60, 0.40, 0.02), + ), + "bottle": _FakeEntity( + "bottle", + _pose(0.10, 0.20, 0.78), + _rect_vertices(0.02, 0.03, 0.10), + ), + } + env = _FakeEnv(entities) + env.agent_robot_profile = "dual_ur10" + env.agent_initial_object_poses = {"bottle": _pose(0.10, 0.20, 0.78)} + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "orient", + "operator": "orient_object", + "object": "bottle", + "actor": {"mode": "auto"}, + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + "support_object": "table", + }, + "depends_on": [], + } + ) + ) + ) + step = program.semantic_steps[0] + pickup_edge = next( + edge + for edge in program.edges + if edge.actions[0]["atomic_action_class"] == "PickUp" + ) + final_edge = next( + edge + for edge in program.edges + if edge.actions[0]["target_binding"].get("phase") == "final" + ) + grounder = ActionGrounder( + program, + env, + lambda uid: ObjectSemantics( + affordance=Affordance(), + geometry={}, + label=uid, + entity=entities[uid], + ), + ) + + pickup = grounder.ground( + pickup_edge.actions[0], + step, + arm="left_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + final = grounder.ground( + final_edge.actions[0], + step, + arm="left_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + + table_top = 0.72 + bottle_half_height = 0.10 + expected_z = table_top + bottle_half_height + 0.05 + assert torch.equal( + pickup.motion_policy["obj_upright_direction"], + torch.tensor([0.0, 0.0, 1.0]), + ) + assert pickup.motion_policy["rotate_upright"] == pytest.approx(torch.pi / 4) + assert pickup.motion_policy["upright_yaw_samples"] == 8 + assert final.target_object_pose[0, 2, 3] == pytest.approx(expected_z) + assert final.motion_policy["upright_local_axis"] == "long_axis" + assert final.motion_policy["upright_yaw_samples"] == 8 + + +def test_orient_verification_requires_upright_pose_near_initial_xy() -> None: + bottle = _FakeEntity( + "bottle", + _pose(0.10, 0.20, 0.823), + _rect_vertices(0.02, 0.03, 0.10), + ) + env = _FakeEnv({"bottle": bottle}) + env.agent_initial_object_poses = {"bottle": _pose(0.10, 0.20, 0.78)} + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "orient", + "operator": "orient_object", + "object": "bottle", + "actor": {"mode": "auto"}, + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + }, + "depends_on": [], + } + ) + ) + ), + env, + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + executor._policies[step.id] = { + "upright_max_tilt": torch.pi / 12, + "upright_xy_tolerance": 0.05, + "upright_local_axis": "z", + } + + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + assert bool(success[0]) + assert not bool(failed[0]) + + bottle._pose[:, :3, :3] = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + assert not bool(success[0]) + assert bool(failed[0]) + + +def test_orient_verification_accepts_grounded_live_xy_anchor() -> None: + bottle = _FakeEntity( + "bottle", + _pose(0.15, -0.10, 0.823), + _rect_vertices(0.02, 0.03, 0.10), + ) + env = _FakeEnv({"bottle": bottle}) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "orient", + "operator": "orient_object", + "object": "bottle", + "actor": {"mode": "auto"}, + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + "position_anchor": "live_xy", + }, + "depends_on": [], + } + ) + ) + ), + env, + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + executor._targets[step.id] = torch.tensor([[0.15, -0.10, 0.823]]) + executor._policies[step.id] = { + "upright_max_tilt": torch.pi / 12, + "upright_xy_tolerance": 0.05, + "upright_local_axis": "z", + } + + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + + assert bool(success[0]) + assert not bool(failed[0]) + + +def test_long_axis_upright_is_undirected_but_explicit_axis_is_not() -> None: + pose = _pose(0.0, 0.0, 0.75) + pose[:, :3, 1] = torch.tensor([0.0, 0.0, -1.0]) + pose[:, :3, 2] = torch.tensor([0.0, 1.0, 0.0]) + entity = _FakeEntity("can", pose, _rect_vertices(0.03, 0.10, 0.03)) + env = _FakeEnv({"can": entity}) + + assert bool( + evaluate_predicate( + env, + { + "type": "object_upright", + "object": "can", + "local_axis": "long_axis", + }, + )[0] + ) + assert not bool( + evaluate_predicate( + env, + { + "type": "object_upright", + "object": "can", + "local_axis": "long_axis", + "directed": True, + }, + )[0] + ) + assert not bool( + evaluate_predicate( + env, + { + "type": "object_upright", + "object": "can", + "local_axis": "y", + }, + )[0] + ) + + +def test_orient_object_uses_solver_roots_when_control_groups_share_root() -> None: + entities = { + "left_object": _FakeEntity( + "left_object", + _pose(0.0, -0.20, 0.8), + _rect_vertices(0.02, 0.02, 0.08), + ), + "right_object": _FakeEntity( + "right_object", + _pose(0.0, 0.20, 0.8), + _rect_vertices(0.02, 0.02, 0.08), + ), + } + env = _FakeEnv(entities) + env.agent_initial_object_poses = { + uid: entity.get_local_pose(to_matrix=True) for uid, entity in entities.items() + } + execution = compile_task_agent( + _task_agent( + *[ + { + "id": uid, + "operator": "orient_object", + "object": uid, + "actor": {"mode": "auto"}, + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + }, + "depends_on": [], + } + for uid in ("left_object", "right_object") + ] + ) + ) + executor = ProgramExecutor( + load_execution_program(execution), env, record_runtime=False + ) + + torch.testing.assert_close( + env.robot.get_control_part_base_pose(name="physical_left_arm", to_matrix=True), + env.robot.get_control_part_base_pose(name="physical_right_arm", to_matrix=True), + ) + assert executor._preferred_in_place_arm(executor.steps["left_object"], 0) == ( + "left_arm" + ) + assert executor._preferred_in_place_arm(executor.steps["right_object"], 0) == ( + "right_arm" + ) + + +def test_orient_object_arm_preference_follows_translated_robot_and_table() -> None: + entities = { + "table": _FakeEntity( + "table", + _pose(1.50, -0.70, 0.70), + _rect_vertices(0.50, 0.40, 0.02), + ), + "left_object": _FakeEntity( + "left_object", + _pose(1.70, -0.70, 0.80), + _rect_vertices(0.02, 0.02, 0.08), + ), + "right_object": _FakeEntity( + "right_object", + _pose(1.30, -0.70, 0.80), + _rect_vertices(0.02, 0.02, 0.08), + ), + } + env = _FakeEnv(entities) + env.robot.get_link_pose = lambda *, link_name, to_matrix: _pose( + 1.80 if link_name == "physical_left_base" else 1.20, + -0.70, + 0.0, + ) + env.agent_initial_object_poses = { + uid: entity.get_local_pose(to_matrix=True) for uid, entity in entities.items() + } + execution = compile_task_agent( + _task_agent( + *[ + { + "id": uid, + "operator": "orient_object", + "object": uid, + "actor": {"mode": "auto"}, + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + }, + "depends_on": [], + } + for uid in ("left_object", "right_object") + ] + ) + ) + executor = ProgramExecutor( + load_execution_program(execution), env, record_runtime=False + ) + + center, _, lateral = executor._arm_selection_workspace( + executor.steps["left_object"] + ) + torch.testing.assert_close(center, torch.tensor([[1.50, -0.70]])) + torch.testing.assert_close(lateral, torch.tensor([[1.0, 0.0]])) + assert executor._preferred_in_place_arm(executor.steps["left_object"], 0) == ( + "left_arm" + ) + assert executor._preferred_in_place_arm(executor.steps["right_object"], 0) == ( + "right_arm" + ) + + +def test_coordinated_placement_uses_live_typed_target_and_profile_parts() -> None: + entities = { + "placing": _FakeEntity( + "placing", + _pose(0.0, 0.1, 0.75), + _box_vertices(0.04), + ), + "support": _FakeEntity( + "support", + _pose(0.0, -0.1, 0.75), + _box_vertices(0.06), + ), + } + env = _FakeEnv(entities) + compiled = compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "coordinated_place", + "object": "placing", + "actor": { + "mode": "coordinated", + "arms": ["left_arm", "right_arm"], + }, + "goal": { + "support_object": "support", + "relation": "on", + "release": True, + }, + "depends_on": [], + } + ) + ) + program = load_execution_program(compiled) + + def semantics(uid: str) -> ObjectSemantics: + return ObjectSemantics( + affordance=Affordance(), + geometry={}, + label=uid, + entity=entities[uid], + ) + + step = program.semantic_steps[0] + assert [action["atomic_action_class"] for action in program.edges[0].actions] == [ + "PickUp", + "PickUp", + ] + edge = next( + edge + for edge in program.edges + if edge.actions[0]["atomic_action_class"] == "CoordinatedPlacement" + ) + grounder = ActionGrounder(program, env, semantics) + grounded = grounder.ground( + edge.actions[0], + step, + arm="coordinated", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + + assert isinstance(grounded.target, CoordinatedPlacementGoal) + assert grounded.target.release is True + assert torch.allclose( + grounded.target.support_object_target_pose, + entities["support"].get_local_pose(to_matrix=True), + ) + + adapter = AtomicActionAdapter(env) + cfg = adapter._build_config(grounded, CoordinatedPlacementOptions) + bound_endpoints: dict[str, dict[str, str]] = {} + + class _BindingEngine: + binding_owner_id = "runtime-contract-test" + + def bind_control_parts( + self, + _skill_id: str, + endpoints: dict[str, dict[str, str]], + ) -> ActionBinding: + bound_endpoints.update(deepcopy(endpoints)) + return ActionBinding(owner_id=self.binding_owner_id) + + adapter._atomic_engine = _BindingEngine() + binding = adapter._binding( + grounded, + adapter.capabilities.get("CoordinatedPlacement"), + ) + assert binding.owner_id == "runtime-contract-test" + assert bound_endpoints == { + "placing": { + "motion": "physical_left_arm", + "grasp": "physical_left_eef", + }, + "support": { + "motion": "physical_right_arm", + "grasp": "physical_right_eef", + }, + } + assert cfg.release is True + + +def test_online_environment_preserves_result_and_disables_terminations( + monkeypatch: Any, +) -> None: + installed: list[Any] = [] + initialization_order: list[tuple[str, Any]] = [] + + def fake_super_init(self: Any, cfg: Any, **kwargs: Any) -> None: + del kwargs + initialization_order.append(("super", cfg.robot)) + self.cfg = cfg + self.robot = object() + self.ignore_terminations_during_agent = True + + def fake_repair(robot_cfg: Any) -> int: + initialization_order.append(("repair", robot_cfg)) + return 1 + + def fake_install(robot: Any) -> int: + initialization_order.append(("install", robot)) + installed.append(robot) + return 1 + + monkeypatch.setattr(EmbodiedEnv, "__init__", fake_super_init) + monkeypatch.setattr( + env_module, + "repair_action_engine_ur5_solver_cfg", + fake_repair, + ) + monkeypatch.setattr( + env_module, + "install_action_engine_solver_compat", + fake_install, + ) + monkeypatch.setattr( + env_module.ActionEngineEnv, + "_capture_runtime_state", + lambda self: None, + ) + robot_cfg = object() + cfg = SimpleNamespace(ignore_terminations=False, robot=robot_cfg) + env = env_module.ActionEngineEnv( + cfg, + agent_config={"schema_version": "action_engine_config_v2"}, + task_name="task", + agent_config_path="/tmp/agent_config.json", + ) + result = ExecutionResult( + actions=[], + success=torch.tensor([True]), + semantic_success={}, + ) + + assert cfg.ignore_terminations is True + assert installed == [env.robot] + assert [name for name, _ in initialization_order] == [ + "repair", + "super", + "install", + ] + assert initialization_order[0][1] is robot_cfg + assert env._normalize_demo_action_list(result) is result + + +def test_solver_compat_repairs_only_stale_action_engine_ur_dh_defaults() -> None: + stale_ur5 = URSolverCfg() + stale_ur5.ur_type = "ur5" + custom_ur5 = URSolverCfg(ur_type="ur5") + custom_ur5.d1 = 0.1 + ur10 = URSolverCfg() + robot_cfg = SimpleNamespace( + solver_cfg={ + "left": stale_ur5, + "left_alias": stale_ur5, + "custom": custom_ur5, + "right": ur10, + } + ) + expected = URSolverCfg(ur_type="ur5") + dh_fields = ("d1", "a2", "a3", "d4", "d5", "d6") + + assert solver_compat.repair_action_engine_ur5_solver_cfg(robot_cfg) == 1 + assert tuple(getattr(stale_ur5, name) for name in dh_fields) == pytest.approx( + tuple(getattr(expected, name) for name in dh_fields) + ) + assert custom_ur5.d1 == pytest.approx(0.1) + assert ur10.ur_type == "ur10" + assert solver_compat.repair_action_engine_ur5_solver_cfg(robot_cfg) == 0 + + +def test_solver_compat_uses_true_tcp_inverse_and_restores_solver( + monkeypatch: Any, +) -> None: + class FakeSolver: + def __init__(self) -> None: + self.device = torch.device("cpu") + self.tcp_xpos = np.array( + [ + [0.0, -1.0, 0.0, 0.1], + [1.0, 0.0, 0.0, 0.2], + [0.0, 0.0, 1.0, 0.3], + [0.0, 0.0, 0.0, 1.0], + ], + dtype=np.float32, + ) + self.received: torch.Tensor | None = None + self.saw_identity = False + + def get_ik(self, target_xpos: torch.Tensor, **kwargs: Any) -> str: + del kwargs + self.received = target_xpos + self.saw_identity = np.allclose(self.tcp_xpos, np.eye(4)) + return "ok" + + monkeypatch.setattr(solver_compat, "PytorchSolver", FakeSolver) + solver = FakeSolver() + original_tcp = solver.tcp_xpos.copy() + robot = SimpleNamespace(_solvers={"left": solver, "alias": solver}) + target = torch.eye(4).unsqueeze(0) + + assert solver_compat.install_pytorch_solver_tcp_compat(robot) == 1 + assert solver.get_ik(target_xpos=target) == "ok" + assert solver.saw_identity + assert torch.allclose( + solver.received, + target @ torch.linalg.inv(torch.as_tensor(original_tcp)), + ) + assert np.allclose(solver.tcp_xpos, original_tcp) + assert solver_compat.install_pytorch_solver_tcp_compat(robot) == 0 + + +def test_solver_compat_aligns_ur5_analytic_ik_with_urdf_ee_frame( + monkeypatch: Any, +) -> None: + class FakeSolver: + def __init__(self, ur_type: str) -> None: + self.cfg = SimpleNamespace(ur_type=ur_type) + self.device = torch.device("cpu") + self.tcp_xpos = np.array( + [ + [0.0, -1.0, 0.0, 0.0], + [1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.2], + [0.0, 0.0, 0.0, 1.0], + ], + dtype=np.float32, + ) + self.received: torch.Tensor | None = None + + def get_ik( + self, + target_xpos: torch.Tensor, + qpos_seed: torch.Tensor | None = None, + **kwargs: Any, + ) -> str: + del kwargs, qpos_seed + self.received = target_xpos + return "ok" + + monkeypatch.setattr(solver_compat, "URSolver", FakeSolver) + ur5 = FakeSolver("ur5") + ur10 = FakeSolver("ur10") + robot = SimpleNamespace(_solvers={"left": ur5, "alias": ur5, "right": ur10}) + target = torch.eye(4).unsqueeze(0) + target[:, :3, 3] = torch.tensor([0.3, -0.2, 0.8]) + qpos_seed = torch.zeros((1, 6)) + + assert solver_compat.install_ur5_solver_frame_compat(robot) == 1 + assert ur5.get_ik(target, qpos_seed) == "ok" + + tcp = torch.as_tensor(ur5.tcp_xpos) + analytic_to_urdf = torch.eye(4) + analytic_to_urdf[0, 3] = -0.01 + expected = target @ torch.linalg.inv(tcp) @ torch.linalg.inv(analytic_to_urdf) @ tcp + assert torch.allclose(ur5.received, expected) + assert ur10.received is None + assert solver_compat.install_ur5_solver_frame_compat(robot) == 0 diff --git a/tests/gen_sim/action_engine/test_agent.py b/tests/gen_sim/action_engine/test_agent.py new file mode 100644 index 000000000..3e137973e --- /dev/null +++ b/tests/gen_sim/action_engine/test_agent.py @@ -0,0 +1,194 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Action Agent compilation, preflight, and report boundary tests.""" + +from __future__ import annotations + +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +import embodichain.gen_sim.action_engine.agent as module +from embodichain.gen_sim.action_engine.agent import ActionAgent +from embodichain.gen_sim.action_engine.domain import seed_graph_hash +from embodichain.gen_sim.action_engine.runtime import ExecutionResult +from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph + +from .task_fixtures import make_task_spec + + +def _bindings(requirements: dict) -> dict[str, str]: + return { + item["role_id"]: f"scene_{item['role_id']}" for item in requirements["objects"] + } + + +def _task_of_type(task_type: str) -> tuple[dict, dict]: + return make_task_spec(task_type) + + +def test_plan_hash_matches_direct_seed_graph_instantiation(monkeypatch) -> None: + task, requirements = make_task_spec("E1") + bindings = _bindings(requirements) + grounded_plan = { + "task_spec": task, + "role_bindings": {"role_bindings": bindings}, + } + monkeypatch.setattr( + module, + "_validate_grounded_plan", + lambda value: dict(value), + ) + + graph = ActionAgent().plan(grounded_plan) + direct = instantiate_seed_graph(task, bindings) + + assert seed_graph_hash(graph) == seed_graph_hash(direct) + + +def test_planning_only_graph_is_rejected_before_executor_construction() -> None: + task, requirements = _task_of_type("E6") + bindings = _bindings(requirements) + graph = instantiate_seed_graph(task, bindings) + constructed = False + + def executor_factory(*args, **kwargs): + nonlocal constructed + constructed = True + raise AssertionError("preflight must reject before executor construction") + + report = ActionAgent(executor_factory=executor_factory).execute( + graph, + SimpleNamespace(num_envs=2), + known_uids=set(bindings.values()), + run_id="preflight-test", + ) + + assert report.status == "rejected" + assert report.action_count == 0 + assert "planning-only" in (report.error or "") + assert not constructed + + +def test_execution_report_is_strictly_json_serializable(tmp_path: Path) -> None: + task, requirements = make_task_spec("E1") + bindings = _bindings(requirements) + graph = instantiate_seed_graph(task, bindings) + + class FakeExecutor: + def __init__(self, program, env, **kwargs) -> None: + self.program = program + self.env = env + + def run(self, **kwargs) -> ExecutionResult: + return ExecutionResult( + actions=[torch.ones((2, 3), dtype=torch.float32)], + success=torch.tensor([True, False]), + semantic_success={ + "task_01": torch.tensor([True, False]), + }, + record_dir=str(tmp_path), + retry_count=1, + retry_counts=[0, 1], + failure_events=[ + { + "failure_type": "plan_failed", + "env_ids": torch.tensor([1]), + } + ], + ) + + report = ActionAgent(executor_factory=FakeExecutor).execute( + graph, + SimpleNamespace(num_envs=2), + known_uids=set(bindings.values()), + run_id="json-test", + ) + payload = report.as_mapping() + + assert report.status == "failed" + assert payload["environments"][0]["semantic_success"] == {"task_01": True} + assert payload["environments"][1]["semantic_success"] == {"task_01": False} + assert [item["retry_count"] for item in payload["environments"]] == [0, 1] + assert "actions" not in payload + json.dumps(payload, allow_nan=False) + assert ( + json.loads((tmp_path / "execution_report.json").read_text(encoding="utf-8")) + == payload + ) + trajectory = torch.load(tmp_path / "executed_trajectory.pt", weights_only=True) + assert torch.equal(trajectory["actions"][0], torch.ones((2, 3))) + trajectory_manifest = json.loads( + (tmp_path / "executed_trajectory.json").read_text(encoding="utf-8") + ) + assert trajectory_manifest["actions"][0]["shape"] == [2, 3] + + +def test_existing_execution_result_can_be_reported_without_reexecution() -> None: + task, requirements = make_task_spec("E1") + bindings = _bindings(requirements) + graph = instantiate_seed_graph(task, bindings) + result = ExecutionResult( + actions=[torch.zeros((1, 2), dtype=torch.float32)], + success=torch.tensor([True]), + semantic_success={"task_01": torch.tensor([True])}, + ) + + report = ActionAgent().report_execution_result( + result, + action_graph=graph, + run_id="legacy-run", + episode_index=3, + ) + + assert report.status == "succeeded" + assert report.episode_id == "3" + assert report.action_count == 1 + + +def test_runtime_exception_is_reported_as_aborted() -> None: + task, requirements = make_task_spec("E1") + bindings = _bindings(requirements) + graph = instantiate_seed_graph(task, bindings) + + def fail_executor(*_args, **_kwargs): + raise RuntimeError("simulator stopped") + + report = ActionAgent(executor_factory=fail_executor).execute( + graph, + SimpleNamespace(num_envs=1), + known_uids=set(bindings.values()), + run_id="aborted-test", + ) + + assert report.status == "aborted" + assert report.action_count == 0 + assert report.error == "RuntimeError: simulator stopped" + + +def test_preflight_raises_for_planning_only_graph() -> None: + task, requirements = _task_of_type("E8") + bindings = _bindings(requirements) + + with pytest.raises(ValueError, match="planning-only"): + ActionAgent().preflight( + instantiate_seed_graph(task, bindings), + known_uids=set(bindings.values()), + ) diff --git a/tests/gen_sim/action_engine/test_architecture.py b/tests/gen_sim/action_engine/test_architecture.py new file mode 100644 index 000000000..818822124 --- /dev/null +++ b/tests/gen_sim/action_engine/test_architecture.py @@ -0,0 +1,149 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Guard the migration boundaries that make the rewrite meaningful.""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import embodichain.gen_sim.action_engine as action_engine_package +from embodichain.gen_sim.action_engine.capabilities import ( + build_atomic_capability_registry, + build_default_registry, +) +from embodichain.gen_sim.action_engine.protocol import ( + ACTION_ENGINE_CONFIG_SCHEMA, + ACTION_ENGINE_ENV_ID, + EXECUTION_PROGRAM_FILENAME, + SCENE_REQUIREMENTS_FILENAME, + SCENE_REQUIREMENTS_SCHEMA, + SEED_GRAPH_SCHEMA, + TASK_SPEC_FILENAME, + TASK_SPEC_SCHEMA, +) + +_PACKAGE_ROOT = Path(action_engine_package.__file__).resolve().parent +_LEGACY_PACKAGE = "embodichain.gen_sim.action_agent_pipeline" + + +def _production_python_files() -> list[Path]: + return sorted( + path + for path in _PACKAGE_ROOT.rglob("*.py") + if "tests" not in path.relative_to(_PACKAGE_ROOT).parts + ) + + +def test_production_code_has_no_legacy_pipeline_imports() -> None: + offenders: list[str] = [] + for path in _production_python_files(): + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + names = [alias.name for alias in node.names] + elif isinstance(node, ast.ImportFrom): + names = [node.module or ""] + else: + continue + if any(name.startswith(_LEGACY_PACKAGE) for name in names): + offenders.append(path.relative_to(_PACKAGE_ROOT).as_posix()) + break + assert offenders == [] + + +def test_protocol_identifiers_are_new_and_stable() -> None: + assert ACTION_ENGINE_ENV_ID == "ActionEngine-v1" + assert ACTION_ENGINE_CONFIG_SCHEMA == "action_engine_config_v2" + assert SEED_GRAPH_SCHEMA == "action_engine_seed_graph_v3" + assert TASK_SPEC_SCHEMA == "action_engine_task_spec_v2" + assert SCENE_REQUIREMENTS_SCHEMA == "action_engine_scene_requirements_v2" + assert EXECUTION_PROGRAM_FILENAME == "seed_task_graph.json" + assert TASK_SPEC_FILENAME == "task_spec.json" + assert SCENE_REQUIREMENTS_FILENAME == "scene_requirements.json" + + +def test_planner_exposes_exactly_the_first_phase_skill_catalog() -> None: + assert set(build_default_registry().operator_names()) == { + "arrange_line", + "build_stack", + "coordinated_transport", + "orient_object", + "place_relative", + } + + +def test_atomic_actions_have_one_runtime_capability_catalog() -> None: + registry = build_atomic_capability_registry() + assert set(registry.names()) == { + "CoordinatedPickment", + "CoordinatedPlacement", + "HandOver", + "MoveEndEffector", + "MoveHeldObject", + "MoveJoints", + "PickUp", + "Place", + "Pour", + "Press", + "PullArticulatedPart", + "PushArticulatedPart", + "TurnKnob", + } + assert set(registry.executable_names()) == { + "CoordinatedPickment", + "CoordinatedPlacement", + "HandOver", + "MoveEndEffector", + "MoveHeldObject", + "MoveJoints", + "PickUp", + "Place", + "Press", + } + + +def test_action_class_dispatch_is_not_duplicated_across_runtime_layers() -> None: + offenders = [] + for path in _production_python_files(): + if path.name == "atomic.py" and path.parent.name == "capabilities": + continue + source = path.read_text(encoding="utf-8") + if "_ACTION_TYPES" in source: + offenders.append(path.relative_to(_PACKAGE_ROOT).as_posix()) + assert offenders == [] + + +def test_runtime_core_has_no_action_name_dispatch_branches() -> None: + action_names = set(build_atomic_capability_registry().executable_names()) + offenders = {} + for relative in ( + "runtime/actions.py", + "runtime/executor.py", + "runtime/grounding.py", + ): + path = _PACKAGE_ROOT / relative + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + literals = { + node.value + for node in ast.walk(tree) + if isinstance(node, ast.Constant) and isinstance(node.value, str) + } + duplicated = sorted(literals & action_names) + if duplicated: + offenders[relative] = duplicated + assert offenders == {} diff --git a/tests/gen_sim/action_engine/test_unbound.py b/tests/gen_sim/action_engine/test_unbound.py new file mode 100644 index 000000000..5aa3980dc --- /dev/null +++ b/tests/gen_sim/action_engine/test_unbound.py @@ -0,0 +1,111 @@ +# ---------------------------------------------------------------------------- +# 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 copy import deepcopy + +import pytest + +from embodichain.gen_sim.action_engine.agent import ActionAgent +import embodichain.gen_sim.action_engine.agent as action_agent_module +from embodichain.gen_sim.action_engine.unbound import validate_unbound_action_plan + + +def _selector(reference: str) -> dict: + return { + "kind": "scene_ref", + "step_id": "", + "reference": reference, + "quantifier": "one", + "count": 0, + } + + +def _candidate() -> dict: + return { + "candidate_id": "candidate_01", + "draft": { + "task_id": "place_can", + "instruction": "Place the can on the table.", + "steps": [ + { + "id": "place", + "task_type": "E1", + "object": _selector("the can"), + "target": _selector("the table"), + "depends_on": [], + } + ], + }, + } + + +def test_action_agent_drafts_without_scene_uids() -> None: + candidate = _candidate() + original = deepcopy(candidate) + + draft = ActionAgent(registry=object()).draft(candidate) + + assert draft["candidate_id"] == "candidate_01" + assert draft["steps"][0]["object"]["reference"] == "the can" + assert "uid" not in str(draft).lower() + assert candidate == original + + +def test_unbound_plan_rejects_noncanonical_action_recipe() -> None: + draft = ActionAgent(registry=object()).draft(_candidate()) + draft["steps"][0]["actions"] = ["UnknownAction"] + + with pytest.raises(ValueError, match="task contract"): + validate_unbound_action_plan(draft) + + +def test_action_agent_rejects_missing_atomic_action_during_draft() -> None: + class Registry: + def names(self): + return () + + def executable_names(self): + return () + + with pytest.raises(ValueError, match="AtomicAction is not registered"): + ActionAgent(registry=Registry()).draft(_candidate()) + + +def test_bind_and_plan_requires_the_exact_unbound_candidate( + monkeypatch: pytest.MonkeyPatch, +) -> None: + agent = ActionAgent(registry=object()) + unbound = agent.draft(_candidate()) + grounded = { + "selected_candidate_id": "candidate_01", + "task_draft": deepcopy(_candidate()["draft"]), + } + monkeypatch.setattr( + action_agent_module, + "_validate_grounded_plan", + lambda value: deepcopy(value), + ) + monkeypatch.setattr(agent, "plan", lambda value: {"task": value["task_draft"]}) + + graph = agent.bind_and_plan(unbound, grounded) + assert graph["task"] == grounded["task_draft"] + + altered = deepcopy(unbound) + altered["instruction"] = "A different instruction." + with pytest.raises(ValueError, match="does not match"): + agent.bind_and_plan(altered, grounded)