diff --git a/embodichain/__main__.py b/embodichain/__main__.py index 3b897a3fa..0447031de 100644 --- a/embodichain/__main__.py +++ b/embodichain/__main__.py @@ -56,6 +56,11 @@ class Command: target="embodichain.gen_sim.scene_engine.cli.start:main", help="Generate a scene export from an input image using gen_sim/.env.", ), + Command( + name="task-engine", + target="embodichain.gen_sim.task_engine.cli:main", + help="Run a complete cross-engine task workflow.", + ), Command( name="preview-scene", target="embodichain.gen_sim.scene_engine.cli.preview:main", diff --git a/embodichain/gen_sim/action_engine/ARCHITECTURE.md b/embodichain/gen_sim/action_engine/ARCHITECTURE.md new file mode 100644 index 000000000..6ead4afa6 --- /dev/null +++ b/embodichain/gen_sim/action_engine/ARCHITECTURE.md @@ -0,0 +1,398 @@ +# Action Engine v2 Architecture + +Action Engine v2 uses a task-first protocol and executes a direct +`AtomicAction` graph. The persisted graph is symbolic and coordinate-free; +simulator geometry is resolved immediately before each action executes. + +The Task Engine entry point wraps that existing pipeline with three narrow +owners: + +1. `TaskAgent` produces three scene-independent `TaskDraft` candidates and + deterministically derives each `SceneRequest` and `SuccessSpec`. +2. `SceneAdapter` binds one verified candidate to a read-only existing scene, + producing `SceneManifest`, `RoleBindings`, and a complete `BindingReport`. +3. `ActionAgent` lowers the selected `GroundedTaskPlan` to the existing + `action_engine_seed_graph_v3`, performs executable capability preflight, + runs it through `ProgramExecutor`, and emits a tensor-free + `ExecutionReport`. Report v2 records the episode seed, package and Python + versions, Git commit/dirty state when available, and structured runtime + arguments alongside the existing plan and graph hashes. + +The public CLI is `python -m embodichain.gen_sim.task_engine --mode ...` with +strict `image`, `image-edit`, `scene`, and `scene-edit` input modes. Every +invocation runs the complete workflow through real trajectory acceptance and +publishes an isolated, timestamped child under `--output-root`. Source projects +are referenced in place and integrity-hashed; Task Engine does not modify them +or copy them into a scene package store. + +## Package Ownership + +The cross-engine workflow is owned by Task Engine rather than nested under +Action Engine: + +- `embodichain.gen_sim.task_engine` owns scene-independent interpretation, + E1-E9 semantic ontology, `TaskDraft`, `SceneRequest`, `SuccessSpec`, and + `TaskAgent`. +- `embodichain.gen_sim.scene_engine` remains the scene generation subsystem and + exposes auditable image-understanding, materialization, edit-understanding, + and edit-materialization stages. +- `embodichain.gen_sim.task_engine.scene` owns Scene Engine adaptation, the + richer static manifest, and deterministic scene/action feasibility reports. +- `embodichain.gen_sim.action_engine.agent` owns `ActionAgent`; Action Engine's + existing `domain`, `planning`, and `runtime` packages remain authoritative + for graph compilation and execution. +- `embodichain.gen_sim.task_engine.orchestration` owns cross-engine contracts, + read-only source references, scene adaptation, orchestration, and artifacts. + `embodichain.gen_sim.task_engine.cli` owns the unified CLI. + +The former `scene_bridge`, `collaboration`, and +`action_engine.collaboration` namespaces were removed; there are no import +bridges or fallback entry points. + +## Data Flow + +1. `TaskAgent` or a caller creates a validated `TaskSpec`. +2. Action Engine emits `SceneRequirements` for the external Scene Engine. +3. After scene generation, Task Engine's scene adapter preserves geometry, physics, + articulation, affordance evidence, and provenance in a versioned + `StaticSceneManifest` while the existing redacted manifest remains compatible. +4. `FeasibilityBroker` intersects the selected task, role bindings, static scene, + robot profile, and executable capability catalog without repairing unknowns. +5. Offline recipes and the online planner independently create complete + `SeedGraph` candidates whose nodes are `AtomicAction` calls. +6. Runtime preflight checks the capability catalog and rejects planning-only + actions before simulator motion starts. +7. `ActionGrounder` reads live robot, object, articulation, and camera state and + materializes the typed goal and immutable action options just in time. +8. `ProgramExecutor` schedules the DAG, executes vectorized action masks, and + verifies semantic postconditions from live state. + +There is no persisted semantic task graph between `TaskSpec` and `SeedGraph`. +The standalone TaskAgent v1 planner and compiler remain available to their +existing callers, but the generation pipeline neither accepts nor publishes +TaskAgent v1 JSON. + +## Protocols + +### TaskSpec + +`TaskSpec` owns the level, public instruction, E1-E9 task instances, +dependencies, path-independent success conditions, and a private oracle. The +online planner receives `public_task_spec(...)`, which removes the oracle and, +for L4, the hidden reference task instances. Online L4 TaskGroups are inferred +from the instruction and observations rather than matched to an oracle path. + +Levels classify how the task is specified, not action count: + +- L1: one E instance. +- L2: two or more instances of the same E type. +- L3: two or more different E types explicitly composed. +- L4: an abstract instruction that requires memory, visual semantics, pattern, + logic, common-sense, or constraint reasoning. + +Free-language L1-L3 generation uses two structured model calls. The first sees +only the instruction and E1-E9 catalog and emits typed steps whose scene +selectors are open natural-language references. The second sees those +references plus a coordinate-free semantic inventory and may return only +existing scene UIDs, status, and confidence. Local validation enforces complete +request coverage, candidate roles, cardinality, confidence, and non-self +targets; unresolved or ambiguous references fail instead of being guessed. +Each structured model stage gets at most one bounded repair attempt after an +invalid response. If repair still fails, or the model call itself fails, +generation stops before recipe expansion and artifact publication. It never +switches to keyword or rule-based instruction parsing. + +The older public `planning.plan_task` adapter still produces a standalone +`TaskAgent` from structured LLM output, but it does not reinterpret the +instruction after that output exists. Axis, orientation, and arm-allocation +fields come only from the validated model result. It has no keyword fallback. + +Generator callers without a configured LLM must provide a validated `TaskSpec` +with explicit role bindings (or a matching `SceneRequirements` sidecar). This +path is fully offline and never imports or calls an LLM client. + +### SceneRequirements + +`SceneRequirements` is the JSON hand-off to the external Scene Engine. It +declares object roles, counts, categories, affordances, initial states, spatial +constraints, camera requirements, and distractors. Scene results are never +silently repaired. Structural contradictions and explicit affordance +contradictions invalidate the task instance; an absent affordance declaration +remains unknown and is deferred to runtime physical validation. + +The current tabletop importer has one deliberately narrow structural contract: +exactly one `background` object is the support surface and receives runtime UID +`table`; every movable object is assumed to begin on that surface. Zero or +multiple backgrounds are rejected rather than resolved from position, UID, or +description text. Semantic `category`, `color`, and `attributes` come only from +their explicit scene fields. Physics `attrs` are not semantic metadata. + +For task-first inputs, explicit role bindings are authoritative unless they +contradict metadata that the scene actually declares. Automatic role binding +requires a unique match with complete structured category, attribute, state, +and affordance evidence. Object names and descriptions remain available to the +LLM grounding call, but deterministic validation never searches them for +semantic substrings. + +### StaticSceneManifest And FeasibilityReport + +`StaticSceneManifest` is an additive Scene Bridge artifact. It keeps the legacy +scene manifest stable while recording initial poses, geometry hashes, physics, +articulation payloads, structured affordance evidence, and provenance. Legacy +affordance strings become `declared` evidence; only structural facts derived by +the adapter are marked `verified`. + +`FeasibilityReport` classifies each check as `proven`, `runtime_probe`, `unknown`, +or `contradicted`. Missing evidence remains unknown. Declared geometric +affordances require a runtime probe, and unavailable AtomicActions are explicit +contradictions. A contradicted report publishes an `infeasible` audit result and +does not invoke graph or bundle generation. Executable preflight remains a +second authoritative gate before bundle generation. + +Scene Bridge reports arm-layout and whole-task pickup, handover, +target-interaction, and safety-clearance phases as `runtime_probe` evidence. +Arm-side compatibility is not claimed without live arm-base poses and workspace +geometry. + +### SeedGraph + +Every node directly names an `atomic_action`, scene `object_uid`, symbolic +`target_binding`, actor, control, dependencies, resources, pre/postconditions, +motion policy, E type, `task_instance_id`, and an Action Contract v2 +`failure_policy`. `task_required` and `safety_required` failures invalidate the +candidate; `best_effort` failures remain observable but do not erase an already +verified task and safety result. `TaskGroup` groups all nodes of one E instance +with `role=primary|recovery`; it is metadata over the same DAG, not a second +graph. + +Validation guarantees: + +- node and TaskGroup dependencies are DAGs; +- every node belongs to exactly one TaskGroup; +- E groups contain their required core actions; +- concurrent nodes do not claim the same exclusive arm/object resource; +- object references resolve to scene UIDs; +- world poses, qpos, trajectories, grasp poses, and waypoints are rejected + recursively; +- hashes use canonical strict JSON and are stable across processes. + +The production loader accepts v3 graphs only. An older graph, whether supplied as +JSON or an in-memory mapping, receives an explicit regeneration error rather +than an implicit migration. + +## Capability Boundary + +`AtomicCapabilityRegistry` is the single runtime catalog. A descriptor declares +the action/option types, accepted symbolic bindings and controls, resource mode, +held-object state effect, target and config materializers, verifier, failure +classifier, retry mode, and runtime availability. + +The executable catalog currently contains: + +- `PickUp`, `MoveHeldObject`, `MoveEndEffector`, `MoveJoints`, and `Place` +- `Press` +- `CoordinatedPickment` and `CoordinatedPlacement` +- `HandOver` + +`Pour`, `PullArticulatedPart`, `PushArticulatedPart`, and `TurnKnob` are +planning-only until matching lower-level implementations exist. They can be +generated and statically checked, but preflight fails before any motion with +the descriptor's unavailable reason. + +Adding an executable skill consists of registering its descriptor and reusable +materializer/verifier hooks plus focused tests. Planner and executor dispatch +do not maintain a parallel action-class table. + +## Offline And Online Planning + +Offline recipes deterministically instantiate E1-E9 task instances. Current +task mappings are: + +- `place_relative -> E1` +- `orient_object -> E2` +- `coordinated_transport -> E5` +- every member of `build_stack` and `arrange_line` -> one E1 instance + +E5 uses `coordinated_transport` only as the semantic task-group operator. Its +motion graph contains one `CoordinatedPickment`; a `place` terminal behavior +adds synchronized left/right `MoveJoints(gripper_open)` nodes. The executor +clears coordinated hold state only after both grippers are observed open. + +The online path first extracts auditable visual facts from multi-view RGB and, +when available, depth and camera calibration. Facts contain only known UIDs, +normalized bboxes/keypoints, canonical spatial relations, task predicates, and +confidence. Spatial relations use a shared ontology and fixed participant +order. Task-level judgments such as visual or pattern completion are accepted +only when the current `TaskSpec.success` explicitly requests them. A second +structured call produces a complete direct `AtomicAction` graph. Prompts +request facts and graph JSON only; hidden chain-of-thought is neither requested +nor stored. + +Image-space constraints may use normalized keypoints, masks, bboxes, and +relative relations. The Grounder uses live depth and camera calibration to +convert them to world targets. The SeedGraph never stores that result. + +## JIT Grounding + +Each action is grounded again immediately before planning/execution. Grounding +therefore observes object displacement, current qpos, current held-object +ownership, articulation state, and fresh camera measurements. Coordinated and +handover actions are grounded as synchronized execution units. Automatic arm +selection, collision checks, live arrangement slots, and current predicate +semantics remain deterministic runtime responsibilities. + +Arm allocation uses the live right-to-left arm-base axis and the live table +center. The preference therefore follows translated and rotated robot +workspaces; live motion planning remains authoritative for reachability. + +Placement support is a relation, not an entity category. Static adaptation +accepts rigid `physical_entity` targets without requiring a `support_surface` +affordance. Articulations require a link-level runtime target interface and are +rejected until that interface is available. Runtime evaluates +`object_supported_by(payload, support, pose)` from live geometry and center of +mass, applies the requested `orientation_goal`, and requires low motion across a +bounded stability window. Successful relations form a per-environment support +graph that is checked for cycles and revalidated at task completion. + +Orientation is compiled into hard `align_axis` or `match_rotation` terms plus a +separate minimum-rotation planning preference. An omitted orientation request +adds no hard acceptance term; `preserve` remains an explicit full-rotation +contract for persisted bundles, while `upright` constrains only the requested +local axis and declares whether that axis is directed. Grounding and runtime +verification consume the same compiled contract so reachability search cannot +silently relax a required terminal orientation. With no hard term, a live +upright state may still select upright-preserving yaw candidates as a planning +preference; this follows current state and automatically stops after that state +is invalidated rather than becoming a sticky success requirement. + +Grasp generation keeps support-plane collision filtering as its strict first +pass. If diagnostics show that this heuristic alone exhausted otherwise +object-collision-free candidates, Action Engine retries without the heuristic; +the relaxed candidates still pass through the live robot and scene collision +planner before execution. This avoids treating a local support-plane proxy as +a proof of scene-level infeasibility, including for objects already held above +the support surface. + +Grounding samples bounded support-relative placement poses. Planning failures +try the next pose before release; instability after release requires a fresh +grasp and an unused pose. The recovery keeps the original actor contract, and +its edges and failure provenance are recorded separately from the primary +attempt. + +## Mainline Planning Contract + +The runtime keeps only an Action Engine-local `ExecutionState` for full-robot +qpos and held-object relations. Each plan converts that state to the mainline +`PlanningContext` (`RobotObservation`, `TaskState`, and `SceneSnapshot`) and +submits an `ActionInvocation` to `AtomicActionEngine`. The returned +`StateDelta` remains speculative until physical and semantic verification; only +verified vectorized rows are committed. + +`AtomicActionAdapter` accepts a shared `SceneProvider` and otherwise builds a +`RigidObjectSceneProvider` from live simulation entities. Planning snapshots now +carry monotonic timestamps and material-change scene/collision revisions. The +adapter also exposes `start_session(...)` for callers adopting +`ExecutionSession`; the existing compound and per-arm merged trajectory +scheduler remains as the compatibility execution path. + +Single-arm arm motion uses cuRobo `motion_gen` by default. Hand-only and +coordinated dual-arm actions use `ik_interp`, because mainline coordinated +primitives do not support cuRobo motion generation. A failed single-arm cuRobo +row may fall back to `ik_interp` without replacing successful rows. Generated +background objects form the static cuRobo collision world; dynamic obstacles +are an explicit runtime-policy opt-in. + +Generated mesh objects carry V-HACD settings in both the current shape-level +schema and legacy top-level fields. Before antipodal grasp construction, the +runtime prepares a checksummed V-HACD payload at the shared collision-checker +cache path so the unchanged mainline checker does not silently recompute CoACD. +Grasp generation samples multiple deviated approach directions and filters them +through the existing gripper collision model. Safety retreat planning searches +a bounded set of live height and baseward targets instead of treating one exact +height as a geometric reachability certificate. + +## A/B Evaluation + +Test mode retains both candidates. `run_strict_ab` creates distinct offline and +online environments with the same task, scene configuration, seed, Grounder, +verifiers, and retry policy. Both environments reset before execution and a +digest over robot qpos and object state must match exactly; a mismatch aborts +before either branch executes. + +Artifacts are written under `offline/` and `online/`, with a shared +`comparison.json`. The comparison records graph hashes/differences, action and +path lengths, success, retries, recoveries, revisions, latency, record paths, +and planner/VLM metadata supplied by each candidate. + +L4 A/B runs must supply a private-oracle evaluator. The built-in evaluator +checks memory reconstruction, visual completion, pattern completion, numeric +selection, functional placement, and stable/unobstructed goals from the final +state only. The comparison labels whether success came from runtime step +postconditions or the private oracle. + +## Dynamic Recovery + +The persisted `SeedGraph` is immutable. `RuntimeGraph` keeps a detached working +copy and an ordered revision log. One failed `AtomicAction` can be freshly +grounded and retried twice, for three total attempts, and only while its live +precondition remains true. + +Failures use the bounded taxonomy `search_exhausted`, `plan_failed`, +`grasp_missed`, `object_fallen`, `object_dropped`, and +`postcondition_failed`. `search_exhausted` records the blocking edge, planning +stage, strategy, finite budget, and observed evidence; it does not claim that a +target is geometrically unreachable. Known recoverable states can insert a +complete `role=recovery` TaskGroup, such as an E2 upright group. Recovery keeps +the failed TaskGroup's actor contract, and primary, recovery, and replay events +are recorded separately. After recovery, the selected route replans only the +unfinished suffix. Offline and online dynamic replanners are explicit, +separate modes. Revision, recovery-action, transition, and retry budgets bound +every loop. + +## Selection And Fusion + +Product mode statically scores offline and online candidates using schema +validity, capability availability, UID validity, task coverage, visual +confidence, and estimated action cost. Exact mature-template matches favor the +offline route; L4 visual tasks favor sufficiently confident online results. + +Fusion is conservative. It may choose only complete `TaskGroup` units, rewires +dependencies at group boundaries, and rejects unordered state changes to the +same object. It never splits one E instance across candidates. + +## Artifacts + +A normal generated bundle contains: + +- `task_spec.json` +- `scene_requirements.json` +- `seed_task_graph.json` +- `seed_task_graph.png` +- `agent_config.json` +- `fast_gym_config.json` + +Strict A/B adds branch-local graph/result artifacts and `comparison.json`. +Review graphs, runtime records, and videos never become execution inputs. + +`prepare` lowers and preflights resolved semantic candidates in selection order. +A candidate-local lowering, symbolic planning, or preflight error rejects only +that candidate. If no resolved candidate is executable, the transaction +publishes `preparation_failure.json` with each attempted draft, verified +bindings, available grounded plan, failure stage, and exception instead of +leaving an older successful bundle in place. + +## Invariants + +- SeedGraph nodes are direct AtomicActions, not E-level operators. +- Lowering uses original instruction-step order as the stable tie-break among + dependency-ready steps. Independent steps remain independent; the contract + linker serializes only actual resource conflicts. +- E labels are subgraph grouping semantics only. +- Planning artifacts contain no grounded motion coordinates. +- Online planning never receives the private oracle. +- Runtime uses one capability registry for preflight, Grounding, config + construction, execution, verification policy, and recovery policy. +- Required arms are never silently replaced. +- Failed or inactive vectorized rows preserve their last valid state. +- Current five task families preserve their v1 AtomicAction topology and live + Grounding behavior after regeneration. 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/capabilities/__init__.py b/embodichain/gen_sim/action_engine/capabilities/__init__.py new file mode 100644 index 000000000..5e7b6935e --- /dev/null +++ b/embodichain/gen_sim/action_engine/capabilities/__init__.py @@ -0,0 +1,57 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Semantic operator and atomic-action capability registry.""" + +from __future__ import annotations + +from .atomic import ( + ACTION_CONTRACT_VERSION, + AtomicCapability, + AtomicCapabilityRegistry, + ResolvedActionContract, + ResourceClaim, + StateAtom, + StateEffect, + build_atomic_capability_registry, + capability_precondition, +) +from .builtins import build_default_registry +from .registry import ( + ActionCapability, + ActionTemplate, + CapabilityRegistry, + OperatorCapability, + PhaseTemplate, +) + +__all__ = [ + "ACTION_CONTRACT_VERSION", + "ActionCapability", + "ActionTemplate", + "AtomicCapability", + "AtomicCapabilityRegistry", + "CapabilityRegistry", + "OperatorCapability", + "PhaseTemplate", + "ResolvedActionContract", + "ResourceClaim", + "StateAtom", + "StateEffect", + "build_atomic_capability_registry", + "build_default_registry", + "capability_precondition", +] diff --git a/embodichain/gen_sim/action_engine/capabilities/atomic.py b/embodichain/gen_sim/action_engine/capabilities/atomic.py new file mode 100644 index 000000000..2e2697816 --- /dev/null +++ b/embodichain/gen_sim/action_engine/capabilities/atomic.py @@ -0,0 +1,915 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Single-source AtomicAction capability descriptors.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +import hashlib +import json +from typing import Any + +import torch + +__all__ = [ + "ACTION_CONTRACT_VERSION", + "AtomicCapability", + "AtomicCapabilityRegistry", + "ResolvedActionContract", + "ResourceClaim", + "StateAtom", + "StateEffect", + "build_atomic_capability_registry", + "capability_precondition", +] + +_RETRY_MODES = frozenset({"direct", "recover_then_retry", "non_retryable"}) +ACTION_CONTRACT_VERSION = "action_contract_v2" +_PREDICATES = frozenset( + { + "arm_free", + "object_free", + "object_held", + "object_coordinated_held", + "handover_complete", + "arm_clear", + "arm_home", + } +) +_EFFECT_OPERATIONS = frozenset({"add", "delete"}) +_RESOURCE_ACCESS = frozenset({"shared_read", "exclusive"}) +_RESOURCE_LIFETIMES = frozenset({"action", "until_release"}) +_COMPLETION_MODES = frozenset({"ordinary", "cleanup", "terminal_barrier"}) +_FAILURE_POLICIES = frozenset({"task_required", "safety_required", "best_effort"}) + + +@dataclass(frozen=True) +class StateAtom: + """One symbolic state fact used by an Action Contract.""" + + predicate: str + object_uid: str | None = None + arm: str | None = None + + def __post_init__(self) -> None: + if self.predicate not in _PREDICATES: + raise ValueError(f"Unknown Action Contract predicate {self.predicate!r}.") + if self.object_uid is not None and not self.object_uid: + raise ValueError("StateAtom.object_uid must not be empty.") + if self.arm is not None and not self.arm: + raise ValueError("StateAtom.arm must not be empty.") + + def as_mapping(self) -> dict[str, str]: + """Return the stable JSON representation of this fact.""" + result = {"predicate": self.predicate} + if self.object_uid is not None: + result["object_uid"] = self.object_uid + if self.arm is not None: + result["arm"] = self.arm + return result + + +@dataclass(frozen=True) +class StateEffect: + """Add or delete one symbolic state fact.""" + + op: str + atom: StateAtom + + def __post_init__(self) -> None: + if self.op not in _EFFECT_OPERATIONS: + raise ValueError(f"Unknown Action Contract effect operation {self.op!r}.") + + def as_mapping(self) -> dict[str, Any]: + """Return the stable JSON representation of this effect.""" + return {"op": self.op, "atom": self.atom.as_mapping()} + + +@dataclass(frozen=True) +class ResourceClaim: + """One resource access claim made by an AtomicAction.""" + + resource: str + access: str = "exclusive" + lifetime: str = "action" + + def __post_init__(self) -> None: + if not self.resource: + raise ValueError("ResourceClaim.resource must not be empty.") + if self.access not in _RESOURCE_ACCESS: + raise ValueError(f"Unknown resource access mode {self.access!r}.") + if self.lifetime not in _RESOURCE_LIFETIMES: + raise ValueError(f"Unknown resource lifetime {self.lifetime!r}.") + + def as_mapping(self) -> dict[str, str]: + """Return the stable JSON representation of this claim.""" + return { + "resource": self.resource, + "access": self.access, + "lifetime": self.lifetime, + } + + +@dataclass(frozen=True) +class ResolvedActionContract: + """Fully resolved, serializable contract for one action node.""" + + requires: tuple[StateAtom, ...] = () + effects: tuple[StateEffect, ...] = () + claims: tuple[ResourceClaim, ...] = () + completion: str = "ordinary" + failure_policy: str = "task_required" + version: str = ACTION_CONTRACT_VERSION + + def __post_init__(self) -> None: + if self.version != ACTION_CONTRACT_VERSION: + raise ValueError( + f"Unsupported Action Contract version {self.version!r}; " + f"expected {ACTION_CONTRACT_VERSION!r}." + ) + if self.completion not in _COMPLETION_MODES: + raise ValueError(f"Unknown Action Contract completion {self.completion!r}.") + if self.failure_policy not in _FAILURE_POLICIES: + raise ValueError( + f"Unknown Action Contract failure policy {self.failure_policy!r}." + ) + + def as_mapping(self) -> dict[str, Any]: + """Return the stable JSON representation persisted in SeedGraph v3.""" + return { + "version": self.version, + "requires": [atom.as_mapping() for atom in self.requires], + "effects": [effect.as_mapping() for effect in self.effects], + "claims": [claim.as_mapping() for claim in self.claims], + "completion": self.completion, + "failure_policy": self.failure_policy, + } + + +@dataclass(frozen=True) +class AtomicCapability: + """Describe planning, grounding, execution, and recovery for one skill.""" + + name: str + action_type: type | None + config_type: type | None + binding_kinds: frozenset[str] + controls: frozenset[str] + resource_mode: str + state_effect: str + target_materializer: str + motion_base: str | None = None + config_materializer: str = "single_arm" + verifier: str = "postcondition" + failure_classifier: str = "default" + retry_mode: str = "direct" + runtime_available: bool = True + unavailable_reason: str | None = None + target_materializer_hook: Callable[..., Any] | None = None + config_materializer_hook: Callable[..., Any] | None = None + verifier_hook: Callable[..., Any] | None = None + failure_classifier_hook: Callable[..., str] | None = None + contract_resolver_hook: ( + Callable[[Mapping[str, Any]], ResolvedActionContract] | None + ) = None + allows_target_contact: bool = False + """Whether motion planning may temporarily exclude the action target.""" + + def __post_init__(self) -> None: + if not self.name: + raise ValueError("AtomicCapability.name must not be empty.") + if self.motion_base is not None and not self.motion_base: + raise ValueError( + f"AtomicCapability {self.name!r} motion_base must not be empty." + ) + if not self.binding_kinds or not self.controls: + raise ValueError( + f"AtomicCapability {self.name!r} requires bindings and controls." + ) + if self.retry_mode not in _RETRY_MODES: + raise ValueError( + f"AtomicCapability {self.name!r} has invalid retry_mode {self.retry_mode!r}." + ) + if not isinstance(self.allows_target_contact, bool): + raise TypeError("allows_target_contact must be a boolean.") + if self.runtime_available: + if self.action_type is None or self.config_type is None: + raise ValueError( + f"Executable AtomicCapability {self.name!r} requires action/config types." + ) + if self.unavailable_reason is not None: + raise ValueError( + f"Executable AtomicCapability {self.name!r} cannot have an unavailable reason." + ) + elif not self.unavailable_reason: + raise ValueError( + f"Planning-only AtomicCapability {self.name!r} requires unavailable_reason." + ) + for field_name in ( + "target_materializer_hook", + "config_materializer_hook", + "verifier_hook", + "failure_classifier_hook", + "contract_resolver_hook", + ): + value = getattr(self, field_name) + if value is not None and not callable(value): + raise TypeError( + f"AtomicCapability {self.name!r} {field_name} must be callable." + ) + + def resolve_contract(self, node: Mapping[str, Any]) -> ResolvedActionContract: + """Resolve the deterministic Action Contract for one bound node.""" + if self.contract_resolver_hook is not None: + contract = self.contract_resolver_hook(node) + if not isinstance(contract, ResolvedActionContract): + raise TypeError( + f"AtomicCapability {self.name!r} contract resolver must return " + "ResolvedActionContract." + ) + return contract + return _resolve_default_contract(self, node) + + def as_catalog_entry(self) -> dict[str, Any]: + """Return the stable, JSON-safe planning view of this capability.""" + return { + "name": self.name, + "binding_kinds": sorted(self.binding_kinds), + "controls": sorted(self.controls), + "resource_mode": self.resource_mode, + "state_effect": self.state_effect, + "target_materializer": self.target_materializer, + "motion_base": self.motion_base or self.name, + "config_materializer": self.config_materializer, + "verifier": self.verifier, + "failure_classifier": self.failure_classifier, + "retry_mode": self.retry_mode, + "runtime_available": self.runtime_available, + "unavailable_reason": self.unavailable_reason, + "allows_target_contact": self.allows_target_contact, + "custom_target_materializer": _callable_name(self.target_materializer_hook), + "custom_config_materializer": _callable_name(self.config_materializer_hook), + "custom_verifier": _callable_name(self.verifier_hook), + "custom_failure_classifier": _callable_name(self.failure_classifier_hook), + "contract_version": ACTION_CONTRACT_VERSION, + "contract_resolver": _callable_name(self.contract_resolver_hook) + or f"{__name__}._resolve_default_contract", + } + + +class AtomicCapabilityRegistry: + """Strict registry shared by planners, validators, grounders, and runtime.""" + + def __init__(self) -> None: + self._capabilities: dict[str, AtomicCapability] = {} + + def register(self, capability: AtomicCapability) -> None: + if capability.name in self._capabilities: + raise ValueError( + f"AtomicCapability {capability.name!r} is already registered." + ) + self._capabilities[capability.name] = capability + + def get(self, name: str) -> AtomicCapability: + try: + return self._capabilities[name] + except KeyError as exc: + raise ValueError( + f"Unknown AtomicAction {name!r}; available actions are {list(self.names())}." + ) from exc + + def require_executable(self, name: str) -> AtomicCapability: + capability = self.get(name) + if not capability.runtime_available: + raise ValueError( + f"AtomicAction {name!r} is planning-only and cannot be executed: " + f"{capability.unavailable_reason}" + ) + return capability + + def validate_binding(self, action: Mapping[str, Any]) -> None: + name = str(action.get("atomic_action", action.get("atomic_action_class", ""))) + capability = self.get(name) + binding = action.get("target_binding") + if not isinstance(binding, Mapping): + raise ValueError( + f"AtomicAction {name!r} requires a target_binding mapping." + ) + kind = str(binding.get("kind", "")) + if kind not in capability.binding_kinds: + raise ValueError( + f"AtomicAction {name!r} does not accept binding kind {kind!r}; " + f"expected one of {sorted(capability.binding_kinds)}." + ) + control = str(action.get("control", "arm")) + if control not in capability.controls: + raise ValueError( + f"AtomicAction {name!r} does not support control {control!r}; " + f"expected one of {sorted(capability.controls)}." + ) + + def names(self) -> tuple[str, ...]: + return tuple(sorted(self._capabilities)) + + def executable_names(self) -> tuple[str, ...]: + return tuple( + name for name in self.names() if self._capabilities[name].runtime_available + ) + + def catalog(self) -> dict[str, dict[str, Any]]: + return { + name: self._capabilities[name].as_catalog_entry() for name in self.names() + } + + def catalog_hash(self) -> str: + payload = json.dumps( + self.catalog(), sort_keys=True, separators=(",", ":"), allow_nan=False + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def build_atomic_capability_registry() -> AtomicCapabilityRegistry: + """Build the default catalog, including explicit planning-only skills.""" + from embodichain.lab.sim.atomic_actions import ( + CoordinatedPickment, + CoordinatedPickmentOptions, + CoordinatedPlacement, + CoordinatedPlacementOptions, + HandOver, + HandOverOptions, + MoveEndEffector, + MoveEndEffectorOptions, + MoveHeldObject, + MoveHeldObjectOptions, + MoveJoints, + MoveJointsOptions, + PickUp, + PickUpOptions, + Place, + PlaceOptions, + Press, + PressOptions, + ) + + registry = AtomicCapabilityRegistry() + definitions = ( + AtomicCapability( + "PickUp", + PickUp, + PickUpOptions, + frozenset({"object"}), + frozenset({"arm"}), + "single_arm_object", + "hold", + "object_grasp", + verifier="held_object", + failure_classifier="grasp", + allows_target_contact=True, + ), + AtomicCapability( + "MoveHeldObject", + MoveHeldObject, + MoveHeldObjectOptions, + frozenset({"semantic_goal", "visual_constraint", "handover_staging"}), + frozenset({"arm"}), + "single_arm_object", + "preserve_hold", + "semantic_held_object", + ), + AtomicCapability( + "MoveEndEffector", + MoveEndEffector, + MoveEndEffectorOptions, + frozenset({"policy_pose", "visual_constraint"}), + frozenset({"arm"}), + "single_arm", + "preserve", + "eef_pose", + verifier_hook=_verify_arm_clearance, + contract_resolver_hook=_resolve_end_effector_contract, + ), + AtomicCapability( + "MoveJoints", + MoveJoints, + MoveJointsOptions, + frozenset({"joint_state"}), + frozenset({"arm", "hand"}), + "control_part", + "preserve", + "joint_state", + contract_resolver_hook=_resolve_joints_contract, + ), + AtomicCapability( + "Place", + Place, + PlaceOptions, + frozenset({"current_held_pose"}), + frozenset({"arm"}), + "single_arm_object", + "release", + "current_held_pose", + ), + AtomicCapability( + "Press", + Press, + PressOptions, + frozenset({"object", "semantic_goal"}), + frozenset({"arm"}), + "single_arm_object", + "preserve", + "press", + verifier="pressed", + allows_target_contact=True, + ), + AtomicCapability( + "CoordinatedPickment", + CoordinatedPickment, + CoordinatedPickmentOptions, + frozenset({"object", "coordinated_goal"}), + frozenset({"coordinated"}), + "coordinated_object", + "coordinated_hold", + "coordinated_pickment", + config_materializer="coordinated_pickment", + verifier="coordinated_hold", + failure_classifier="grasp", + ), + AtomicCapability( + "CoordinatedPlacement", + CoordinatedPlacement, + CoordinatedPlacementOptions, + frozenset({"coordinated_placement_goal"}), + frozenset({"coordinated"}), + "coordinated_object", + "coordinated_release", + "coordinated_placement", + config_materializer="coordinated_placement", + ), + AtomicCapability( + "HandOver", + HandOver, + HandOverOptions, + frozenset({"handover_goal"}), + frozenset({"coordinated"}), + "coordinated_object", + "transfer_hold", + "handover", + config_materializer="handover", + verifier="receiver_holds", + failure_classifier="handover", + retry_mode="recover_then_retry", + ), + ) + for capability in definitions: + registry.register(capability) + + unavailable = "The corresponding AtomicAction is not implemented in embodichain.lab.sim.atomic_actions." + for name, binding, effect in ( + ("Pour", "pour_goal", "preserve_hold"), + ("PullArticulatedPart", "articulation_goal", "articulation_change"), + ("PushArticulatedPart", "articulation_goal", "articulation_change"), + ("TurnKnob", "articulation_goal", "articulation_change"), + ): + registry.register( + AtomicCapability( + name, + None, + None, + frozenset({binding}), + frozenset({"arm"}), + "single_arm_object", + effect, + binding, + retry_mode="non_retryable", + runtime_available=False, + unavailable_reason=unavailable, + ) + ) + return registry + + +def capability_precondition( + capability: AtomicCapability, + *, + object_uid: str, + actor: Mapping[str, Any], + target_binding: Mapping[str, Any], +) -> dict[str, Any]: + """Build the generic live precondition used to authorize a retry.""" + if target_binding.get("coordinated_release_role") is not None: + # Opening a gripper is idempotent. A retry must remain legal when one + # hand opened on the first attempt and the physical dual-hold predicate + # therefore no longer holds. + return {} + if capability.state_effect == "coordinated_release": + return {"type": "held_by_both_grippers", "object": object_uid} + if capability.state_effect in {"preserve_hold", "release", "transfer_hold"}: + result = {"type": "object_held", "object": object_uid} + arm = target_binding.get("transfer_arm") + if arm is None and actor.get("mode") in {"required", "preferred"}: + arm = actor.get("arm") + if isinstance(arm, str) and arm: + result["arm"] = arm + return result + return {} + + +def _resolve_default_contract( + capability: AtomicCapability, node: Mapping[str, Any] +) -> ResolvedActionContract: + object_uid = _required_string(node.get("object_uid"), "node.object_uid") + actor = node.get("actor", {}) + if not isinstance(actor, Mapping): + raise ValueError("Action Contract resolution requires a mapping actor.") + binding = node.get("target_binding", {}) + if not isinstance(binding, Mapping): + raise ValueError("Action Contract resolution requires a target_binding.") + arms = _actor_arms(actor) + arm = arms[0] if len(arms) == 1 else None + arm_claims = tuple(ResourceClaim(f"arm:{item}") for item in arms) + object_claim = ResourceClaim(f"object:{object_uid}") + payload_claims = _payload_resource_claims(binding, object_uid) + + if capability.name == "PickUp": + required_arm = _required_arm(arm, capability.name) + return ResolvedActionContract( + requires=( + StateAtom("arm_free", arm=required_arm), + StateAtom("object_free", object_uid=object_uid), + ), + effects=( + StateEffect("delete", StateAtom("arm_free", arm=required_arm)), + StateEffect("delete", StateAtom("object_free", object_uid=object_uid)), + StateEffect( + "add", + StateAtom("object_held", object_uid=object_uid, arm=required_arm), + ), + ), + claims=( + ResourceClaim(f"arm:{required_arm}", lifetime="until_release"), + ResourceClaim(f"object:{object_uid}", lifetime="until_release"), + ) + + payload_claims, + ) + if capability.name == "MoveHeldObject": + required_arm = _required_arm(arm, capability.name) + return ResolvedActionContract( + requires=( + StateAtom("object_held", object_uid=object_uid, arm=required_arm), + ), + claims=( + ResourceClaim(f"arm:{required_arm}", lifetime="until_release"), + ResourceClaim(f"object:{object_uid}", lifetime="until_release"), + ) + + payload_claims, + ) + if capability.name == "Place": + required_arm = _required_arm(arm, capability.name) + return ResolvedActionContract( + requires=( + StateAtom("object_held", object_uid=object_uid, arm=required_arm), + ), + effects=( + StateEffect( + "delete", + StateAtom("object_held", object_uid=object_uid, arm=required_arm), + ), + StateEffect("add", StateAtom("arm_free", arm=required_arm)), + StateEffect("add", StateAtom("object_free", object_uid=object_uid)), + ), + claims=arm_claims + (object_claim,) + payload_claims, + ) + if capability.name == "HandOver": + transfer = _required_string( + binding.get("transfer_arm"), "target_binding.transfer_arm" + ) + receive = _required_string( + binding.get("receive_arm"), "target_binding.receive_arm" + ) + if transfer == receive: + raise ValueError("HandOver requires distinct transfer and receive arms.") + return ResolvedActionContract( + requires=( + StateAtom("object_held", object_uid=object_uid, arm=transfer), + StateAtom("arm_free", arm=receive), + ), + effects=( + StateEffect( + "delete", + StateAtom("object_held", object_uid=object_uid, arm=transfer), + ), + StateEffect("delete", StateAtom("arm_free", arm=receive)), + StateEffect("add", StateAtom("arm_free", arm=transfer)), + StateEffect( + "add", + StateAtom("object_held", object_uid=object_uid, arm=receive), + ), + StateEffect( + "add", StateAtom("handover_complete", object_uid=object_uid) + ), + ), + claims=( + ResourceClaim(f"arm:{transfer}"), + ResourceClaim(f"arm:{receive}", lifetime="until_release"), + ResourceClaim(f"object:{object_uid}", lifetime="until_release"), + ), + ) + if capability.name == "CoordinatedPickment": + coordinated_arms = _coordinated_arms(arms, capability.name) + requires = tuple(StateAtom("arm_free", arm=item) for item in coordinated_arms) + effects = tuple( + StateEffect("delete", StateAtom("arm_free", arm=item)) + for item in coordinated_arms + ) + ( + StateEffect("delete", StateAtom("object_free", object_uid=object_uid)), + StateEffect( + "add", StateAtom("object_coordinated_held", object_uid=object_uid) + ), + ) + claims = tuple( + ResourceClaim(f"arm:{item}", lifetime="until_release") + for item in coordinated_arms + ) + ( + ResourceClaim(f"object:{object_uid}", lifetime="until_release"), + ) + return ResolvedActionContract( + requires=requires + (StateAtom("object_free", object_uid=object_uid),), + effects=effects, + claims=claims + payload_claims, + ) + if capability.name == "CoordinatedPlacement": + coordinated_arms = _coordinated_arms(arms, capability.name) + effects = ( + StateEffect( + "delete", StateAtom("object_coordinated_held", object_uid=object_uid) + ), + StateEffect("add", StateAtom("object_free", object_uid=object_uid)), + ) + tuple( + StateEffect("add", StateAtom("arm_free", arm=item)) + for item in coordinated_arms + ) + claims = tuple(ResourceClaim(f"arm:{item}") for item in coordinated_arms) + ( + object_claim, + ) + return ResolvedActionContract( + requires=(StateAtom("object_coordinated_held", object_uid=object_uid),), + effects=effects, + claims=claims + payload_claims, + ) + + requirements: tuple[StateAtom, ...] = () + if capability.state_effect == "coordinated_hold": + requirements = (StateAtom("object_free", object_uid=object_uid),) + elif capability.state_effect == "coordinated_release": + requirements = (StateAtom("object_coordinated_held", object_uid=object_uid),) + elif capability.resource_mode in {"single_arm", "single_arm_object"}: + required_arm = _required_arm(arm, capability.name) + requirements = (StateAtom("arm_free", arm=required_arm),) + claims = arm_claims + if "object" in capability.resource_mode: + claims += (object_claim,) + return ResolvedActionContract( + requires=requirements, + claims=claims + payload_claims, + ) + + +def _payload_resource_claims( + binding: Mapping[str, Any], object_uid: str +) -> tuple[ResourceClaim, ...]: + """Resolve exclusive claims for objects physically carried by a carrier.""" + raw_payloads = binding.get("payloads", ()) + if not isinstance(raw_payloads, Sequence) or isinstance( + raw_payloads, (str, bytes, bytearray) + ): + raise ValueError("target_binding.payloads must be a list.") + payload_uids: list[str] = [] + for index, raw_payload in enumerate(raw_payloads): + value = ( + raw_payload.get("object") + if isinstance(raw_payload, Mapping) + else raw_payload + ) + if not isinstance(value, str) or not value: + raise ValueError( + f"target_binding.payloads[{index}] requires an object UID." + ) + if value == object_uid: + raise ValueError("An AtomicAction carrier cannot be its own payload.") + payload_uids.append(value) + if len(payload_uids) != len(set(payload_uids)): + raise ValueError("target_binding payload objects must be unique.") + return tuple(ResourceClaim(f"object:{uid}") for uid in payload_uids) + + +def _resolve_end_effector_contract( + node: Mapping[str, Any], +) -> ResolvedActionContract: + object_uid = _required_string(node.get("object_uid"), "node.object_uid") + actor = node.get("actor", {}) + binding = node.get("target_binding", {}) + if not isinstance(actor, Mapping) or not isinstance(binding, Mapping): + raise ValueError("MoveEndEffector contract requires actor and target_binding.") + arm = _required_arm(_actor_arms(actor)[0], "MoveEndEffector") + if binding.get("operation") == "retreat" or node.get("role") == "cleanup": + requires = [StateAtom("arm_free", arm=arm)] + if binding.get("source") == "handover": + requires.append(StateAtom("handover_complete", object_uid=object_uid)) + return ResolvedActionContract( + requires=tuple(requires), + effects=(StateEffect("add", StateAtom("arm_clear", arm=arm)),), + claims=(ResourceClaim(f"arm:{arm}"),), + completion="cleanup", + failure_policy="safety_required", + ) + return ResolvedActionContract( + requires=(StateAtom("arm_free", arm=arm),), + claims=(ResourceClaim(f"arm:{arm}"),), + ) + + +def _resolve_joints_contract(node: Mapping[str, Any]) -> ResolvedActionContract: + object_uid = _required_string(node.get("object_uid"), "node.object_uid") + actor = node.get("actor", {}) + if not isinstance(actor, Mapping): + raise ValueError("MoveJoints contract requires an actor mapping.") + arm = _required_arm(_actor_arms(actor)[0], "MoveJoints") + binding = node.get("target_binding", {}) + if not isinstance(binding, Mapping): + raise ValueError("MoveJoints contract requires a target_binding mapping.") + release_role = binding.get("coordinated_release_role") + if release_role is not None: + if ( + node.get("task_type") != "E5" + or node.get("control") != "hand" + or binding.get("source") != "gripper_open" + or not node.get("sync_group") + ): + raise ValueError( + "Coordinated MoveJoints release requires an E5 synchronized " + "hand action targeting gripper_open." + ) + if release_role not in {"participant", "commit"}: + raise ValueError( + "coordinated_release_role must be 'participant' or 'commit'." + ) + claims = ( + ResourceClaim(f"arm:{arm}", lifetime="until_release"), + ResourceClaim(f"object:{object_uid}", lifetime="until_release"), + ) + if release_role == "participant": + return ResolvedActionContract( + requires=(StateAtom("object_coordinated_held", object_uid=object_uid),), + claims=claims, + ) + return ResolvedActionContract( + requires=(StateAtom("object_coordinated_held", object_uid=object_uid),), + effects=( + StateEffect( + "delete", + StateAtom("object_coordinated_held", object_uid=object_uid), + ), + StateEffect("add", StateAtom("object_free", object_uid=object_uid)), + StateEffect("add", StateAtom("arm_free", arm="left_arm")), + StateEffect("add", StateAtom("arm_free", arm="right_arm")), + ), + claims=claims, + ) + if node.get("control") == "hand": + return ResolvedActionContract( + requires=(StateAtom("object_held", object_uid=object_uid, arm=arm),), + claims=( + ResourceClaim(f"arm:{arm}", lifetime="until_release"), + ResourceClaim(f"object:{object_uid}", lifetime="until_release"), + ), + ) + if node.get("role") == "cleanup": + return ResolvedActionContract( + requires=(StateAtom("arm_clear", arm=arm),), + effects=( + StateEffect("add", StateAtom("arm_home", arm=arm)), + StateEffect("add", StateAtom("arm_free", arm=arm)), + ), + claims=(ResourceClaim(f"arm:{arm}"),), + completion="terminal_barrier", + failure_policy="best_effort", + ) + return ResolvedActionContract( + requires=(StateAtom("arm_free", arm=arm),), + claims=(ResourceClaim(f"arm:{arm}"),), + ) + + +def _verify_arm_clearance( + *, + executor: Any, + step: Any, + arm: str, + outcome: Any, + attempted: torch.Tensor, +) -> torch.Tensor: + """Verify a released TCP is clear, plus the transfer side for handover.""" + policy = outcome.grounded.motion_policy + object_uid = policy.get("clearance_object_uid") + if not isinstance(object_uid, str) or not object_uid: + return attempted + transfer_arm = str(policy.get("transfer_arm", arm)) + if transfer_arm not in {"left_arm", "right_arm"}: + return torch.zeros_like(attempted) + entity = executor.env.sim.get_rigid_object(object_uid) + getter = getattr(executor.env, "get_current_xpos_agent", None) + if entity is None or not callable(getter): + return torch.zeros_like(attempted) + left, right = getter() + eef = torch.as_tensor( + left if transfer_arm == "left_arm" else right, + dtype=torch.float32, + device=executor.env.device, + ) + if eef.ndim == 2: + eef = eef.unsqueeze(0).repeat(int(executor.env.num_envs), 1, 1) + object_pose = torch.as_tensor( + entity.get_local_pose(to_matrix=True), + dtype=torch.float32, + device=executor.env.device, + ) + if object_pose.ndim == 2: + object_pose = object_pose.unsqueeze(0).repeat(int(executor.env.num_envs), 1, 1) + offset = eef[:, :3, 3] - object_pose[:, :3, 3] + distance = torch.linalg.vector_norm(offset, dim=1) + minimum_clearance = policy.get( + "minimum_clearance", + policy.get("minimum_transfer_clearance", 0.10), + ) + clear = distance >= float(minimum_clearance) + role_axis = policy.get("transfer_role_axis") + if role_axis is None: + return attempted & clear + role_axis = torch.as_tensor( + role_axis, + dtype=offset.dtype, + device=offset.device, + ) + if role_axis.ndim == 1: + role_axis = role_axis.unsqueeze(0).repeat(int(executor.env.num_envs), 1) + lateral = torch.sum(offset * role_axis, dim=1) + clear &= lateral >= float(policy.get("minimum_transfer_lateral_clearance", 0.06)) + return attempted & clear + + +def _actor_arms(actor: Mapping[str, Any]) -> tuple[str, ...]: + mode = str(actor.get("mode", "auto")) + if mode == "coordinated": + arms = actor.get("arms", ()) + if not isinstance(arms, (list, tuple)): + raise ValueError("Coordinated actor arms must be a sequence.") + result = tuple(str(item) for item in arms) + if len(result) < 2 or any(not item for item in result): + raise ValueError("Coordinated actor requires at least two named arms.") + return result + if mode in {"required", "preferred"}: + return (_required_string(actor.get("arm"), "actor.arm"),) + return ("auto",) + + +def _coordinated_arms(arms: tuple[str, ...], action: str) -> tuple[str, ...]: + if len(arms) < 2: + raise ValueError(f"{action} requires a coordinated actor.") + return arms + + +def _required_arm(arm: str | None, action: str) -> str: + if arm is None: + raise ValueError(f"{action} requires exactly one arm.") + return arm + + +def _required_string(value: Any, context: str) -> str: + if not isinstance(value, str) or not value: + raise ValueError(f"{context} must be a non-empty string.") + return value + + +def _callable_name(value: Callable[..., Any] | None) -> str | None: + if value is None: + return None + module = getattr(value, "__module__", "") + name = getattr( + value, "__qualname__", getattr(value, "__name__", type(value).__name__) + ) + return f"{module}.{name}" if module else str(name) diff --git a/embodichain/gen_sim/action_engine/capabilities/builtins.py b/embodichain/gen_sim/action_engine/capabilities/builtins.py new file mode 100644 index 000000000..839797e37 --- /dev/null +++ b/embodichain/gen_sim/action_engine/capabilities/builtins.py @@ -0,0 +1,1018 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Built-in semantic operators lowered to public atomic-action contracts.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +from typing import Any + +from embodichain.gen_sim.action_engine.domain.motion import ( + motion_policy as build_motion_policy, +) +from embodichain.gen_sim.action_engine.domain.task_contracts import ( + PLACEMENT_RELATIONS, + TERMINAL_BEHAVIORS, + TRANSPORT_DIRECTIONS, + normalize_placement_relation, +) +from embodichain.gen_sim.action_engine.orientation import ( + compile_orientation_constraint, +) + +from .registry import ( + ActionCapability, + ActionTemplate, + CapabilityRegistry, + OperatorCapability, + PhaseTemplate, +) + +__all__ = ["build_default_registry"] + +_SINGLE_ARM_PHASE_OPERATORS = frozenset( + {"build_stack", "hold_hover", "orient_object", "place_relative"} +) + + +def build_default_registry() -> CapabilityRegistry: + """Build a fresh registry containing all Action Engine v1 capabilities.""" + registry = CapabilityRegistry() + from .atomic import build_atomic_capability_registry + + for capability in build_atomic_capability_registry().catalog().values(): + registry.register_action( + ActionCapability( + str(capability["name"]), + frozenset(capability["binding_kinds"]), + frozenset(capability["controls"]), + ) + ) + + definitions = ( + OperatorCapability( + "arrange_line", + "Arrange two or more movable objects into one live-grounded line.", + _expand_arrange_line, + _build_arrange_line_phases, + expansion_topology="parallel_children", + ), + OperatorCapability( + "build_stack", + "Build one ordered vertical or nested stack.", + _expand_build_stack, + _build_single_arm_phases, + ), + OperatorCapability( + "place_relative", + "Place one object at a symbolic relation to another object.", + _expand_place_relative, + _build_single_arm_phases, + ), + OperatorCapability( + "orient_object", + "Reorient one object in place and release it in a stable pose.", + _expand_orient_object, + _build_orient_object_phases, + ), + OperatorCapability( + "coordinated_transport", + "Use both arms to pick and transport one shared object.", + _expand_coordinated_transport, + _build_coordinated_transport_phases, + ), + # These internal operators preserve runtime characterization coverage + # for public Atomic Actions. They are intentionally absent from the + # planner catalog during the five-skill first phase. + OperatorCapability( + "hold_hover", + "Internal terminal-hold compatibility operator.", + _expand_hold_hover, + _build_single_arm_phases, + lifecycle="terminal_hold", + planner_visible=False, + ), + OperatorCapability( + "press", + "Internal press compatibility operator.", + _expand_press, + _build_press_phases, + planner_visible=False, + ), + OperatorCapability( + "coordinated_place", + "Internal coordinated-placement compatibility operator.", + _expand_coordinated_place, + _build_coordinated_place_phases, + planner_visible=False, + ), + ) + for definition in definitions: + registry.register_operator(definition) + return registry + + +def _expand_arrange_line(step: Mapping[str, Any]) -> list[dict[str, Any]]: + objects = _collective_objects(step, "arrange_line", minimum=2) + goal = _goal( + step, + allowed={ + "anchor", + "axis", + "order_by", + "order_constraint", + "order_direction", + "orientation_axis", + "orientation_constraint", + "orientation_directed", + "orientation_goal", + "participation", + }, + ) + orientation_goal, orientation_axis = _orientation(goal, "arrange_line") + axis = str(goal.get("axis", "world_y")) + if axis not in {"world_x", "world_y", "table_long_axis"}: + raise ValueError("arrange_line goal.axis must be a symbolic table axis.") + anchor = str(goal.get("anchor", "table_center")) + if anchor != "table_center": + raise ValueError("arrange_line currently requires anchor='table_center'.") + order_constraint = str(goal.get("order_constraint", "free")) + if order_constraint not in {"free", "ordered"}: + raise ValueError( + "arrange_line goal.order_constraint must be 'free' or 'ordered'." + ) + order_by = str(goal.get("order_by", "explicit")) + if order_by not in {"explicit", "size", "color"}: + raise ValueError("arrange_line order_by must be explicit, size, or color.") + order_direction = str(goal.get("order_direction", "given")) + if order_direction not in {"given", "ascending", "descending"}: + raise ValueError( + "arrange_line order_direction must be given, ascending, or descending." + ) + participation = str(goal.get("participation", "auto")) + if participation not in {"auto", "both_arms"}: + raise ValueError("arrange_line participation must be auto or both_arms.") + + common_goal = { + "layout": "line", + "objects": objects, + "axis": axis, + "anchor": anchor, + "order_by": order_by, + "order_direction": order_direction, + "order_constraint": order_constraint, + "participation": participation, + "orientation_goal": orientation_goal, + "orientation_axis": orientation_axis, + **_orientation_extensions(goal), + } + actor = _single_arm_actor(step) + expanded: list[dict[str, Any]] = [] + for slot_index, object_uid in enumerate(objects): + child_goal = { + **deepcopy(common_goal), + "nominal_slot_index": slot_index, + "slot_constraint": ( + "required" if order_constraint == "ordered" else "free_reassignable" + ), + } + expanded.append( + _execution_step( + step, + child_id=f"{step['id']}__{slot_index + 1:02d}", + object_uid=object_uid, + actor=( + { + **actor, + "allocation_group": f"{step['id']}_both_arms", + } + if participation == "both_arms" and slot_index < 2 + else actor + ), + goal=child_goal, + postcondition={ + "type": "line_member_placed", + "nominal_slot_index": slot_index, + "slot_constraint": child_goal["slot_constraint"], + "order_constraint": order_constraint, + }, + ) + ) + return expanded + + +def _expand_build_stack(step: Mapping[str, Any]) -> list[dict[str, Any]]: + objects = _collective_objects(step, "build_stack", minimum=1) + goal = _goal( + step, + allowed={ + "anchor", + "orientation_axis", + "orientation_constraint", + "orientation_directed", + "orientation_goal", + "stack_mode", + }, + ) + orientation_goal, orientation_axis = _orientation(goal, "build_stack") + stack_mode = str(goal.get("stack_mode", "on_top")) + if stack_mode not in {"on_top", "nested"}: + raise ValueError("build_stack goal.stack_mode must be 'on_top' or 'nested'.") + anchor = goal.get("anchor", "table_center") + if not isinstance(anchor, str) or not anchor: + raise ValueError("build_stack goal.anchor must be an object or table_center.") + + actor = _single_arm_actor(step) + expanded: list[dict[str, Any]] = [] + for layer_index, object_uid in enumerate(objects): + reference = objects[layer_index - 1] if layer_index else anchor + support_reference = "table" if reference == "table_center" else reference + child_goal: dict[str, Any] = { + "relation": ( + "inside" if stack_mode == "nested" and layer_index > 0 else "on" + ), + "reference_object": support_reference, + "reference_state": "live", + "layer_index": layer_index, + "stack_mode": stack_mode, + "orientation_goal": orientation_goal, + "orientation_axis": orientation_axis, + **_orientation_extensions(goal), + } + expanded.append( + _execution_step( + step, + child_id=f"{step['id']}__{layer_index + 1:02d}", + object_uid=object_uid, + actor=actor, + goal=child_goal, + postcondition={ + "type": "stack_layer_supported", + "layer_index": layer_index, + "reference_object": support_reference, + }, + ) + ) + return expanded + + +def _expand_place_relative(step: Mapping[str, Any]) -> list[dict[str, Any]]: + object_uid = _single_object(step, "place_relative") + goal = _goal( + step, + allowed={ + "orientation_axis", + "orientation_constraint", + "orientation_directed", + "orientation_goal", + "orientation_reference_object", + "payloads", + "reference_object", + "reference_state", + "relation", + "slot", + }, + ) + orientation_goal, orientation_axis = _orientation(goal, "place_relative") + reference = _required_string(goal, "reference_object", "place_relative") + relation = normalize_placement_relation(goal.get("relation", "on")) + normalized_goal = { + "reference_object": reference, + "reference_state": str(goal.get("reference_state", "live")), + "relation": relation, + "orientation_goal": orientation_goal, + "orientation_axis": orientation_axis, + **_orientation_extensions(goal), + "slot": str(goal.get("slot", "auto")), + } + if normalized_goal["reference_state"] not in {"initial", "live"}: + raise ValueError("place_relative reference_state must be 'initial' or 'live'.") + if normalized_goal["slot"] not in {"auto", "left", "center", "right"}: + raise ValueError("place_relative slot must be left, center, right, or auto.") + if "orientation_reference_object" in goal: + normalized_goal["orientation_reference_object"] = goal[ + "orientation_reference_object" + ] + payloads = _normalize_payloads( + goal.get("payloads", []), + object_uid, + "place_relative", + ) + if payloads: + normalized_goal["payloads"] = payloads + return [ + _execution_step( + step, + object_uid=object_uid, + actor=_single_arm_actor(step), + goal=normalized_goal, + postcondition={ + "type": "semantic_goal", + "relation": relation, + "reference_object": reference, + }, + ) + ] + + +def _expand_hold_hover(step: Mapping[str, Any]) -> list[dict[str, Any]]: + object_uid = _single_object(step, "hold_hover") + goal = _goal( + step, + allowed={ + "orientation_axis", + "orientation_constraint", + "orientation_directed", + "orientation_goal", + "reference_object", + "reference_state", + }, + ) + orientation_goal, orientation_axis = _orientation( + goal, + "hold_hover", + allow_change=False, + ) + reference = str(goal.get("reference_object", object_uid)) + return [ + _execution_step( + step, + object_uid=object_uid, + actor=_single_arm_actor(step), + goal={ + "relation": "held_above_initial", + "reference_object": reference, + "reference_state": str(goal.get("reference_state", "initial")), + "orientation_goal": orientation_goal, + "orientation_axis": orientation_axis, + **_orientation_extensions(goal), + }, + postcondition={"type": "object_held", "object": object_uid}, + ) + ] + + +def _expand_orient_object(step: Mapping[str, Any]) -> list[dict[str, Any]]: + """Normalize an in-place orientation request into one executable step. + + Keeping the target position symbolic is important: runtime observes the + object's live position immediately before grounding, so prior independent + operations and simulator settling cannot make this plan stale. + """ + object_uid = _single_object(step, "orient_object") + goal = _goal( + step, + allowed={ + "orientation_axis", + "orientation_constraint", + "orientation_directed", + "orientation_goal", + "position_anchor", + "support_object", + "upright_local_axis", + }, + ) + orientation_goal, orientation_axis = _orientation(goal, "orient_object") + if orientation_goal not in {"upright", "lay_flat", "axis_align"}: + raise ValueError( + "orient_object requires upright, lay_flat, or axis_align orientation." + ) + position_anchor = str(goal.get("position_anchor", "initial_xy")) + if position_anchor not in {"initial_xy", "live_xy"}: + raise ValueError( + "orient_object position_anchor must be 'initial_xy' or 'live_xy'." + ) + upright_local_axis = str(goal.get("upright_local_axis", "auto")) + if upright_local_axis not in {"auto", "long_axis", "x", "y", "z"}: + raise ValueError( + "orient_object upright_local_axis must be auto, long_axis, x, y, or z." + ) + support_object = str(goal.get("support_object", "table")) + if not support_object: + raise ValueError("orient_object support_object must be a non-empty string.") + return [ + _execution_step( + step, + object_uid=object_uid, + actor=_single_arm_actor(step), + goal={ + "relation": "none", + "reference_state": "live", + "orientation_goal": orientation_goal, + "orientation_axis": orientation_axis, + **_orientation_extensions(goal), + "position_anchor": position_anchor, + "support_object": support_object, + "upright_local_axis": upright_local_axis, + }, + postcondition={ + "type": "semantic_goal", + "relation": "none", + "orientation_goal": orientation_goal, + }, + ) + ] + + +def _expand_coordinated_transport( + step: Mapping[str, Any], +) -> list[dict[str, Any]]: + object_uid = _single_object(step, "coordinated_transport") + goal = _goal( + step, + allowed={ + "direction", + "orientation_axis", + "orientation_constraint", + "orientation_directed", + "orientation_goal", + "payloads", + "reference_object", + "relation", + "terminal_behavior", + }, + ) + orientation_goal, orientation_axis = _orientation( + goal, + "coordinated_transport", + ) + terminal_behavior = str(goal.get("terminal_behavior", "hold")) + if terminal_behavior not in TERMINAL_BEHAVIORS - {"none"}: + raise ValueError( + "coordinated_transport terminal_behavior must be 'hold' or 'place'." + ) + direction = str(goal.get("direction", "none")) + if direction not in TRANSPORT_DIRECTIONS: + raise ValueError( + f"coordinated_transport direction {direction!r} is unsupported." + ) + relation = goal.get("relation") + if relation is not None and str(relation) not in PLACEMENT_RELATIONS: + raise ValueError( + f"coordinated_transport relation {str(relation)!r} is unsupported." + ) + normalized_goal = { + "direction": direction, + "terminal_behavior": terminal_behavior, + "orientation_goal": orientation_goal, + "orientation_axis": orientation_axis, + **_orientation_extensions(goal), + } + normalized_payloads = _normalize_payloads( + goal.get("payloads", []), + object_uid, + "coordinated_transport", + ) + if normalized_payloads: + normalized_goal["payloads"] = normalized_payloads + for key in ("reference_object", "relation"): + if key in goal: + normalized_goal[key] = goal[key] + postcondition = ( + {"type": "semantic_goal", "relation": normalized_goal.get("relation", "at")} + if terminal_behavior == "place" + else {"type": "held_by_both_grippers", "object": object_uid} + ) + return [ + _execution_step( + step, + object_uid=object_uid, + actor={"mode": "coordinated", "arms": ["left_arm", "right_arm"]}, + goal=normalized_goal, + postcondition=postcondition, + ) + ] + + +def _expand_press(step: Mapping[str, Any]) -> list[dict[str, Any]]: + object_uid = _single_object(step, "press") + goal = _goal( + step, + allowed={"interaction", "reference_object", "terminal_state"}, + ) + return [ + _execution_step( + step, + object_uid=object_uid, + actor=_single_arm_actor(step), + goal={ + "interaction": str(goal.get("interaction", "press")), + "terminal_state": str(goal.get("terminal_state", "activated")), + **( + {"reference_object": goal["reference_object"]} + if "reference_object" in goal + else {} + ), + }, + postcondition={ + "type": "pressed", + "object": object_uid, + "terminal_state": str(goal.get("terminal_state", "activated")), + }, + ) + ] + + +def _expand_coordinated_place(step: Mapping[str, Any]) -> list[dict[str, Any]]: + object_uid = _single_object(step, "coordinated_place") + goal = _goal( + step, + allowed={"relation", "release", "support_object"}, + ) + support_object = _required_string(goal, "support_object", "coordinated_place") + if support_object == object_uid: + raise ValueError("coordinated_place requires two distinct objects.") + relation = str(goal.get("relation", "on")) + if relation not in {"on", "inside"}: + raise ValueError("coordinated_place relation must be 'on' or 'inside'.") + release = goal.get("release", True) + if not isinstance(release, bool): + raise ValueError("coordinated_place goal.release must be boolean.") + return [ + _execution_step( + step, + object_uid=object_uid, + actor={"mode": "coordinated", "arms": ["left_arm", "right_arm"]}, + goal={ + "support_object": support_object, + "relation": relation, + "release": release, + }, + postcondition={ + "type": "coordinated_placed", + "object": object_uid, + "support_object": support_object, + "relation": relation, + }, + ) + ] + + +def _build_arrange_line_phases( + step: Mapping[str, Any], +) -> tuple[PhaseTemplate, ...]: + return ( + _pickup_phase(step), + _move_phase(step, "staging"), + _move_phase(step, "final"), + *_release_retreat_home(step), + ) + + +def _build_single_arm_phases( + step: Mapping[str, Any], +) -> tuple[PhaseTemplate, ...]: + if step["operator"] not in _SINGLE_ARM_PHASE_OPERATORS: + raise ValueError(f"Unexpected single-arm operator {step['operator']!r}.") + phases: tuple[PhaseTemplate, ...] = ( + _pickup_phase(step), + _move_phase(step), + ) + if step["operator"] == "hold_hover": + return phases + ( + PhaseTemplate( + name="keep_holding", + state_semantic=f"`{step['object']}` remains held", + actions=( + ActionTemplate( + "MoveJoints", + {"kind": "joint_state", "source": "gripper_closed"}, + build_motion_policy(), + control="hand", + ), + ), + ), + ) + return phases + _release_retreat_home(step) + + +def _build_orient_object_phases( + step: Mapping[str, Any], +) -> tuple[PhaseTemplate, ...]: + """Rotate at a clearance waypoint before descending to the support.""" + upright = build_motion_policy(("orientation", "upright")) + return ( + _pickup_phase(step, motion_policy=upright), + _move_phase( + step, + "staging", + motion_policy=upright, + ), + _move_phase( + step, + "final", + motion_policy=upright, + ), + *_release_retreat_home( + step, + release_policy=upright, + retreat_policy=upright, + ), + ) + + +def _build_coordinated_transport_phases( + step: Mapping[str, Any], +) -> tuple[PhaseTemplate, ...]: + phases: tuple[PhaseTemplate, ...] = ( + PhaseTemplate( + name="coordinated_transport", + state_semantic=f"`{step['object']}` reaches its coordinated goal", + actions=( + ActionTemplate( + "CoordinatedPickment", + { + "kind": "coordinated_goal", + "semantic_step": step["id"], + "object": step["object"], + "payloads": deepcopy(step["goal"].get("payloads", [])), + }, + build_motion_policy(), + control="coordinated", + ), + ), + ), + ) + if step["goal"]["terminal_behavior"] != "place": + return phases + return phases + ( + PhaseTemplate( + name="dual_release", + state_semantic="Both grippers release the transported object", + actions=tuple( + ActionTemplate( + "MoveJoints", + { + "kind": "joint_state", + "source": "gripper_open", + "coordinated_release_role": release_role, + }, + build_motion_policy(), + control="hand", + actor={"mode": "required", "arm": arm}, + ) + for arm, release_role in ( + ("left_arm", "participant"), + ("right_arm", "commit"), + ) + ), + ), + ) + + +def _build_press_phases(step: Mapping[str, Any]) -> tuple[PhaseTemplate, ...]: + return ( + PhaseTemplate( + name="press", + state_semantic=f"`{step['object']}` has been pressed", + actions=( + ActionTemplate( + "Press", + { + "kind": "semantic_goal", + "semantic_step": step["id"], + "object": step["object"], + "interaction": "press", + }, + build_motion_policy(), + ), + ), + ), + ) + + +def _build_coordinated_place_phases( + step: Mapping[str, Any], +) -> tuple[PhaseTemplate, ...]: + return ( + PhaseTemplate( + name="dual_pick_up", + state_semantic=( + f"`{step['object']}` is held by the left arm and " + f"`{step['goal']['support_object']}` is held by the right arm" + ), + actions=( + ActionTemplate( + "PickUp", + { + "kind": "object", + "object": step["object"], + "affordance": "antipodal", + }, + build_motion_policy(), + actor={"mode": "required", "arm": "left_arm"}, + ), + ActionTemplate( + "PickUp", + { + "kind": "object", + "object": step["goal"]["support_object"], + "affordance": "antipodal", + }, + build_motion_policy(), + actor={"mode": "required", "arm": "right_arm"}, + ), + ), + ), + PhaseTemplate( + name="coordinated_place", + state_semantic=( + f"`{step['object']}` is coordinated with " + f"`{step['goal']['support_object']}`" + ), + actions=( + ActionTemplate( + "CoordinatedPlacement", + { + "kind": "coordinated_placement_goal", + "semantic_step": step["id"], + "placing_object": step["object"], + "support_object": step["goal"]["support_object"], + }, + build_motion_policy(), + control="coordinated", + ), + ), + ), + ) + + +def _pickup_phase( + step: Mapping[str, Any], + *, + motion_policy: Mapping[str, Any] | None = None, +) -> PhaseTemplate: + payloads = deepcopy(step["goal"].get("payloads", [])) + return PhaseTemplate( + name="pick_up", + state_semantic=f"Holding `{step['object']}`", + actions=( + ActionTemplate( + "PickUp", + { + "kind": "object", + "object": step["object"], + "affordance": "antipodal", + **({"payloads": payloads} if payloads else {}), + }, + motion_policy or build_motion_policy(), + ), + ), + ) + + +def _move_phase( + step: Mapping[str, Any], + phase: str | None = None, + *, + motion_policy: Mapping[str, Any] | None = None, +) -> PhaseTemplate: + target_binding = { + "kind": "semantic_goal", + "semantic_step": step["id"], + } + if phase is not None: + target_binding["phase"] = phase + payloads = deepcopy(step["goal"].get("payloads", [])) + if payloads: + target_binding["payloads"] = payloads + return PhaseTemplate( + name=f"move_to_{phase or 'semantic_goal'}", + state_semantic=f"`{step['object']}` is held at {phase or 'its semantic goal'}", + actions=( + ActionTemplate( + "MoveHeldObject", + target_binding, + motion_policy or build_motion_policy(), + ), + ), + ) + + +def _release_retreat_home( + step: Mapping[str, Any], + *, + release_policy: Mapping[str, Any] | None = None, + retreat_policy: Mapping[str, Any] | None = None, +) -> tuple[PhaseTemplate, ...]: + payloads = deepcopy(step["goal"].get("payloads", [])) + return ( + PhaseTemplate( + name="release", + state_semantic=f"`{step['object']}` is released at its semantic goal", + actions=( + ActionTemplate( + "Place", + { + "kind": "current_held_pose", + **({"payloads": payloads} if payloads else {}), + }, + release_policy or build_motion_policy(), + ), + ), + ), + PhaseTemplate( + name="retreat", + state_semantic=f"The end effector retreats from `{step['object']}`", + actions=( + ActionTemplate( + "MoveEndEffector", + { + "kind": "policy_pose", + "source": "release", + "operation": "retreat", + }, + retreat_policy or build_motion_policy(), + ), + ), + ), + PhaseTemplate( + name="home", + state_semantic="The selected arm returns to its initial state", + actions=( + ActionTemplate( + "MoveJoints", + {"kind": "joint_state", "source": "initial"}, + build_motion_policy(), + ), + ), + ), + ) + + +def _dual_arm_phase( + name: str, + state_semantic: str, + action_class: str, + target_binding: Mapping[str, Any], + motion_policy: Mapping[str, Any], + *, + control: str = "arm", +) -> PhaseTemplate: + return PhaseTemplate( + name=name, + state_semantic=state_semantic, + actions=tuple( + ActionTemplate( + action_class, + target_binding, + motion_policy, + control=control, + actor={"mode": "required", "arm": arm}, + ) + for arm in ("left_arm", "right_arm") + ), + ) + + +def _execution_step( + parent: Mapping[str, Any], + *, + object_uid: str, + actor: Mapping[str, Any], + goal: Mapping[str, Any], + postcondition: Mapping[str, Any], + child_id: str | None = None, +) -> dict[str, Any]: + return { + "id": child_id or parent["id"], + "parent_step_id": parent["id"], + "operator": parent["operator"], + "object": object_uid, + "actor": deepcopy(dict(actor)), + "goal": deepcopy(dict(goal)), + "depends_on": [], + "postcondition": deepcopy(dict(postcondition)), + "edge_ids": [], + } + + +def _single_object(step: Mapping[str, Any], operator: str) -> str: + if "object" not in step: + raise ValueError(f"{operator} requires one 'object', not 'objects'.") + return str(step["object"]) + + +def _collective_objects( + step: Mapping[str, Any], + operator: str, + *, + minimum: int, +) -> list[str]: + if "objects" not in step: + raise ValueError(f"{operator} requires an 'objects' list.") + objects = [str(value) for value in step["objects"]] + if len(objects) < minimum: + raise ValueError(f"{operator} requires at least {minimum} object(s).") + return objects + + +def _single_arm_actor(step: Mapping[str, Any]) -> dict[str, Any]: + actor = deepcopy(dict(step["actor"])) + if actor["mode"] == "coordinated": + raise ValueError(f"{step['operator']} requires one arm, not coordinated arms.") + if actor["mode"] == "required": + arm = str(actor["arm"]) + if arm in {"left", "right"}: + actor["arm"] = f"{arm}_arm" + return actor + + +def _goal(step: Mapping[str, Any], *, allowed: set[str]) -> dict[str, Any]: + goal = deepcopy(dict(step["goal"])) + unknown = sorted(set(goal) - allowed) + if unknown: + raise ValueError( + f"{step['operator']} goal contains unsupported fields: {unknown}." + ) + return goal + + +def _normalize_payloads( + value: Any, + carrier_uid: str, + operator: str, +) -> list[dict[str, str]]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes, bytearray)): + raise ValueError(f"{operator} payloads must be a list.") + if len(value) > 4: + raise ValueError(f"{operator} supports at most four payloads.") + result = [] + for index, payload in enumerate(value): + item = {"object": payload} if isinstance(payload, str) else dict(payload) + uid = item.get("object") + slot = str(item.get("slot", "auto")) + if not isinstance(uid, str) or not uid: + raise ValueError(f"payloads[{index}] requires an object UID.") + if uid == carrier_uid: + raise ValueError(f"A {operator} carrier cannot be its own payload.") + if slot not in {"left", "right", "center", "auto"}: + raise ValueError(f"Unsupported payload slot {slot!r}.") + result.append({"object": uid, "slot": slot}) + payload_uids = [item["object"] for item in result] + if len(payload_uids) != len(set(payload_uids)): + raise ValueError(f"{operator} payload objects must be unique.") + return result + + +def _required_string( + value: Mapping[str, Any], + key: str, + operator: str, +) -> str: + result = value.get(key) + if not isinstance(result, str) or not result.strip(): + raise ValueError(f"{operator} goal.{key} must be a non-empty string.") + return result + + +def _orientation( + goal: Mapping[str, Any], + operator: str, + *, + allow_change: bool = True, +) -> tuple[str, str]: + default_goal = "none" if allow_change else "preserve" + orientation_goal = str(goal.get("orientation_goal", default_goal)) + orientation_axis = str(goal.get("orientation_axis", "none")) + allowed_goals = ( + {"none", "preserve", "upright", "lay_flat", "axis_align"} + if allow_change + else {"preserve"} + ) + if orientation_goal not in allowed_goals: + raise ValueError( + f"{operator} orientation_goal {orientation_goal!r} is unsupported." + ) + if orientation_axis not in {"none", "x", "y", "long_axis", "short_axis"}: + raise ValueError( + f"{operator} orientation_axis {orientation_axis!r} is unsupported." + ) + if orientation_goal == "axis_align" and orientation_axis == "none": + raise ValueError(f"{operator} axis_align requires an orientation_axis.") + compile_orientation_constraint(goal) + return orientation_goal, orientation_axis + + +def _orientation_extensions(goal: Mapping[str, Any]) -> dict[str, Any]: + """Copy optional composable fields after operator-level validation.""" + return { + key: deepcopy(goal[key]) + for key in ("orientation_constraint", "orientation_directed") + if key in goal + } diff --git a/embodichain/gen_sim/action_engine/capabilities/registry.py b/embodichain/gen_sim/action_engine/capabilities/registry.py new file mode 100644 index 000000000..83ade2d3a --- /dev/null +++ b/embodichain/gen_sim/action_engine/capabilities/registry.py @@ -0,0 +1,183 @@ +# ---------------------------------------------------------------------------- +# 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 registry shared by planning metadata and compilation.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any + +from embodichain.gen_sim.action_engine.domain.motion import validate_motion_policy + +__all__ = [ + "ActionCapability", + "ActionTemplate", + "CapabilityRegistry", + "OperatorCapability", + "PhaseTemplate", +] + + +@dataclass(frozen=True) +class ActionCapability: + """Describe one public AtomicAction class exposed to the compiler.""" + + class_name: str + target_binding_kinds: frozenset[str] + controls: frozenset[str] + + +@dataclass(frozen=True) +class ActionTemplate: + """Describe one symbolic atomic action before actor materialization.""" + + atomic_action_class: str + target_binding: Mapping[str, Any] + motion_policy: Mapping[str, Any] + control: str = "arm" + actor: Mapping[str, Any] | None = None + + def __post_init__(self) -> None: + object.__setattr__( + self, + "target_binding", + MappingProxyType(dict(self.target_binding)), + ) + object.__setattr__( + self, + "motion_policy", + MappingProxyType(validate_motion_policy(self.motion_policy)), + ) + if self.actor is not None: + object.__setattr__(self, "actor", MappingProxyType(dict(self.actor))) + + +@dataclass(frozen=True) +class PhaseTemplate: + """Group atomic actions that execute on one graph edge.""" + + name: str + state_semantic: str + actions: tuple[ActionTemplate, ...] + + +ExpandOperator = Callable[[Mapping[str, Any]], list[dict[str, Any]]] +BuildPhases = Callable[[Mapping[str, Any]], Sequence[PhaseTemplate]] + + +@dataclass(frozen=True) +class OperatorCapability: + """Bind a semantic operator to deterministic expansion and lowering.""" + + name: str + description: str + expand: ExpandOperator + build_phases: BuildPhases + expansion_topology: str = "serial" + lifecycle: str = "release" + planner_visible: bool = True + + def __post_init__(self) -> None: + if self.expansion_topology not in {"serial", "parallel_children"}: + raise ValueError( + "Operator expansion_topology must be 'serial' or " + "'parallel_children'." + ) + if self.lifecycle not in {"release", "terminal_hold"}: + raise ValueError("Operator lifecycle must be 'release' or 'terminal_hold'.") + + +class CapabilityRegistry: + """Store explicit operator and atomic-action capabilities. + + Registration is intentionally strict. Replacing a capability by accident + would silently change compilation semantics, so callers must construct a + new registry when they need a different definition. + """ + + def __init__(self) -> None: + self._operators: dict[str, OperatorCapability] = {} + self._actions: dict[str, ActionCapability] = {} + + def register_operator(self, capability: OperatorCapability) -> None: + """Register one semantic operator.""" + if capability.name in self._operators: + raise ValueError(f"Operator {capability.name!r} is already registered.") + self._operators[capability.name] = capability + + def register_action(self, capability: ActionCapability) -> None: + """Register one public AtomicAction contract.""" + if capability.class_name in self._actions: + raise ValueError( + f"Atomic action {capability.class_name!r} is already registered." + ) + self._actions[capability.class_name] = capability + + def operator(self, name: str) -> OperatorCapability: + """Return an operator or raise a capability-focused error.""" + try: + return self._operators[name] + except KeyError as exc: + raise ValueError( + f"Unknown semantic operator {name!r}; available operators are " + f"{sorted(self._operators)}." + ) from exc + + def action(self, class_name: str) -> ActionCapability: + """Return an atomic-action contract or raise a focused error.""" + try: + return self._actions[class_name] + except KeyError as exc: + raise ValueError( + f"Unknown atomic action {class_name!r}; available actions are " + f"{sorted(self._actions)}." + ) from exc + + def operator_names(self) -> tuple[str, ...]: + """Return only the semantic skills exposed to the LLM planner.""" + return tuple( + sorted( + name + for name, capability in self._operators.items() + if capability.planner_visible + ) + ) + + def operator_descriptions(self) -> dict[str, str]: + """Return JSON-safe operator descriptions.""" + return { + name: self._operators[name].description for name in self.operator_names() + } + + def validate_action_template(self, template: ActionTemplate) -> None: + """Validate one compiler-produced action against its registered API.""" + capability = self.action(template.atomic_action_class) + kind = template.target_binding.get("kind") + if kind not in capability.target_binding_kinds: + raise ValueError( + f"{template.atomic_action_class} does not accept target binding " + f"kind {kind!r}; expected one of " + f"{sorted(capability.target_binding_kinds)}." + ) + if template.control not in capability.controls: + raise ValueError( + f"{template.atomic_action_class} does not support control " + f"{template.control!r}; expected one of " + f"{sorted(capability.controls)}." + ) diff --git a/embodichain/gen_sim/action_engine/cli/__init__.py b/embodichain/gen_sim/action_engine/cli/__init__.py new file mode 100644 index 000000000..564654c85 --- /dev/null +++ b/embodichain/gen_sim/action_engine/cli/__init__.py @@ -0,0 +1,21 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Command-line entry points for Action Engine generation and execution.""" + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/embodichain/gen_sim/action_engine/cli/generate_action_agent_config.py b/embodichain/gen_sim/action_engine/cli/generate_action_agent_config.py new file mode 100644 index 000000000..542dc1c70 --- /dev/null +++ b/embodichain/gen_sim/action_engine/cli/generate_action_agent_config.py @@ -0,0 +1,240 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""CLI for generating Action Engine configs from a Prompt2Scene gym export.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from embodichain.gen_sim.action_engine.config import generation_defaults +from embodichain.gen_sim.action_engine.generation import ( + generate_action_engine_config, +) + +__all__ = ["build_parser", "cli"] + +_GENERATION_DEFAULTS = generation_defaults() +_TASK_DEFAULTS = _GENERATION_DEFAULTS["task"] +_SCENE_DEFAULTS = _GENERATION_DEFAULTS["scene"] + +_ROBOT_PROFILE_CHOICES = ( + "ur5", + "ur10", + "dual_ur5", + "dual_ur10", + "franka", + "dual_franka", +) + + +def build_parser() -> argparse.ArgumentParser: + """Build the standalone config-generation argument parser.""" + parser = argparse.ArgumentParser( + description=( + "Plan and compile an Action Engine task from an exported tabletop " + "gym project." + ) + ) + parser.add_argument( + "--gym_project", + "--gym-project", + required=True, + help=( + "Prompt2Scene task/export directory or gym_config.json/" + "scene_config.json path." + ), + ) + parser.add_argument( + "--output_dir", + "--output-dir", + required=True, + help="Directory receiving canonical JSON artifacts and the Seed PNG.", + ) + parser.add_argument( + "--task_name", + "--task-name", + required=True, + help="Stable task identifier stored in both programs.", + ) + parser.add_argument( + "--task_description", + "--task-description", + help="Natural-language goal passed to structured LLM interpretation.", + ) + parser.add_argument( + "--task_file", + "--task-file", + help="Optional UTF-8 file containing the natural-language goal.", + ) + parser.add_argument( + "--task-spec", + "--task_spec", + dest="task_spec", + help=( + "Optional existing Action Engine v2 TaskSpec JSON; bypasses text " + "LLM interpretation and uses its role_bindings hand-off." + ), + ) + parser.add_argument( + "--robot-profile", + "--robot_profile", + choices=_ROBOT_PROFILE_CHOICES, + default=str(_TASK_DEFAULTS["default_robot_profile"]), + help="Robot template used in fast_gym_config.json.", + ) + parser.add_argument( + "--llm_model", + "--llm-model", + default=None, + help="Optional planner model override.", + ) + parser.add_argument( + "--vlm_model", + "--vlm-model", + default=None, + help="Optional online visual/planner model override stored for A/B runs.", + ) + parser.add_argument( + "--planning-mode", + "--planning_mode", + choices=("offline", "ab"), + default="offline", + help="Generate one offline bundle or an offline/online A/B bundle.", + ) + parser.add_argument( + "--source_scene_z_rotation_degrees", + "--source-scene-z-rotation-degrees", + type=float, + default=None, + help=( + "World-frame scene rotation. Prompt2Scene exports default to -90 " + "degrees; other inputs default to zero." + ), + ) + parser.add_argument( + "--body-scale-policy", + choices=("preserve", "multiply", "absolute"), + default=str(_SCENE_DEFAULTS["body_scale_policy"]), + help="How the requested xyz scale combines with source body_scale.", + ) + parser.add_argument( + "--body-scale", + type=float, + nargs=3, + default=tuple(float(value) for value in _SCENE_DEFAULTS["body_scale"]), + metavar=("X", "Y", "Z"), + help="Positive xyz scale used by multiply or absolute policy.", + ) + parser.add_argument( + "--max_episodes", + "--max-episodes", + type=int, + default=int(_TASK_DEFAULTS["max_episodes"]), + help="Episode count written to fast_gym_config.json.", + ) + parser.add_argument( + "--max_episode_steps", + "--max-episode-steps", + type=int, + default=int(_TASK_DEFAULTS["max_episode_steps"]), + help="Per-episode step limit written to fast_gym_config.json.", + ) + parser.add_argument( + "--overwrite", + action="store_true", + help="Replace existing canonical artifacts in the output directory.", + ) + parser.add_argument( + "--randomize-scene", + action="store_true", + help="Randomize rigid-object poses and table height on every reset.", + ) + parser.add_argument( + "--randomize-table-material", + action="store_true", + help="Randomize the table material independently on every reset.", + ) + return parser + + +def cli() -> None: + """Generate and report the canonical Action Engine artifact bundle.""" + args = build_parser().parse_args() + task_description = _resolve_task_description(args) + paths = generate_action_engine_config( + args.gym_project, + args.output_dir, + task_name=args.task_name, + task_description=task_description, + task_spec=args.task_spec, + robot_profile=args.robot_profile, + llm_model=args.llm_model, + source_scene_z_rotation_degrees=args.source_scene_z_rotation_degrees, + body_scale_policy=args.body_scale_policy, + body_scale=args.body_scale, + overwrite=args.overwrite, + max_episodes=args.max_episodes, + max_episode_steps=args.max_episode_steps, + randomize_scene=args.randomize_scene, + randomize_table_material=args.randomize_table_material, + planning_mode=args.planning_mode, + vlm_model=args.vlm_model, + ) + + print(f"Generated gym config: {paths.gym_config}") + print(f"Generated agent config: {paths.agent_config}") + print(f"Generated TaskSpec: {paths.task_spec}") + print(f"Generated SceneRequirements: {paths.scene_requirements}") + print(f"Generated SeedGraph: {paths.seed_task_graph}") + print(f"Generated Seed graph PNG: {paths.seed_task_graph_png}") + print( + "Run with:\n" + "python -m embodichain.gen_sim.action_engine.cli.run_agent " + f"--task_name {args.task_name} " + f'--gym_config "{paths.gym_config}" ' + f'--agent_config "{paths.agent_config}" ' + "--regenerate" + ) + + +def _resolve_task_description(args: argparse.Namespace) -> str: + task_spec = getattr(args, "task_spec", None) + if task_spec: + if args.task_description or args.task_file: + raise ValueError( + "--task-spec cannot be combined with --task_description or " + "--task_file." + ) + return "" + if args.task_description and args.task_file: + raise ValueError("Use either --task_description or --task_file, not both.") + if args.task_file: + description = ( + Path(args.task_file).expanduser().read_text(encoding="utf-8").strip() + ) + else: + description = str(args.task_description or "").strip() + if not description: + raise ValueError( + "--task_description (or --task_file) must provide a non-empty goal." + ) + return description + + +if __name__ == "__main__": + cli() 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/compiler/__init__.py b/embodichain/gen_sim/action_engine/compiler/__init__.py new file mode 100644 index 000000000..fbd36a858 --- /dev/null +++ b/embodichain/gen_sim/action_engine/compiler/__init__.py @@ -0,0 +1,33 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Stable deterministic compiler API.""" + +from __future__ import annotations + +from .core import compile_task_agent +from .v2 import ( + compile_task_agent_v2, + execution_program_to_seed_graph, + seed_graph_to_execution_program, +) + +__all__ = [ + "compile_task_agent", + "compile_task_agent_v2", + "execution_program_to_seed_graph", + "seed_graph_to_execution_program", +] diff --git a/embodichain/gen_sim/action_engine/compiler/core.py b/embodichain/gen_sim/action_engine/compiler/core.py new file mode 100644 index 000000000..260e613fc --- /dev/null +++ b/embodichain/gen_sim/action_engine/compiler/core.py @@ -0,0 +1,590 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Deterministically lower a route-free TaskAgent into an action DAG.""" + +from __future__ import annotations + +import re +from collections import deque +from collections.abc import Collection, Mapping, Sequence +from copy import deepcopy +from typing import Any + +from embodichain.gen_sim.action_engine.capabilities import ( + ActionTemplate, + CapabilityRegistry, + PhaseTemplate, + build_default_registry, +) +from embodichain.gen_sim.action_engine.domain import ( + EXECUTION_PROGRAM_SCHEMA, + MOTION_POLICY_VERSION, + validate_execution_program, + validate_task_agent, +) + +__all__ = ["compile_task_agent"] + +_UNSAFE_ID_RE = re.compile(r"[^0-9a-z]+") + + +def compile_task_agent( + program: Mapping[str, Any], + *, + registry: CapabilityRegistry | None = None, + known_objects: Collection[str] | None = None, +) -> dict[str, Any]: + """Compile semantic steps into a complete coordinate-free action DAG. + + Compilation never reads simulator state and never calls an LLM. Collective + operators such as ``arrange_line`` and ``build_stack`` expand into one + execution semantic step per object, while dependencies are rewritten to + point at the terminal expanded step of each parent operation. + + Args: + program: Valid or validation-ready TaskAgent mapping. + registry: Optional capability registry for controlled extensions. + known_objects: Optional runtime scene UIDs used for pre-simulator + object-reference validation. + + Returns: + A validated ``action_engine_execution_graph_v1`` mapping. + """ + task_agent = validate_task_agent(program, known_objects=known_objects) + capabilities = registry or build_default_registry() + ordered_task_steps = _stable_topological_steps(task_agent["semantic_steps"]) + + expanded_by_parent: dict[str, list[dict[str, Any]]] = {} + all_expanded_ids: set[str] = set() + for task_step in ordered_task_steps: + definition = capabilities.operator(task_step["operator"]) + expanded = definition.expand(task_step) + if not expanded: + raise ValueError( + f"Operator {task_step['operator']!r} produced no execution steps." + ) + for child in expanded: + child_id = str(child.get("id", "")) + if not child_id or child_id in all_expanded_ids: + raise ValueError( + f"Operator {task_step['operator']!r} produced duplicate or " + f"empty execution step ID {child_id!r}." + ) + all_expanded_ids.add(child_id) + expanded_by_parent[task_step["id"]] = expanded + + # Operator expansion validates each step's shape first, so held-state + # diagnostics never mask a more direct capability-contract error. + _validate_held_state_contract(ordered_task_steps) + + terminal_children: dict[str, list[str]] = {} + for task_step in ordered_task_steps: + definition = capabilities.operator(task_step["operator"]) + children = expanded_by_parent[task_step["id"]] + terminal_children[task_step["id"]] = ( + [child["id"] for child in children] + if definition.expansion_topology == "parallel_children" + else [children[-1]["id"]] + ) + expanded_steps: list[dict[str, Any]] = [] + for task_step in ordered_task_steps: + definition = capabilities.operator(task_step["operator"]) + children = expanded_by_parent[task_step["id"]] + parent_dependencies = [ + child_id + for parent_id in task_step["depends_on"] + for child_id in terminal_children[parent_id] + ] + for index, child in enumerate(children): + child["depends_on"] = ( + parent_dependencies + if index == 0 or definition.expansion_topology == "parallel_children" + else [children[index - 1]["id"]] + ) + expanded_steps.append(child) + + phases_by_step: dict[str, tuple[PhaseTemplate, ...]] = {} + for step in expanded_steps: + definition = capabilities.operator(step["operator"]) + phases = tuple(definition.build_phases(step)) + if not phases or any(not phase.actions for phase in phases): + raise ValueError( + f"Operator {step['operator']!r} produced an empty action phase." + ) + for phase in phases: + for action in phase.actions: + capabilities.validate_action_template(action) + phases_by_step[step["id"]] = phases + + graph = _build_graph( + task=task_agent["task"], + goal_description=task_agent["goal"], + semantic_steps=expanded_steps, + phases_by_step=phases_by_step, + ) + graph["allocation_groups"] = _merge_allocation_groups( + _compile_task_allocation_groups( + task_agent["allocation_groups"], + expanded_by_parent, + ), + _derive_allocation_groups( + expanded_steps, + phases_by_step, + ), + ) + return validate_execution_program(graph) + + +def _compile_task_allocation_groups( + groups: Sequence[Mapping[str, Any]], + expanded_by_parent: Mapping[str, Sequence[Mapping[str, Any]]], +) -> list[dict[str, Any]]: + result: list[dict[str, Any]] = [] + for group in groups: + members = [ + expanded_by_parent[parent_id][0]["id"] + for parent_id in group["semantic_step_ids"] + ] + result.append( + { + "id": group["id"], + "semantic_step_ids": members, + "arm_constraint": "distinct_arms", + "execution_policy": "parallel_if_feasible", + "parallel_action_classes": ["PickUp"], + "workspace_policy": "shared_target_serial", + } + ) + return result + + +def _merge_allocation_groups( + explicit: Sequence[Mapping[str, Any]], + derived: Sequence[Mapping[str, Any]], +) -> list[dict[str, Any]]: + result = [deepcopy(dict(group)) for group in explicit] + assigned = {step_id for group in result for step_id in group["semantic_step_ids"]} + used_ids = {group["id"] for group in result} + for group in derived: + if set(group["semantic_step_ids"]) & assigned: + continue + candidate = deepcopy(dict(group)) + base_id = candidate["id"] + suffix = 2 + while candidate["id"] in used_ids: + candidate["id"] = f"{base_id}_{suffix}" + suffix += 1 + result.append(candidate) + used_ids.add(candidate["id"]) + assigned.update(candidate["semantic_step_ids"]) + return result + + +def _build_graph( + *, + task: str, + goal_description: str, + semantic_steps: list[dict[str, Any]], + phases_by_step: Mapping[str, tuple[PhaseTemplate, ...]], +) -> dict[str, Any]: + start_id = "v0_start" + goal_id = "v_goal" + dependents: dict[str, list[str]] = {step["id"]: [] for step in semantic_steps} + for step in semantic_steps: + for dependency in step["depends_on"]: + dependents[dependency].append(step["id"]) + + terminal_node = { + step["id"]: ( + f"v_{_slug(step['id'])}_done" if dependents[step["id"]] else goal_id + ) + for step in semantic_steps + } + nodes: list[dict[str, str]] = [ + { + "id": start_id, + "semantic": "Initial state before executing the semantic action DAG", + } + ] + node_ids = {start_id} + edges: list[dict[str, Any]] = [] + final_edge_by_step: dict[str, str] = {} + + def add_node(node_id: str, semantic: str) -> None: + if node_id in node_ids or node_id == goal_id: + return + node_ids.add(node_id) + nodes.append({"id": node_id, "semantic": semantic}) + + for step in semantic_steps: + phases = phases_by_step[step["id"]] + if step["depends_on"]: + source_id = terminal_node[step["depends_on"][0]] + else: + source_id = start_id + add_node( + source_id, + f"Dependencies for semantic step `{step['id']}` are complete", + ) + + step_edge_ids: list[str] = [] + previous_edge_id: str | None = None + for phase_index, phase in enumerate(phases, start=1): + is_last = phase_index == len(phases) + target_id = ( + terminal_node[step["id"]] + if is_last + else f"v_{_slug(step['id'])}_{phase_index:02d}_{_slug(phase.name)}" + ) + add_node(target_id, phase.state_semantic) + edge_id = f"e{len(edges) + 1:03d}_{_slug(step['id'])}_{_slug(phase.name)}" + edge_dependencies = ( + [final_edge_by_step[item] for item in step["depends_on"]] + if previous_edge_id is None + else [previous_edge_id] + ) + actions = [ + _materialize_action(action, default_actor=step["actor"]) + for action in phase.actions + ] + edges.append( + { + "id": edge_id, + "source": source_id, + "target": target_id, + "semantic_step_id": step["id"], + "actions": actions, + "depends_on": edge_dependencies, + "resources": _edge_resources(step, actions), + } + ) + step_edge_ids.append(edge_id) + previous_edge_id = edge_id + source_id = target_id + step["edge_ids"] = step_edge_ids + final_edge_by_step[step["id"]] = step_edge_ids[-1] + + nodes.append( + { + "id": goal_id, + "semantic": "All required semantic steps have reached their postconditions", + } + ) + return { + "schema_version": EXECUTION_PROGRAM_SCHEMA, + "task": task, + "goal_description": goal_description, + "start": start_id, + "goal": goal_id, + "nodes": nodes, + "edges": edges, + "semantic_steps": semantic_steps, + "allocation_groups": [], + "motion_policy_version": MOTION_POLICY_VERSION, + } + + +def _materialize_action( + template: ActionTemplate, + *, + default_actor: Mapping[str, Any], +) -> dict[str, Any]: + actor = template.actor if template.actor is not None else default_actor + return { + "atomic_action_class": template.atomic_action_class, + "actor": deepcopy(dict(actor)), + "control": template.control, + "target_binding": deepcopy(dict(template.target_binding)), + "motion_policy": deepcopy(dict(template.motion_policy)), + } + + +def _edge_resources( + step: Mapping[str, Any], + actions: Sequence[Mapping[str, Any]], +) -> list[str]: + resources = {f"object:{step['object']}"} + reference = step["goal"].get("reference_object") + support = step["goal"].get("support_object") + + for action in actions: + actor = action["actor"] + if actor["mode"] == "auto": + resources.add("arm:auto") + elif actor["mode"] == "required": + resources.add(f"arm:{actor['arm']}") + else: + resources.update(f"arm:{arm}" for arm in actor["arms"]) + + binding = action["target_binding"] + for key in ("object", "placing_object", "support_object"): + object_uid = binding.get(key) + if isinstance(object_uid, str) and object_uid: + resources.add(f"object:{object_uid}") + for payload in binding.get("payloads", []): + object_uid = ( + payload.get("object") if isinstance(payload, Mapping) else payload + ) + if isinstance(object_uid, str) and object_uid: + resources.add(f"object:{object_uid}") + + action_classes = {action["atomic_action_class"] for action in actions} + uses_goal_workspace = bool( + action_classes + & { + "MoveHeldObject", + "MoveEndEffector", + "Place", + "CoordinatedPickment", + "CoordinatedPlacement", + "Press", + } + ) + if isinstance(reference, str) and reference and uses_goal_workspace: + resources.add(f"workspace:{reference}") + elif isinstance(support, str) and support: + # Passive supports such as a table may be shared by independent + # pickups. Only a coordinated placement manipulates and owns its + # support object throughout the semantic step. + if step["operator"] == "coordinated_place": + resources.add(f"object:{support}") + if uses_goal_workspace: + resources.add(f"workspace:{support}") + + if action_classes & { + "MoveHeldObject", + "Place", + "CoordinatedPickment", + "CoordinatedPlacement", + "Press", + }: + if step["operator"] == "arrange_line": + resources.add("workspace:table") + elif reference is None and support is None: + resources.add("workspace:world") + return sorted(resources) + + +def _derive_allocation_groups( + semantic_steps: Sequence[Mapping[str, Any]], + phases_by_step: Mapping[str, tuple[PhaseTemplate, ...]], +) -> list[dict[str, Any]]: + """Declare only explicit, independent distinct-arm pickup pairs.""" + groups: list[dict[str, Any]] = [] + ancestor_ids = _ancestor_sets(semantic_steps) + used_steps: set[str] = set() + for index, first in enumerate(semantic_steps): + if first["id"] in used_steps or not _starts_with_pickup( + phases_by_step[first["id"]] + ): + continue + for second in semantic_steps[index + 1 :]: + if second["id"] in used_steps or not _starts_with_pickup( + phases_by_step[second["id"]] + ): + continue + if not _actors_request_distinct_arms( + first["actor"], + second["actor"], + ): + continue + if ( + second["id"] in ancestor_ids[first["id"]] + or first["id"] in ancestor_ids[second["id"]] + ): + continue + if first["object"] == second["object"]: + continue + groups.append( + { + "id": f"g{len(groups) + 1:02d}_distinct_arms", + "semantic_step_ids": [first["id"], second["id"]], + "arm_constraint": "distinct_arms", + "execution_policy": "parallel_if_feasible", + "parallel_action_classes": ["PickUp"], + "workspace_policy": "shared_target_serial", + } + ) + used_steps.update({first["id"], second["id"]}) + break + return groups + + +def _validate_held_state_contract( + semantic_steps: Sequence[Mapping[str, Any]], +) -> None: + """Validate persistent object ownership and required-arm reservations. + + ``hold_hover`` is terminal behavior for its object and reserves the + selected arm through task completion. Unrelated downstream work remains + legal because runtime can assign it to another free arm. Action Engine v1 + does not expose a "continue with currently held object" operator, however, + so any second step that references the held object would imply an unsafe + pickup, handover, or use of a moving reference. Planner-produced + hold/place pairs are fused before this boundary. + """ + ancestors = _ancestor_sets(semantic_steps) + for hold in semantic_steps: + terminal_coordinated = ( + hold["operator"] == "coordinated_transport" + and hold["goal"].get("terminal_behavior", "hold") == "hold" + ) + if hold["operator"] != "hold_hover" and not terminal_coordinated: + continue + hold_id = hold["id"] + held_object = hold["object"] + for other in semantic_steps: + other_id = other["id"] + if other_id == hold_id or other_id in ancestors[hold_id]: + continue + if held_object in _step_object_references(other): + raise ValueError( + f"hold_hover step {hold_id!r} reserves object " + f"{held_object!r} through task completion, but step " + f"{other_id!r} also references it." + ) + hold_actor = hold["actor"] + other_actor = other["actor"] + if terminal_coordinated: + raise ValueError( + f"Terminal coordinated step {hold_id!r} reserves both arms, " + f"but step {other_id!r} is not an ancestor." + ) + if hold_actor["mode"] != "required": + continue + reserved_arm = _canonical_arm(hold_actor["arm"]) + conflicts = other_actor["mode"] == "coordinated" or ( + other_actor["mode"] == "required" + and _canonical_arm(other_actor["arm"]) == reserved_arm + ) + if conflicts: + raise ValueError( + f"hold_hover step {hold_id!r} reserves arm " + f"{reserved_arm!r}, but non-ancestor step {other_id!r} " + "also requires it." + ) + + +def _step_object_references(step: Mapping[str, Any]) -> set[str]: + """Return object UIDs whose ownership or workspace a step may require.""" + result = {step["object"]} if "object" in step else set(step.get("objects", ())) + goal = step["goal"] + for key in ( + "anchor", + "orientation_reference_object", + "reference_object", + "support_object", + ): + value = goal.get(key) + if isinstance(value, str): + result.add(value) + for payload in goal.get("payloads", []): + value = payload.get("object") if isinstance(payload, Mapping) else payload + if isinstance(value, str): + result.add(value) + return result + + +def _ancestor_sets( + semantic_steps: Sequence[Mapping[str, Any]], +) -> dict[str, set[str]]: + direct = {step["id"]: set(step["depends_on"]) for step in semantic_steps} + ancestors: dict[str, set[str]] = {} + for step in semantic_steps: + pending = list(direct[step["id"]]) + result: set[str] = set() + while pending: + dependency = pending.pop() + if dependency in result: + continue + result.add(dependency) + pending.extend(direct[dependency]) + ancestors[step["id"]] = result + return ancestors + + +def _starts_with_pickup(phases: Sequence[PhaseTemplate]) -> bool: + return bool( + phases + and phases[0].actions + and phases[0].actions[0].atomic_action_class == "PickUp" + ) + + +def _actors_request_distinct_arms( + first: Mapping[str, Any], + second: Mapping[str, Any], +) -> bool: + """Return whether actors explicitly request a distinct-arm assignment.""" + first_group = first.get("allocation_group") + same_group = first_group is not None and first_group == second.get( + "allocation_group" + ) + required_opposite = ( + first["mode"] == "required" + and second["mode"] == "required" + and _canonical_arm(first["arm"]) != _canonical_arm(second["arm"]) + ) + if same_group and not required_opposite: + both_required = first["mode"] == second["mode"] == "required" + if both_required: + raise ValueError( + f"Allocation group {first_group!r} requires distinct arms, " + "but both steps require the same arm." + ) + return same_group or required_opposite + + +def _canonical_arm(value: Any) -> str: + arm = str(value) + return f"{arm}_arm" if arm in {"left", "right"} else arm + + +def _stable_topological_steps( + semantic_steps: Sequence[Mapping[str, Any]], +) -> list[dict[str, Any]]: + original = [deepcopy(dict(step)) for step in semantic_steps] + order = {step["id"]: index for index, step in enumerate(original)} + by_id = {step["id"]: step for step in original} + indegree = {step["id"]: len(step["depends_on"]) for step in original} + dependents: dict[str, list[str]] = {step["id"]: [] for step in original} + for step in original: + for dependency in step["depends_on"]: + dependents[dependency].append(step["id"]) + + ready = deque( + sorted( + (step_id for step_id, degree in indegree.items() if degree == 0), + key=order.__getitem__, + ) + ) + result: list[dict[str, Any]] = [] + while ready: + step_id = ready.popleft() + result.append(by_id[step_id]) + newly_ready: list[str] = [] + for dependent in dependents[step_id]: + indegree[dependent] -= 1 + if indegree[dependent] == 0: + newly_ready.append(dependent) + ready.extend(sorted(newly_ready, key=order.__getitem__)) + return result + + +def _slug(value: Any) -> str: + slug = _UNSAFE_ID_RE.sub("_", str(value).lower()).strip("_") + return slug[:64].rstrip("_") or "step" diff --git a/embodichain/gen_sim/action_engine/compiler/v2.py b/embodichain/gen_sim/action_engine/compiler/v2.py new file mode 100644 index 000000000..13942afe7 --- /dev/null +++ b/embodichain/gen_sim/action_engine/compiler/v2.py @@ -0,0 +1,424 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Bridge mature v1 task recipes to the direct AtomicAction SeedGraph v3.""" + +from __future__ import annotations + +from collections import deque +from collections.abc import Collection, Mapping, Sequence +from copy import deepcopy +import re +from typing import Any + +from embodichain.gen_sim.action_engine.capabilities import ( + AtomicCapabilityRegistry, + build_atomic_capability_registry, + capability_precondition, +) +from embodichain.gen_sim.action_engine.domain import ( + EXECUTION_PROGRAM_SCHEMA, + MOTION_POLICY_VERSION, + validate_execution_program, + validate_seed_graph, +) +from embodichain.gen_sim.action_engine.protocol import SEED_GRAPH_SCHEMA +from embodichain.gen_sim.action_engine.planning.linker import ( + link_seed_graph, + validate_persisted_contracts, +) + +__all__ = [ + "compile_task_agent_v2", + "execution_program_to_seed_graph", + "seed_graph_to_execution_program", +] + +_UNSAFE_ID_RE = re.compile(r"[^0-9a-z]+") +_OPERATOR_TASK_TYPES = { + "arrange_line": "E1", + "build_stack": "E1", + "coordinated_place": "E5", + "coordinated_transport": "E5", + "hold_hover": "E1", + "orient_object": "E2", + "place_in_line": "E1", + "place_relative": "E1", + "press": "E9", +} + + +def compile_task_agent_v2( + program: Mapping[str, Any], + *, + known_objects: Collection[str] | None = None, + registry: AtomicCapabilityRegistry | None = None, +) -> dict[str, Any]: + """Compile a mature semantic recipe directly to the persisted v3 graph.""" + from .core import compile_task_agent + + legacy = compile_task_agent(program, known_objects=known_objects) + return execution_program_to_seed_graph( + legacy, + known_objects=known_objects, + registry=registry, + ) + + +def execution_program_to_seed_graph( + program: Mapping[str, Any], + *, + planner_route: str = "offline", + known_objects: Collection[str] | None = None, + registry: AtomicCapabilityRegistry | None = None, +) -> dict[str, Any]: + """Convert a mature v1 result without changing its AtomicAction topology.""" + legacy = validate_execution_program(program) + capabilities = registry or build_atomic_capability_registry() + steps = {str(step["id"]): step for step in legacy["semantic_steps"]} + node_ids_by_edge: dict[str, list[str]] = {} + nodes: list[dict[str, Any]] = [] + + for edge in legacy["edges"]: + edge_id = str(edge["id"]) + step = steps[str(edge["semantic_step_id"])] + task_type = _task_type(str(step["operator"])) + dependencies = [ + node_id + for dependency in edge.get("depends_on", []) + for node_id in node_ids_by_edge[str(dependency)] + ] + edge_nodes: list[str] = [] + actions = list(edge["actions"]) + for action_index, action in enumerate(actions): + action_name = str(action["atomic_action_class"]) + descriptor_view = { + "atomic_action": action_name, + "control": action.get("control", "arm"), + "target_binding": action["target_binding"], + } + capabilities.validate_binding(descriptor_view) + capability = capabilities.get(action_name) + node_id = _node_id(edge_id, action_name, action_index, len(actions)) + postcondition = ( + deepcopy(step["postcondition"]) + if edge_id == step["edge_ids"][-1] + else {} + ) + node = { + "id": node_id, + "atomic_action": action_name, + "object_uid": str(step["object"]), + "actor": _v2_actor(action["actor"]), + "control": str(action.get("control", "arm")), + "target_binding": deepcopy(dict(action["target_binding"])), + "depends_on": list(dict.fromkeys(dependencies)), + "task_instance_id": str(step["id"]), + "task_type": task_type, + "role": _node_role(action_name, action["target_binding"]), + "precondition": capability_precondition( + capability, + object_uid=str(step["object"]), + actor=_v2_actor(action["actor"]), + target_binding=action["target_binding"], + ), + "postcondition": postcondition, + "motion_policy": deepcopy(dict(action["motion_policy"])), + } + if len(actions) > 1: + node["sync_group"] = edge_id + nodes.append(node) + edge_nodes.append(node_id) + node_ids_by_edge[edge_id] = edge_nodes + + groups = [] + for step in legacy["semantic_steps"]: + group_node_ids = [ + node_id + for edge_id in step["edge_ids"] + for node_id in node_ids_by_edge[str(edge_id)] + ] + groups.append( + { + "id": str(step["id"]), + "task_type": _task_type(str(step["operator"])), + "role": "primary", + "operator": str(step["operator"]), + "object_uid": str(step["object"]), + "actor": _v2_actor(step["actor"]), + "goal": deepcopy(dict(step.get("goal", {}))), + "depends_on": list(step.get("depends_on", [])), + "parent_task_instance_id": str(step.get("parent_step_id", step["id"])), + "node_ids": group_node_ids, + "success": deepcopy(dict(step["postcondition"])), + } + ) + + level = _level(groups) + graph = { + "schema_version": SEED_GRAPH_SCHEMA, + "task_id": str(legacy["task"]), + "instruction": str(legacy["goal_description"]), + "level": level, + "reasoning_type": "none", + "planner_route": planner_route, + "nodes": nodes, + "task_groups": groups, + "success": { + "op": "all", + "terms": [deepcopy(group["success"]) for group in groups], + }, + "capability_catalog_hash": capabilities.catalog_hash(), + "metadata": { + "source_schema": EXECUTION_PROGRAM_SCHEMA, + "legacy_allocation_groups": deepcopy(legacy.get("allocation_groups", [])), + "planning_latency_seconds": 0.0, + "vlm_call_count": 0, + }, + } + return link_seed_graph( + graph, + registry=capabilities, + task_order=[str(step["id"]) for step in legacy["semantic_steps"]], + known_objects=known_objects, + ) + + +def seed_graph_to_execution_program( + graph: Mapping[str, Any], + *, + known_objects: Collection[str] | None = None, + registry: AtomicCapabilityRegistry | None = None, + require_executable: bool = True, +) -> dict[str, Any]: + """Materialize the v3 DAG as the existing in-memory runtime view.""" + capabilities = registry or build_atomic_capability_registry() + seed = validate_seed_graph( + graph, + known_objects=known_objects, + known_actions=capabilities.names(), + executable_actions=capabilities.executable_names(), + require_executable=require_executable, + ) + if seed["capability_catalog_hash"] != capabilities.catalog_hash(): + raise ValueError( + "SeedGraph capability_catalog_hash does not match the runtime catalog." + ) + validate_persisted_contracts(seed, capabilities) + for node in seed["nodes"]: + capabilities.validate_binding(node) + + node_by_id = {str(node["id"]): node for node in seed["nodes"]} + unit_by_node, units = _execution_units(seed["nodes"]) + ordered_units = _topological_units(units) + edge_id_by_unit = {unit_id: f"edge_{_slug(unit_id)}" for unit_id in ordered_units} + target_by_unit = { + unit_id: f"state_{index + 1:04d}_{_slug(unit_id)}" + for index, unit_id in enumerate(ordered_units) + } + start = "state_start" + edges = [] + graph_nodes = [{"id": start, "semantic": "Initial live simulator state"}] + for unit_id in ordered_units: + unit = units[unit_id] + dependencies = sorted(unit["depends_on"]) + source = start if not dependencies else target_by_unit[dependencies[0]] + target = target_by_unit[unit_id] + graph_nodes.append( + { + "id": target, + "semantic": f"Completed AtomicAction unit {unit_id}", + } + ) + unit_nodes = [node_by_id[node_id] for node_id in unit["node_ids"]] + edges.append( + { + "id": edge_id_by_unit[unit_id], + "source": source, + "target": target, + "semantic_step_id": str(unit_nodes[0]["task_instance_id"]), + "actions": [ + { + "atomic_action_class": node["atomic_action"], + "actor": deepcopy(node["actor"]), + "control": node["control"], + "target_binding": deepcopy(node["target_binding"]), + "motion_policy": node["motion_policy"], + "seed_node_id": node["id"], + "failure_policy": node["contract"]["failure_policy"], + } + for node in unit_nodes + ], + "depends_on": [edge_id_by_unit[item] for item in dependencies], + "resources": sorted( + { + str(claim["resource"]) + for node in unit_nodes + for claim in node["contract"]["claims"] + } + ), + } + ) + + group_by_id = {str(group["id"]): group for group in seed["task_groups"]} + semantic_steps = [] + for group_id in _topological_groups(seed["task_groups"]): + group = group_by_id[group_id] + group_units = [ + unit_id + for unit_id in ordered_units + if any( + node_by_id[node_id]["task_instance_id"] == group_id + for node_id in units[unit_id]["node_ids"] + ) + ] + semantic_steps.append( + { + "id": group_id, + "parent_step_id": str(group.get("parent_task_instance_id", group_id)), + "operator": str(group["operator"]), + "object": str(group["object_uid"]), + "actor": deepcopy(group["actor"]), + "goal": deepcopy(group["goal"]), + "depends_on": list(group["depends_on"]), + "postcondition": deepcopy(group["success"]), + "edge_ids": [edge_id_by_unit[unit_id] for unit_id in group_units], + } + ) + + metadata = seed.get("metadata", {}) + allocation_groups = ( + deepcopy(metadata.get("legacy_allocation_groups", [])) + if isinstance(metadata, Mapping) + else [] + ) + program = { + "schema_version": EXECUTION_PROGRAM_SCHEMA, + "task": seed["task_id"], + "goal_description": seed["instruction"], + "start": start, + "goal": target_by_unit[ordered_units[-1]], + "nodes": graph_nodes, + "edges": edges, + "semantic_steps": semantic_steps, + "allocation_groups": allocation_groups, + "motion_policy_version": MOTION_POLICY_VERSION, + } + return validate_execution_program(program) + + +def _execution_units( + nodes: Sequence[Mapping[str, Any]], +) -> tuple[dict[str, str], dict[str, dict[str, Any]]]: + unit_by_node = { + str(node["id"]): str(node.get("sync_group", node["id"])) for node in nodes + } + units: dict[str, dict[str, Any]] = {} + for node in nodes: + node_id = str(node["id"]) + unit_id = unit_by_node[node_id] + unit = units.setdefault(unit_id, {"node_ids": [], "depends_on": set()}) + unit["node_ids"].append(node_id) + for dependency in node["depends_on"]: + dependency_unit = unit_by_node[str(dependency)] + if dependency_unit == unit_id: + raise ValueError( + f"Synchronized unit {unit_id!r} has an internal dependency." + ) + unit["depends_on"].add(dependency_unit) + for unit_id, unit in units.items(): + groups = { + str( + next(node for node in nodes if node["id"] == node_id)[ + "task_instance_id" + ] + ) + for node_id in unit["node_ids"] + } + if len(groups) != 1: + raise ValueError(f"Synchronized unit {unit_id!r} crosses task groups.") + return unit_by_node, units + + +def _topological_units(units: Mapping[str, Mapping[str, Any]]) -> list[str]: + return _topological_ids( + {unit_id: list(unit["depends_on"]) for unit_id, unit in units.items()} + ) + + +def _topological_groups(groups: Sequence[Mapping[str, Any]]) -> list[str]: + return _topological_ids( + {str(group["id"]): list(group["depends_on"]) for group in groups} + ) + + +def _topological_ids(dependencies: Mapping[str, Sequence[str]]) -> list[str]: + outgoing = {item_id: [] for item_id in dependencies} + indegree = {item_id: 0 for item_id in dependencies} + for item_id, parents in dependencies.items(): + for parent in parents: + outgoing[parent].append(item_id) + indegree[item_id] += 1 + ready = deque( + sorted(item_id for item_id, degree in indegree.items() if degree == 0) + ) + ordered = [] + while ready: + item_id = ready.popleft() + ordered.append(item_id) + for child in sorted(outgoing[item_id]): + indegree[child] -= 1 + if indegree[child] == 0: + ready.append(child) + if len(ordered) != len(dependencies): + raise ValueError("Graph contains a dependency cycle.") + return ordered + + +def _v2_actor(value: Mapping[str, Any]) -> dict[str, Any]: + actor = deepcopy(dict(value)) + actor.pop("allocation_group", None) + if actor.get("mode") == "required" and actor.get("arm") in {"left", "right"}: + actor["arm"] = f"{actor['arm']}_arm" + return actor + + +def _task_type(operator: str) -> str: + return _OPERATOR_TASK_TYPES.get(operator, "E1") + + +def _level(groups: Sequence[Mapping[str, Any]]) -> str: + types = {str(group["task_type"]) for group in groups} + if len(groups) == 1: + return "L1" + return "L2" if len(types) == 1 else "L3" + + +def _node_role(action_name: str, binding: Mapping[str, Any]) -> str: + if ( + action_name == "MoveJoints" and binding.get("source") == "initial" + ) or binding.get("kind") == "policy_pose": + return "cleanup" + return "primary" + + +def _node_id(edge_id: str, action: str, index: int, count: int) -> str: + base = f"{_slug(edge_id)}_{_slug(action)}" + return base if count == 1 else f"{base}_{index + 1}" + + +def _slug(value: str) -> str: + return _UNSAFE_ID_RE.sub("_", value.lower()).strip("_") or "node" diff --git a/embodichain/gen_sim/action_engine/config/__init__.py b/embodichain/gen_sim/action_engine/config/__init__.py new file mode 100644 index 000000000..d9759d702 --- /dev/null +++ b/embodichain/gen_sim/action_engine/config/__init__.py @@ -0,0 +1,41 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Validated package policy for Action Engine generation and runtime.""" + +from __future__ import annotations + +from .runtime_policy import ( + ACTION_ENGINE_DEFAULTS_SCHEMA, + RUNTIME_POLICY_SCHEMA, + ArmSelectionPolicyCfg, + RuntimePolicyCfg, + default_runtime_policy, + generation_defaults, + resolve_agent_runtime_policy, + runtime_policy_hash, +) + +__all__ = [ + "ACTION_ENGINE_DEFAULTS_SCHEMA", + "RUNTIME_POLICY_SCHEMA", + "ArmSelectionPolicyCfg", + "RuntimePolicyCfg", + "default_runtime_policy", + "generation_defaults", + "resolve_agent_runtime_policy", + "runtime_policy_hash", +] diff --git a/embodichain/gen_sim/action_engine/config/defaults.yaml b/embodichain/gen_sim/action_engine/config/defaults.yaml new file mode 100644 index 000000000..809810c5d --- /dev/null +++ b/embodichain/gen_sim/action_engine/config/defaults.yaml @@ -0,0 +1,344 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +schema_version: action_engine_defaults_v1 + +# Generation policy is materialized into fast_gym_config.json. It is not part +# of the coordinate-free Execution Program or the runtime-policy hash. +generation: + task: + default_robot_profile: ur10 + max_episodes: 1 + max_episode_steps: 2000 + environment: + viewer_camera_uid: cam_high + ignore_terminations_during_agent: true + recording: + enabled: true + resolution: [640, 360] + interval_step: 1 + arm_aim_yaw_offset: + left: 0.0 + right: 0.0 + scene: + prompt2scene_z_rotation_degrees: -90.0 + default_tabletop_z: 0.7 + body_scale_policy: preserve + body_scale: [1.0, 1.0, 1.0] + object_length_sample_points: 5000 + physics: + background: + mass: 10.0 + static_friction: 0.95 + dynamic_friction: 0.9 + restitution: 0.01 + max_convex_hull_num: 1 + rigid_object: + mass: 0.1 + static_friction: 0.95 + dynamic_friction: 0.9 + linear_damping: 0.9 + angular_damping: 0.9 + contact_offset: 0.003 + rest_offset: 0.001 + restitution: 0.05 + max_depenetration_velocity: 0.8 + max_linear_velocity: 5.0 + max_angular_velocity: 5.0 + min_position_iters: 32 + min_velocity_iters: 8 + max_convex_hull_num: 16 + acd_method: vhacd + randomization: + rigid_object_position_range: [[-0.04, -0.04, 0.0], [0.04, 0.04, 0.0]] + rigid_object_rotation_range: [[0.0, 0.0, -30.0], [0.0, 0.0, 30.0]] + table_height_delta_range: [[-0.05], [0.05]] + table_material: + random_texture_prob: 0.0 + base_color_range: [[0.55, 0.55, 0.55], [0.95, 0.95, 0.95]] + metallic_range: [0.0, 0.15] + roughness_range: [0.45, 0.95] + dataset: + control_frequency: 25 + save_failed_episodes: true + use_videos: true + +# Runtime policy is resolved per robot profile, snapshotted in agent_config, +# hash-verified at startup, and recorded with every execution. +runtime: + common: + execution: + max_transitions: 1000 + semantic_step_settle_steps: 10 + max_retries_per_action: 2 + max_graph_revisions: 8 + max_recovery_actions: 12 + support_stability_samples: 3 + support_stability_interval_steps: 5 + support_linear_velocity_tolerance: 0.02 + support_angular_velocity_tolerance: 0.20 + + planner: + backend: curobo + single_arm_strategy: motion_gen + coordinated_strategy: ik_interp + fallback_strategy: ik_interp + allow_fallback: true + dynamic_collision: false + static_obstacle_uids: [] + dynamic_obstacle_uids: [] + curobo: + log_level: error + obstacle_representation: cuboid + multi_env: false + use_cuda_graph: true + preserve_plan_samples: false + max_attempts: 5 + collision_activation_distance: 0.01 + + # Crossing is measured along the live right-to-left arm-base axis so the + # same-side constraint follows translated and rotated robot workspaces. + arm_selection: + crossing_deadband_ratio: 0.08 + allow_cross_side_fallback: false + pickup_crossing_weight: 1.0 + placement_crossing_weight: 1.5 + motion_cost_scale: 3.141592653589793 + fallback_workspace_half_width: 0.5 + orient_object_preferred_arm_deadband: 0.02 + + grounding: + semantic_defaults: + surface_clearance: 0.003 + transport_clearance: 0.10 + staging_lift_height: 0.12 + relation_distance: 0.16 + hover_height: 0.10 + press_depth: 0.004 + retreat_height: 0.10 + maximum_eef_height: 0.80 + arrangement: + slot_margin: 0.08 + minimum_spacing: 0.07 + layout_clearance: 0.025 + row_search_step: 0.025 + row_search_radius: 0.25 + placement: + clearance: 0.012 + candidate_count: 5 + candidate_offset_fraction: 0.50 + support_margin: 0.002 + recovery_attempts: 2 + coordinated_grasp: + inset_fraction: 0.15 + minimum_inset: 0.01 + handover: + retreat_height: 0.10 + retreat_distance: 0.10 + maximum_eef_height: 1.50 + minimum_transfer_clearance: 0.10 + minimum_transfer_lateral_clearance: 0.06 + joint_state: + hand_close_sample_interval: 10 + hand_open_sample_interval: 15 + + grasp: + antipodal_n_sample: 10000 + antipodal_max_angle: 0.2617993877991494 + max_open_length: 0.15 + min_open_length: 0.01 + finger_length: 0.13 + point_sample_dense: 0.012 + max_deviation_angle: 0.3490658503988659 + n_deviated_approach_directions: 4 + viser_port: 11801 + max_decomposition_hulls: 16 + force_grasp_reannotate: false + + motion_defaults: + PickUp: + pre_grasp_distance: 0.08 + lift_height: 0.30 + sample_interval: 45 + MoveHeldObject: + sample_interval: 45 + relation_distance: 0.18 + robot_relative_distance: 0.10 + relation_clearance: 0.02 + exchange_clearance: 0.06 + exchange_candidate_offset: 0.16 + exchange_obstacle_clearance: 0.04 + exchange_gripper_horizontal_envelope: 0.035 + exchange_wrist_horizontal_envelope: 0.055 + exchange_gripper_vertical_envelope: 0.025 + exchange_wrist_vertical_envelope: 0.04 + exchange_minimum_reach: 0.10 + exchange_maximum_reach: 1.00 + exchange_candidate_count: 4 + hover_height: 0.10 + line_spacing: 0.14 + transport_clearance: 0.10 + staging_lift_height: 0.30 + surface_clearance: 0.005 + postcondition_tolerance: 0.08 + line_axis_tolerance: 0.06 + line_perpendicular_tolerance: 0.06 + preserve_orientation_tolerance: 0.2617993877991494 + Place: + sample_interval: 15 + lift_height: 0.0 + post_hold_steps: 0 + cartesian_waypoint_count: 2 + MoveEndEffector: + sample_interval: 20 + retreat_height: 0.30 + minimum_retreat_height: 0.05 + maximum_eef_height: 1.10 + postcondition_tolerance: 0.05 + MoveJoints: + sample_interval: 30 + postcondition_tolerance: 0.05 + Press: + sample_interval: 80 + press_depth: 0.004 + postcondition_tolerance: 0.03 + CoordinatedPickment: + sample_interval: 120 + object_motion_keyframes: 6 + pre_grasp_distance: 0.10 + lift_height: 0.08 + middle_empty_ratio: 0.7 + is_filter_ground_collision: false + postcondition_tolerance: 0.06 + HandOver: + sample_interval: 140 + pre_grasp_distance: 0.08 + lift_height: 0.08 + receiver_hold_joint_tolerance: 0.002 + receive_pick_object_part: bottom + exchange_clearance: 0.06 + exchange_candidate_offset: 0.16 + exchange_obstacle_clearance: 0.04 + exchange_gripper_horizontal_envelope: 0.035 + exchange_wrist_horizontal_envelope: 0.055 + exchange_gripper_vertical_envelope: 0.025 + exchange_wrist_vertical_envelope: 0.04 + exchange_minimum_reach: 0.10 + exchange_maximum_reach: 1.00 + exchange_candidate_count: 4 + held_position_tolerance: 0.03 + hand_interp_steps: 10 + hold_steps: 4 + retreat_steps: 28 + postcondition_tolerance: 0.06 + CoordinatedPlacement: + sample_interval: 100 + hand_interp_steps: 10 + hold_steps: 4 + retreat_steps: 16 + postcondition_tolerance: 0.06 + + motion_modifiers: + orientation: + upright: + PickUp: + rotate_upright: 0.7853981633974483 + upright_yaw_samples: 8 + MoveHeldObject: + staging_lift_height: 0.25 + surface_clearance: 0.05 + upright_yaw_samples: 8 + upright_xy_tolerance: 0.05 + upright_max_tilt: 0.2617993877991494 + Place: + sample_interval: 64 + post_hold_steps: 12 + hand_interp_steps: 12 + MoveEndEffector: + sample_interval: 30 + retreat_height: 0.30 + retreat_distance: 0.10 + maximum_eef_height: 1.50 + handover_role: + transfer: + PickUp: + sample_interval: 80 + hand_interp_steps: 5 + pick_object_part: top + + predicate_fallbacks: + held_position_tolerance: 0.06 + held_gripper_tolerance: 0.01 + position_tolerance: 0.05 + xy_tolerance: 0.05 + container_xy_radius: 0.20 + container_min_z_offset: -0.05 + container_max_z_offset: 0.35 + support_xy_radius: 0.08 + support_com_margin: 0.002 + support_max_vertical_gap: 0.03 + support_max_penetration: 0.01 + support_min_overlap_ratio: 0.25 + not_fallen_max_tilt: 0.7853981633974483 + upright_max_tilt: 0.2617993877991494 + axis_tolerance: 0.03 + collinearity_tolerance: 0.03 + ordering_tolerance: 0.02 + minimum_lift_height: 0.08 + arm_initial_qpos_tolerance: 0.05 + gripper_state_tolerance: 0.001 + gripper_clear_min_distance: 0.08 + line_axis_tolerance: 0.06 + line_perpendicular_tolerance: 0.06 + preserve_orientation_tolerance: 0.2617993877991494 + payload_minimum_upright_cosine: 0.94 + payload_position_tolerance: 0.08 + payload_support_margin: 0.015 + + profiles: + dual_franka: + motion_defaults: + MoveHeldObject: + exchange_maximum_reach: 0.85 + MoveEndEffector: + retreat_height: 0.10 + HandOver: + exchange_maximum_reach: 0.85 + dual_ur3: + motion_defaults: + MoveHeldObject: + exchange_maximum_reach: 0.55 + HandOver: + exchange_maximum_reach: 0.55 + dual_ur5: + motion_defaults: + MoveHeldObject: + exchange_maximum_reach: 0.85 + HandOver: + exchange_maximum_reach: 0.85 + motion_modifiers: + orientation: + upright: + PickUp: + lift_height: 0.12 + MoveHeldObject: + staging_lift_height: 0.12 + dual_ur10: + motion_defaults: + MoveHeldObject: + exchange_maximum_reach: 1.25 + HandOver: + exchange_maximum_reach: 1.25 diff --git a/embodichain/gen_sim/action_engine/config/runtime_policy.py b/embodichain/gen_sim/action_engine/config/runtime_policy.py new file mode 100644 index 000000000..56198211b --- /dev/null +++ b/embodichain/gen_sim/action_engine/config/runtime_policy.py @@ -0,0 +1,775 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Load, resolve, snapshot, and hash package-owned Action Engine defaults.""" + +from __future__ import annotations + +from collections.abc import Mapping +from copy import deepcopy +import hashlib +import json +import math +from pathlib import Path +from typing import Any, Final + +from embodichain.gen_sim.action_engine.domain.motion import MOTION_MODIFIER_MODES +from embodichain.utils import configclass +from embodichain.utils.utility import load_config + +__all__ = [ + "ACTION_ENGINE_DEFAULTS_SCHEMA", + "RUNTIME_POLICY_SCHEMA", + "ArmSelectionPolicyCfg", + "RuntimePolicyCfg", + "default_runtime_policy", + "generation_defaults", + "resolve_agent_runtime_policy", + "runtime_policy_hash", +] + +ACTION_ENGINE_DEFAULTS_SCHEMA: Final = "action_engine_defaults_v1" +RUNTIME_POLICY_SCHEMA: Final = "action_engine_runtime_policy_v6" +_PREVIOUS_RUNTIME_POLICY_SCHEMA: Final = "action_engine_runtime_policy_v5" +_PRE_GRASP_RUNTIME_POLICY_SCHEMA: Final = "action_engine_runtime_policy_v4" +_PRE_PLANNER_RUNTIME_POLICY_SCHEMA: Final = "action_engine_runtime_policy_v3" +_LEGACY_RUNTIME_POLICY_SCHEMA: Final = "action_engine_runtime_policy_v1" +_DEFAULTS_PATH = Path(__file__).with_name("defaults.yaml") +_ARM_SELECTION_KEYS = ( + "crossing_deadband_ratio", + "pickup_crossing_weight", + "placement_crossing_weight", + "motion_cost_scale", + "fallback_workspace_half_width", + "orient_object_preferred_arm_deadband", +) +_ARM_SELECTION_OPTIONAL_KEYS = {"allow_cross_side_fallback"} +_GROUNDING_KEYS = { + "semantic_defaults": { + "surface_clearance", + "transport_clearance", + "staging_lift_height", + "relation_distance", + "hover_height", + "press_depth", + "retreat_height", + "maximum_eef_height", + }, + "arrangement": { + "slot_margin", + "minimum_spacing", + "layout_clearance", + "row_search_step", + "row_search_radius", + }, + "placement": { + "clearance", + "candidate_count", + "candidate_offset_fraction", + "support_margin", + "recovery_attempts", + }, + "coordinated_grasp": {"inset_fraction", "minimum_inset"}, + "handover": { + "retreat_height", + "retreat_distance", + "maximum_eef_height", + "minimum_transfer_clearance", + "minimum_transfer_lateral_clearance", + }, + "joint_state": { + "hand_close_sample_interval", + "hand_open_sample_interval", + }, +} +_GRASP_KEYS = { + "antipodal_n_sample", + "antipodal_max_angle", + "max_open_length", + "min_open_length", + "finger_length", + "point_sample_dense", + "max_deviation_angle", + "n_deviated_approach_directions", + "viser_port", + "max_decomposition_hulls", + "force_grasp_reannotate", +} +_PLANNER_KEYS = { + "backend", + "single_arm_strategy", + "coordinated_strategy", + "fallback_strategy", + "allow_fallback", + "dynamic_collision", + "static_obstacle_uids", + "dynamic_obstacle_uids", + "curobo", +} +_CUROBO_KEYS = { + "log_level", + "obstacle_representation", + "multi_env", + "use_cuda_graph", + "preserve_plan_samples", + "max_attempts", + "collision_activation_distance", +} +_MOTION_DEFAULT_ACTIONS = { + "CoordinatedPickment", + "CoordinatedPlacement", + "HandOver", + "MoveEndEffector", + "MoveHeldObject", + "MoveJoints", + "PickUp", + "Place", + "Press", +} +_PREDICATE_KEYS = { + "held_position_tolerance", + "held_gripper_tolerance", + "position_tolerance", + "xy_tolerance", + "container_xy_radius", + "container_min_z_offset", + "container_max_z_offset", + "support_xy_radius", + "support_com_margin", + "support_max_vertical_gap", + "support_max_penetration", + "support_min_overlap_ratio", + "not_fallen_max_tilt", + "upright_max_tilt", + "axis_tolerance", + "collinearity_tolerance", + "ordering_tolerance", + "minimum_lift_height", + "arm_initial_qpos_tolerance", + "gripper_state_tolerance", + "gripper_clear_min_distance", + "line_axis_tolerance", + "line_perpendicular_tolerance", + "preserve_orientation_tolerance", + "payload_minimum_upright_cosine", + "payload_position_tolerance", + "payload_support_margin", +} +_DEPRECATED_PREDICATE_KEYS = { + "support_min_z_offset", + "support_max_z_offset", +} + + +@configclass +class ArmSelectionPolicyCfg: + """Arm-allocation constraints and costs resolved for one robot profile.""" + + crossing_deadband_ratio: float = 0.08 + allow_cross_side_fallback: bool = False + pickup_crossing_weight: float = 1.0 + placement_crossing_weight: float = 1.5 + motion_cost_scale: float = math.pi + fallback_workspace_half_width: float = 0.5 + orient_object_preferred_arm_deadband: float = 0.02 + + def __post_init__(self) -> None: + if not isinstance(self.allow_cross_side_fallback, bool): + raise TypeError("allow_cross_side_fallback must be a bool.") + for name in _ARM_SELECTION_KEYS: + value = float(getattr(self, name)) + if not math.isfinite(value): + raise ValueError(f"{name} must be finite.") + if not 0.0 <= float(self.crossing_deadband_ratio) < 1.0: + raise ValueError("crossing_deadband_ratio must be in [0, 1).") + for name in _ARM_SELECTION_KEYS[1:3]: + if float(getattr(self, name)) < 0.0: + raise ValueError(f"{name} must be non-negative.") + for name in _ARM_SELECTION_KEYS[3:]: + if float(getattr(self, name)) <= 0.0: + raise ValueError(f"{name} must be positive.") + + @classmethod + def from_mapping(cls, value: Mapping[str, Any]) -> ArmSelectionPolicyCfg: + """Build a strict policy from a JSON/YAML mapping.""" + keys = frozenset(value) + if keys not in { + frozenset(_ARM_SELECTION_KEYS), + frozenset((*_ARM_SELECTION_KEYS, *_ARM_SELECTION_OPTIONAL_KEYS)), + }: + raise ValueError("arm_selection fields do not match the policy schema.") + fields: dict[str, Any] = {key: float(value[key]) for key in _ARM_SELECTION_KEYS} + if "allow_cross_side_fallback" in value: + fields["allow_cross_side_fallback"] = value["allow_cross_side_fallback"] + return cls(**fields) + + def as_mapping(self) -> dict[str, float | bool]: + """Return a stable JSON-compatible representation.""" + return { + "crossing_deadband_ratio": float(self.crossing_deadband_ratio), + "allow_cross_side_fallback": bool(self.allow_cross_side_fallback), + "pickup_crossing_weight": float(self.pickup_crossing_weight), + "placement_crossing_weight": float(self.placement_crossing_weight), + "motion_cost_scale": float(self.motion_cost_scale), + "fallback_workspace_half_width": float(self.fallback_workspace_half_width), + "orient_object_preferred_arm_deadband": float( + self.orient_object_preferred_arm_deadband + ), + } + + +@configclass +class RuntimePolicyCfg: + """Effective runtime policy persisted in generated agent artifacts.""" + + schema_version: str = RUNTIME_POLICY_SCHEMA + arm_selection: ArmSelectionPolicyCfg = ArmSelectionPolicyCfg() + execution: dict[str, Any] = {} + planner: dict[str, Any] = {} + grounding: dict[str, Any] = {} + grasp: dict[str, Any] = {} + motion_defaults: dict[str, dict[str, Any]] = {} + motion_modifiers: dict[str, dict[str, dict[str, dict[str, Any]]]] = {} + predicate_fallbacks: dict[str, Any] = {} + + def __post_init__(self) -> None: + if self.schema_version != RUNTIME_POLICY_SCHEMA: + raise ValueError( + f"Unsupported runtime policy schema {self.schema_version!r}." + ) + if not isinstance(self.arm_selection, ArmSelectionPolicyCfg): + raise TypeError("arm_selection must be an ArmSelectionPolicyCfg.") + for name in ( + "execution", + "planner", + "grounding", + "grasp", + "motion_defaults", + "motion_modifiers", + "predicate_fallbacks", + ): + if not isinstance(getattr(self, name), dict): + raise TypeError(f"{name} must be a mapping.") + _validate_finite_numbers(getattr(self, name), name) + if int(self.execution.get("max_transitions", 0)) <= 0: + raise ValueError("execution.max_transitions must be positive.") + if int(self.execution.get("semantic_step_settle_steps", -1)) < 0: + raise ValueError( + "execution.semantic_step_settle_steps must be non-negative." + ) + for name in ( + "max_retries_per_action", + "max_graph_revisions", + "max_recovery_actions", + "support_stability_interval_steps", + ): + if int(self.execution.get(name, -1)) < 0: + raise ValueError(f"execution.{name} must be non-negative.") + _require_keys( + self.execution, + { + "max_transitions", + "semantic_step_settle_steps", + "max_retries_per_action", + "max_graph_revisions", + "max_recovery_actions", + "support_stability_samples", + "support_stability_interval_steps", + "support_linear_velocity_tolerance", + "support_angular_velocity_tolerance", + }, + "execution", + ) + if int(self.execution["support_stability_samples"]) <= 0: + raise ValueError("execution.support_stability_samples must be positive.") + for name in ( + "support_linear_velocity_tolerance", + "support_angular_velocity_tolerance", + ): + if float(self.execution[name]) < 0.0: + raise ValueError(f"execution.{name} must be non-negative.") + _validate_planner(self.planner) + _require_keys(self.grounding, set(_GROUNDING_KEYS), "grounding") + for name, keys in _GROUNDING_KEYS.items(): + section = self.grounding.get(name) + if not isinstance(section, Mapping): + raise ValueError(f"grounding.{name} must be a mapping.") + _require_keys(section, keys, f"grounding.{name}") + placement = self.grounding["placement"] + if not 1 <= int(placement["candidate_count"]) <= 9: + raise ValueError("grounding.placement.candidate_count must be in [1, 9].") + if int(placement["recovery_attempts"]) < 0: + raise ValueError( + "grounding.placement.recovery_attempts must be non-negative." + ) + if not 0.0 <= float(placement["candidate_offset_fraction"]) <= 1.0: + raise ValueError( + "grounding.placement.candidate_offset_fraction must be in [0, 1]." + ) + if float(placement["support_margin"]) < 0.0: + raise ValueError("grounding.placement.support_margin must be non-negative.") + _require_keys(self.grasp, _GRASP_KEYS, "grasp") + _require_keys( + self.motion_defaults, + _MOTION_DEFAULT_ACTIONS, + "motion_defaults", + ) + if not all( + isinstance(policy, Mapping) and policy + for policy in self.motion_defaults.values() + ): + raise ValueError("Every motion default must be a non-empty mapping.") + _validate_motion_modifiers(self.motion_modifiers) + _require_keys( + self.predicate_fallbacks, + _PREDICATE_KEYS, + "predicate_fallbacks", + ) + if float(self.grasp.get("min_open_length", -1.0)) < 0.0: + raise ValueError("grasp.min_open_length must be non-negative.") + if float(self.grasp.get("max_open_length", 0.0)) <= float( + self.grasp.get("min_open_length", 0.0) + ): + raise ValueError("grasp.max_open_length must exceed min_open_length.") + direction_count = self.grasp.get("n_deviated_approach_directions") + if ( + isinstance(direction_count, bool) + or not isinstance(direction_count, int) + or not 1 <= direction_count <= 16 + ): + raise ValueError("grasp.n_deviated_approach_directions must be in [1, 16].") + + @classmethod + def from_mapping(cls, value: Mapping[str, Any]) -> RuntimePolicyCfg: + """Parse one fully resolved policy snapshot.""" + fields = { + "schema_version", + "execution", + "planner", + "arm_selection", + "grounding", + "grasp", + "motion_defaults", + "motion_modifiers", + "predicate_fallbacks", + } + if set(value) != fields: + raise ValueError("Runtime policy fields do not match the policy schema.") + if value.get("schema_version") != RUNTIME_POLICY_SCHEMA: + raise ValueError("Runtime policy has an unexpected schema_version.") + arm_selection = value.get("arm_selection") + if not isinstance(arm_selection, Mapping): + raise ValueError("Runtime policy requires an arm_selection mapping.") + sections = { + name: value.get(name) + for name in fields + if name not in {"schema_version", "arm_selection"} + } + if not all(isinstance(section, Mapping) for section in sections.values()): + raise ValueError("Runtime policy sections must be mappings.") + resolved_sections = { + name: deepcopy(dict(section)) for name, section in sections.items() + } + predicate_fallbacks = resolved_sections["predicate_fallbacks"] + for key in _DEPRECATED_PREDICATE_KEYS: + predicate_fallbacks.pop(key, None) + return cls( + schema_version=RUNTIME_POLICY_SCHEMA, + arm_selection=ArmSelectionPolicyCfg.from_mapping(arm_selection), + **resolved_sections, + ) + + def as_mapping(self) -> dict[str, Any]: + """Return the canonical artifact snapshot.""" + return { + "schema_version": self.schema_version, + "execution": deepcopy(self.execution), + "planner": deepcopy(self.planner), + "arm_selection": self.arm_selection.as_mapping(), + "grounding": deepcopy(self.grounding), + "grasp": deepcopy(self.grasp), + "motion_defaults": deepcopy(self.motion_defaults), + "motion_modifiers": deepcopy(self.motion_modifiers), + "predicate_fallbacks": deepcopy(self.predicate_fallbacks), + } + + +def default_runtime_policy(robot_profile: str) -> RuntimePolicyCfg: + """Resolve a package policy for one canonical robot profile.""" + document = _load_defaults() + runtime = document.get("runtime") + if not isinstance(runtime, Mapping) or set(runtime) != {"common", "profiles"}: + raise ValueError("Runtime defaults require common and profiles mappings.") + common, profiles = runtime["common"], runtime["profiles"] + if not isinstance(common, Mapping) or not isinstance(profiles, Mapping): + raise ValueError("Runtime common and profiles must be mappings.") + override = profiles.get(str(robot_profile)) + if not isinstance(override, Mapping): + raise ValueError(f"Unknown runtime robot profile {robot_profile!r}.") + resolved = _deep_merge(common, override) + return RuntimePolicyCfg.from_mapping( + { + "schema_version": RUNTIME_POLICY_SCHEMA, + **resolved, + } + ) + + +def generation_defaults() -> dict[str, Any]: + """Return a detached generation-policy mapping.""" + value = _load_defaults().get("generation") + if not isinstance(value, Mapping): + raise ValueError("Action Engine defaults require a generation mapping.") + required = { + "task", + "environment", + "scene", + "physics", + "randomization", + "dataset", + } + if set(value) != required: + raise ValueError("Generation defaults do not match the expected sections.") + return deepcopy(dict(value)) + + +def _load_defaults() -> dict[str, Any]: + document = load_config(_DEFAULTS_PATH) + if not isinstance(document, dict) or set(document) != { + "schema_version", + "generation", + "runtime", + }: + raise ValueError("Action Engine defaults do not match the package schema.") + if document.get("schema_version") != ACTION_ENGINE_DEFAULTS_SCHEMA: + raise ValueError("Action Engine defaults have an unexpected schema_version.") + return document + + +def _deep_merge( + base: Mapping[str, Any], + override: Mapping[str, Any], +) -> dict[str, Any]: + result = deepcopy(dict(base)) + for key, value in override.items(): + current = result.get(key) + result[key] = ( + _deep_merge(current, value) + if isinstance(current, Mapping) and isinstance(value, Mapping) + else deepcopy(value) + ) + return result + + +def _validate_finite_numbers(value: Any, path: str) -> None: + if isinstance(value, Mapping): + for key, item in value.items(): + _validate_finite_numbers(item, f"{path}.{key}") + elif isinstance(value, (list, tuple)): + for index, item in enumerate(value): + _validate_finite_numbers(item, f"{path}[{index}]") + elif isinstance(value, float) and not math.isfinite(value): + raise ValueError(f"{path} must be finite.") + + +def _require_keys( + value: Mapping[str, Any], + expected: set[str], + path: str, +) -> None: + if set(value) != expected: + raise ValueError(f"{path} fields do not match the defaults schema.") + + +def _validate_string_sequence(value: Any, path: str) -> None: + if not isinstance(value, (list, tuple)): + raise ValueError(f"{path} must be a list of object UIDs.") + normalized = [str(item) for item in value] + if any(not item.strip() for item in normalized): + raise ValueError(f"{path} entries must be non-empty strings.") + if any(not isinstance(item, str) for item in value): + raise ValueError(f"{path} entries must be strings.") + if len(set(normalized)) != len(normalized): + raise ValueError(f"{path} must not contain duplicate object UIDs.") + + +def _validate_planner(value: Mapping[str, Any]) -> None: + _require_keys(value, _PLANNER_KEYS, "planner") + backend = value.get("backend") + if backend not in {"curobo", "toppra"}: + raise ValueError("planner.backend must be 'curobo' or 'toppra'.") + for name in ("single_arm_strategy", "coordinated_strategy"): + if value.get(name) not in {"motion_gen", "ik_interp"}: + raise ValueError(f"planner.{name} must be 'motion_gen' or 'ik_interp'.") + if value.get("fallback_strategy") != "ik_interp": + raise ValueError("planner.fallback_strategy must be 'ik_interp'.") + if value.get("coordinated_strategy") == "motion_gen" and backend == "curobo": + raise ValueError( + "planner.coordinated_strategy must be 'ik_interp' with cuRobo." + ) + for name in ("allow_fallback", "dynamic_collision"): + if not isinstance(value.get(name), bool): + raise ValueError(f"planner.{name} must be a boolean.") + if value.get("dynamic_collision") and backend != "curobo": + raise ValueError("planner.dynamic_collision requires the cuRobo backend.") + _validate_string_sequence( + value.get("static_obstacle_uids"), + "planner.static_obstacle_uids", + ) + _validate_string_sequence( + value.get("dynamic_obstacle_uids"), + "planner.dynamic_obstacle_uids", + ) + overlap = set(value["static_obstacle_uids"]) & set(value["dynamic_obstacle_uids"]) + if overlap: + raise ValueError( + "Planner obstacle UIDs cannot be both static and dynamic: " + f"{sorted(overlap)}." + ) + + curobo = value.get("curobo") + if not isinstance(curobo, Mapping): + raise ValueError("planner.curobo must be a mapping.") + _require_keys(curobo, _CUROBO_KEYS, "planner.curobo") + if curobo.get("log_level") not in { + "debug", + "info", + "warning", + "warn", + "error", + }: + raise ValueError("planner.curobo.log_level is unsupported.") + if curobo.get("obstacle_representation") not in {"sphere", "cuboid", "mesh"}: + raise ValueError( + "planner.curobo.obstacle_representation must be sphere, cuboid, or mesh." + ) + for name in ("multi_env", "use_cuda_graph", "preserve_plan_samples"): + if not isinstance(curobo.get(name), bool): + raise ValueError(f"planner.curobo.{name} must be a boolean.") + max_attempts = curobo.get("max_attempts") + if ( + isinstance(max_attempts, bool) + or not isinstance(max_attempts, int) + or max_attempts <= 0 + ): + raise ValueError("planner.curobo.max_attempts must be positive.") + activation_distance = curobo.get("collision_activation_distance") + if ( + isinstance(activation_distance, bool) + or not isinstance(activation_distance, (int, float)) + or float(activation_distance) < 0.0 + ): + raise ValueError( + "planner.curobo.collision_activation_distance must be non-negative." + ) + + +def _validate_motion_modifiers(value: Mapping[str, Any]) -> None: + _require_keys(value, set(MOTION_MODIFIER_MODES), "motion_modifiers") + for modifier_type, modes in MOTION_MODIFIER_MODES.items(): + configured_modes = value.get(modifier_type) + if not isinstance(configured_modes, Mapping): + raise ValueError(f"motion_modifiers.{modifier_type} must be a mapping.") + _require_keys( + configured_modes, + set(modes), + f"motion_modifiers.{modifier_type}", + ) + for mode, patches in configured_modes.items(): + path = f"motion_modifiers.{modifier_type}.{mode}" + if not isinstance(patches, Mapping) or not patches: + raise ValueError(f"{path} must contain action-specific patches.") + unknown_actions = set(patches) - _MOTION_DEFAULT_ACTIONS + if unknown_actions: + raise ValueError( + f"{path} references unknown actions: {sorted(unknown_actions)}." + ) + if not all( + isinstance(patch, Mapping) and patch for patch in patches.values() + ): + raise ValueError(f"Every {path} action patch must be non-empty.") + + +def runtime_policy_hash(policy: RuntimePolicyCfg | Mapping[str, Any]) -> str: + """Hash the canonical effective policy independently of the Seed graph.""" + resolved = ( + policy + if isinstance(policy, RuntimePolicyCfg) + else RuntimePolicyCfg.from_mapping(policy) + ) + return _mapping_hash(resolved.as_mapping()) + + +def resolve_agent_runtime_policy(agent_config: Mapping[str, Any]) -> RuntimePolicyCfg: + """Resolve a generated snapshot or fall back for a legacy v1 artifact.""" + snapshot = agent_config.get("runtime_policy") + expected_hash = agent_config.get("runtime_policy_hash") + if snapshot is None: + if expected_hash is not None: + raise ValueError("runtime_policy_hash requires a runtime_policy snapshot.") + return default_runtime_policy( + str(agent_config.get("robot_profile", "dual_ur10")) + ) + if not isinstance(snapshot, Mapping): + raise ValueError("agent_config.runtime_policy must be a mapping.") + if not isinstance(expected_hash, str) or not expected_hash: + raise ValueError( + "agent_config.runtime_policy requires a non-empty runtime_policy_hash." + ) + if _mapping_hash(snapshot) != expected_hash: + raise ValueError( + "agent_config runtime policy hash does not match its snapshot." + ) + if snapshot.get("schema_version") == _LEGACY_RUNTIME_POLICY_SCHEMA: + if set(snapshot) != {"schema_version", "arm_selection"} or not isinstance( + snapshot.get("arm_selection"), Mapping + ): + raise ValueError("Legacy runtime policy snapshot is malformed.") + policy = default_runtime_policy( + str(agent_config.get("robot_profile", "dual_ur10")) + ) + merged = policy.arm_selection.as_mapping() + merged.update(snapshot["arm_selection"]) + policy.arm_selection = ArmSelectionPolicyCfg.from_mapping(merged) + return policy + if snapshot.get("schema_version") == _PREVIOUS_RUNTIME_POLICY_SCHEMA: + defaults = default_runtime_policy( + str(agent_config.get("robot_profile", "dual_ur10")) + ) + migrated = deepcopy(dict(snapshot)) + migrated["schema_version"] = RUNTIME_POLICY_SCHEMA + migrated_execution = deepcopy(dict(migrated.get("execution", {}))) + for key in ( + "support_stability_samples", + "support_stability_interval_steps", + "support_linear_velocity_tolerance", + "support_angular_velocity_tolerance", + ): + migrated_execution[key] = defaults.execution[key] + migrated["execution"] = migrated_execution + migrated_grounding = deepcopy(dict(migrated.get("grounding", {}))) + migrated_placement = deepcopy(dict(migrated_grounding.get("placement", {}))) + for key, value in defaults.grounding["placement"].items(): + migrated_placement.setdefault(key, value) + migrated_grounding["placement"] = migrated_placement + migrated["grounding"] = migrated_grounding + migrated_predicates = deepcopy(dict(migrated.get("predicate_fallbacks", {}))) + for key in ( + "support_com_margin", + "support_max_vertical_gap", + "support_max_penetration", + "support_min_overlap_ratio", + ): + migrated_predicates[key] = defaults.predicate_fallbacks[key] + migrated["predicate_fallbacks"] = migrated_predicates + return RuntimePolicyCfg.from_mapping(migrated) + if snapshot.get("schema_version") == _PRE_GRASP_RUNTIME_POLICY_SCHEMA: + defaults = default_runtime_policy( + str(agent_config.get("robot_profile", "dual_ur10")) + ) + migrated = deepcopy(dict(snapshot)) + migrated["schema_version"] = RUNTIME_POLICY_SCHEMA + migrated_grasp = deepcopy(dict(migrated.get("grasp", {}))) + migrated_grasp["n_deviated_approach_directions"] = defaults.grasp[ + "n_deviated_approach_directions" + ] + migrated["grasp"] = migrated_grasp + migrated_execution = deepcopy(dict(migrated.get("execution", {}))) + for key in ( + "support_stability_samples", + "support_stability_interval_steps", + "support_linear_velocity_tolerance", + "support_angular_velocity_tolerance", + ): + migrated_execution[key] = defaults.execution[key] + migrated["execution"] = migrated_execution + migrated_grounding = deepcopy(dict(migrated.get("grounding", {}))) + migrated_placement = deepcopy(dict(migrated_grounding.get("placement", {}))) + for key, value in defaults.grounding["placement"].items(): + migrated_placement.setdefault(key, value) + migrated_grounding["placement"] = migrated_placement + migrated["grounding"] = migrated_grounding + migrated_predicates = deepcopy(dict(migrated.get("predicate_fallbacks", {}))) + for key in ( + "support_com_margin", + "support_max_vertical_gap", + "support_max_penetration", + "support_min_overlap_ratio", + ): + migrated_predicates[key] = defaults.predicate_fallbacks[key] + migrated["predicate_fallbacks"] = migrated_predicates + return RuntimePolicyCfg.from_mapping(migrated) + if snapshot.get("schema_version") == _PRE_PLANNER_RUNTIME_POLICY_SCHEMA: + expected_fields = { + "schema_version", + "execution", + "arm_selection", + "grounding", + "grasp", + "motion_defaults", + "motion_modifiers", + "predicate_fallbacks", + } + if set(snapshot) != expected_fields: + raise ValueError("Previous runtime policy snapshot is malformed.") + defaults = default_runtime_policy( + str(agent_config.get("robot_profile", "dual_ur10")) + ) + migrated = deepcopy(dict(snapshot)) + migrated["schema_version"] = RUNTIME_POLICY_SCHEMA + migrated["planner"] = deepcopy(defaults.planner) + migrated_execution = deepcopy(dict(migrated.get("execution", {}))) + for key in ( + "support_stability_samples", + "support_stability_interval_steps", + "support_linear_velocity_tolerance", + "support_angular_velocity_tolerance", + ): + migrated_execution[key] = defaults.execution[key] + migrated["execution"] = migrated_execution + migrated_grounding = deepcopy(dict(migrated.get("grounding", {}))) + migrated_placement = deepcopy(dict(migrated_grounding.get("placement", {}))) + for key, value in defaults.grounding["placement"].items(): + migrated_placement.setdefault(key, value) + migrated_grounding["placement"] = migrated_placement + migrated["grounding"] = migrated_grounding + migrated_grasp = deepcopy(dict(migrated["grasp"])) + migrated_grasp["n_deviated_approach_directions"] = defaults.grasp[ + "n_deviated_approach_directions" + ] + migrated["grasp"] = migrated_grasp + migrated_predicates = deepcopy(dict(migrated["predicate_fallbacks"])) + for key in ( + "support_com_margin", + "support_max_vertical_gap", + "support_max_penetration", + "support_min_overlap_ratio", + ): + migrated_predicates[key] = defaults.predicate_fallbacks[key] + migrated["predicate_fallbacks"] = migrated_predicates + return RuntimePolicyCfg.from_mapping(migrated) + policy = RuntimePolicyCfg.from_mapping(snapshot) + return policy + + +def _mapping_hash(value: Mapping[str, Any]) -> str: + payload = json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() diff --git a/embodichain/gen_sim/action_engine/domain/__init__.py b/embodichain/gen_sim/action_engine/domain/__init__.py new file mode 100644 index 000000000..5f2c303ad --- /dev/null +++ b/embodichain/gen_sim/action_engine/domain/__init__.py @@ -0,0 +1,93 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Stable public contracts for Action Engine programs.""" + +from __future__ import annotations + +from .motion import ( + MOTION_MODIFIER_MODES, + MOTION_POLICY_VERSION, + motion_policy, + validate_motion_policy, +) +from .programs import ( + EXECUTION_PROGRAM_SCHEMA, + TASK_AGENT_SCHEMA, + execution_program_hash, + validate_execution_program, + validate_task_agent, +) +from .task_contracts import ( + PLACEMENT_RELATIONS, + RELATIONS, + TASK_CONTRACTS, + TERMINAL_BEHAVIORS, + TRANSPORT_DIRECTIONS, + TaskContract, + normalize_placement_relation, + task_contract, + task_success_type, +) +from .v2 import ( + REASONING_TYPES, + TASK_LEVELS, + TASK_TYPES, + public_task_spec, + seed_graph_hash, + validate_public_task_spec, + validate_scene_requirements, + validate_seed_graph, + validate_task_spec, +) +from .visual_contracts import ( + OCCLUSION_RELATION, + VISUAL_RELATION_PARTICIPANTS, + requested_visual_task_predicates, +) + +__all__ = [ + "EXECUTION_PROGRAM_SCHEMA", + "MOTION_POLICY_VERSION", + "MOTION_MODIFIER_MODES", + "OCCLUSION_RELATION", + "REASONING_TYPES", + "RELATIONS", + "PLACEMENT_RELATIONS", + "TASK_CONTRACTS", + "TASK_LEVELS", + "TASK_TYPES", + "TASK_AGENT_SCHEMA", + "TERMINAL_BEHAVIORS", + "TRANSPORT_DIRECTIONS", + "VISUAL_RELATION_PARTICIPANTS", + "TaskContract", + "execution_program_hash", + "motion_policy", + "normalize_placement_relation", + "public_task_spec", + "requested_visual_task_predicates", + "seed_graph_hash", + "task_contract", + "task_success_type", + "validate_public_task_spec", + "validate_scene_requirements", + "validate_seed_graph", + "validate_task_spec", + "validate_execution_program", + "validate_motion_policy", + "validate_task_agent", +] diff --git a/embodichain/gen_sim/action_engine/domain/motion.py b/embodichain/gen_sim/action_engine/domain/motion.py new file mode 100644 index 000000000..ae9ced35d --- /dev/null +++ b/embodichain/gen_sim/action_engine/domain/motion.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. +# ---------------------------------------------------------------------------- + +"""Typed, composable motion-policy references persisted in symbolic graphs.""" + +from __future__ import annotations + +from collections.abc import Mapping +from copy import deepcopy +from typing import Any, Final + +__all__ = [ + "MOTION_MODIFIER_MODES", + "MOTION_POLICY_VERSION", + "motion_policy", + "validate_motion_policy", +] + +MOTION_POLICY_VERSION: Final = "action_engine_motion_policy_v3" +MOTION_MODIFIER_MODES: Final = { + "orientation": frozenset({"upright"}), + "handover_role": frozenset({"transfer"}), +} + +_POLICY_KEYS = frozenset({"modifiers"}) +_MODIFIER_KEYS = frozenset({"type", "mode"}) + + +def motion_policy(*modifiers: tuple[str, str]) -> dict[str, Any]: + """Build one canonical policy reference from typed modifier pairs.""" + return validate_motion_policy( + { + "modifiers": [ + {"type": modifier_type, "mode": mode} + for modifier_type, mode in modifiers + ] + } + ) + + +def validate_motion_policy( + value: Any, + context: str = "motion_policy", +) -> dict[str, Any]: + """Validate and detach one symbolic motion-policy reference.""" + if not isinstance(value, Mapping): + raise ValueError( + f"{context} must be a mapping with typed modifiers; named string " + "policies are no longer supported. Regenerate the graph." + ) + if set(value) != _POLICY_KEYS: + raise ValueError(f"{context} fields must be {sorted(_POLICY_KEYS)}.") + raw_modifiers = value.get("modifiers") + if not isinstance(raw_modifiers, (list, tuple)): + raise ValueError(f"{context}.modifiers must be a sequence.") + + modifiers: list[dict[str, str]] = [] + seen: set[tuple[str, str]] = set() + seen_types: set[str] = set() + for index, raw_modifier in enumerate(raw_modifiers): + modifier_context = f"{context}.modifiers[{index}]" + if not isinstance(raw_modifier, Mapping): + raise ValueError(f"{modifier_context} must be a mapping.") + if set(raw_modifier) != _MODIFIER_KEYS: + raise ValueError( + f"{modifier_context} fields must be {sorted(_MODIFIER_KEYS)}." + ) + modifier_type = raw_modifier.get("type") + mode = raw_modifier.get("mode") + if not isinstance(modifier_type, str) or not modifier_type: + raise ValueError(f"{modifier_context}.type must be a non-empty string.") + if modifier_type not in MOTION_MODIFIER_MODES: + raise ValueError( + f"{modifier_context}.type {modifier_type!r} is unsupported; " + f"expected one of {sorted(MOTION_MODIFIER_MODES)}." + ) + if ( + not isinstance(mode, str) + or mode not in MOTION_MODIFIER_MODES[modifier_type] + ): + raise ValueError( + f"{modifier_context}.mode {mode!r} is unsupported for " + f"{modifier_type!r}; expected one of " + f"{sorted(MOTION_MODIFIER_MODES[modifier_type])}." + ) + key = (modifier_type, mode) + if key in seen: + raise ValueError(f"{modifier_context} duplicates modifier {key!r}.") + if modifier_type in seen_types: + raise ValueError( + f"{context} may select only one mode for modifier type " + f"{modifier_type!r}." + ) + seen.add(key) + seen_types.add(modifier_type) + modifiers.append({"type": modifier_type, "mode": mode}) + + return {"modifiers": deepcopy(modifiers)} diff --git a/embodichain/gen_sim/action_engine/domain/programs.py b/embodichain/gen_sim/action_engine/domain/programs.py new file mode 100644 index 000000000..a1243f1ea --- /dev/null +++ b/embodichain/gen_sim/action_engine/domain/programs.py @@ -0,0 +1,848 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Coordinate-free task and execution program contracts. + +The task agent is the only structure an LLM is allowed to influence. The +execution program is produced deterministically and contains the complete +symbolic action DAG consumed by runtime. Neither representation may contain +poses, trajectories, joint values, or other environment-specific geometry. +""" + +from __future__ import annotations + +import hashlib +import json +from collections import deque +from collections.abc import Collection, Mapping, Sequence +from copy import deepcopy +from typing import Any + +from embodichain.gen_sim.action_engine.protocol import ( + EXECUTION_PROGRAM_SCHEMA, + TASK_AGENT_SCHEMA, +) +from .motion import MOTION_POLICY_VERSION, validate_motion_policy + +__all__ = [ + "EXECUTION_PROGRAM_SCHEMA", + "MOTION_POLICY_VERSION", + "TASK_AGENT_SCHEMA", + "execution_program_hash", + "validate_execution_program", + "validate_task_agent", +] + +_ACTOR_MODES = frozenset({"auto", "required", "coordinated"}) +_CONTROL_MODES = frozenset({"arm", "hand", "coordinated"}) +_TASK_KEYS = frozenset( + {"schema_version", "task", "goal", "semantic_steps", "allocation_groups"} +) +_TASK_STEP_KEYS = frozenset( + {"id", "operator", "object", "objects", "actor", "goal", "depends_on"} +) +_EXECUTION_KEYS = frozenset( + { + "schema_version", + "task", + "goal_description", + "start", + "goal", + "nodes", + "edges", + "semantic_steps", + "allocation_groups", + "motion_policy_version", + } +) +_EXECUTION_STEP_KEYS = frozenset( + { + "id", + "parent_step_id", + "operator", + "object", + "actor", + "goal", + "depends_on", + "postcondition", + "edge_ids", + } +) +_EDGE_KEYS = frozenset( + { + "id", + "source", + "target", + "semantic_step_id", + "actions", + "depends_on", + "resources", + } +) +_ACTION_KEYS = frozenset( + { + "atomic_action_class", + "actor", + "control", + "target_binding", + "motion_policy", + "seed_node_id", + "failure_policy", + } +) +_FAILURE_POLICIES = frozenset({"task_required", "safety_required", "best_effort"}) +_TASK_ALLOCATION_GROUP_KEYS = frozenset({"id", "semantic_step_ids", "arm_constraint"}) +_BINDING_REQUIREMENTS = { + "articulation_goal": frozenset({"object"}), + "coordinated_goal": frozenset({"object"}), + "coordinated_placement_goal": frozenset({"placing_object", "support_object"}), + "current_held_pose": frozenset(), + "handover_goal": frozenset({"object"}), + "handover_staging": frozenset({"object"}), + "joint_state": frozenset({"source"}), + "object": frozenset({"object"}), + "policy_pose": frozenset(), + "pour_goal": frozenset({"object", "reference_object"}), + "semantic_goal": frozenset({"semantic_step"}), + "visual_constraint": frozenset({"camera_uid", "normalized_keypoint"}), +} +_POSTCONDITION_TYPES = frozenset( + { + "both_arms_at_initial_qpos", + "both_grippers_open", + "coordinated_placed", + "grippers_clear_of_object", + "held_by_both_grippers", + "line_member_placed", + "object_axis_near", + "object_axis_offset_near", + "object_held", + "object_held_by_both_grippers", + "object_held_by_gripper", + "object_in_container", + "object_lifted", + "object_not_fallen", + "object_on_object", + "object_supported_by", + "object_position_near", + "object_upright", + "object_xy_near", + "objects_collinear", + "objects_ordered", + "pressed", + "poured", + "articulation_joint_near", + "handover_complete", + "visual_relation", + "stable_unobstructed", + "sum_equals", + "semantic_goal", + "stack_layer_supported", + } +) +_OBJECT_REFERENCE_KEYS = frozenset( + { + "anchor", + "container", + "object", + "object_uid", + "orientation_reference_object", + "placing_object", + "reference", + "reference_object", + "support", + "support_object", + } +) + +# These fields indicate that planning-time or runtime geometry leaked into a +# symbolic program. Integers such as slot and layer indices remain valid. +_GROUNDED_FIELD_NAMES = frozenset( + { + "absolute_position", + "coordinates", + "joint_positions", + "object_target_pose", + "position", + "positions", + "pose", + "qpos", + "release_position", + "staging_position", + "target_pose", + "trajectory", + "waypoints", + "xpos", + } +) + + +def validate_task_agent( + program: Mapping[str, Any], + *, + known_objects: Collection[str] | None = None, +) -> dict[str, Any]: + """Validate and return a detached canonical TaskAgent mapping. + + Defaults are added only for structural fields that have one unambiguous + meaning: ``actor={"mode": "auto"}``, an empty goal, and no dependencies. + Operator-specific semantics are validated by the capability registry. + + Args: + program: Candidate route-free task agent. + known_objects: Optional runtime scene UIDs. When supplied, every object + reference is validated before compilation. + + Returns: + A deep-copied, canonical mapping safe for compilation. + + Raises: + ValueError: If the program violates the TaskAgent contract. + """ + value = _mapping_copy(program, "TaskAgent") + _reject_unknown_keys(value, _TASK_KEYS, "TaskAgent") + _require_schema(value, TASK_AGENT_SCHEMA, "TaskAgent") + _require_nonempty_string(value.get("task"), "TaskAgent.task") + _require_nonempty_string(value.get("goal"), "TaskAgent.goal") + + raw_steps = _sequence(value.get("semantic_steps"), "TaskAgent.semantic_steps") + if not raw_steps: + raise ValueError("TaskAgent.semantic_steps must not be empty.") + + steps: list[dict[str, Any]] = [] + for index, raw_step in enumerate(raw_steps): + context = f"TaskAgent.semantic_steps[{index}]" + step = _mapping_copy(raw_step, context) + _reject_unknown_keys(step, _TASK_STEP_KEYS, context) + _require_nonempty_string(step.get("id"), f"{context}.id") + _require_nonempty_string(step.get("operator"), f"{context}.operator") + + has_object = "object" in step + has_objects = "objects" in step + if has_object == has_objects: + raise ValueError( + f"{context} must contain exactly one of 'object' or 'objects'." + ) + if has_object: + _require_nonempty_string(step["object"], f"{context}.object") + else: + objects = _string_list(step["objects"], f"{context}.objects") + if not objects: + raise ValueError(f"{context}.objects must not be empty.") + _require_unique(objects, f"{context}.objects") + step["objects"] = objects + + step["actor"] = _validate_actor( + step.get("actor", {"mode": "auto"}), + f"{context}.actor", + ) + step["goal"] = _mapping_copy(step.get("goal", {}), f"{context}.goal") + step["depends_on"] = _string_list( + step.get("depends_on", []), + f"{context}.depends_on", + ) + _require_unique(step["depends_on"], f"{context}.depends_on") + steps.append(step) + + step_ids = [step["id"] for step in steps] + _require_unique(step_ids, "TaskAgent semantic step IDs") + dependencies = {step["id"]: step["depends_on"] for step in steps} + _validate_dependency_dag(dependencies, "TaskAgent semantic steps") + value["semantic_steps"] = steps + value["allocation_groups"] = _validate_task_allocation_groups( + value.get("allocation_groups", []), + set(step_ids), + ) + if known_objects is not None: + _validate_known_objects(value, known_objects) + _reject_grounded_values(value) + return value + + +def validate_execution_program(program: Mapping[str, Any]) -> dict[str, Any]: + """Validate and return a detached canonical ExecutionProgram mapping. + + Validation covers both dependency DAGs: semantic-step dependencies and + executable edge dependencies. It also proves node reachability, edge + ownership, resource declarations, and the symbolic action envelope. + + Args: + program: Candidate deterministic execution program. + + Returns: + A deep-copied mapping safe for runtime consumption or hashing. + + Raises: + ValueError: If the program violates the ExecutionProgram contract. + """ + value = _mapping_copy(program, "ExecutionProgram") + _reject_unknown_keys(value, _EXECUTION_KEYS, "ExecutionProgram") + _require_schema(value, EXECUTION_PROGRAM_SCHEMA, "ExecutionProgram") + _require_nonempty_string(value.get("task"), "ExecutionProgram.task") + _require_nonempty_string( + value.get("goal_description"), + "ExecutionProgram.goal_description", + ) + _require_nonempty_string(value.get("start"), "ExecutionProgram.start") + _require_nonempty_string(value.get("goal"), "ExecutionProgram.goal") + if value.get("motion_policy_version") != MOTION_POLICY_VERSION: + raise ValueError( + "ExecutionProgram.motion_policy_version must be " + f"{MOTION_POLICY_VERSION!r}." + ) + + nodes = _validate_nodes(value.get("nodes")) + node_ids = {node["id"] for node in nodes} + if value["start"] not in node_ids or value["goal"] not in node_ids: + raise ValueError("ExecutionProgram start and goal must reference nodes.") + + edges = _validate_edges(value.get("edges"), node_ids) + edge_by_id = {edge["id"]: edge for edge in edges} + _validate_dependency_dag( + {edge_id: edge["depends_on"] for edge_id, edge in edge_by_id.items()}, + "ExecutionProgram edges", + ) + _validate_node_reachability( + start=value["start"], + goal=value["goal"], + node_ids=node_ids, + edges=edges, + ) + + semantic_steps = _validate_execution_steps( + value.get("semantic_steps"), + edge_by_id, + ) + step_ids = {step["id"] for step in semantic_steps} + _validate_allocation_groups(value.get("allocation_groups"), step_ids) + + value["nodes"] = nodes + value["edges"] = edges + value["semantic_steps"] = semantic_steps + value["allocation_groups"] = deepcopy(list(value.get("allocation_groups", []))) + _reject_grounded_values(value) + return value + + +def execution_program_hash(program: Mapping[str, Any]) -> str: + """Return the stable SHA-256 hash of a validated ExecutionProgram.""" + canonical = validate_execution_program(program) + try: + payload = json.dumps( + canonical, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + except (TypeError, ValueError) as exc: + raise ValueError( + "ExecutionProgram must contain JSON-serializable values." + ) from exc + return hashlib.sha256(payload).hexdigest() + + +def _validate_nodes(value: Any) -> list[dict[str, Any]]: + raw_nodes = _sequence(value, "ExecutionProgram.nodes") + if len(raw_nodes) < 2: + raise ValueError("ExecutionProgram.nodes must contain start and goal nodes.") + nodes: list[dict[str, Any]] = [] + for index, raw_node in enumerate(raw_nodes): + context = f"ExecutionProgram.nodes[{index}]" + node = _mapping_copy(raw_node, context) + _reject_unknown_keys(node, frozenset({"id", "semantic"}), context) + _require_nonempty_string(node.get("id"), f"{context}.id") + _require_nonempty_string(node.get("semantic"), f"{context}.semantic") + nodes.append(node) + _require_unique([node["id"] for node in nodes], "ExecutionProgram node IDs") + return nodes + + +def _validate_edges(value: Any, node_ids: set[str]) -> list[dict[str, Any]]: + raw_edges = _sequence(value, "ExecutionProgram.edges") + if not raw_edges: + raise ValueError("ExecutionProgram.edges must not be empty.") + edges: list[dict[str, Any]] = [] + for index, raw_edge in enumerate(raw_edges): + context = f"ExecutionProgram.edges[{index}]" + edge = _mapping_copy(raw_edge, context) + _reject_unknown_keys(edge, _EDGE_KEYS, context) + for key in ("id", "source", "target", "semantic_step_id"): + _require_nonempty_string(edge.get(key), f"{context}.{key}") + if edge["source"] not in node_ids or edge["target"] not in node_ids: + raise ValueError(f"{context} references an unknown graph node.") + + edge["depends_on"] = _string_list( + edge.get("depends_on", []), + f"{context}.depends_on", + ) + edge["resources"] = _string_list( + edge.get("resources", []), + f"{context}.resources", + ) + _require_unique(edge["depends_on"], f"{context}.depends_on") + _require_unique(edge["resources"], f"{context}.resources") + edge["actions"] = _validate_actions(edge.get("actions"), context) + edges.append(edge) + + edge_ids = [edge["id"] for edge in edges] + _require_unique(edge_ids, "ExecutionProgram edge IDs") + known_edges = set(edge_ids) + for edge in edges: + unknown = set(edge["depends_on"]) - known_edges + if unknown: + raise ValueError( + f"Edge {edge['id']!r} depends on unknown edges: {sorted(unknown)}." + ) + return edges + + +def _validate_actions(value: Any, edge_context: str) -> list[dict[str, Any]]: + raw_actions = _sequence(value, f"{edge_context}.actions") + if not raw_actions: + raise ValueError(f"{edge_context}.actions must not be empty.") + actions: list[dict[str, Any]] = [] + for index, raw_action in enumerate(raw_actions): + context = f"{edge_context}.actions[{index}]" + action = _mapping_copy(raw_action, context) + _reject_unknown_keys(action, _ACTION_KEYS, context) + _require_nonempty_string( + action.get("atomic_action_class"), + f"{context}.atomic_action_class", + ) + action["actor"] = _validate_actor(action.get("actor"), f"{context}.actor") + control = _require_nonempty_string(action.get("control"), f"{context}.control") + if control not in _CONTROL_MODES: + raise ValueError( + f"{context}.control must be one of {sorted(_CONTROL_MODES)}." + ) + binding = _mapping_copy( + action.get("target_binding"), + f"{context}.target_binding", + ) + _require_nonempty_string( + binding.get("kind"), + f"{context}.target_binding.kind", + ) + kind = binding["kind"] + required = _BINDING_REQUIREMENTS.get(kind) + if required is None: + raise ValueError(f"{context}.target_binding.kind {kind!r} is unsupported.") + missing = sorted( + key for key in required if not _is_present_binding_value(binding.get(key)) + ) + if missing: + raise ValueError( + f"{context}.target_binding is missing required fields: {missing}." + ) + action["target_binding"] = binding + action["motion_policy"] = validate_motion_policy( + action.get("motion_policy"), + f"{context}.motion_policy", + ) + if "seed_node_id" in action: + _require_nonempty_string( + action.get("seed_node_id"), + f"{context}.seed_node_id", + ) + failure_policy = action.get("failure_policy", "task_required") + if failure_policy not in _FAILURE_POLICIES: + raise ValueError( + f"{context}.failure_policy must be one of " + f"{sorted(_FAILURE_POLICIES)}." + ) + action["failure_policy"] = str(failure_policy) + actions.append(action) + return actions + + +def _validate_execution_steps( + value: Any, + edge_by_id: Mapping[str, Mapping[str, Any]], +) -> list[dict[str, Any]]: + raw_steps = _sequence(value, "ExecutionProgram.semantic_steps") + if not raw_steps: + raise ValueError("ExecutionProgram.semantic_steps must not be empty.") + steps: list[dict[str, Any]] = [] + for index, raw_step in enumerate(raw_steps): + context = f"ExecutionProgram.semantic_steps[{index}]" + step = _mapping_copy(raw_step, context) + _reject_unknown_keys(step, _EXECUTION_STEP_KEYS, context) + for key in ("id", "parent_step_id", "operator", "object"): + _require_nonempty_string(step.get(key), f"{context}.{key}") + step["actor"] = _validate_actor(step.get("actor"), f"{context}.actor") + step["goal"] = _mapping_copy(step.get("goal"), f"{context}.goal") + step["postcondition"] = _mapping_copy( + step.get("postcondition"), + f"{context}.postcondition", + ) + _require_nonempty_string( + step["postcondition"].get("type"), + f"{context}.postcondition.type", + ) + if step["postcondition"]["type"] not in _POSTCONDITION_TYPES: + raise ValueError( + f"{context}.postcondition.type " + f"{step['postcondition']['type']!r} is unsupported." + ) + step["depends_on"] = _string_list( + step.get("depends_on", []), + f"{context}.depends_on", + ) + step["edge_ids"] = _string_list( + step.get("edge_ids"), + f"{context}.edge_ids", + ) + if not step["edge_ids"]: + raise ValueError(f"{context}.edge_ids must not be empty.") + _require_unique(step["depends_on"], f"{context}.depends_on") + _require_unique(step["edge_ids"], f"{context}.edge_ids") + steps.append(step) + + step_ids = [step["id"] for step in steps] + _require_unique(step_ids, "ExecutionProgram semantic step IDs") + _validate_dependency_dag( + {step["id"]: step["depends_on"] for step in steps}, + "ExecutionProgram semantic steps", + ) + + covered_edges: list[str] = [] + for step in steps: + for edge_id in step["edge_ids"]: + edge = edge_by_id.get(edge_id) + if edge is None: + raise ValueError( + f"Semantic step {step['id']!r} owns unknown edge {edge_id!r}." + ) + if edge["semantic_step_id"] != step["id"]: + raise ValueError( + f"Edge {edge_id!r} is assigned to {edge['semantic_step_id']!r}, " + f"not {step['id']!r}." + ) + covered_edges.append(edge_id) + _require_unique(covered_edges, "ExecutionProgram semantic edge ownership") + if set(covered_edges) != set(edge_by_id): + missing = sorted(set(edge_by_id) - set(covered_edges)) + raise ValueError(f"ExecutionProgram has unowned edges: {missing}.") + return steps + + +def _validate_allocation_groups(value: Any, step_ids: set[str]) -> None: + groups = _sequence(value, "ExecutionProgram.allocation_groups") + group_ids: list[str] = [] + for index, raw_group in enumerate(groups): + context = f"ExecutionProgram.allocation_groups[{index}]" + group = _mapping_copy(raw_group, context) + allowed = frozenset( + { + "id", + "semantic_step_ids", + "arm_constraint", + "execution_policy", + "parallel_action_classes", + "workspace_policy", + } + ) + _reject_unknown_keys(group, allowed, context) + group_ids.append(_require_nonempty_string(group.get("id"), f"{context}.id")) + members = _string_list( + group.get("semantic_step_ids"), + f"{context}.semantic_step_ids", + ) + if len(members) < 2: + raise ValueError(f"{context} must contain at least two semantic steps.") + _require_unique(members, f"{context}.semantic_step_ids") + unknown = set(members) - step_ids + if unknown: + raise ValueError(f"{context} references unknown steps: {sorted(unknown)}.") + for key in ("arm_constraint", "execution_policy", "workspace_policy"): + _require_nonempty_string(group.get(key), f"{context}.{key}") + if group["arm_constraint"] != "distinct_arms": + raise ValueError(f"{context}.arm_constraint must be 'distinct_arms'.") + if group["execution_policy"] != "parallel_if_feasible": + raise ValueError( + f"{context}.execution_policy must be 'parallel_if_feasible'." + ) + if group["workspace_policy"] != "shared_target_serial": + raise ValueError( + f"{context}.workspace_policy must be 'shared_target_serial'." + ) + action_classes = _string_list( + group.get("parallel_action_classes"), + f"{context}.parallel_action_classes", + ) + if not action_classes: + raise ValueError(f"{context}.parallel_action_classes must not be empty.") + _require_unique(group_ids, "ExecutionProgram allocation group IDs") + + +def _validate_task_allocation_groups( + value: Any, + step_ids: set[str], +) -> list[dict[str, Any]]: + groups = _sequence(value, "TaskAgent.allocation_groups") + result: list[dict[str, Any]] = [] + ids: list[str] = [] + members_seen: set[str] = set() + for index, raw_group in enumerate(groups): + context = f"TaskAgent.allocation_groups[{index}]" + group = _mapping_copy(raw_group, context) + _reject_unknown_keys(group, _TASK_ALLOCATION_GROUP_KEYS, context) + group_id = _require_nonempty_string(group.get("id"), f"{context}.id") + members = _string_list( + group.get("semantic_step_ids"), + f"{context}.semantic_step_ids", + ) + if len(members) < 2: + raise ValueError(f"{context} must contain at least two semantic steps.") + _require_unique(members, f"{context}.semantic_step_ids") + unknown = set(members) - step_ids + if unknown: + raise ValueError(f"{context} references unknown steps: {sorted(unknown)}.") + overlap = set(members) & members_seen + if overlap: + raise ValueError( + f"TaskAgent allocation groups overlap at steps: {sorted(overlap)}." + ) + constraint = _require_nonempty_string( + group.get("arm_constraint"), + f"{context}.arm_constraint", + ) + if constraint != "distinct_arms": + raise ValueError(f"{context}.arm_constraint must be 'distinct_arms'.") + ids.append(group_id) + members_seen.update(members) + result.append( + { + "id": group_id, + "semantic_step_ids": members, + "arm_constraint": constraint, + } + ) + _require_unique(ids, "TaskAgent allocation group IDs") + return result + + +def _validate_known_objects( + program: Mapping[str, Any], + known_objects: Collection[str], +) -> None: + known = {str(uid) for uid in known_objects} + if not known: + raise ValueError("known_objects must not be empty when supplied.") + allowed_sentinels = {"self", "table", "table_center"} + references: list[tuple[str, str]] = [] + for index, step in enumerate(program["semantic_steps"]): + if "object" in step: + references.append((f"semantic_steps[{index}].object", step["object"])) + for item_index, uid in enumerate(step.get("objects", [])): + references.append((f"semantic_steps[{index}].objects[{item_index}]", uid)) + _collect_object_references( + step["goal"], + f"semantic_steps[{index}].goal", + references, + ) + unknown = [ + f"{path}={uid!r}" + for path, uid in references + if uid not in known and uid not in allowed_sentinels + ] + if unknown: + raise ValueError( + "TaskAgent references objects not present in the scene: " + + ", ".join(unknown) + + "." + ) + + +def _collect_object_references( + value: Any, + path: str, + output: list[tuple[str, str]], +) -> None: + if isinstance(value, Mapping): + for key, child in value.items(): + child_path = f"{path}.{key}" + if key in _OBJECT_REFERENCE_KEYS and isinstance(child, str): + output.append((child_path, child)) + elif ( + key in {"objects", "object_uids", "payloads"} + and isinstance(child, Sequence) + and not isinstance(child, (str, bytes, bytearray)) + ): + for index, item in enumerate(child): + uid = item.get("object") if isinstance(item, Mapping) else item + if isinstance(uid, str): + output.append((f"{child_path}[{index}]", uid)) + else: + _collect_object_references(child, child_path, output) + + +def _is_present_binding_value(value: Any) -> bool: + return value is not None and value != "" and value != [] + + +def _validate_actor(value: Any, context: str) -> dict[str, Any]: + actor = _mapping_copy(value, context) + mode = _require_nonempty_string(actor.get("mode"), f"{context}.mode") + if mode not in _ACTOR_MODES: + raise ValueError(f"{context}.mode must be one of {sorted(_ACTOR_MODES)}.") + if mode == "auto": + _reject_unknown_keys(actor, frozenset({"mode", "allocation_group"}), context) + elif mode == "required": + _reject_unknown_keys( + actor, + frozenset({"mode", "arm", "allocation_group"}), + context, + ) + _require_nonempty_string(actor.get("arm"), f"{context}.arm") + else: + _reject_unknown_keys(actor, frozenset({"mode", "arms"}), context) + arms = _string_list(actor.get("arms"), f"{context}.arms") + if len(arms) < 2: + raise ValueError(f"{context}.arms must contain at least two arms.") + _require_unique(arms, f"{context}.arms") + actor["arms"] = arms + if "allocation_group" in actor: + _require_nonempty_string( + actor["allocation_group"], + f"{context}.allocation_group", + ) + return actor + + +def _validate_dependency_dag( + dependencies: Mapping[str, Sequence[str]], + context: str, +) -> None: + known = set(dependencies) + outgoing: dict[str, list[str]] = {item_id: [] for item_id in known} + indegree = {item_id: 0 for item_id in known} + for item_id, required_ids in dependencies.items(): + unknown = set(required_ids) - known + if unknown: + raise ValueError(f"{context} reference unknown IDs: {sorted(unknown)}.") + if item_id in required_ids: + raise ValueError(f"{context} contain a self-dependency at {item_id!r}.") + for required_id in required_ids: + outgoing[required_id].append(item_id) + indegree[item_id] += 1 + + ready = deque( + sorted(item_id for item_id, degree in indegree.items() if degree == 0) + ) + visited = 0 + while ready: + item_id = ready.popleft() + visited += 1 + for dependent_id in sorted(outgoing[item_id]): + indegree[dependent_id] -= 1 + if indegree[dependent_id] == 0: + ready.append(dependent_id) + if visited != len(known): + cyclic = sorted(item_id for item_id, degree in indegree.items() if degree > 0) + raise ValueError(f"{context} contain a dependency cycle: {cyclic}.") + + +def _validate_node_reachability( + *, + start: str, + goal: str, + node_ids: set[str], + edges: Sequence[Mapping[str, Any]], +) -> None: + outgoing: dict[str, list[str]] = {node_id: [] for node_id in node_ids} + for edge in edges: + outgoing[edge["source"]].append(edge["target"]) + reachable = {start} + ready = deque([start]) + while ready: + node_id = ready.popleft() + for target_id in outgoing[node_id]: + if target_id not in reachable: + reachable.add(target_id) + ready.append(target_id) + if goal not in reachable: + raise ValueError("ExecutionProgram goal is unreachable from start.") + unreachable = sorted(node_ids - reachable) + if unreachable: + raise ValueError(f"ExecutionProgram contains unreachable nodes: {unreachable}.") + + +def _reject_grounded_values(value: Any, path: str = "program") -> None: + if isinstance(value, Mapping): + for key, child in value.items(): + normalized = str(key).strip().lower() + if normalized in _GROUNDED_FIELD_NAMES: + raise ValueError( + f"{path}.{key} is grounded runtime data and is not allowed." + ) + _reject_grounded_values(child, f"{path}.{key}") + return + if isinstance(value, list): + for index, child in enumerate(value): + _reject_grounded_values(child, f"{path}[{index}]") + return + if isinstance(value, float): + raise ValueError( + f"{path} contains a floating-point runtime value; use a named " + "motion policy or symbolic relation instead." + ) + + +def _mapping_copy(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, bytearray)): + raise ValueError(f"{context} must be a list.") + return list(value) + + +def _string_list(value: Any, context: str) -> list[str]: + items = _sequence(value, context) + result: list[str] = [] + for index, item in enumerate(items): + result.append(_require_nonempty_string(item, f"{context}[{index}]")) + return result + + +def _require_schema(value: Mapping[str, Any], expected: str, context: str) -> None: + if value.get("schema_version") != expected: + raise ValueError(f"{context}.schema_version must be {expected!r}.") + + +def _require_nonempty_string(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 + + +def _require_unique(values: Sequence[str], context: str) -> None: + if len(values) != len(set(values)): + raise ValueError(f"{context} must not contain duplicates.") + + +def _reject_unknown_keys( + value: Mapping[str, Any], + allowed: frozenset[str], + context: str, +) -> None: + unknown = sorted(set(value) - allowed) + if unknown: + raise ValueError(f"{context} contains unknown fields: {unknown}.") diff --git a/embodichain/gen_sim/action_engine/domain/task_contracts.py b/embodichain/gen_sim/action_engine/domain/task_contracts.py new file mode 100644 index 000000000..c5bc7ece2 --- /dev/null +++ b/embodichain/gen_sim/action_engine/domain/task_contracts.py @@ -0,0 +1,125 @@ +# ---------------------------------------------------------------------------- +# 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 Engine recipes layered over the Task Engine semantic ontology.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any + +from embodichain.gen_sim.task_engine.ontology import ( + RELATIONS, + TERMINAL_BEHAVIORS, + TRANSPORT_DIRECTIONS, + TaskContract as SemanticTaskContract, + task_success_type, +) +from embodichain.gen_sim.task_engine.ontology import ( + TASK_CONTRACTS as SEMANTIC_TASK_CONTRACTS, +) + +__all__ = [ + "PLACEMENT_RELATIONS", + "RELATIONS", + "TASK_CONTRACTS", + "TERMINAL_BEHAVIORS", + "TRANSPORT_DIRECTIONS", + "TaskContract", + "task_contract", + "task_success_type", + "normalize_placement_relation", +] + +PLACEMENT_RELATIONS = RELATIONS - {"none"} +_SUPPORTED_PLACEMENT_ALIASES = frozenset({"above", "on_top", "on_top_of"}) + + +def normalize_placement_relation(value: Any) -> str: + """Lower task-language relations to physically executable release goals. + + A released object cannot remain freely hovering. Task-language ``above`` + therefore lowers to the supported ``on`` relation for placement operators; + non-placement operators such as pouring retain their distinct ``above`` + semantics. + """ + relation = str(value) + if ( + relation not in PLACEMENT_RELATIONS + and relation not in _SUPPORTED_PLACEMENT_ALIASES + ): + raise ValueError(f"Unsupported placement relation {relation!r}.") + return "on" if relation in _SUPPORTED_PLACEMENT_ALIASES else relation + + +_CORE_ACTIONS: Mapping[str, tuple[str, ...]] = MappingProxyType( + { + "E1": ("PickUp", "MoveHeldObject", "Place"), + "E2": ("PickUp", "MoveHeldObject", "Place"), + "E3": ("Pour",), + "E4": ("PickUp", "MoveHeldObject", "HandOver"), + "E5": ("CoordinatedPickment",), + "E6": ("PullArticulatedPart",), + "E7": ("PushArticulatedPart",), + "E8": ("TurnKnob",), + "E9": ("Press",), + } +) + + +@dataclass(frozen=True, slots=True) +class TaskContract: + """Action-facing view of one Task Engine semantic contract.""" + + task_type: str + semantics: str + core_actions: tuple[str, ...] + applicable_intent_fields: frozenset[str] + source_structure: str + required_affordances: frozenset[str] + success_type: str + scene_affordances: frozenset[str] + + +def _action_contract(value: SemanticTaskContract) -> TaskContract: + return TaskContract( + task_type=value.task_type, + semantics=value.semantics, + core_actions=_CORE_ACTIONS[value.task_type], + applicable_intent_fields=value.applicable_intent_fields, + source_structure=value.source_structure, + required_affordances=value.required_affordances, + success_type=value.success_type, + scene_affordances=value.scene_affordances, + ) + + +TASK_CONTRACTS: Mapping[str, TaskContract] = MappingProxyType( + { + task_type: _action_contract(contract) + for task_type, contract in SEMANTIC_TASK_CONTRACTS.items() + } +) + + +def task_contract(task_type: str) -> TaskContract: + """Return the Action Engine view of one canonical task contract.""" + try: + return TASK_CONTRACTS[str(task_type)] + except KeyError as exc: + raise ValueError(f"Unsupported task type {task_type!r}.") from exc diff --git a/embodichain/gen_sim/action_engine/domain/v2.py b/embodichain/gen_sim/action_engine/domain/v2.py new file mode 100644 index 000000000..9b15fe8ed --- /dev/null +++ b/embodichain/gen_sim/action_engine/domain/v2.py @@ -0,0 +1,1291 @@ +# ---------------------------------------------------------------------------- +# 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 coordinate-free contracts for Action Engine SeedGraph v3.""" + +from __future__ import annotations + +from collections import deque +from collections.abc import Collection, Mapping, Sequence +from copy import deepcopy +import hashlib +import json +import math +from typing import Any + +from embodichain.gen_sim.action_engine.protocol import ( + SCENE_REQUIREMENTS_SCHEMA, + SEED_GRAPH_SCHEMA, + TASK_SPEC_SCHEMA, +) +from .motion import validate_motion_policy + +__all__ = [ + "REASONING_TYPES", + "TASK_LEVELS", + "TASK_TYPES", + "public_task_spec", + "seed_graph_hash", + "validate_public_task_spec", + "validate_scene_requirements", + "validate_seed_graph", + "validate_task_spec", +] + +TASK_LEVELS = frozenset({"L1", "L2", "L3", "L4"}) +TASK_TYPES = frozenset({f"E{index}" for index in range(1, 10)}) +REASONING_TYPES = frozenset( + { + "none", + "memory", + "visual_semantics", + "pattern", + "logic", + "common_sense", + "constraint", + } +) + +_TASK_SPEC_KEYS = frozenset( + { + "schema_version", + "task_id", + "level", + "instruction", + "reasoning_type", + "task_instances", + "success", + "oracle", + "metadata", + } +) +_TASK_INSTANCE_KEYS = frozenset({"id", "task_type", "params", "depends_on", "role"}) +_SCENE_REQUIREMENTS_KEYS = frozenset( + { + "schema_version", + "task_id", + "objects", + "cameras", + "spatial_constraints", + "distractor_count", + "metadata", + } +) +_OBJECT_REQUIREMENT_KEYS = frozenset( + { + "role_id", + "category", + "count", + "affordances", + "initial_state", + "attributes", + } +) +_SEED_GRAPH_KEYS = frozenset( + { + "schema_version", + "task_id", + "instruction", + "level", + "reasoning_type", + "planner_route", + "nodes", + "task_groups", + "success", + "capability_catalog_hash", + "metadata", + } +) +_ACTION_NODE_KEYS = frozenset( + { + "id", + "atomic_action", + "object_uid", + "actor", + "control", + "target_binding", + "depends_on", + "contract", + "task_instance_id", + "task_type", + "role", + "precondition", + "postcondition", + "motion_policy", + "sync_group", + } +) +_TASK_GROUP_KEYS = frozenset( + { + "id", + "task_type", + "role", + "operator", + "object_uid", + "actor", + "goal", + "depends_on", + "parent_task_instance_id", + "node_ids", + "success", + "contract", + } +) +_ACTOR_MODES = frozenset({"auto", "required", "preferred", "coordinated"}) +_NODE_ROLES = frozenset({"primary", "recovery", "cleanup"}) +_GROUP_ROLES = frozenset({"primary", "recovery"}) +_PLANNER_ROUTES = frozenset({"offline", "online", "selected", "fused"}) +_ACTION_CONTRACT_KEYS = frozenset( + { + "version", + "requires", + "effects", + "claims", + "completion", + "failure_policy", + } +) +_TASK_GROUP_CONTRACT_KEYS = frozenset( + { + "entry_requires", + "exit_effects", + "claims", + "entry_node_ids", + "terminal_node_ids", + "completion", + } +) +_STATE_ATOM_KEYS = frozenset({"predicate", "object_uid", "arm"}) +_STATE_PREDICATES = frozenset( + { + "arm_free", + "object_free", + "object_held", + "object_coordinated_held", + "handover_complete", + "arm_clear", + "arm_home", + } +) +_EFFECT_KEYS = frozenset({"op", "atom"}) +_EFFECT_OPERATIONS = frozenset({"add", "delete"}) +_CLAIM_KEYS = frozenset({"resource", "access", "lifetime"}) +_CLAIM_ACCESS = frozenset({"shared_read", "exclusive"}) +_CLAIM_LIFETIMES = frozenset({"action", "until_release"}) +_ACTION_COMPLETION = frozenset({"ordinary", "cleanup", "terminal_barrier"}) +_FAILURE_POLICIES = frozenset({"task_required", "safety_required", "best_effort"}) +_GROUP_COMPLETION = frozenset({"ordinary", "terminal_barrier"}) +_OBJECT_REFERENCE_KEYS = frozenset( + { + "anchor", + "container", + "object", + "object_uid", + "placing_object", + "reference", + "reference_object", + "support", + "support_object", + } +) +_GROUNDED_FIELD_NAMES = frozenset( + { + "absolute_position", + "coordinates", + "grasp_pose", + "joint_positions", + "object_target_pose", + "position", + "positions", + "pose", + "qpos", + "release_position", + "staging_position", + "target_pose", + "trajectory", + "waypoints", + "xpos", + } +) + + +def validate_task_spec(value: Mapping[str, Any]) -> dict[str, Any]: + """Validate one task-first, scene-independent task specification.""" + result = _mapping(value, "TaskSpec") + _keys(result, _TASK_SPEC_KEYS, "TaskSpec") + _schema(result, TASK_SPEC_SCHEMA, "TaskSpec") + _string(result.get("task_id"), "TaskSpec.task_id") + level = _enum(result.get("level"), TASK_LEVELS, "TaskSpec.level") + _string(result.get("instruction"), "TaskSpec.instruction") + reasoning = _enum( + result.get("reasoning_type", "none"), + REASONING_TYPES, + "TaskSpec.reasoning_type", + ) + if level == "L4" and reasoning == "none": + raise ValueError("TaskSpec L4 tasks require a non-'none' reasoning_type.") + if level != "L4" and reasoning != "none": + raise ValueError("Only TaskSpec L4 tasks may declare reasoning_type.") + + instances: list[dict[str, Any]] = [] + for index, item in enumerate( + _sequence(result.get("task_instances"), "TaskSpec.task_instances") + ): + context = f"TaskSpec.task_instances[{index}]" + instance = _mapping(item, context) + _keys(instance, _TASK_INSTANCE_KEYS, context) + _string(instance.get("id"), f"{context}.id") + _enum(instance.get("task_type"), TASK_TYPES, f"{context}.task_type") + instance["params"] = _mapping(instance.get("params", {}), f"{context}.params") + instance["depends_on"] = _strings( + instance.get("depends_on", []), f"{context}.depends_on" + ) + instance["role"] = _enum( + instance.get("role", "primary"), _GROUP_ROLES, f"{context}.role" + ) + instances.append(instance) + if not instances: + raise ValueError("TaskSpec.task_instances must not be empty.") + _unique([item["id"] for item in instances], "TaskSpec task instance IDs") + _dag( + {item["id"]: item["depends_on"] for item in instances}, + "TaskSpec task instances", + ) + _validate_level_shape(level, instances) + result["task_instances"] = instances + result["success"] = _mapping(result.get("success"), "TaskSpec.success") + result["oracle"] = _mapping(result.get("oracle", {}), "TaskSpec.oracle") + result["metadata"] = _mapping(result.get("metadata", {}), "TaskSpec.metadata") + _reject_grounded(result) + _finite(result) + return result + + +def public_task_spec(value: Mapping[str, Any]) -> dict[str, Any]: + """Return the validated TaskSpec view safe to expose to an online agent.""" + if "task_instances" not in value and value.get("level") == "L4": + result = dict(value) + _strip_public_private_metadata(result) + return validate_public_task_spec(result) + result = validate_task_spec(value) + result.pop("oracle", None) + if result["level"] == "L4": + # L4 task instances are the hidden reference plan, not public intent. + result.pop("task_instances", None) + _strip_public_private_metadata(result) + return validate_public_task_spec(result) + + +def _strip_public_private_metadata(value: dict[str, Any]) -> None: + """Remove role/UID bindings that would turn the public view into an oracle.""" + metadata = value.get("metadata") + if not isinstance(metadata, Mapping): + return + private_keys = { + "role_bindings", + "uid_map", + "source_uid_map", + "reference_seed_graph", + "oracle", + } + + def strip(child: Any) -> Any: + if isinstance(child, Mapping): + return { + key: strip(nested) + for key, nested in child.items() + if str(key).lower() not in private_keys + } + if isinstance(child, list): + return [strip(item) for item in child] + if isinstance(child, tuple): + return [strip(item) for item in child] + return deepcopy(child) + + value["metadata"] = strip(metadata) + + +def validate_public_task_spec(value: Mapping[str, Any]) -> dict[str, Any]: + """Validate the oracle-free TaskSpec projection consumed online.""" + result = _mapping(value, "PublicTaskSpec") + allowed = _TASK_SPEC_KEYS - {"oracle"} + _keys(result, allowed, "PublicTaskSpec") + if "oracle" in result: + raise ValueError("PublicTaskSpec must not contain oracle data.") + _schema(result, TASK_SPEC_SCHEMA, "PublicTaskSpec") + _string(result.get("task_id"), "PublicTaskSpec.task_id") + level = _enum(result.get("level"), TASK_LEVELS, "PublicTaskSpec.level") + _string(result.get("instruction"), "PublicTaskSpec.instruction") + reasoning = _enum( + result.get("reasoning_type", "none"), + REASONING_TYPES, + "PublicTaskSpec.reasoning_type", + ) + if (level == "L4") != (reasoning != "none"): + raise ValueError( + "PublicTaskSpec reasoning_type must be non-'none' exactly for L4." + ) + if level == "L4": + if "task_instances" in result: + raise ValueError("Public L4 TaskSpec must hide reference task instances.") + else: + # Reuse the complete structural validator for explicit L1-L3 tasks. + normalized = validate_task_spec({**result, "oracle": {}}) + normalized.pop("oracle", None) + return normalized + result["success"] = _mapping(result.get("success"), "PublicTaskSpec.success") + result["metadata"] = _mapping(result.get("metadata", {}), "PublicTaskSpec.metadata") + _reject_grounded(result) + _finite(result) + return result + + +def validate_scene_requirements(value: Mapping[str, Any]) -> dict[str, Any]: + """Validate the structured hand-off contract consumed by a Scene Engine.""" + result = _mapping(value, "SceneRequirements") + _keys(result, _SCENE_REQUIREMENTS_KEYS, "SceneRequirements") + _schema(result, SCENE_REQUIREMENTS_SCHEMA, "SceneRequirements") + _string(result.get("task_id"), "SceneRequirements.task_id") + objects: list[dict[str, Any]] = [] + for index, item in enumerate( + _sequence(result.get("objects"), "SceneRequirements.objects") + ): + context = f"SceneRequirements.objects[{index}]" + requirement = _mapping(item, context) + _keys(requirement, _OBJECT_REQUIREMENT_KEYS, context) + _string(requirement.get("role_id"), f"{context}.role_id") + _string(requirement.get("category"), f"{context}.category") + count = requirement.get("count", 1) + if not isinstance(count, int) or isinstance(count, bool) or count < 1: + raise ValueError(f"{context}.count must be a positive integer.") + requirement["count"] = count + requirement["affordances"] = _strings( + requirement.get("affordances", []), f"{context}.affordances" + ) + requirement["initial_state"] = _mapping( + requirement.get("initial_state", {}), f"{context}.initial_state" + ) + requirement["attributes"] = _mapping( + requirement.get("attributes", {}), f"{context}.attributes" + ) + objects.append(requirement) + if not objects: + raise ValueError("SceneRequirements.objects must not be empty.") + _unique([item["role_id"] for item in objects], "SceneRequirements role IDs") + result["objects"] = objects + result["cameras"] = [ + _mapping(item, f"SceneRequirements.cameras[{index}]") + for index, item in enumerate( + _sequence(result.get("cameras", []), "SceneRequirements.cameras") + ) + ] + result["spatial_constraints"] = [ + _mapping(item, f"SceneRequirements.spatial_constraints[{index}]") + for index, item in enumerate( + _sequence( + result.get("spatial_constraints", []), + "SceneRequirements.spatial_constraints", + ) + ) + ] + distractors = result.get("distractor_count", 0) + if ( + not isinstance(distractors, int) + or isinstance(distractors, bool) + or distractors < 0 + ): + raise ValueError("SceneRequirements.distractor_count must be non-negative.") + result["distractor_count"] = distractors + result["metadata"] = _mapping( + result.get("metadata", {}), "SceneRequirements.metadata" + ) + _finite(result) + return result + + +def validate_seed_graph( + value: Mapping[str, Any], + *, + known_objects: Collection[str] | None = None, + known_actions: Collection[str] | None = None, + executable_actions: Collection[str] | None = None, + require_executable: bool = False, +) -> dict[str, Any]: + """Validate a direct, coordinate-free AtomicAction DAG.""" + result = _mapping(value, "SeedGraph") + _keys(result, _SEED_GRAPH_KEYS, "SeedGraph") + _schema(result, SEED_GRAPH_SCHEMA, "SeedGraph") + _string(result.get("task_id"), "SeedGraph.task_id") + _string(result.get("instruction"), "SeedGraph.instruction") + level = _enum(result.get("level"), TASK_LEVELS, "SeedGraph.level") + reasoning = _enum( + result.get("reasoning_type", "none"), + REASONING_TYPES, + "SeedGraph.reasoning_type", + ) + if (level == "L4") != (reasoning != "none"): + raise ValueError("SeedGraph reasoning_type must be non-'none' exactly for L4.") + result["planner_route"] = _enum( + result.get("planner_route"), _PLANNER_ROUTES, "SeedGraph.planner_route" + ) + _string(result.get("capability_catalog_hash"), "SeedGraph.capability_catalog_hash") + + nodes: list[dict[str, Any]] = [] + for index, item in enumerate(_sequence(result.get("nodes"), "SeedGraph.nodes")): + context = f"SeedGraph.nodes[{index}]" + node = _mapping(item, context) + _keys(node, _ACTION_NODE_KEYS, context) + _string(node.get("id"), f"{context}.id") + action = _string(node.get("atomic_action"), f"{context}.atomic_action") + if known_actions is not None and action not in set(known_actions): + raise ValueError(f"{context} references unknown AtomicAction {action!r}.") + if ( + require_executable + and executable_actions is not None + and action not in set(executable_actions) + ): + raise ValueError( + f"AtomicAction {action!r} is planning-only and cannot be executed." + ) + object_uid = _string(node.get("object_uid"), f"{context}.object_uid") + if known_objects is not None and object_uid not in set(known_objects): + raise ValueError(f"{context} references unknown object {object_uid!r}.") + node["actor"] = _actor(node.get("actor", {"mode": "auto"}), f"{context}.actor") + node["control"] = _string(node.get("control", "arm"), f"{context}.control") + binding = _mapping(node.get("target_binding"), f"{context}.target_binding") + _string(binding.get("kind"), f"{context}.target_binding.kind") + node["target_binding"] = binding + node["depends_on"] = _strings( + node.get("depends_on", []), f"{context}.depends_on" + ) + node["contract"] = _action_contract(node.get("contract"), f"{context}.contract") + node["task_instance_id"] = _string( + node.get("task_instance_id"), f"{context}.task_instance_id" + ) + node["task_type"] = _enum( + node.get("task_type"), TASK_TYPES, f"{context}.task_type" + ) + node["role"] = _enum( + node.get("role", "primary"), _NODE_ROLES, f"{context}.role" + ) + node["precondition"] = _mapping( + node.get("precondition", {}), f"{context}.precondition" + ) + node["postcondition"] = _mapping( + node.get("postcondition", {}), f"{context}.postcondition" + ) + node["motion_policy"] = validate_motion_policy( + node.get("motion_policy"), f"{context}.motion_policy" + ) + if "sync_group" in node: + node["sync_group"] = _string(node["sync_group"], f"{context}.sync_group") + nodes.append(node) + if not nodes: + raise ValueError("SeedGraph.nodes must not be empty.") + node_ids = [node["id"] for node in nodes] + _unique(node_ids, "SeedGraph node IDs") + _dag({node["id"]: node["depends_on"] for node in nodes}, "SeedGraph nodes") + + groups: list[dict[str, Any]] = [] + for index, item in enumerate( + _sequence(result.get("task_groups"), "SeedGraph.task_groups") + ): + context = f"SeedGraph.task_groups[{index}]" + group = _mapping(item, context) + _keys(group, _TASK_GROUP_KEYS, context) + group["id"] = _string(group.get("id"), f"{context}.id") + group["task_type"] = _enum( + group.get("task_type"), TASK_TYPES, f"{context}.task_type" + ) + group["role"] = _enum( + group.get("role", "primary"), _GROUP_ROLES, f"{context}.role" + ) + group["operator"] = _string(group.get("operator"), f"{context}.operator") + group["object_uid"] = _string(group.get("object_uid"), f"{context}.object_uid") + group["actor"] = _actor( + group.get("actor", {"mode": "auto"}), f"{context}.actor" + ) + group["goal"] = _mapping(group.get("goal", {}), f"{context}.goal") + group["depends_on"] = _strings( + group.get("depends_on", []), f"{context}.depends_on" + ) + if "parent_task_instance_id" in group: + group["parent_task_instance_id"] = _string( + group["parent_task_instance_id"], + f"{context}.parent_task_instance_id", + ) + group["node_ids"] = _strings(group.get("node_ids"), f"{context}.node_ids") + if not group["node_ids"]: + raise ValueError(f"{context}.node_ids must not be empty.") + group["success"] = _mapping(group.get("success"), f"{context}.success") + group["contract"] = _task_group_contract( + group.get("contract"), f"{context}.contract" + ) + groups.append(group) + if not groups: + raise ValueError("SeedGraph.task_groups must not be empty.") + _validate_groups(nodes, groups) + _validate_group_contract_topology(nodes, groups) + _validate_cleanup_barriers(nodes, groups) + _validate_task_group_semantics(nodes, groups) + _validate_ownership_transitions(nodes, groups) + _validate_resource_conflicts(nodes, groups, result.get("metadata", {})) + _dag( + {group["id"]: group["depends_on"] for group in groups}, + "SeedGraph task groups", + ) + _validate_group_dependency_alignment(nodes, groups) + result["nodes"] = nodes + result["task_groups"] = groups + result["success"] = _mapping(result.get("success"), "SeedGraph.success") + result["metadata"] = _mapping(result.get("metadata", {}), "SeedGraph.metadata") + _reject_grounded(result) + _finite(result) + if known_objects is not None: + _known_object_references(result, set(known_objects)) + return result + + +def seed_graph_hash(value: Mapping[str, Any]) -> str: + """Return a stable SHA-256 hash for one validated SeedGraph.""" + canonical = validate_seed_graph(value) + payload = json.dumps( + canonical, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def _validate_level_shape(level: str, instances: Sequence[Mapping[str, Any]]) -> None: + primary = [item for item in instances if item["role"] == "primary"] + types = {str(item["task_type"]) for item in primary} + if level == "L1" and len(primary) != 1: + raise ValueError("L1 requires exactly one primary task instance.") + if level == "L2" and (len(primary) < 2 or len(types) != 1): + raise ValueError("L2 requires at least two primary instances of one E type.") + if level == "L3" and (len(primary) < 2 or len(types) < 2): + raise ValueError( + "L3 requires at least two primary instances of different E types." + ) + + +def _validate_groups( + nodes: Sequence[Mapping[str, Any]], groups: Sequence[Mapping[str, Any]] +) -> None: + node_by_id = {str(node["id"]): node for node in nodes} + _unique([str(group["id"]) for group in groups], "SeedGraph task group IDs") + memberships: dict[str, str] = {} + for group in groups: + group_id = str(group["id"]) + for node_id in group["node_ids"]: + if node_id not in node_by_id: + raise ValueError( + f"SeedGraph task group {group_id!r} references unknown node {node_id!r}." + ) + if node_id in memberships: + raise ValueError( + f"SeedGraph node {node_id!r} belongs to multiple task groups." + ) + node = node_by_id[node_id] + if node["task_instance_id"] != group_id: + raise ValueError( + f"SeedGraph node {node_id!r} task_instance_id does not match {group_id!r}." + ) + if node["task_type"] != group["task_type"]: + raise ValueError( + f"SeedGraph node {node_id!r} task_type does not match its group." + ) + memberships[node_id] = group_id + missing = sorted(set(node_by_id) - set(memberships)) + if missing: + raise ValueError( + f"SeedGraph nodes are missing task group membership: {missing}." + ) + + +def _validate_task_group_semantics( + nodes: Sequence[Mapping[str, Any]], + groups: Sequence[Mapping[str, Any]], +) -> None: + node_by_id = {str(node["id"]): node for node in nodes} + required_actions = { + "E1": set(), + "E2": {"MoveHeldObject", "Place"}, + "E3": {"Pour"}, + "E4": {"HandOver"}, + "E5": set(), + "E6": {"PullArticulatedPart"}, + "E7": {"PushArticulatedPart"}, + "E8": {"TurnKnob"}, + "E9": {"Press"}, + } + for group in groups: + task_type = str(group["task_type"]) + group_nodes = [node_by_id[node_id] for node_id in group["node_ids"]] + actions = {str(node["atomic_action"]) for node in group_nodes} + missing = required_actions[task_type] - actions + if task_type == "E1": + if not actions.intersection({"MoveHeldObject", "Place"}): + missing = {"MoveHeldObject|Place"} + elif "PickUp" not in actions: + first = group_nodes[0] + precondition = first.get("precondition", {}) + if precondition.get("type") != "object_held": + missing = {"PickUp|object_held precondition"} + if task_type == "E2" and "PickUp" not in actions: + first = group_nodes[0] + precondition = first.get("precondition", {}) + if precondition.get("type") != "object_held": + missing = {"PickUp|object_held precondition"} + # Recovery may explicitly preserve a verified downstream hold. Ordinary + # E2 groups always complete their supported world state with Place. + if ( + task_type == "E2" + and "Place" not in actions + and group.get("goal", {}).get("terminal_behavior") == "hold" + and "MoveHeldObject" in actions + and group.get("role") == "recovery" + ): + missing.discard("Place") + if task_type == "E5" and not actions.intersection( + {"CoordinatedPickment", "CoordinatedPlacement"} + ): + missing = {"CoordinatedPickment|CoordinatedPlacement"} + if missing: + raise ValueError( + f"SeedGraph TaskGroup {group['id']!r} is missing {task_type} " + f"core actions: {sorted(missing)}." + ) + + +def _validate_ownership_transitions( + nodes: Sequence[Mapping[str, Any]], + groups: Sequence[Mapping[str, Any]], +) -> None: + """Check release/reacquire and explicit single-arm hold transitions. + + An ordinary E2 -> E4 transition persists the supported upright state, ends + the predecessor resource lease, and lets E4 acquire a fresh transfer grasp. + E4 -> E1 keeps receiver ownership because the exchanged object is not yet + supported. Recovery groups may preserve an explicitly requested hold. + """ + node_by_id = {str(node["id"]): node for node in nodes} + group_by_id = {str(group["id"]): group for group in groups} + nodes_by_group = { + group_id: [node_by_id[node_id] for node_id in group["node_ids"]] + for group_id, group in group_by_id.items() + } + + def direct_predecessor( + group: Mapping[str, Any], task_type: str + ) -> Mapping[str, Any] | None: + for dependency in group.get("depends_on", []): + candidate = group_by_id.get(str(dependency)) + if candidate is not None and candidate.get("task_type") == task_type: + return candidate + return None + + def held_arm(node: Mapping[str, Any]) -> str | None: + precondition = node.get("precondition", {}) + if ( + isinstance(precondition, Mapping) + and precondition.get("type") == "object_held" + ): + arm = str(precondition.get("arm", "")) + if arm in {"left_arm", "right_arm"}: + return arm + actor = node.get("actor", {}) + if isinstance(actor, Mapping) and actor.get("mode") == "required": + arm = str(actor.get("arm", "")) + if arm in {"left_arm", "right_arm"}: + return arm + return None + + for group_id, group in group_by_id.items(): + task_type = str(group.get("task_type")) + group_nodes = nodes_by_group[group_id] + actions = [str(node.get("atomic_action")) for node in group_nodes] + object_uid = str(group.get("object_uid")) + + if task_type == "E2": + handover = next( + ( + candidate + for candidate in groups + if candidate.get("task_type") == "E4" + and group_id + in {str(item) for item in candidate.get("depends_on", [])} + and str(candidate.get("object_uid")) == object_uid + ), + None, + ) + if handover is None: + continue + if ( + group.get("goal", {}).get("terminal_behavior") == "hold" + and group.get("role") != "recovery" + ): + raise ValueError( + f"SeedGraph E2 group {group_id!r} may not preserve a holder " + "across an ordinary E2->E4 TaskGroup boundary." + ) + if group.get("role") != "recovery" and "Place" not in actions: + raise ValueError( + f"SeedGraph E2 group {group_id!r} must release its supported " + "object before E4 reacquires it." + ) + + if task_type == "E4": + predecessor = direct_predecessor(group, "E2") + if ( + predecessor is not None + and str(predecessor.get("object_uid")) == object_uid + ): + predecessor_nodes = nodes_by_group[str(predecessor["id"])] + preserves_hold = ( + predecessor.get("role") == "recovery" + and predecessor.get("goal", {}).get("terminal_behavior") == "hold" + ) + if preserves_hold: + if "PickUp" in actions or not group_nodes: + raise ValueError( + f"SeedGraph E4 group {group_id!r} must consume the " + "recovery-held object without PickUp." + ) + first = group_nodes[0] + holder_arm = next( + ( + held_arm(node) + for node in reversed(predecessor_nodes) + if held_arm(node) is not None + ), + None, + ) + if ( + str(first.get("atomic_action")) != "MoveHeldObject" + or held_arm(first) is None + or held_arm(first) != holder_arm + ): + raise ValueError( + f"SeedGraph E2->E4 recovery holder mismatch for object " + f"{object_uid!r}." + ) + else: + predecessor_actions = { + str(node.get("atomic_action")) for node in predecessor_nodes + } + if "Place" not in predecessor_actions: + raise ValueError( + f"SeedGraph E2 predecessor {predecessor['id']!r} must " + "release its object before E4." + ) + if ( + not group_nodes + or str(group_nodes[0].get("atomic_action")) != "PickUp" + ): + raise ValueError( + f"SeedGraph E4 group {group_id!r} must reacquire the " + "supported E2 object with PickUp." + ) + + if task_type == "E1": + predecessor = direct_predecessor(group, "E4") + if ( + predecessor is not None + and str(predecessor.get("object_uid")) == object_uid + ): + if "PickUp" in actions or not group_nodes: + raise ValueError( + f"SeedGraph E1 group {group_id!r} must preserve the E4 receiver hold " + "without PickUp." + ) + first = group_nodes[0] + if ( + str(first.get("atomic_action")) != "MoveHeldObject" + or held_arm(first) is None + ): + raise ValueError( + f"SeedGraph E1 group {group_id!r} must start with MoveHeldObject " + "from the receiver hold." + ) + + +def _validate_group_dependency_alignment( + nodes: Sequence[Mapping[str, Any]], + groups: Sequence[Mapping[str, Any]], +) -> None: + group_by_node = { + str(node_id): str(group["id"]) + for group in groups + for node_id in group["node_ids"] + } + group_dependencies = { + str(group["id"]): set(str(parent) for parent in group["depends_on"]) + for group in groups + } + + def group_reaches(child: str, parent: str) -> bool: + pending = list(group_dependencies[child]) + visited = set() + while pending: + current = pending.pop() + if current == parent: + return True + if current not in visited: + visited.add(current) + pending.extend(group_dependencies[current]) + return False + + for node in nodes: + child_group = group_by_node[str(node["id"])] + for dependency in node["depends_on"]: + parent_group = group_by_node[str(dependency)] + if parent_group != child_group and not group_reaches( + child_group, parent_group + ): + raise ValueError( + f"SeedGraph node {node['id']!r} depends on TaskGroup " + f"{parent_group!r}, but TaskGroup {child_group!r} does not." + ) + + +def _validate_resource_conflicts( + nodes: Sequence[Mapping[str, Any]], + groups: Sequence[Mapping[str, Any]], + metadata: Any, +) -> None: + dependencies = { + str(node["id"]): set(str(item) for item in node["depends_on"]) for node in nodes + } + + def reaches(child: str, parent: str) -> bool: + pending = list(dependencies[child]) + visited = set() + while pending: + current = pending.pop() + if current == parent: + return True + if current not in visited: + visited.add(current) + pending.extend(dependencies[current]) + return False + + group_by_node = { + str(node_id): str(group["id"]) + for group in groups + for node_id in group["node_ids"] + } + distinct_arm_pairs = _distinct_arm_pairs(metadata) + for index, first in enumerate(nodes): + for second in nodes[index + 1 :]: + first_id = str(first["id"]) + second_id = str(second["id"]) + if reaches(first_id, second_id) or reaches(second_id, first_id): + continue + if ( + first.get("sync_group") == second.get("sync_group") + and first.get("sync_group") is not None + ): + continue + first_claims = { + str(claim["resource"]): str(claim["access"]) + for claim in first["contract"]["claims"] + } + second_claims = { + str(claim["resource"]): str(claim["access"]) + for claim in second["contract"]["claims"] + } + conflicts = sorted( + resource + for resource in set(first_claims) & set(second_claims) + if "exclusive" in {first_claims[resource], second_claims[resource]} + ) + if ( + frozenset({group_by_node[first_id], group_by_node[second_id]}) + in distinct_arm_pairs + ): + conflicts = [item for item in conflicts if item != "arm:auto"] + if conflicts: + raise ValueError( + f"SeedGraph concurrent nodes {first_id!r} and {second_id!r} " + f"have resource conflicts: {conflicts}." + ) + + +def _distinct_arm_pairs(value: Any) -> set[frozenset[str]]: + if not isinstance(value, Mapping): + return set() + groups = value.get("legacy_allocation_groups", value.get("allocation_groups", ())) + if not isinstance(groups, Sequence) or isinstance(groups, (str, bytes, bytearray)): + return set() + result: set[frozenset[str]] = set() + for group in groups: + if ( + not isinstance(group, Mapping) + or group.get("arm_constraint") != "distinct_arms" + ): + continue + members = group.get("semantic_step_ids", group.get("task_instance_ids", ())) + if not isinstance(members, Sequence) or isinstance( + members, (str, bytes, bytearray) + ): + continue + member_ids = [str(item) for item in members] + for index, first in enumerate(member_ids): + for second in member_ids[index + 1 :]: + result.add(frozenset({first, second})) + return result + + +def _action_contract(value: Any, context: str) -> dict[str, Any]: + contract = _mapping(value, context) + _keys(contract, _ACTION_CONTRACT_KEYS, context) + if set(contract) != _ACTION_CONTRACT_KEYS: + missing = sorted(_ACTION_CONTRACT_KEYS - set(contract)) + raise ValueError(f"{context} is missing required fields: {missing}.") + if contract["version"] != "action_contract_v2": + raise ValueError(f"{context}.version must be 'action_contract_v2'.") + contract["requires"] = [ + _state_atom(item, f"{context}.requires[{index}]") + for index, item in enumerate( + _sequence(contract["requires"], f"{context}.requires") + ) + ] + contract["effects"] = [ + _state_effect(item, f"{context}.effects[{index}]") + for index, item in enumerate( + _sequence(contract["effects"], f"{context}.effects") + ) + ] + contract["claims"] = [ + _resource_claim(item, f"{context}.claims[{index}]") + for index, item in enumerate(_sequence(contract["claims"], f"{context}.claims")) + ] + _unique( + [str(item["resource"]) for item in contract["claims"]], + f"{context}.claims resources", + ) + contract["completion"] = _enum( + contract["completion"], _ACTION_COMPLETION, f"{context}.completion" + ) + contract["failure_policy"] = _enum( + contract["failure_policy"], + _FAILURE_POLICIES, + f"{context}.failure_policy", + ) + return contract + + +def _task_group_contract(value: Any, context: str) -> dict[str, Any]: + contract = _mapping(value, context) + _keys(contract, _TASK_GROUP_CONTRACT_KEYS, context) + if set(contract) != _TASK_GROUP_CONTRACT_KEYS: + missing = sorted(_TASK_GROUP_CONTRACT_KEYS - set(contract)) + raise ValueError(f"{context} is missing required fields: {missing}.") + contract["entry_requires"] = [ + _state_atom(item, f"{context}.entry_requires[{index}]") + for index, item in enumerate( + _sequence(contract["entry_requires"], f"{context}.entry_requires") + ) + ] + contract["exit_effects"] = [ + _state_effect(item, f"{context}.exit_effects[{index}]") + for index, item in enumerate( + _sequence(contract["exit_effects"], f"{context}.exit_effects") + ) + ] + contract["claims"] = [ + _resource_claim(item, f"{context}.claims[{index}]") + for index, item in enumerate(_sequence(contract["claims"], f"{context}.claims")) + ] + _unique( + [str(item["resource"]) for item in contract["claims"]], + f"{context}.claims resources", + ) + contract["entry_node_ids"] = _strings( + contract["entry_node_ids"], f"{context}.entry_node_ids" + ) + contract["terminal_node_ids"] = _strings( + contract["terminal_node_ids"], f"{context}.terminal_node_ids" + ) + if not contract["entry_node_ids"] or not contract["terminal_node_ids"]: + raise ValueError(f"{context} requires entry and terminal node IDs.") + contract["completion"] = _enum( + contract["completion"], _GROUP_COMPLETION, f"{context}.completion" + ) + return contract + + +def _state_atom(value: Any, context: str) -> dict[str, str]: + atom = _mapping(value, context) + _keys(atom, _STATE_ATOM_KEYS, context) + predicate = _enum(atom.get("predicate"), _STATE_PREDICATES, f"{context}.predicate") + required = { + "arm_free": {"arm"}, + "object_free": {"object_uid"}, + "object_held": {"object_uid", "arm"}, + "object_coordinated_held": {"object_uid"}, + "handover_complete": {"object_uid"}, + "arm_clear": {"arm"}, + "arm_home": {"arm"}, + }[predicate] + present = set(atom) - {"predicate"} + if present != required: + raise ValueError( + f"{context} predicate {predicate!r} requires exactly {sorted(required)}." + ) + for field in required: + atom[field] = _string(atom.get(field), f"{context}.{field}") + return atom + + +def _state_effect(value: Any, context: str) -> dict[str, Any]: + effect = _mapping(value, context) + _keys(effect, _EFFECT_KEYS, context) + if set(effect) != _EFFECT_KEYS: + raise ValueError(f"{context} requires op and atom.") + effect["op"] = _enum(effect["op"], _EFFECT_OPERATIONS, f"{context}.op") + effect["atom"] = _state_atom(effect["atom"], f"{context}.atom") + return effect + + +def _resource_claim(value: Any, context: str) -> dict[str, str]: + claim = _mapping(value, context) + _keys(claim, _CLAIM_KEYS, context) + if set(claim) != _CLAIM_KEYS: + raise ValueError(f"{context} requires resource, access, and lifetime.") + claim["resource"] = _string(claim["resource"], f"{context}.resource") + claim["access"] = _enum(claim["access"], _CLAIM_ACCESS, f"{context}.access") + claim["lifetime"] = _enum( + claim["lifetime"], _CLAIM_LIFETIMES, f"{context}.lifetime" + ) + return claim + + +def _validate_group_contract_topology( + nodes: Sequence[Mapping[str, Any]], groups: Sequence[Mapping[str, Any]] +) -> None: + node_by_id = {str(node["id"]): node for node in nodes} + children: dict[str, set[str]] = {node_id: set() for node_id in node_by_id} + for node in nodes: + for dependency in node["depends_on"]: + children[str(dependency)].add(str(node["id"])) + for group in groups: + group_id = str(group["id"]) + node_ids = {str(item) for item in group["node_ids"]} + expected_entries = { + node_id + for node_id in node_ids + if not any( + str(parent) in node_ids for parent in node_by_id[node_id]["depends_on"] + ) + } + expected_terminals = { + node_id for node_id in node_ids if not (children[node_id] & node_ids) + } + contract = group["contract"] + if set(contract["entry_node_ids"]) != expected_entries: + raise ValueError( + f"SeedGraph TaskGroup {group_id!r} contract entry_node_ids do not " + "match its internal topology." + ) + if set(contract["terminal_node_ids"]) != expected_terminals: + raise ValueError( + f"SeedGraph TaskGroup {group_id!r} contract terminal_node_ids do not " + "match its internal topology." + ) + if contract["completion"] == "terminal_barrier" and not all( + node_by_id[node_id]["contract"]["completion"] == "terminal_barrier" + for node_id in expected_terminals + ): + raise ValueError( + f"SeedGraph TaskGroup {group_id!r} terminal barrier must end in " + "terminal_barrier AtomicActions." + ) + + +def _validate_cleanup_barriers( + nodes: Sequence[Mapping[str, Any]], groups: Sequence[Mapping[str, Any]] +) -> None: + node_by_id = {str(node["id"]): node for node in nodes} + for group in groups: + group_id = str(group["id"]) + group_nodes = [node_by_id[str(node_id)] for node_id in group["node_ids"]] + cleanup = [ + node for node in group_nodes if node["contract"]["completion"] == "cleanup" + ] + has_handover = any(node["atomic_action"] == "HandOver" for node in group_nodes) + recovery_successor = any( + candidate.get("role") == "recovery" + and candidate.get("parent_task_instance_id") == group_id + and group_id in candidate.get("depends_on", ()) + for candidate in groups + ) + if has_handover and not cleanup and recovery_successor: + continue + if has_handover and not cleanup: + raise ValueError( + f"SeedGraph HandOver TaskGroup {group_id!r} is missing retreat cleanup." + ) + if not cleanup and not has_handover: + continue + terminal_ids = group["contract"]["terminal_node_ids"] + if group["contract"]["completion"] != "terminal_barrier" or not all( + node_by_id[node_id]["contract"]["completion"] == "terminal_barrier" + for node_id in terminal_ids + ): + raise ValueError( + f"SeedGraph TaskGroup {group_id!r} cleanup must end at a home " + "terminal barrier." + ) + + +def _actor(value: Any, context: str) -> dict[str, Any]: + actor = _mapping(value, context) + mode = _enum(actor.get("mode"), _ACTOR_MODES, f"{context}.mode") + allowed = {"mode"} + if mode in {"required", "preferred"}: + allowed.add("arm") + _string(actor.get("arm"), f"{context}.arm") + elif mode == "coordinated": + allowed.add("arms") + arms = _strings(actor.get("arms"), f"{context}.arms") + if len(arms) < 2: + raise ValueError(f"{context}.arms must contain at least two arms.") + actor["arms"] = arms + _keys(actor, frozenset(allowed), context) + return actor + + +def _known_object_references( + value: Any, known: set[str], path: str = "SeedGraph" +) -> None: + if isinstance(value, Mapping): + for key, child in value.items(): + child_path = f"{path}.{key}" + if key in _OBJECT_REFERENCE_KEYS and isinstance(child, str): + if child not in known and child not in {"table_center", "world"}: + raise ValueError( + f"{child_path} references unknown object {child!r}." + ) + _known_object_references(child, known, child_path) + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + for index, child in enumerate(value): + _known_object_references(child, known, f"{path}[{index}]") + + +def _reject_grounded(value: Any, path: str = "document") -> None: + if isinstance(value, Mapping): + for key, child in value.items(): + if str(key).lower() in _GROUNDED_FIELD_NAMES: + raise ValueError(f"{path}.{key} contains grounded motion data.") + _reject_grounded(child, f"{path}.{key}") + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + for index, child in enumerate(value): + _reject_grounded(child, f"{path}[{index}]") + + +def _finite(value: Any, path: str = "document") -> None: + if isinstance(value, Mapping): + for key, child in value.items(): + _finite(child, f"{path}.{key}") + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + for index, child in enumerate(value): + _finite(child, f"{path}[{index}]") + elif isinstance(value, float) and not math.isfinite(value): + raise ValueError(f"{path} must be finite.") + + +def _dag(dependencies: Mapping[str, Sequence[str]], context: str) -> None: + known = set(dependencies) + outgoing = {item_id: [] for item_id in known} + indegree = {item_id: 0 for item_id in known} + for item_id, required in dependencies.items(): + unknown = set(required) - known + if unknown: + raise ValueError(f"{context} reference unknown IDs: {sorted(unknown)}.") + if item_id in required: + raise ValueError(f"{context} contain a self-dependency at {item_id!r}.") + for parent in required: + outgoing[parent].append(item_id) + indegree[item_id] += 1 + ready = deque( + sorted(item_id for item_id, degree in indegree.items() if degree == 0) + ) + visited = 0 + while ready: + item_id = ready.popleft() + visited += 1 + for child in sorted(outgoing[item_id]): + indegree[child] -= 1 + if indegree[child] == 0: + ready.append(child) + if visited != len(known): + cyclic = sorted(item_id for item_id, degree in indegree.items() if degree) + raise ValueError(f"{context} contain a dependency cycle: {cyclic}.") + + +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, bytearray)): + raise TypeError(f"{context} must be a list.") + return list(value) + + +def _strings(value: Any, context: str) -> list[str]: + result = [ + _string(item, f"{context}[{index}]") + for index, item in enumerate(_sequence(value, context)) + ] + _unique(result, context) + return result + + +def _string(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() + + +def _enum(value: Any, allowed: Collection[str], context: str) -> str: + result = _string(value, context) + if result not in allowed: + raise ValueError(f"{context} must be one of {sorted(allowed)}.") + return result + + +def _unique(values: Sequence[str], context: str) -> None: + if len(values) != len(set(values)): + raise ValueError(f"{context} must be unique.") + + +def _keys(value: Mapping[str, Any], allowed: frozenset[str], context: str) -> None: + unknown = sorted(set(value) - allowed) + if unknown: + raise ValueError(f"{context} contains unsupported fields: {unknown}.") + + +def _schema(value: Mapping[str, Any], expected: str, context: str) -> None: + if value.get("schema_version") != expected: + raise ValueError(f"{context}.schema_version must be {expected!r}.") diff --git a/embodichain/gen_sim/action_engine/domain/visual_contracts.py b/embodichain/gen_sim/action_engine/domain/visual_contracts.py new file mode 100644 index 000000000..78ca4a0f1 --- /dev/null +++ b/embodichain/gen_sim/action_engine/domain/visual_contracts.py @@ -0,0 +1,60 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Canonical visual-fact contracts shared by planning and evaluation.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Any + +__all__ = [ + "OCCLUSION_RELATION", + "VISUAL_RELATION_PARTICIPANTS", + "requested_visual_task_predicates", +] + + +OCCLUSION_RELATION = "occludes" + +# Participant order is semantic. For ``occludes`` it is +# ``[occluder_uid, occluded_uid]``. +VISUAL_RELATION_PARTICIPANTS: Mapping[str, tuple[str, ...]] = MappingProxyType( + {OCCLUSION_RELATION: ("occluder", "occluded")} +) + + +def requested_visual_task_predicates(task_spec: Mapping[str, Any]) -> frozenset[str]: + """Return task-level visual predicates explicitly requested by a TaskSpec.""" + result: set[str] = set() + + def collect(value: Any) -> None: + if isinstance(value, Mapping): + if value.get("type") == "visual_relation": + relation = value.get("relation") + if isinstance(relation, str) and relation: + result.add(relation) + for child in value.values(): + collect(child) + elif isinstance(value, Sequence) and not isinstance( + value, (str, bytes, bytearray) + ): + for child in value: + collect(child) + + collect(task_spec.get("success", {})) + return frozenset(result) 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/generation/__init__.py b/embodichain/gen_sim/action_engine/generation/__init__.py new file mode 100644 index 000000000..4446010d6 --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/__init__.py @@ -0,0 +1,33 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Independent config generation for Action Engine.""" + +from __future__ import annotations + +from .config_builder import VLM_CAMERA_UIDS, canonical_robot_profile +from .assets import normalize_scene_assets +from .generator import generate_action_engine_config +from .models import GeneratedConfigPaths, PreparedScene + +__all__ = [ + "GeneratedConfigPaths", + "PreparedScene", + "VLM_CAMERA_UIDS", + "canonical_robot_profile", + "generate_action_engine_config", + "normalize_scene_assets", +] diff --git a/embodichain/gen_sim/action_engine/generation/artifacts.py b/embodichain/gen_sim/action_engine/generation/artifacts.py new file mode 100644 index 000000000..0af961051 --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/artifacts.py @@ -0,0 +1,141 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Publish canonical generation artifacts without intermediate copies.""" + +from __future__ import annotations + +from collections.abc import Mapping +import json +import os +from pathlib import Path +import tempfile +from typing import Any + +from embodichain.gen_sim.action_engine.protocol import ( + AGENT_CONFIG_FILENAME, + EXECUTION_PROGRAM_FILENAME, + FAST_GYM_CONFIG_FILENAME, + SCENE_REQUIREMENTS_FILENAME, + SEED_TASK_GRAPH_PNG_FILENAME, + TASK_SPEC_FILENAME, +) + +from .models import GeneratedConfigPaths + +__all__ = ["artifact_paths", "write_generation_artifacts"] + + +def artifact_paths( + output_dir: str | Path, + *, + planning_mode: str = "offline", +) -> GeneratedConfigPaths: + """Return canonical resolved paths for one output directory.""" + directory = Path(output_dir).expanduser().resolve() + _validate_planning_mode(planning_mode) + graph_directory = directory if planning_mode == "offline" else directory / "offline" + return GeneratedConfigPaths( + gym_config=directory / FAST_GYM_CONFIG_FILENAME, + agent_config=directory / AGENT_CONFIG_FILENAME, + task_spec=directory / TASK_SPEC_FILENAME, + scene_requirements=directory / SCENE_REQUIREMENTS_FILENAME, + seed_task_graph=graph_directory / EXECUTION_PROGRAM_FILENAME, + seed_task_graph_png=graph_directory / SEED_TASK_GRAPH_PNG_FILENAME, + planning_mode=planning_mode, + ) + + +def write_generation_artifacts( + output_dir: str | Path, + *, + gym_config: Mapping[str, Any], + agent_config: Mapping[str, Any], + task_spec: Mapping[str, Any], + scene_requirements: Mapping[str, Any], + seed_task_graph: Mapping[str, Any], + seed_task_graph_png: bytes, + overwrite: bool, + planning_mode: str = "offline", +) -> GeneratedConfigPaths: + """Serialize validated artifacts and replace their destinations atomically.""" + paths = artifact_paths(output_dir, planning_mode=planning_mode) + if not isinstance(seed_task_graph_png, (bytes, bytearray)): + raise TypeError("seed_task_graph_png must be bytes.") + payloads = { + paths.gym_config: _serialize_json(gym_config), + paths.agent_config: _serialize_json(agent_config), + paths.task_spec: _serialize_json(task_spec), + paths.scene_requirements: _serialize_json(scene_requirements), + paths.seed_task_graph: _serialize_json(seed_task_graph), + paths.seed_task_graph_png: bytes(seed_task_graph_png), + } + existing = sorted(path for path in payloads if path.exists()) + if existing and not overwrite: + names = ", ".join(path.name for path in existing) + raise FileExistsError( + f"Generated artifacts already exist in {paths.gym_config.parent}: " + f"{names}. Pass --overwrite to replace them." + ) + + paths.gym_config.parent.mkdir(parents=True, exist_ok=True) + temporary: dict[Path, Path] = {} + try: + for destination, payload in payloads.items(): + destination.parent.mkdir(parents=True, exist_ok=True) + temporary[destination] = _write_temporary(destination.parent, payload) + for destination, temporary_path in temporary.items(): + os.replace(temporary_path, destination) + finally: + for temporary_path in temporary.values(): + temporary_path.unlink(missing_ok=True) + return paths + + +def _serialize_json(value: Mapping[str, Any]) -> str: + try: + return ( + json.dumps( + dict(value), + ensure_ascii=False, + indent=2, + sort_keys=False, + allow_nan=False, + ) + + "\n" + ) + except (TypeError, ValueError) as exc: + raise ValueError("Generated artifact is not strict JSON data.") from exc + + +def _write_temporary(directory: Path, payload: str | bytes) -> Path: + data = payload if isinstance(payload, bytes) else payload.encode("utf-8") + with tempfile.NamedTemporaryFile( + mode="wb", + dir=directory, + prefix=".action_engine_", + suffix=".tmp", + delete=False, + ) as stream: + stream.write(data) + stream.flush() + os.fsync(stream.fileno()) + return Path(stream.name) + + +def _validate_planning_mode(value: Any) -> None: + if value not in {"offline", "ab"}: + raise ValueError("planning_mode must be 'offline' or 'ab'.") diff --git a/embodichain/gen_sim/action_engine/generation/assets.py b/embodichain/gen_sim/action_engine/generation/assets.py new file mode 100644 index 000000000..d1b4da830 --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/assets.py @@ -0,0 +1,158 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Normalize GLB node transforms and body scale into reusable runtime assets.""" + +from __future__ import annotations + +from copy import deepcopy +from dataclasses import replace +import hashlib +import json +from pathlib import Path +from typing import Any + +import numpy as np + +from .models import PreparedScene + +__all__ = ["normalize_scene_assets"] + +_POLICY = "action_engine_glb_geometry_v2" + + +def normalize_scene_assets( + scene: PreparedScene, + output_dir: str | Path, +) -> PreparedScene: + """Return a scene whose valid GLB meshes have flattened runtime geometry. + + Source files are never modified. Cache names derive from source bytes, + object scale, and the normalization policy, so repeated generation reuses + identical assets. + """ + sections = { + "background": [deepcopy(value) for value in scene.background], + "rigid_object": [deepcopy(value) for value in scene.rigid_objects], + "articulation": [deepcopy(value) for value in scene.articulations], + } + cache_dir = Path(output_dir).expanduser().resolve() / "mesh_assets" / "normalized" + reports: list[dict[str, Any]] = [] + hashes = dict(scene.asset_hashes) + normalized_by_uid: dict[str, dict[str, Any]] = {} + for section in ("background", "rigid_object"): + for config in sections[section]: + report = _normalize_object(config, cache_dir) + if report is not None: + reports.append(report) + hashes[str(config["uid"])] = str(report["runtime_sha256"]) + normalized_by_uid[str(config["uid"])] = config + + planner = [deepcopy(value) for value in scene.planner_objects] + for item in planner: + runtime = normalized_by_uid.get(str(item["runtime_uid"])) + if runtime is None: + continue + item["shape"] = deepcopy(runtime.get("shape", {})) + item["body_scale"] = list(runtime.get("body_scale", [1.0, 1.0, 1.0])) + return replace( + scene, + planner_objects=tuple(planner), + background=tuple(sections["background"]), + rigid_objects=tuple(sections["rigid_object"]), + articulations=tuple(sections["articulation"]), + asset_hashes=hashes, + asset_provenance=tuple(reports), + ) + + +def _normalize_object( + config: dict[str, Any], + cache_dir: Path, +) -> dict[str, Any] | None: + shape = config.get("shape") + if not isinstance(shape, dict) or not shape.get("fpath"): + return None + source = Path(str(shape["fpath"])).expanduser().resolve() + if source.suffix.lower() not in {".glb", ".gltf"}: + return None + source_hash = _file_hash(source) + scale = [float(value) for value in config.get("body_scale", [1.0, 1.0, 1.0])] + key = hashlib.sha256( + json.dumps( + {"source": source_hash, "scale": scale, "policy": _POLICY}, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + ).hexdigest() + destination = cache_dir / f"{source.stem[:32]}_{key[:16]}.glb" + status = "reused" if destination.is_file() else "generated" + if status == "generated": + try: + _bake_glb(source, destination, scale) + except Exception as exc: + return { + "uid": str(config.get("uid", "")), + "source_path": source.as_posix(), + "source_sha256": source_hash, + "runtime_path": source.as_posix(), + "runtime_sha256": source_hash, + "body_scale": scale, + "status": "preserved_invalid_source", + "error": f"{type(exc).__name__}: {exc}", + "policy_version": _POLICY, + } + shape["fpath"] = destination.as_posix() + config["body_scale"] = [1.0, 1.0, 1.0] + return { + "uid": str(config.get("uid", "")), + "source_path": source.as_posix(), + "source_sha256": source_hash, + "runtime_path": destination.as_posix(), + "runtime_sha256": _file_hash(destination), + "body_scale": scale, + "status": status, + "policy_version": _POLICY, + } + + +def _bake_glb(source: Path, destination: Path, sim_scale: list[float]) -> None: + import trimesh + + source_scene = trimesh.load(source.as_posix(), force="scene") + baked = trimesh.Scene() + scale = np.diag([sim_scale[0], sim_scale[2], sim_scale[1], 1.0]) + for node_name in source_scene.graph.nodes_geometry: + node_transform, geometry_name = source_scene.graph.get(node_name) + mesh = source_scene.geometry[geometry_name].copy() + mesh.apply_transform(scale @ node_transform) + baked.add_geometry( + mesh, + node_name=str(node_name), + geom_name=f"geometry_{len(baked.geometry)}", + ) + if not baked.geometry: + raise ValueError(f"GLB contains no mesh geometry: {source}") + destination.parent.mkdir(parents=True, exist_ok=True) + baked.export(destination.as_posix(), file_type="glb") + + +def _file_hash(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() diff --git a/embodichain/gen_sim/action_engine/generation/config_builder.py b/embodichain/gen_sim/action_engine/generation/config_builder.py new file mode 100644 index 000000000..51e53648e --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/config_builder.py @@ -0,0 +1,814 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Build the simulator and Action Engine artifact manifests.""" + +from __future__ import annotations + +from collections.abc import Sequence +from copy import deepcopy +from functools import lru_cache +import json +import math +from pathlib import Path +from typing import Any + +from embodichain.gen_sim.action_engine.config import ( + ACTION_ENGINE_DEFAULTS_SCHEMA, + RuntimePolicyCfg, + default_runtime_policy, + generation_defaults, + runtime_policy_hash, +) +from embodichain.gen_sim.action_engine.protocol import ( + ACTION_ENGINE_CONFIG_SCHEMA, + ACTION_ENGINE_ENV_ID, + EXECUTION_PROGRAM_FILENAME, + SCENE_REQUIREMENTS_FILENAME, + TASK_SPEC_FILENAME, +) + +from .models import PreparedScene + +__all__ = [ + "build_agent_config", + "build_fast_gym_config", + "canonical_robot_profile", + "VLM_CAMERA_UIDS", + "validate_fast_gym_config", +] + +_TEMPLATE_DIR = Path(__file__).resolve().parent / "templates" +_GENERATION_DEFAULTS = generation_defaults() +_DEFAULT_TABLETOP_Z = float(_GENERATION_DEFAULTS["scene"]["default_tabletop_z"]) + +_ARM_SLOTS = { + "left": {"arm": "left_arm", "eef": "left_eef"}, + "right": {"arm": "right_arm", "eef": "right_eef"}, +} + +# These IDs are part of the A/B runtime contract. Keep the order stable so +# visual-fact payloads and comparison reports are reproducible across runs. +VLM_CAMERA_UIDS = ( + "vlm_front", + "vlm_left", + "vlm_rear", + "vlm_right", +) + + +def canonical_robot_profile(profile: str) -> str: + """Normalize the supported CLI aliases to one runtime profile ID.""" + normalized = str(profile).strip().lower().replace("-", "_") + profiles = _robot_profiles() + if normalized in profiles: + return normalized + for profile_id, value in profiles.items(): + if normalized in value["aliases"]: + return profile_id + raise ValueError( + f"Unsupported robot profile {profile!r}; expected one of: " + f"{', '.join(sorted(profiles))}" + ) + + +def build_agent_config( + *, + task_name: str, + robot_profile: str, + execution_program_hash: str, + source_config_path: Path, + uid_map: dict[str, str], + static_obstacle_uids: Sequence[str] | None = None, + dynamic_obstacle_uids: Sequence[str] | None = None, + table_top_z: float | None = None, + planning_mode: str = "offline", + seed_task_graph_path: str | Path | None = EXECUTION_PROGRAM_FILENAME, + vlm_model: str | None = None, + vlm_camera_uids: Sequence[str] | None = None, +) -> dict[str, Any]: + """Build the small manifest consumed by ``run_agent``.""" + profile = canonical_robot_profile(robot_profile) + runtime_policy = default_runtime_policy(profile) + if ( + static_obstacle_uids is not None + or dynamic_obstacle_uids is not None + or table_top_z is not None + ): + policy = runtime_policy.as_mapping() + planner = policy["planner"] + if static_obstacle_uids is not None: + planner["static_obstacle_uids"] = [str(uid) for uid in static_obstacle_uids] + if dynamic_obstacle_uids is not None: + planner["dynamic_obstacle_uids"] = [ + str(uid) for uid in dynamic_obstacle_uids + ] + planner["dynamic_collision"] = bool(dynamic_obstacle_uids) + if table_top_z is not None: + tabletop = float(table_top_z) + if not math.isfinite(tabletop): + raise ValueError("table_top_z must be finite when provided.") + height_offset = tabletop - _DEFAULT_TABLETOP_Z + height_policies = ( + policy["grounding"]["semantic_defaults"], + policy["grounding"]["handover"], + policy["motion_defaults"]["MoveEndEffector"], + policy["motion_modifiers"]["orientation"]["upright"]["MoveEndEffector"], + ) + for height_policy in height_policies: + height_policy["maximum_eef_height"] = round( + float(height_policy["maximum_eef_height"]) + height_offset, + 6, + ) + runtime_policy = RuntimePolicyCfg.from_mapping(policy) + _validate_planning_mode(planning_mode) + graph_path = _validate_seed_graph_path(seed_task_graph_path) + if planning_mode == "ab" and graph_path == EXECUTION_PROGRAM_FILENAME: + graph_path = f"offline/{EXECUTION_PROGRAM_FILENAME}" + result = { + "schema_version": ACTION_ENGINE_CONFIG_SCHEMA, + "task_name": task_name, + "robot_profile": profile, + "planning_mode": planning_mode, + "task_spec": TASK_SPEC_FILENAME, + "scene_requirements": SCENE_REQUIREMENTS_FILENAME, + "seed_task_graph": graph_path, + "seed_task_graph_hash": execution_program_hash, + "runtime_policy": runtime_policy.as_mapping(), + "runtime_policy_hash": runtime_policy_hash(runtime_policy), + "source": { + "gym_config": source_config_path.as_posix(), + "uid_map": dict(sorted(uid_map.items())), + }, + } + if planning_mode == "ab": + camera_uids = _normalize_vlm_camera_uids(vlm_camera_uids) + configured_model = _optional_model(vlm_model) + # Retain concise top-level aliases for early A/B bundles while keeping + # the nested section as the canonical runtime namespace. + result["offline_seed_task_graph"] = graph_path + result["vlm_model"] = configured_model + result["vlm_camera_uids"] = list(camera_uids) + result["online_planning"] = { + # Model names are deliberately persisted only when explicitly + # supplied by the generator. Runtime resolution can then apply + # the documented ACTION_ENGINE_VLM_MODEL/OPENAI_MODEL fallback. + "vlm_model": configured_model, + "camera_uids": camera_uids, + } + return result + + +def build_fast_gym_config( + scene: PreparedScene, + *, + task_name: str, + task_description: str, + robot_profile: str, + execution_program_hash: str, + max_episodes: int, + max_episode_steps: int, + randomize_scene: bool = False, + randomize_table_material: bool = False, + planning_mode: str = "offline", + seed_task_graph_path: str | Path | None = EXECUTION_PROGRAM_FILENAME, +) -> dict[str, Any]: + """Build a runnable EmbodiChain gym config from a prepared source scene.""" + if max_episodes < 1: + raise ValueError("max_episodes must be at least 1.") + if max_episode_steps < 1: + raise ValueError("max_episode_steps must be at least 1.") + _validate_planning_mode(planning_mode) + graph_path = _validate_seed_graph_path(seed_task_graph_path) + if planning_mode == "ab" and graph_path == EXECUTION_PROGRAM_FILENAME: + graph_path = f"offline/{EXECUTION_PROGRAM_FILENAME}" + profile = canonical_robot_profile(robot_profile) + + profile_config = _profile(profile) + robot = _make_robot(profile, profile_config, scene.table_top_z) + observations = _make_observations(robot) + # These two template fields describe serialization order to generation, not + # RobotCfg. Remove them after deriving observation IDs to avoid parser noise. + robot.pop("observation_joint_parts", None) + robot.pop("qpos_control_part_order", None) + sensors = _load_template("default_sensors.json") + if not isinstance(sensors, list) or not sensors: + raise ValueError("Default sensor template must define at least one camera.") + environment_policy = _GENERATION_DEFAULTS["environment"] + viewer_camera_uid = str(environment_policy["viewer_camera_uid"]) + sensors[0]["uid"] = viewer_camera_uid + if planning_mode == "ab": + vlm_sensors = _load_template("vlm_sensors.json") + if not isinstance(vlm_sensors, list) or len(vlm_sensors) != len( + VLM_CAMERA_UIDS + ): + raise ValueError("A/B planning requires exactly four VLM cameras.") + _validate_vlm_sensors(vlm_sensors) + _anchor_vlm_sensors(vlm_sensors, scene) + sensors.extend(vlm_sensors) + light = _load_template("default_lights.json") + + rigid_uids = [str(config["uid"]) for config in scene.rigid_objects] + background_uids = [str(config["uid"]) for config in scene.background] + engine_extension = { + "schema_version": "action_engine_runtime_v2", + "defaults_schema_version": ACTION_ENGINE_DEFAULTS_SCHEMA, + "task_name": task_name, + "robot_profile": profile, + "planning_mode": planning_mode, + "task_spec": TASK_SPEC_FILENAME, + "scene_requirements": SCENE_REQUIREMENTS_FILENAME, + "seed_task_graph": graph_path, + "seed_task_graph_hash": execution_program_hash, + "source_gym_config": scene.source_config_path.as_posix(), + "source_scene_z_rotation_degrees": scene.z_rotation_degrees, + "body_scale_policy": scene.body_scale_policy, + "body_scale": list(scene.body_scale), + "asset_hashes": dict(sorted(scene.asset_hashes.items())), + "asset_provenance": [deepcopy(value) for value in scene.asset_provenance], + "uid_map": dict(sorted(scene.uid_map.items())), + } + extensions = { + "action_engine": engine_extension, + "agent_robot_profile": profile, + "agent_arm_slots": deepcopy(_ARM_SLOTS), + "agent_static_obstacle_uids": background_uids, + "agent_dynamic_obstacle_uids": rigid_uids, + "gripper_open_state": list(profile_config["gripper_open_state"]), + "gripper_close_state": list(profile_config["gripper_close_state"]), + "arm_aim_yaw_offset": deepcopy(environment_policy["arm_aim_yaw_offset"]), + "ignore_terminations_during_agent": bool( + environment_policy["ignore_terminations_during_agent"] + ), + "viewer_camera_uid": viewer_camera_uid, + } + + config: dict[str, Any] = { + "id": ACTION_ENGINE_ENV_ID, + "max_episodes": int(max_episodes), + "max_episode_steps": int(max_episode_steps), + "env": { + "extensions": extensions, + "events": _make_events( + sensors[0], + rigid_uids, + planning_mode=planning_mode, + randomize_scene=randomize_scene, + randomize_table_material=randomize_table_material, + ), + "observations": observations, + "dataset": _make_dataset( + task_name=task_name, + task_description=task_description, + source_config_path=scene.source_config_path, + robot_type=str(robot["uid"]), + ), + }, + "robot": robot, + "sensor": sensors, + "light": light, + "background": [deepcopy(obj_config) for obj_config in scene.background], + "rigid_object": [deepcopy(obj_config) for obj_config in scene.rigid_objects], + } + if scene.articulations: + config["articulation"] = [ + deepcopy(articulation) for articulation in scene.articulations + ] + validate_fast_gym_config(config) + return config + + +def validate_fast_gym_config(config: dict[str, Any]) -> None: + """Check the cross-file and simulator-facing invariants generation owns.""" + if config.get("id") != ACTION_ENGINE_ENV_ID: + raise ValueError(f"Gym config id must be {ACTION_ENGINE_ENV_ID!r}.") + if not isinstance(config.get("robot"), dict) or not config["robot"].get("uid"): + raise ValueError("Gym config requires a concrete robot template.") + if not config.get("sensor"): + raise ValueError("Gym config requires at least one sensor.") + if not all(isinstance(sensor, dict) for sensor in config["sensor"]): + raise ValueError("Generated sensors must be object mappings.") + sensor_uids = [str(sensor.get("uid", "")) for sensor in config["sensor"]] + if not all(sensor_uids) or len(sensor_uids) != len(set(sensor_uids)): + raise ValueError("Generated sensor UIDs must be non-empty and unique.") + if not config.get("background"): + raise ValueError("Gym config requires at least one background object.") + + objects = [ + *config.get("background", []), + *config.get("rigid_object", []), + *config.get("articulation", []), + ] + uids = [str(obj.get("uid", "")) for obj in objects] + if not all(uids) or len(uids) != len(set(uids)): + raise ValueError("Generated scene object UIDs must be non-empty and unique.") + if "table" not in uids: + raise ValueError("Generated tabletop scene must expose runtime UID 'table'.") + + for obj in objects: + shape = obj.get("shape") + fpath = shape.get("fpath") if isinstance(shape, dict) else obj.get("fpath") + if fpath is None: + continue + path = Path(str(fpath)) + if not path.is_absolute() or not path.is_file(): + raise ValueError( + f"Generated asset path for {obj.get('uid')!r} is not an " + f"existing absolute file: {path}" + ) + + action_engine = config.get("env", {}).get("extensions", {}).get("action_engine", {}) + if action_engine.get("defaults_schema_version") != ACTION_ENGINE_DEFAULTS_SCHEMA: + raise ValueError("Gym config has an unexpected defaults schema version.") + if action_engine.get("task_spec") != TASK_SPEC_FILENAME: + raise ValueError("Gym config points to an unexpected TaskSpec artifact.") + if action_engine.get("scene_requirements") != SCENE_REQUIREMENTS_FILENAME: + raise ValueError("Gym config points to unexpected SceneRequirements.") + graph_path = action_engine.get("seed_task_graph") + if ( + not isinstance(graph_path, str) + or Path(graph_path).name != EXECUTION_PROGRAM_FILENAME + ): + raise ValueError("Gym config points to an unexpected SeedGraph artifact.") + + planning_mode = action_engine.get("planning_mode", "offline") + _validate_planning_mode(planning_mode) + if planning_mode == "ab": + sensors = config["sensor"] + vlm_sensors = [ + sensor + for sensor in sensors + if isinstance(sensor, dict) + and str(sensor.get("uid", "")).startswith("vlm_") + ] + _validate_vlm_sensors(vlm_sensors) + + registered = { + entry.get("entity_cfg", {}).get("uid") + for entry in ( + config.get("env", {}) + .get("events", {}) + .get("register_info_to_env", {}) + .get("params", {}) + .get("registry", []) + ) + } + rigid_uids = {obj["uid"] for obj in config.get("rigid_object", [])} + if registered != rigid_uids: + raise ValueError("Every rigid object must have one live-pose registry entry.") + + +def _make_robot( + profile_id: str, + profile: dict[str, Any], + table_top_z: float | None, +) -> dict[str, Any]: + robot = _load_template(str(profile["template"])) + tabletop_z = _DEFAULT_TABLETOP_Z if table_top_z is None else float(table_top_z) + robot["init_pos"][2] = round( + tabletop_z + + float(profile["tabletop_clearance"]) + - float(profile["arm_component_z"]), + 6, + ) + family = str(profile["robot_family"]) + if family.startswith("ur"): + display = family.upper() + urdf_dir = display + robot["uid"] = f"Dual{display}" + robot["urdf_cfg"]["fname"] = f"dual_{family}_robotiq_arg2f_140_basket" + for component in robot["urdf_cfg"]["components"]: + if str(component.get("component_type", "")).endswith("_arm"): + component["urdf_path"] = f"UniversalRobots/{urdf_dir}/{urdf_dir}.urdf" + component["transform"][0][3] = float(profile["arm_base_x"]) + component["transform"][2][3] = float(profile["arm_component_z"]) + for arm in ("left_arm", "right_arm"): + robot["solver_cfg"][arm]["ur_type"] = family + robot["drive_pros"]["max_effort"][arm] = float(profile["max_effort"]) + robot["qpos_control_part_order"] = [ + "left_arm", + "right_arm", + "left_eef", + "right_eef", + ] + robot["observation_joint_parts"] = ["left_eef", "right_eef"] + if profile_id != canonical_robot_profile(profile_id): + raise ValueError(f"Invalid canonical robot profile {profile_id!r}.") + return robot + + +@lru_cache(maxsize=1) +def _robot_profiles() -> dict[str, dict[str, Any]]: + value = _read_template("robot_profiles.json") + if not isinstance(value, dict) or not value: + raise ValueError("robot_profiles.json must contain a non-empty object.") + return value + + +def _profile(profile_id: str) -> dict[str, Any]: + profile = deepcopy(_robot_profiles()[profile_id]) + required = { + "aliases", + "template", + "robot_family", + "tabletop_clearance", + "arm_component_z", + "gripper_open_state", + "gripper_close_state", + } + missing = sorted(required - set(profile)) + if missing: + raise ValueError(f"Robot profile {profile_id!r} is missing fields: {missing}.") + return profile + + +def _make_events( + camera: dict[str, Any], + rigid_uids: list[str], + *, + planning_mode: str, + randomize_scene: bool = False, + randomize_table_material: bool = False, +) -> dict[str, Any]: + extrinsics = camera["extrinsics"] + eye = list(extrinsics["eye"]) + target = list(extrinsics["target"]) + # The recording view mirrors the interactive viewer around its target. + audience_eye = [ + 2.0 * float(target[0]) - float(eye[0]), + 2.0 * float(target[1]) - float(eye[1]), + float(eye[2]), + ] + recording_enabled, recording_resolution, recording_interval = _recording_policy( + planning_mode + ) + source_width = int(camera["width"]) + source_height = int(camera["height"]) + if source_width <= 0 or source_height <= 0: + raise ValueError("Recording source camera resolution must be positive.") + intrinsics = camera.get("intrinsics") + if ( + not isinstance(intrinsics, Sequence) + or isinstance(intrinsics, (str, bytes, bytearray)) + or len(intrinsics) != 4 + ): + raise ValueError("Recording source camera intrinsics must be a 4-vector.") + scale_x = recording_resolution[0] / source_width + scale_y = recording_resolution[1] / source_height + recording_intrinsics = [ + float(intrinsics[0]) * scale_x, + float(intrinsics[1]) * scale_y, + float(intrinsics[2]) * scale_x, + float(intrinsics[3]) * scale_y, + ] + events = { + "record_camera": { + "func": "record_camera_data", + "mode": "interval", + "interval_step": recording_interval, + "params": { + "name": "record_cam_audience_view", + "resolution": list(recording_resolution), + "intrinsics": recording_intrinsics, + "eye": audience_eye, + "target": target, + "up": [ + -float(extrinsics["up"][0]), + -float(extrinsics["up"][1]), + float(extrinsics["up"][2]), + ], + }, + }, + "validation_cameras": { + "func": "validation_cameras", + "mode": "trigger", + "params": {}, + }, + "prepare_extra_attr": { + "func": "prepare_extra_attr", + "mode": "reset", + "params": { + "attrs": [ + { + "name": "object_lengths", + "mode": "callable", + "entity_uids": "all_objects", + "func_name": "compute_object_length", + "func_kwargs": { + "is_svd_frame": True, + "sample_points": int( + _GENERATION_DEFAULTS["scene"][ + "object_length_sample_points" + ] + ), + }, + } + ] + }, + }, + "register_info_to_env": { + "func": "register_info_to_env", + "mode": "reset", + "params": { + "registry": [ + { + "entity_cfg": {"uid": uid}, + "pose_register_params": { + "compute_relative": False, + "compute_pose_object_to_arena": True, + "to_matrix": True, + }, + } + for uid in sorted(rigid_uids) + ], + "registration": "affordance_datas", + "sim_update": True, + }, + }, + } + if not recording_enabled: + events.pop("record_camera") + if randomize_table_material: + material = _GENERATION_DEFAULTS["randomization"]["table_material"] + events["randomize_table_material"] = { + "func": "randomize_visual_material", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "table"}, + "random_texture_prob": float(material["random_texture_prob"]), + "base_color_range": deepcopy(material["base_color_range"]), + "metallic_range": list(material["metallic_range"]), + "roughness_range": list(material["roughness_range"]), + }, + } + if randomize_scene: + randomization = _GENERATION_DEFAULTS["randomization"] + for uid in sorted(rigid_uids): + events[f"randomize_{uid}_pose"] = { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": uid}, + "position_range": deepcopy( + randomization["rigid_object_position_range"] + ), + "rotation_range": deepcopy( + randomization["rigid_object_rotation_range"] + ), + "relative_position": True, + "relative_rotation": True, + }, + } + events["randomize_table_height"] = { + "func": "randomize_anchor_height", + "mode": "reset", + "params": { + "anchor_uid": "table", + "height_delta_range": deepcopy( + randomization["table_height_delta_range"] + ), + }, + } + return events + + +def _recording_policy(planning_mode: str) -> tuple[bool, tuple[int, int], int]: + """Resolve the bounded GenSim audience-recording policy.""" + value = _GENERATION_DEFAULTS["environment"].get("recording") + required = {"enabled", "resolution", "interval_step"} + if not isinstance(value, dict) or set(value) != required: + raise ValueError( + "generation.environment.recording must define enabled, resolution, " + "and interval_step." + ) + enabled = value["enabled"] + if not isinstance(enabled, bool): + raise ValueError("generation.environment.recording.enabled must be a boolean.") + resolution = value["resolution"] + if ( + not isinstance(resolution, Sequence) + or isinstance(resolution, (str, bytes, bytearray)) + or len(resolution) != 2 + or any( + isinstance(item, bool) or not isinstance(item, int) for item in resolution + ) + or any(int(item) <= 0 for item in resolution) + ): + raise ValueError( + "generation.environment.recording.resolution must contain two " + "positive integers." + ) + interval_step = value["interval_step"] + if ( + isinstance(interval_step, bool) + or not isinstance(interval_step, int) + or interval_step <= 0 + ): + raise ValueError( + "generation.environment.recording.interval_step must be positive." + ) + return ( + bool(enabled or planning_mode == "ab"), + (int(resolution[0]), int(resolution[1])), + int(interval_step), + ) + + +def _make_observations(robot: dict[str, Any]) -> dict[str, Any]: + control_parts = robot["control_parts"] + qpos_order = robot["qpos_control_part_order"] + observed_parts = set(robot["observation_joint_parts"]) + offset = 0 + joint_ids: list[int] = [] + for part in qpos_order: + count = len(control_parts[part]) + if part in observed_parts: + joint_ids.extend(range(offset, offset + count)) + offset += count + return { + "norm_robot_eef_joint": { + "func": "normalize_robot_joint_data", + "mode": "modify", + "name": "robot/qpos", + "params": {"joint_ids": joint_ids}, + } + } + + +def _make_dataset( + *, + task_name: str, + task_description: str, + source_config_path: Path, + robot_type: str, +) -> dict[str, Any]: + dataset_policy = _GENERATION_DEFAULTS["dataset"] + return { + "lerobot": { + "func": "LeRobotRecorder", + "mode": "save", + "save_failed_episodes": bool(dataset_policy["save_failed_episodes"]), + "params": { + "robot_meta": { + "robot_type": robot_type, + "control_freq": int(dataset_policy["control_frequency"]), + }, + "instruction": {"lang": task_description}, + "extra": { + "scene_type": source_config_path.parent.name, + "task_name": task_name, + # LeRobotRecorder uses this legacy field as a directory label. + "task_description": task_name, + "data_type": "sim", + }, + "use_videos": bool(dataset_policy["use_videos"]), + }, + } + } + + +def _load_template(name: str) -> Any: + return deepcopy(_read_template(name)) + + +@lru_cache(maxsize=None) +def _read_template(name: str) -> Any: + path = _TEMPLATE_DIR / name + if not path.is_file(): + raise FileNotFoundError(f"Action Engine template not found: {path}") + return json.loads(path.read_text(encoding="utf-8")) + + +def _validate_planning_mode(value: Any) -> str: + """Validate and return the two supported generation/runtime modes.""" + if value not in {"offline", "ab"}: + raise ValueError("planning_mode must be 'offline' or 'ab'.") + return str(value) + + +def _validate_seed_graph_path(value: str | Path | None) -> str: + """Validate a relative or absolute path while preserving caller spelling.""" + if value is None: + return EXECUTION_PROGRAM_FILENAME + if not isinstance(value, (str, Path)): + raise ValueError("seed_task_graph_path must be a non-empty path string.") + path = str(value).strip() + if not path: + raise ValueError("seed_task_graph_path must be a non-empty path string.") + if Path(path).name != EXECUTION_PROGRAM_FILENAME: + raise ValueError("seed_task_graph_path must point to seed_task_graph.json.") + return path + + +def _optional_model(value: Any) -> str | None: + """Normalize optional model names without serializing blank strings.""" + if value is None: + return None + if not isinstance(value, str): + raise TypeError("Model name must be a string or None.") + normalized = value.strip() + return normalized or None + + +def _normalize_vlm_camera_uids(value: Sequence[str] | None) -> list[str]: + """Return the canonical four-camera list used by A/B execution.""" + if value is None: + return list(VLM_CAMERA_UIDS) + if isinstance(value, (str, bytes, bytearray)) or not isinstance(value, Sequence): + raise TypeError("vlm_camera_uids must be a list of strings.") + if not all(isinstance(item, str) for item in value): + raise TypeError("vlm_camera_uids must be a list of strings.") + normalized = [item.strip() for item in value] + if normalized != list(VLM_CAMERA_UIDS): + raise ValueError( + "A/B planning requires VLM cameras in canonical order: " + f"{list(VLM_CAMERA_UIDS)}." + ) + return normalized + + +def _validate_vlm_sensors(value: list[dict[str, Any]]) -> None: + """Validate camera template fields needed by visual fact extraction.""" + if len(value) != len(VLM_CAMERA_UIDS): + raise ValueError("A/B planning requires exactly four VLM cameras.") + if not all(isinstance(sensor, dict) for sensor in value): + raise ValueError("VLM sensors must be object mappings.") + uids = [str(sensor.get("uid", "")) for sensor in value] + if uids != list(VLM_CAMERA_UIDS): + raise ValueError("VLM camera UIDs must be exactly " f"{list(VLM_CAMERA_UIDS)}.") + for sensor in value: + if sensor.get("sensor_type", "Camera") != "Camera": + raise ValueError(f"VLM sensor {sensor.get('uid')!r} must be a Camera.") + if int(sensor.get("width", 0)) != 640 or int(sensor.get("height", 0)) != 480: + raise ValueError("VLM cameras must use 640x480 resolution.") + if not bool(sensor.get("enable_color")) or not bool(sensor.get("enable_depth")): + raise ValueError("VLM cameras must enable RGB and depth.") + extrinsics = sensor.get("extrinsics") + if not isinstance(extrinsics, dict) or not all( + key in extrinsics for key in ("eye", "target", "up") + ): + raise ValueError( + f"VLM sensor {sensor.get('uid')!r} requires eye/target/up extrinsics." + ) + for name in ("eye", "target", "up"): + vector = extrinsics[name] + if ( + not isinstance(vector, Sequence) + or isinstance(vector, (str, bytes, bytearray)) + or len(vector) != 3 + ): + raise ValueError( + f"VLM sensor {sensor.get('uid')!r} {name} must be a 3-vector." + ) + try: + values = [float(item) for item in vector] + except (TypeError, ValueError) as exc: + raise ValueError( + f"VLM sensor {sensor.get('uid')!r} {name} must be numeric." + ) from exc + if not all(math.isfinite(item) for item in values): + raise ValueError( + f"VLM sensor {sensor.get('uid')!r} {name} must be finite." + ) + + +def _anchor_vlm_sensors(sensors: list[dict[str, Any]], scene: PreparedScene) -> None: + """Aim the fixed high views at the normalized tabletop center.""" + table = next( + ( + item + for item in scene.background + if isinstance(item, dict) and str(item.get("uid")) == "table" + ), + None, + ) + init_pos = table.get("init_pos", [0.0, 0.0, 0.0]) if table else [0.0, 0.0, 0.0] + if not isinstance(init_pos, Sequence) or len(init_pos) != 3: + init_pos = [0.0, 0.0, 0.0] + center = [ + float(init_pos[0]), + float(init_pos[1]), + float(scene.table_top_z if scene.table_top_z is not None else 0.75), + ] + for sensor in sensors: + extrinsics = sensor["extrinsics"] + eye = [float(value) for value in extrinsics["eye"]] + target = [float(value) for value in extrinsics["target"]] + offset = [target[index] - 0.0 for index in range(3)] + extrinsics["target"] = list(center) + extrinsics["eye"] = [ + center[index] + eye[index] - offset[index] for index in range(3) + ] diff --git a/embodichain/gen_sim/action_engine/generation/generator.py b/embodichain/gen_sim/action_engine/generation/generator.py new file mode 100644 index 000000000..ecd389d52 --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/generator.py @@ -0,0 +1,866 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Orchestrate source-scene preparation, planning, compilation, and publication.""" + +from __future__ import annotations + +import json +from collections.abc import Mapping +from collections.abc import Sequence +from copy import deepcopy +from pathlib import Path +from typing import Any + +from embodichain.gen_sim.action_engine.config import ( + generation_defaults, + resolve_agent_runtime_policy, +) +from embodichain.gen_sim.action_engine.protocol import ( + ACTION_ENGINE_CONFIG_SCHEMA, + EXECUTION_PROGRAM_FILENAME, + SCENE_REQUIREMENTS_FILENAME, + SCENE_REQUIREMENTS_SCHEMA, + TASK_SPEC_FILENAME, +) + +from .artifacts import artifact_paths, write_generation_artifacts +from .assets import normalize_scene_assets +from .config_builder import ( + VLM_CAMERA_UIDS, + build_agent_config, + build_fast_gym_config, +) +from .models import GeneratedConfigPaths +from .source_scene import prepare_scene + +__all__ = ["generate_action_engine_config"] + +_GENERATION_DEFAULTS = generation_defaults() +_TASK_DEFAULTS = _GENERATION_DEFAULTS["task"] +_SCENE_DEFAULTS = _GENERATION_DEFAULTS["scene"] +_DEFAULT_BODY_SCALE = tuple(float(value) for value in _SCENE_DEFAULTS["body_scale"]) + + +def generate_action_engine_config( + gym_project: str | Path, + output_dir: str | Path, + *, + task_name: str, + task_description: str | None = None, + task_spec: Mapping[str, Any] | str | Path | None = None, + robot_profile: str = str(_TASK_DEFAULTS["default_robot_profile"]), + llm_model: str | None = None, + source_scene_z_rotation_degrees: float | None = None, + body_scale_policy: str = str(_SCENE_DEFAULTS["body_scale_policy"]), + body_scale: Sequence[float] = _DEFAULT_BODY_SCALE, + overwrite: bool = False, + max_episodes: int = int(_TASK_DEFAULTS["max_episodes"]), + max_episode_steps: int = int(_TASK_DEFAULTS["max_episode_steps"]), + randomize_scene: bool = False, + randomize_table_material: bool = False, + planning_mode: str = "offline", + vlm_model: str | None = None, +) -> GeneratedConfigPaths: + """Generate the complete Action Engine input bundle. + + Natural-language input is interpreted and grounded by the structured LLM + path. Callers may instead provide an already grounded v2 TaskSpec; that + path never invokes a text model. + """ + task_name = str(task_name).strip() + task_description = "" if task_description is None else str(task_description).strip() + if not task_name: + raise ValueError("task_name must be a non-empty string.") + if task_spec is not None and task_description: + raise ValueError("task_spec cannot be combined with task_description.") + if task_spec is None and not task_description: + raise ValueError("task_description is required when task_spec is not supplied.") + if planning_mode not in {"offline", "ab"}: + raise ValueError("planning_mode must be 'offline' or 'ab'.") + _raise_if_outputs_exist( + output_dir, + overwrite=overwrite, + planning_mode=planning_mode, + ) + + scene = prepare_scene( + gym_project, + z_rotation_degrees=source_scene_z_rotation_degrees, + body_scale_policy=body_scale_policy, + body_scale=body_scale, + ) + + # Delayed imports keep scene/config tooling lightweight and avoid importing + # an LLM client when callers only inspect exported projects. + from embodichain.gen_sim.action_engine.capabilities import ( + build_atomic_capability_registry, + ) + from embodichain.gen_sim.action_engine.domain import ( + seed_graph_hash, + validate_seed_graph, + validate_scene_requirements, + validate_task_spec, + ) + from embodichain.gen_sim.action_engine.tasks import ( + interpret_and_ground_task_spec, + instantiate_seed_graph, + ) + + known_objects = [str(item["runtime_uid"]) for item in scene.planner_objects] + if task_spec is not None: + supplied_task_spec, source_path = _read_task_spec(task_spec) + task_spec = _validated_mapping( + supplied_task_spec, + validator=validate_task_spec, + label="TaskSpec", + ) + _require_matching_task_spec(task_spec, task_name) + task_description = str(task_spec["instruction"]) + supplied_requirements = _read_sibling_scene_requirements( + source_path, + task_name, + ) + if supplied_requirements is not None: + supplied_requirements = _validated_mapping( + supplied_requirements, + validator=validate_scene_requirements, + label="SceneRequirements", + ) + role_bindings = _task_spec_role_bindings( + task_spec, + known_objects, + scene_requirements=supplied_requirements, + scene_objects=scene.planner_objects, + robot_profile=robot_profile, + ) + task_spec = _with_role_bindings(task_spec, role_bindings) + if supplied_requirements is None: + scene_requirements = _scene_requirements_from_bindings( + task_name, + scene.planner_objects, + role_bindings, + ) + else: + scene_requirements = supplied_requirements + _validate_requirement_roles(scene_requirements, role_bindings) + compiled = instantiate_seed_graph(task_spec, role_bindings) + else: + planned = interpret_and_ground_task_spec( + task_name=task_name, + task_description=task_description, + scene_objects=[deepcopy(obj) for obj in scene.planner_objects], + robot_profile=robot_profile, + model=llm_model, + ) + task_spec = _validated_mapping( + planned.task_spec, + validator=validate_task_spec, + label="TaskSpec", + ) + # Persist the validated Scene-Engine hand-off alongside the shared + # semantic TaskSpec. The binding is not an oracle for online planning, + # but it is required for ``--regenerate`` and runtime-only loading. + task_spec = _with_role_bindings(task_spec, planned.role_bindings) + scene_requirements = _validated_mapping( + planned.scene_requirements, + validator=validate_scene_requirements, + label="SceneRequirements", + ) + compiled = instantiate_seed_graph( + task_spec, + planned.role_bindings, + ) + if planning_mode == "ab": + scene_requirements = _add_ab_camera_requirements(scene_requirements) + capabilities = build_atomic_capability_registry() + execution_program = _validated_mapping( + compiled, + validator=lambda value: validate_seed_graph( + value, + known_objects=known_objects, + known_actions=capabilities.names(), + ), + label="SeedGraph", + ) + if execution_program.get("task_id") != task_name: + raise ValueError("SeedGraph task_id does not match requested task_name.") + program_hash = str(seed_graph_hash(execution_program)) + if not program_hash: + raise ValueError("SeedGraph hash must be non-empty.") + + # Validate planning before materializing normalized meshes in output_dir so + # an ambiguous instruction cannot leave a half-generated bundle behind. + scene = normalize_scene_assets(scene, output_dir) + + # Rendering consumes the exact validated in-memory program that runtime + # consumes. The PNG is review-only and never appears in agent input fields. + from embodichain.gen_sim.action_engine.graph_visualization import ( + render_seed_task_graph_png, + ) + + seed_task_graph_png = render_seed_task_graph_png(execution_program) + if not isinstance(seed_task_graph_png, bytes): + raise TypeError("render_seed_task_graph_png must return bytes.") + + paths = artifact_paths(output_dir, planning_mode=planning_mode) + graph_relative_path = paths.seed_task_graph.relative_to( + paths.agent_config.parent + ).as_posix() + vlm_camera_uids = list(VLM_CAMERA_UIDS) + agent_config = build_agent_config( + task_name=task_name, + robot_profile=robot_profile, + execution_program_hash=program_hash, + source_config_path=scene.source_config_path, + uid_map=scene.uid_map, + static_obstacle_uids=[str(config["uid"]) for config in scene.background], + dynamic_obstacle_uids=[str(config["uid"]) for config in scene.rigid_objects], + table_top_z=scene.table_top_z, + planning_mode=planning_mode, + seed_task_graph_path=graph_relative_path, + vlm_model=vlm_model, + vlm_camera_uids=vlm_camera_uids, + ) + gym_config = build_fast_gym_config( + scene, + task_name=task_name, + task_description=task_description, + robot_profile=robot_profile, + execution_program_hash=program_hash, + max_episodes=max_episodes, + max_episode_steps=max_episode_steps, + randomize_scene=randomize_scene, + randomize_table_material=randomize_table_material, + planning_mode=planning_mode, + seed_task_graph_path=graph_relative_path, + ) + if planning_mode == "ab": + output_root = Path(output_dir).expanduser().resolve() + gym_config["env"]["events"]["record_camera"]["params"]["save_path"] = ( + output_root / ".ab_video_staging" + ).as_posix() + gym_config["env"]["dataset"]["lerobot"]["params"]["save_path"] = ( + output_root / ".ab_datasets" + ).as_posix() + _validate_agent_config(agent_config) + return write_generation_artifacts( + output_dir, + gym_config=gym_config, + agent_config=agent_config, + task_spec=task_spec, + scene_requirements=scene_requirements, + seed_task_graph=execution_program, + seed_task_graph_png=seed_task_graph_png, + overwrite=overwrite, + planning_mode=planning_mode, + ) + + +def _read_task_spec( + source: Mapping[str, Any] | str | Path, +) -> tuple[dict[str, Any], Path | None]: + """Read one existing v2 TaskSpec without invoking a text planner.""" + if isinstance(source, Mapping): + return deepcopy(dict(source)), None + path = Path(source).expanduser().resolve() + return _read_json_mapping(path, label="TaskSpec"), path + + +def _read_json_mapping(path: Path, *, label: str) -> dict[str, Any]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except OSError as exc: + raise ValueError(f"Unable to read {label} at {path}: {exc}") from exc + except json.JSONDecodeError as exc: + raise ValueError(f"{label} at {path} is not valid JSON: {exc}") from exc + if not isinstance(value, Mapping): + raise ValueError(f"{label} JSON must contain an object.") + return deepcopy(dict(value)) + + +def _read_sibling_scene_requirements( + task_spec_path: Path | None, + task_name: str, +) -> dict[str, Any] | None: + """Load the canonical sidecar when a task-first batch supplied one.""" + if task_spec_path is None: + return None + candidate = task_spec_path.parent / SCENE_REQUIREMENTS_FILENAME + if not candidate.is_file(): + return None + requirements = _read_json_mapping(candidate, label="SceneRequirements") + if requirements.get("task_id") != task_name: + raise ValueError( + "Sibling SceneRequirements task_id does not match the requested " + "task_name." + ) + return requirements + + +def _require_matching_task_spec(task_spec: Mapping[str, Any], task_name: str) -> None: + if task_spec.get("task_id") != task_name: + raise ValueError( + f"TaskSpec task_id {task_spec.get('task_id')!r} does not match " + f"requested task_name {task_name!r}." + ) + + +def _task_spec_role_bindings( + task_spec: Mapping[str, Any], + known_objects: Sequence[str], + *, + scene_requirements: Mapping[str, Any] | None = None, + scene_objects: Sequence[Mapping[str, Any]] | None = None, + robot_profile: str = "dual_ur10", +) -> dict[str, str]: + """Resolve v2 roles from explicit hand-off data or a strict sidecar match. + + Task-first artifacts may contain abstract role IDs rather than scene UIDs. + When their sibling SceneRequirements is available, match + every still-unbound role against the source scene's static category, + attributes, state, and affordance metadata. This is a deterministic + Scene-Engine hand-off, not a text-model fallback: missing or ambiguous + evidence remains an error. + """ + metadata = task_spec.get("metadata", {}) + if not isinstance(metadata, Mapping): + raise ValueError("TaskSpec.metadata must be a mapping.") + metadata_bindings = metadata.get("role_bindings", {}) + if not isinstance(metadata_bindings, Mapping): + raise ValueError("TaskSpec.metadata.role_bindings must be a mapping.") + candidates: list[tuple[str, Mapping[str, Any]]] = [] + if metadata_bindings: + candidates.append(("TaskSpec.metadata", metadata_bindings)) + + # Older grounded v2 TaskSpecs kept this private hand-off in ``oracle`` + # rather than metadata. Accept that representation while publishing the + # normalized binding in metadata for runtime regeneration. + oracle = task_spec.get("oracle", {}) + if isinstance(oracle, Mapping) and oracle.get("role_bindings"): + oracle_bindings = oracle["role_bindings"] + if not isinstance(oracle_bindings, Mapping): + raise ValueError("TaskSpec.oracle.role_bindings must be a mapping.") + candidates.append(("TaskSpec.oracle", oracle_bindings)) + if isinstance(oracle, Mapping): + reference = oracle.get("reference_seed_graph") + if isinstance(reference, Mapping): + graph_metadata = reference.get("metadata", {}) + if isinstance(graph_metadata, Mapping) and graph_metadata.get( + "role_bindings" + ): + graph_bindings = graph_metadata["role_bindings"] + if not isinstance(graph_bindings, Mapping): + raise ValueError( + "SeedGraph.metadata.role_bindings must be a mapping." + ) + candidates.append(("SeedGraph.metadata", graph_bindings)) + + if scene_requirements is not None: + requirement_metadata = scene_requirements.get("metadata", {}) + if isinstance(requirement_metadata, Mapping) and requirement_metadata.get( + "role_bindings" + ): + requirement_bindings = requirement_metadata["role_bindings"] + if not isinstance(requirement_bindings, Mapping): + raise ValueError( + "SceneRequirements.metadata.role_bindings must be a mapping." + ) + candidates.append(("SceneRequirements.metadata", requirement_bindings)) + + supplied: dict[str, Any] = {} + supplied_sources: dict[str, str] = {} + for source, candidate in candidates: + for raw_role, uid in candidate.items(): + if not isinstance(raw_role, str) or not raw_role.strip(): + raise ValueError(f"{source}.role_bindings must use non-empty role IDs.") + role = raw_role.strip() + if role in supplied and supplied[role] != uid: + raise ValueError( + "Conflicting role_bindings were supplied for " + f"{role!r} by {supplied_sources[role]} and {source}." + ) + supplied[role] = uid + supplied_sources[role] = source + + known = {str(uid) for uid in known_objects} + required = _task_spec_role_references(task_spec.get("task_instances", [])) + required.discard("table") + if not required: + raise ValueError("TaskSpec must reference at least one non-table object role.") + + bindings: dict[str, str] = {} + missing: list[str] = [] + for role in sorted(required): + raw_uid = supplied.get(role, role if role in known else None) + if raw_uid is None: + missing.append(role) + continue + if not isinstance(raw_uid, str) or not raw_uid.strip(): + raise ValueError( + "TaskSpec.metadata.role_bindings must map role IDs to non-empty " + "runtime UIDs." + ) + uid = raw_uid.strip() + if uid not in known: + raise ValueError(f"TaskSpec role {role!r} binds unknown scene UID {uid!r}.") + bindings[role] = uid + if missing and scene_requirements is not None and scene_objects is not None: + bindings.update( + _infer_role_bindings_from_scene_requirements( + missing, + known_objects=known, + scene_objects=scene_objects, + scene_requirements=scene_requirements, + existing_bindings=bindings, + robot_profile=robot_profile, + ) + ) + missing = [role for role in missing if role not in bindings] + if missing: + raise ValueError( + "TaskSpec requires explicit role_bindings or an unambiguous sibling " + f"SceneRequirements match for roles {missing}; a task-first spec must " + "be grounded by a Scene Engine before it can be compiled for this gym " + "project." + ) + if len(bindings.values()) != len(set(bindings.values())): + raise ValueError("TaskSpec role bindings must resolve to unique scene UIDs.") + if scene_requirements is not None and scene_objects is not None: + _validate_bound_role_requirements( + bindings, + scene_requirements=scene_requirements, + scene_objects=scene_objects, + robot_profile=robot_profile, + ) + return bindings + + +def _infer_role_bindings_from_scene_requirements( + roles: Sequence[str], + *, + known_objects: set[str], + scene_objects: Sequence[Mapping[str, Any]], + scene_requirements: Mapping[str, Any], + existing_bindings: Mapping[str, str], + robot_profile: str, +) -> dict[str, str]: + """Bind abstract task roles only when static evidence is unique.""" + from embodichain.gen_sim.action_engine.tasks.assembly import SceneInventory + + requirements = _requirements_by_role(scene_requirements) + inventory = SceneInventory(scene_objects, robot_profile=robot_profile) + entities = [entity for entity in inventory.entities if entity.uid in known_objects] + used_uids = set(existing_bindings.values()) + inferred: dict[str, str] = {} + for role in sorted(roles): + requirement = requirements.get(role) + if requirement is None: + raise ValueError( + "Sibling SceneRequirements is missing TaskSpec role " f"{role!r}." + ) + count = requirement.get("count", 1) + if count != 1: + raise ValueError( + f"TaskSpec role {role!r} has count={count}; direct SeedGraph " + "binding requires exactly one concrete scene UID." + ) + matches = [ + entity + for entity in entities + if entity.uid not in used_uids + and _entity_matches_requirement( + entity, + requirement, + require_complete_static_evidence=True, + ) + ] + if len(matches) != 1: + raise ValueError( + "TaskSpec role " + f"{role!r} requires one unambiguous scene match, found " + f"{[entity.uid for entity in matches]}." + ) + uid = matches[0].uid + inferred[role] = uid + used_uids.add(uid) + return inferred + + +def _validate_bound_role_requirements( + bindings: Mapping[str, str], + *, + scene_requirements: Mapping[str, Any], + scene_objects: Sequence[Mapping[str, Any]], + robot_profile: str, +) -> None: + """Ensure an explicit binding does not contradict its static sidecar.""" + from embodichain.gen_sim.action_engine.tasks.assembly import SceneInventory + + requirements = _requirements_by_role(scene_requirements) + entities = SceneInventory(scene_objects, robot_profile=robot_profile).by_uid + for role, uid in bindings.items(): + requirement = requirements.get(role) + if requirement is None: + raise ValueError( + "Sibling SceneRequirements is missing TaskSpec role " f"{role!r}." + ) + entity = entities.get(uid) + if entity is None: + raise ValueError( + f"TaskSpec role {role!r} binds unavailable scene UID {uid!r}." + ) + if not _entity_matches_requirement( + entity, + requirement, + require_complete_static_evidence=False, + ): + raise ValueError( + f"TaskSpec role {role!r} binding {uid!r} conflicts with its " + "SceneRequirements category, attributes, state, or affordances." + ) + + +def _requirements_by_role( + scene_requirements: Mapping[str, Any], +) -> dict[str, Mapping[str, Any]]: + objects = scene_requirements.get("objects", []) + if not isinstance(objects, Sequence) or isinstance(objects, (str, bytes)): + raise ValueError("SceneRequirements.objects must be a list.") + result: dict[str, Mapping[str, Any]] = {} + for requirement in objects: + if not isinstance(requirement, Mapping): + raise ValueError("SceneRequirements.objects must contain mappings.") + role = requirement.get("role_id") + if not isinstance(role, str) or not role: + raise ValueError("SceneRequirements role_id must be a non-empty string.") + result[role] = requirement + return result + + +def _entity_matches_requirement( + entity: Any, + requirement: Mapping[str, Any], + *, + require_complete_static_evidence: bool, +) -> bool: + """Match explicit metadata; UID inference requires complete evidence.""" + category = requirement.get("category") + expected_category = category.strip().casefold() if isinstance(category, str) else "" + actual_category = str(entity.category).strip().casefold() + if expected_category: + if not actual_category: + if require_complete_static_evidence: + return False + elif expected_category != actual_category: + return False + required_affordances = requirement.get("affordances", []) + if not isinstance(required_affordances, Sequence) or isinstance( + required_affordances, (str, bytes) + ): + return False + expected_affordances = { + str(value).strip().casefold() for value in required_affordances + } + if ( + expected_affordances + and (require_complete_static_evidence or entity.affordances) + and not expected_affordances.issubset(entity.affordances) + ): + return False + expected_attributes = requirement.get("attributes", {}) + if not isinstance(expected_attributes, Mapping): + return False + for name, expected in expected_attributes.items(): + if not _static_attribute_matches( + entity, + str(name), + expected, + require_complete_static_evidence=require_complete_static_evidence, + ): + return False + expected_state = requirement.get("initial_state", {}) + if not isinstance(expected_state, Mapping): + return False + missing = object() + for name, expected in expected_state.items(): + actual = entity.initial_state.get(str(name), missing) + if actual is missing: + if require_complete_static_evidence: + return False + elif actual != expected: + return False + return True + + +def _static_attribute_matches( + entity: Any, + name: str, + expected: Any, + *, + require_complete_static_evidence: bool, +) -> bool: + """Compare one requirement against explicit exported metadata only.""" + marker = object() + actual = entity.color if name == "color" else entity.attributes.get(name, marker) + if actual is marker or actual is None or actual == "": + return not require_complete_static_evidence + if name == "color" and isinstance(actual, str) and isinstance(expected, str): + return actual.strip().casefold() == expected.strip().casefold() + return actual == expected + + +def _task_spec_role_references(value: Any, key: str = "") -> set[str]: + if isinstance(value, Mapping): + return { + role + for child_key, child in value.items() + for role in _task_spec_role_references(child, str(child_key)) + } + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + return { + role for child in value for role in _task_spec_role_references(child, key) + } + if isinstance(value, str) and (key.endswith("_role") or key.endswith("_roles")): + return {value} + return set() + + +def _with_role_bindings( + task_spec: Mapping[str, Any], + role_bindings: Mapping[str, str], +) -> dict[str, Any]: + result = deepcopy(dict(task_spec)) + metadata = result.setdefault("metadata", {}) + if not isinstance(metadata, dict): + raise ValueError("TaskSpec.metadata must be a mapping.") + metadata["role_bindings"] = dict(sorted(role_bindings.items())) + return result + + +def _validate_requirement_roles( + requirements: Mapping[str, Any], + role_bindings: Mapping[str, str], +) -> None: + requirement_roles = { + str(item["role_id"]) + for item in requirements["objects"] + if isinstance(item, Mapping) + } + missing = sorted(set(role_bindings) - requirement_roles) + if missing: + raise ValueError( + "SceneRequirements is missing TaskSpec role bindings for " f"{missing}." + ) + + +def _scene_requirements_from_bindings( + task_id: str, + planner_objects: Sequence[Mapping[str, Any]], + role_bindings: Mapping[str, str], +) -> dict[str, Any]: + """Derive a minimal concrete SceneRequirements view for grounded roles.""" + source = _scene_requirements_from_scene(task_id, planner_objects) + by_uid = {str(item["role_id"]): item for item in source["objects"]} + objects = [] + for role, uid in sorted(role_bindings.items()): + requirement = by_uid.get(uid) + if requirement is None: + raise ValueError( + f"TaskSpec role {role!r} binds UID {uid!r}, which has no " + "source-scene requirement." + ) + resolved = deepcopy(requirement) + resolved["role_id"] = role + objects.append(resolved) + return { + "schema_version": SCENE_REQUIREMENTS_SCHEMA, + "task_id": task_id, + "objects": objects, + "cameras": [], + "spatial_constraints": [{"type": "preserve_source_scene"}], + "distractor_count": 0, + "metadata": {"source": "task_spec_role_bindings"}, + } + + +def _validated_mapping( + value: Any, + *, + validator: Any, + label: str, +) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise TypeError( + f"{label} producer returned {type(value).__name__}, not a mapping." + ) + candidate = deepcopy(dict(value)) + validated = validator(candidate) + if validated is None: + # Validators may either return a normalized mapping or validate in place. + validated = candidate + if not isinstance(validated, Mapping): + raise TypeError(f"{label} validator must return a mapping or None.") + return deepcopy(dict(validated)) + + +def _validate_agent_config(config: Mapping[str, Any]) -> None: + if config.get("schema_version") != ACTION_ENGINE_CONFIG_SCHEMA: + raise ValueError("Agent config has an unexpected schema_version.") + if config.get("task_spec") != TASK_SPEC_FILENAME: + raise ValueError("Agent config must point to the canonical TaskSpec.") + if config.get("scene_requirements") != SCENE_REQUIREMENTS_FILENAME: + raise ValueError("Agent config must point to canonical SceneRequirements.") + graph_path = config.get("seed_task_graph") + if ( + not isinstance(graph_path, str) + or Path(graph_path).name != EXECUTION_PROGRAM_FILENAME + ): + raise ValueError("Agent config must point to the canonical SeedGraph.") + planning_mode = config.get("planning_mode", "offline") + if planning_mode not in {"offline", "ab"}: + raise ValueError("Agent config planning_mode must be 'offline' or 'ab'.") + if planning_mode == "ab": + online = config.get("online_planning") + if not isinstance(online, Mapping): + raise ValueError("A/B agent config requires online_planning settings.") + camera_uids = online.get("camera_uids") + if camera_uids != list(VLM_CAMERA_UIDS): + raise ValueError( + "A/B agent config must list the canonical four VLM cameras." + ) + model = online.get("vlm_model") + if model is not None and (not isinstance(model, str) or not model.strip()): + raise ValueError("online_planning.vlm_model must be a string or null.") + if config.get("offline_seed_task_graph") != graph_path: + raise ValueError( + "A/B agent config offline_seed_task_graph must match seed_task_graph." + ) + if config.get("vlm_camera_uids") != camera_uids: + raise ValueError( + "A/B agent config vlm_camera_uids must match online_planning." + ) + if config.get("vlm_model") != model: + raise ValueError("A/B agent config vlm_model must match online_planning.") + resolve_agent_runtime_policy(config) + + +def _raise_if_outputs_exist( + output_dir: str | Path, + *, + overwrite: bool, + planning_mode: str = "offline", +) -> None: + if overwrite: + return + paths = artifact_paths(output_dir, planning_mode=planning_mode) + existing = [ + path + for path in ( + paths.gym_config, + paths.agent_config, + paths.task_spec, + paths.scene_requirements, + paths.seed_task_graph, + paths.seed_task_graph_png, + ) + if path.exists() + ] + if existing: + names = ", ".join(path.name for path in existing) + raise FileExistsError( + f"Generated artifacts already exist in {paths.gym_config.parent}: " + f"{names}. Pass --overwrite to replace them." + ) + + +def _scene_requirements_from_scene( + task_id: str, + planner_objects: Sequence[Mapping[str, Any]], +) -> dict[str, Any]: + objects = [] + for item in planner_objects: + uid = str(item.get("runtime_uid", item.get("uid", ""))).strip() + if not uid: + raise ValueError("Planner scene object is missing a runtime UID.") + role = str(item.get("role", "object")).strip().lower() + raw_category = item.get("category", item.get("object_category", "")) + category = str(raw_category).strip().lower() or role or "object" + raw_attributes = item.get("attributes", {}) + attributes = ( + deepcopy(dict(raw_attributes)) + if isinstance(raw_attributes, Mapping) + else {} + ) + color = item.get("color") + if color not in (None, ""): + attributes.setdefault("color", color) + objects.append( + { + "role_id": uid, + "category": category, + "count": 1, + "affordances": [], + "initial_state": {}, + "attributes": attributes, + } + ) + return { + "schema_version": SCENE_REQUIREMENTS_SCHEMA, + "task_id": task_id, + "objects": objects, + "cameras": [], + "spatial_constraints": [{"type": "preserve_source_scene"}], + "distractor_count": 0, + "metadata": {"source": "existing_gym_project"}, + } + + +def _add_ab_camera_requirements( + requirements: Mapping[str, Any], +) -> dict[str, Any]: + """Declare fixed multi-view inputs in the shared A/B hand-off.""" + from embodichain.gen_sim.action_engine.domain import validate_scene_requirements + + result = deepcopy(dict(requirements)) + cameras = result.get("cameras", []) + if not isinstance(cameras, list): + raise ValueError("SceneRequirements.cameras must be a list.") + existing_uids = { + str(item.get("uid")) + for item in cameras + if isinstance(item, Mapping) and item.get("uid") + } + for uid in VLM_CAMERA_UIDS: + if uid in existing_uids: + continue + cameras.append( + { + "uid": uid, + "role": "vlm_view", + "modalities": ["rgb", "depth"], + "coverage": "all_interaction_objects", + "resolution": [640, 480], + } + ) + result["cameras"] = cameras + metadata = result.setdefault("metadata", {}) + if not isinstance(metadata, dict): + metadata = {} + result["metadata"] = metadata + metadata["planning_mode"] = "ab" + metadata["vlm_camera_uids"] = list(VLM_CAMERA_UIDS) + return validate_scene_requirements(result) diff --git a/embodichain/gen_sim/action_engine/generation/models.py b/embodichain/gen_sim/action_engine/generation/models.py new file mode 100644 index 000000000..ba35566f2 --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/models.py @@ -0,0 +1,82 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Small value objects used by Action Engine config generation.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +__all__ = ["GeneratedConfigPaths", "PreparedScene"] + + +@dataclass(frozen=True) +class GeneratedConfigPaths: + """Paths written by one successful generation transaction.""" + + gym_config: Path + agent_config: Path + task_spec: Path + scene_requirements: Path + seed_task_graph: Path + seed_task_graph_png: Path + planning_mode: str = "offline" + + @property + def execution_program(self) -> Path: + """Retain the Python API alias for callers migrating to SeedGraph v3.""" + return self.seed_task_graph + + @property + def offline_seed_task_graph(self) -> Path: + """Explicit A/B alias for the immutable offline SeedGraph artifact.""" + return self.seed_task_graph + + @property + def seed_task_graph_path(self) -> Path: + """Path-style alias used by runtime config loaders.""" + return self.seed_task_graph + + @property + def offline_seed_task_graph_path(self) -> Path: + """Verbose alias for callers that distinguish A/B graph branches.""" + return self.seed_task_graph + + @property + def offline_seed_task_graph_png(self) -> Path: + """Explicit A/B alias for the review rendering of the offline graph.""" + return self.seed_task_graph_png + + +@dataclass(frozen=True) +class PreparedScene: + """A source scene normalized for both planning and simulator loading.""" + + source_config_path: Path + scene_dir: Path + planner_objects: tuple[dict[str, Any], ...] + background: tuple[dict[str, Any], ...] + rigid_objects: tuple[dict[str, Any], ...] + articulations: tuple[dict[str, Any], ...] + uid_map: dict[str, str] + table_top_z: float | None + z_rotation_degrees: float + body_scale_policy: str + body_scale: tuple[float, float, float] + asset_hashes: dict[str, str] + asset_provenance: tuple[dict[str, Any], ...] = () diff --git a/embodichain/gen_sim/action_engine/generation/source_scene.py b/embodichain/gen_sim/action_engine/generation/source_scene.py new file mode 100644 index 000000000..699acde40 --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/source_scene.py @@ -0,0 +1,600 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Read and normalize an exported Prompt2Scene source scene. + +The source scene remains the authority for object geometry and initial poses. +Generation only makes asset paths absolute, gives runtime objects stable UIDs, +applies one explicit world-frame rotation, and adds conservative physics values +needed by manipulation tasks. +""" + +from __future__ import annotations + +import hashlib +from collections.abc import Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass +import json +import math +from pathlib import Path +import re +from typing import Any +import warnings + +from embodichain.gen_sim.action_engine.config import generation_defaults + +from .models import PreparedScene + +__all__ = [ + "ResolvedSceneSource", + "is_prompt2scene_export", + "prepare_scene", + "resolve_gym_config_path", + "resolve_source_scene", +] + +_LEGACY_CONFIG_FILENAMES = ("gym_config_merged.json", "gym_config.json") +_SCENE_CONFIG_FILENAME = "scene_config.json" +_CONFIG_FILENAMES = (*_LEGACY_CONFIG_FILENAMES, _SCENE_CONFIG_FILENAME) +_EXPORT_DIRECTORY_NAMES = ("gym_export", "scene_export") +_LEGACY_GYM_FORMAT = "legacy_gym_config" +_SCENE_EXPORT_FORMAT = "embodichain.scene-export/v1" +_SCENE_SECTIONS = ("background", "rigid_object", "articulation") +_UID_SUFFIX_RE = re.compile(r"_0$") +_UID_INVALID_RE = re.compile(r"[^0-9A-Za-z_.-]+") + +_GENERATION_DEFAULTS = generation_defaults() +_SCENE_DEFAULTS = _GENERATION_DEFAULTS["scene"] +_PHYSICS_DEFAULTS = _GENERATION_DEFAULTS["physics"] +_BACKGROUND_POLICY = _PHYSICS_DEFAULTS["background"] +_RIGID_POLICY = _PHYSICS_DEFAULTS["rigid_object"] +_BACKGROUND_ATTRS = { + key: value + for key, value in _BACKGROUND_POLICY.items() + if key != "max_convex_hull_num" +} +_RIGID_ATTRS = { + key: value + for key, value in _RIGID_POLICY.items() + if key not in {"max_convex_hull_num", "acd_method"} +} +_DEFAULT_BODY_SCALE = tuple(float(value) for value in _SCENE_DEFAULTS["body_scale"]) + + +@dataclass(frozen=True) +class ResolvedSceneSource: + """One validated source-scene config selected from an export layout. + + Attributes: + path: Absolute path to the selected source configuration. + source_format: Stable identifier for the detected source schema. + is_prompt2scene: Whether Prompt2Scene world alignment should be applied. + """ + + path: Path + source_format: str + is_prompt2scene: bool + + +def resolve_source_scene(gym_project: str | Path) -> ResolvedSceneSource: + """Resolve and classify one supported source-scene configuration. + + Args: + gym_project: Task root, export directory, or explicit configuration path. + + Returns: + The selected path together with its source format and provenance. + + Raises: + FileNotFoundError: If no supported source configuration exists. + ValueError: If a config is unsupported or recursive discovery is ambiguous. + """ + input_path = Path(gym_project).expanduser().resolve() + if input_path.is_file(): + return _classify_source_config(input_path) + if not input_path.is_dir(): + raise FileNotFoundError(f"Scene project does not exist: {input_path}") + + for directory in ( + input_path, + *(input_path / name for name in _EXPORT_DIRECTORY_NAMES), + ): + preferred = _preferred_config(directory) + if preferred is not None: + return _classify_source_config(preferred) + + matches = sorted( + { + candidate.parent + for filename in _CONFIG_FILENAMES + for candidate in input_path.rglob(filename) + } + ) + preferred = [ + config + for directory in matches + if (config := _preferred_config(directory)) is not None + ] + if len(preferred) == 1: + return _classify_source_config(preferred[0]) + if not preferred: + expected = ", ".join(_CONFIG_FILENAMES) + raise FileNotFoundError( + f"No supported scene config ({expected}) found under: {input_path}" + ) + paths = ", ".join(path.as_posix() for path in preferred) + raise ValueError(f"Multiple exported scene configs found: {paths}") + + +def resolve_gym_config_path(gym_project: str | Path) -> Path: + """Return the selected config path for callers using the legacy API name.""" + return resolve_source_scene(gym_project).path + + +def is_prompt2scene_export(gym_project: str | Path) -> bool: + """Return whether the input has Prompt2Scene export provenance.""" + try: + return resolve_source_scene(gym_project).is_prompt2scene + except (FileNotFoundError, ValueError): + return False + + +def prepare_scene( + gym_project: str | Path, + *, + z_rotation_degrees: float | None = None, + body_scale_policy: str = str(_SCENE_DEFAULTS["body_scale_policy"]), + body_scale: Sequence[float] = _DEFAULT_BODY_SCALE, +) -> PreparedScene: + """Load a source config and return planner/runtime views of one scene.""" + scale_policy = str(body_scale_policy).strip().lower() + if scale_policy not in {"preserve", "multiply", "absolute"}: + raise ValueError("body_scale_policy must be preserve, multiply, or absolute.") + requested_scale = _vector3(body_scale) + if any(value <= 0.0 for value in requested_scale): + raise ValueError("body_scale values must be positive.") + resolved_source = resolve_source_scene(gym_project) + source_path = resolved_source.path + source = _read_json_object(source_path) + scene_dir = source_path.parent + source_entries = _collect_source_entries(source) + if not source_entries: + raise ValueError( + "Source scene config has no background, rigid_object, or articulation." + ) + + table_source_uid = _find_table_source_uid(source_entries) + uid_map = _make_uid_map(source_entries, table_source_uid=table_source_uid) + rotation = ( + float(_SCENE_DEFAULTS["prompt2scene_z_rotation_degrees"]) + if z_rotation_degrees is None and resolved_source.is_prompt2scene + else float(z_rotation_degrees or 0.0) + ) + + planner_objects: list[dict[str, Any]] = [] + runtime_sections: dict[str, list[dict[str, Any]]] = { + section: [] for section in _SCENE_SECTIONS + } + asset_hashes: dict[str, str] = {} + for role, source_config in source_entries: + source_uid = _require_uid(source_config, role=role) + normalized = deepcopy(source_config) + normalized["uid"] = uid_map[source_uid] + _make_asset_paths_absolute(normalized, scene_dir=scene_dir, role=role) + _normalize_pose_fields(normalized) + _apply_body_scale_policy( + normalized, + policy=scale_policy, + requested=requested_scale, + ) + _apply_world_z_rotation(normalized, rotation) + shape = normalized.get("shape") + if isinstance(shape, Mapping) and shape.get("fpath"): + asset_hashes[normalized["uid"]] = _file_hash(Path(str(shape["fpath"]))) + + planner_objects.append( + _planner_object( + normalized, + source_uid=source_uid, + role=role, + ) + ) + runtime_sections[role].append(_runtime_object(normalized, role=role)) + + table = next( + (obj for obj in runtime_sections["background"] if obj.get("uid") == "table"), + None, + ) + table_top_z = _estimate_mesh_top_z(table) if table is not None else None + return PreparedScene( + source_config_path=source_path, + scene_dir=scene_dir, + planner_objects=tuple(planner_objects), + background=tuple(runtime_sections["background"]), + rigid_objects=tuple(runtime_sections["rigid_object"]), + articulations=tuple(runtime_sections["articulation"]), + uid_map=uid_map, + table_top_z=table_top_z, + z_rotation_degrees=rotation, + body_scale_policy=scale_policy, + body_scale=tuple(requested_scale), + asset_hashes=asset_hashes, + ) + + +def _preferred_config(directory: Path) -> Path | None: + for filename in _CONFIG_FILENAMES: + candidate = directory / filename + if candidate.is_file(): + return candidate + return None + + +def _classify_source_config(path: Path) -> ResolvedSceneSource: + if path.name not in _CONFIG_FILENAMES: + expected = ", ".join(_CONFIG_FILENAMES) + raise ValueError(f"Expected one of {expected}, got: {path}") + if path.name == _SCENE_CONFIG_FILENAME: + source = _read_json_object(path) + source_format = source.get("format") + if source_format != _SCENE_EXPORT_FORMAT: + raise ValueError( + f"Scene config {path} has unsupported format {source_format!r}; " + f"expected {_SCENE_EXPORT_FORMAT!r}." + ) + return ResolvedSceneSource( + path=path, + source_format=_SCENE_EXPORT_FORMAT, + is_prompt2scene=True, + ) + return ResolvedSceneSource( + path=path, + source_format=_LEGACY_GYM_FORMAT, + is_prompt2scene=( + _has_legacy_prompt2scene_marker(path) or _has_scene_export_companion(path) + ), + ) + + +def _has_legacy_prompt2scene_marker(config_path: Path) -> bool: + config_dir = config_path.parent + directories = [config_dir, config_dir / "gym_export"] + return any( + (directory / "scene_state" / "result.json").is_file() + for directory in directories + ) + + +def _has_scene_export_companion(config_path: Path) -> bool: + config_dir = config_path.parent + candidates = [config_dir / _SCENE_CONFIG_FILENAME] + if config_dir.name == "gym_export": + candidates.append(config_dir.parent / "scene_export" / _SCENE_CONFIG_FILENAME) + else: + candidates.append(config_dir / "scene_export" / _SCENE_CONFIG_FILENAME) + return any(_is_scene_export_v1(candidate) for candidate in candidates) + + +def _is_scene_export_v1(path: Path) -> bool: + if not path.is_file(): + return False + try: + return _read_json_object(path).get("format") == _SCENE_EXPORT_FORMAT + except ValueError: + return False + + +def _read_json_object(path: Path) -> dict[str, Any]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise ValueError(f"Invalid JSON in source scene config {path}: {exc}") from exc + if not isinstance(value, dict): + raise ValueError(f"Source scene config must contain a JSON object: {path}") + return value + + +def _collect_source_entries( + source: Mapping[str, Any], +) -> list[tuple[str, dict[str, Any]]]: + entries: list[tuple[str, dict[str, Any]]] = [] + for section in _SCENE_SECTIONS: + value = source.get(section, []) + if isinstance(value, Mapping): + value = [value] + if not isinstance(value, list): + raise ValueError(f"Source scene section {section!r} must be a list.") + for config in value: + if not isinstance(config, Mapping): + raise ValueError(f"Entries in {section!r} must be JSON objects.") + entries.append((section, dict(config))) + return entries + + +def _find_table_source_uid(entries: Sequence[tuple[str, Mapping[str, Any]]]) -> str: + backgrounds = [config for role, config in entries if role == "background"] + if len(backgrounds) != 1: + raise ValueError( + "A tabletop action scene requires exactly one background object; " + f"found {len(backgrounds)}." + ) + return _require_uid(backgrounds[0], role="background") + + +def _make_uid_map( + entries: Sequence[tuple[str, Mapping[str, Any]]], + *, + table_source_uid: str, +) -> dict[str, str]: + uid_map: dict[str, str] = {} + used: set[str] = set() + for role, config in entries: + source_uid = _require_uid(config, role=role) + if source_uid in uid_map: + raise ValueError(f"Duplicate scene object UID: {source_uid!r}") + candidate = ( + "table" if source_uid == table_source_uid else _normalize_uid(source_uid) + ) + runtime_uid = candidate + suffix = 2 + while runtime_uid in used: + runtime_uid = f"{candidate}_{suffix}" + suffix += 1 + uid_map[source_uid] = runtime_uid + used.add(runtime_uid) + return uid_map + + +def _normalize_uid(source_uid: str) -> str: + candidate = _UID_SUFFIX_RE.sub("", source_uid.strip()) + candidate = _UID_INVALID_RE.sub("_", candidate).strip("._-") + if not candidate: + raise ValueError(f"Cannot derive a runtime UID from {source_uid!r}.") + if candidate[0].isdigit(): + candidate = f"object_{candidate}" + return candidate + + +def _require_uid(config: Mapping[str, Any], *, role: str) -> str: + uid = str(config.get("uid", "")).strip() + if not uid: + raise ValueError(f"Scene object in {role!r} has no UID.") + return uid + + +def _make_asset_paths_absolute( + config: dict[str, Any], + *, + scene_dir: Path, + role: str, +) -> None: + shape = config.get("shape") + if isinstance(shape, Mapping): + normalized_shape = deepcopy(dict(shape)) + fpath = normalized_shape.get("fpath") + if fpath: + normalized_shape["fpath"] = _resolve_asset_path( + scene_dir, str(fpath) + ).as_posix() + config["shape"] = normalized_shape + if role == "articulation" and config.get("fpath"): + config["fpath"] = _resolve_asset_path( + scene_dir, str(config["fpath"]) + ).as_posix() + + +def _resolve_asset_path(scene_dir: Path, fpath: str) -> Path: + raw = Path(fpath).expanduser() + resolved = raw.resolve() if raw.is_absolute() else (scene_dir / raw).resolve() + if not resolved.is_file(): + raise FileNotFoundError(f"Scene asset does not exist: {resolved}") + return resolved + + +def _normalize_pose_fields(config: dict[str, Any]) -> None: + config["init_pos"] = _vector3(config.get("init_pos", [0.0, 0.0, 0.0])) + config["init_rot"] = _vector3(config.get("init_rot", [0.0, 0.0, 0.0])) + if "body_scale" in config: + scale = _vector3(config["body_scale"]) + if any(value <= 0.0 for value in scale): + raise ValueError( + f"Object {config.get('uid')!r} has non-positive body_scale." + ) + config["body_scale"] = scale + + +def _apply_body_scale_policy( + config: dict[str, Any], + *, + policy: str, + requested: Sequence[float], +) -> None: + source = _vector3(config.get("body_scale", [1.0, 1.0, 1.0])) + if policy == "preserve": + result = source + elif policy == "multiply": + result = [left * right for left, right in zip(source, requested)] + else: + result = list(requested) + config["body_scale"] = [_clean_float(value) for value in result] + + +def _file_hash(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _apply_world_z_rotation(config: dict[str, Any], degrees: float) -> None: + if math.isclose(degrees, 0.0, abs_tol=1e-12): + return + theta = math.radians(degrees) + cos_theta, sin_theta = math.cos(theta), math.sin(theta) + x, y, z = _vector3(config["init_pos"]) + config["init_pos"] = [ + _clean_float(x * cos_theta - y * sin_theta), + _clean_float(x * sin_theta + y * cos_theta), + _clean_float(z), + ] + + # EmbodiChain and Prompt2Scene both interpret these values as intrinsic XYZ. + from scipy.spatial.transform import Rotation + + original = Rotation.from_euler("XYZ", config["init_rot"], degrees=True) + world_z = Rotation.from_rotvec([0.0, 0.0, theta]) + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="Gimbal lock detected") + rotated = (world_z * original).as_euler("XYZ", degrees=True) + config["init_rot"] = [_clean_float(value) for value in rotated] + if "init_local_pose" in config: + # Keeping two pose representations risks the stale local matrix + # overriding the rotated Euler pose in ObjectBaseCfg.from_dict. + del config["init_local_pose"] + + +def _planner_object( + config: Mapping[str, Any], + *, + source_uid: str, + role: str, +) -> dict[str, Any]: + description = str(config.get("description", "")).strip() + shape = deepcopy(dict(config.get("shape", {}))) + raw_attributes = config.get("attributes", {}) + if not isinstance(raw_attributes, Mapping): + raw_attributes = {} + raw_initial_state = config.get("initial_state", config.get("state", {})) + if not isinstance(raw_initial_state, Mapping): + raw_initial_state = {} + raw_affordances = config.get("affordances", config.get("capabilities", [])) + affordances = ( + [str(value) for value in raw_affordances] + if isinstance(raw_affordances, Sequence) + and not isinstance(raw_affordances, (str, bytes)) + else [] + ) + return { + "uid": str(config["uid"]), + "runtime_uid": str(config["uid"]), + "source_uid": source_uid, + "role": role, + "name": str(config.get("name", "")).strip(), + "description": description, + "shape": shape, + "init_pos": list(config["init_pos"]), + "init_rot": list(config["init_rot"]), + "body_scale": list(config.get("body_scale", [1.0, 1.0, 1.0])), + "category": config.get("category", config.get("object_category", "")), + "color": config.get("color", raw_attributes.get("color")), + "attributes": deepcopy(dict(raw_attributes)), + "initial_state": deepcopy(dict(raw_initial_state)), + "affordances": affordances, + } + + +def _runtime_object(config: Mapping[str, Any], *, role: str) -> dict[str, Any]: + if role == "articulation": + # Articulation schemas vary by asset; preserve their source fields after + # path and pose normalization instead of guessing a reduced schema. + result = deepcopy(dict(config)) + result.pop("description", None) + return result + + result = { + key: deepcopy(config[key]) + for key in ( + "uid", + "shape", + "init_pos", + "init_rot", + "body_scale", + ) + if key in config + } + result.setdefault("body_scale", [1.0, 1.0, 1.0]) + source_attrs = dict(config.get("attrs", {})) + if role == "background": + result["attrs"] = {**source_attrs, **_BACKGROUND_ATTRS} + result["body_type"] = "kinematic" + result["max_convex_hull_num"] = int(_BACKGROUND_POLICY["max_convex_hull_num"]) + else: + result["attrs"] = {**source_attrs, **_RIGID_ATTRS} + result["body_type"] = "dynamic" + hull_limit = int(_RIGID_POLICY["max_convex_hull_num"]) + max_hulls = max( + 1, + min(int(config.get("max_convex_hull_num", hull_limit)), hull_limit), + ) + result["max_convex_hull_num"] = max_hulls + result["acd_method"] = str(_RIGID_POLICY["acd_method"]) + shape = result.get("shape") + if isinstance(shape, dict): + shape["acd_method"] = str(_RIGID_POLICY["acd_method"]) + shape["max_convex_hull_num"] = max_hulls + return result + + +def _estimate_mesh_top_z(config: Mapping[str, Any]) -> float | None: + shape = config.get("shape", {}) + if not isinstance(shape, Mapping) or not shape.get("fpath"): + return None + try: + import numpy as np + import trimesh + from scipy.spatial.transform import Rotation + + loaded = trimesh.load(str(shape["fpath"]), force="scene") + geometry = ( + loaded.to_geometry() + if hasattr(loaded, "to_geometry") + else loaded.dump(concatenate=True) + ) + vertices = np.asarray(geometry.vertices, dtype=np.float64) + if vertices.size == 0: + return None + # DexSim converts glTF Y-up vertices to its Z-up basis at load time. + sim_vertices = np.column_stack( + (vertices[:, 0], -vertices[:, 2], vertices[:, 1]) + ) + sim_vertices *= np.asarray( + config.get("body_scale", [1.0, 1.0, 1.0]), dtype=np.float64 + ) + rotated = Rotation.from_euler( + "XYZ", config.get("init_rot", [0.0, 0.0, 0.0]), degrees=True + ).apply(sim_vertices) + rotated += np.asarray(config.get("init_pos", [0.0, 0.0, 0.0]), dtype=np.float64) + return float(rotated[:, 2].max()) + except Exception: + # Mesh bounds improve robot placement but are not needed to preserve the + # exported scene. The robot builder has a conservative tabletop fallback. + return None + + +def _vector3(value: Any) -> list[float]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise ValueError(f"Expected a finite xyz vector, got: {value!r}") + values = [float(item) for item in value] + if len(values) != 3 or not all(math.isfinite(item) for item in values): + raise ValueError(f"Expected a finite xyz vector, got: {value!r}") + return values + + +def _clean_float(value: float) -> float: + rounded = round(float(value), 12) + return 0.0 if abs(rounded) < 1e-12 else rounded diff --git a/embodichain/gen_sim/action_engine/generation/templates/default_lights.json b/embodichain/gen_sim/action_engine/generation/templates/default_lights.json new file mode 100644 index 000000000..5ea73ee5b --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/templates/default_lights.json @@ -0,0 +1,3 @@ +{ + "direct": [] +} diff --git a/embodichain/gen_sim/action_engine/generation/templates/default_sensors.json b/embodichain/gen_sim/action_engine/generation/templates/default_sensors.json new file mode 100644 index 000000000..f9ad7aea8 --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/templates/default_sensors.json @@ -0,0 +1,14 @@ +[ + { + "sensor_type": "Camera", + "width": 960, + "height": 540, + "intrinsics": [420, 420, 480, 270], + "extrinsics": { + "pos": [0.4, 0.0, 2.2], + "eye": [-0.6, 0.0, 1.8], + "target": [0.0, 0.0, 0.75], + "up": [1.0, 0.0, 0.0] + } + } +] diff --git a/embodichain/gen_sim/action_engine/generation/templates/dual_franka_robot.json b/embodichain/gen_sim/action_engine/generation/templates/dual_franka_robot.json new file mode 100644 index 000000000..b5709f40d --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/templates/dual_franka_robot.json @@ -0,0 +1,185 @@ +{ + "uid": "DualFrankaPanda", + "urdf_cfg": { + "fname": "dual_franka_panda_basket", + "name_case": { + "joint": "original", + "link": "original" + }, + "components": [ + { + "component_type": "left_arm", + "urdf_path": "Franka/Panda/Panda.urdf", + "transform": [ + [1.0, 0.0, 0.0, -1.25], + [0.0, 1.0, 0.0, 0.3], + [0.0, 0.0, 1.0, 0.4], + [0.0, 0.0, 0.0, 1.0] + ] + }, + { + "component_type": "left_hand", + "urdf_path": "Robotiq/robotiq_arg2f_140/robotiq_arg2f_140.urdf", + "transform": [ + [0.0, 1.0, 0.0, 0.0], + [-1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0] + ] + }, + { + "component_type": "right_arm", + "urdf_path": "Franka/Panda/Panda.urdf", + "transform": [ + [1.0, 0.0, 0.0, -1.25], + [0.0, 1.0, 0.0, -0.3], + [0.0, 0.0, 1.0, 0.4], + [0.0, 0.0, 0.0, 1.0] + ] + }, + { + "component_type": "right_hand", + "urdf_path": "Robotiq/robotiq_arg2f_140/robotiq_arg2f_140.urdf", + "transform": [ + [0.0, 1.0, 0.0, 0.0], + [-1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0] + ] + } + ] + }, + "init_pos": [-0.7, 0.0, 0.0], + "init_rot": [0.0, 0.0, 180.0], + "init_qpos": [ + 0.0, + 0.0, + -0.569, + -0.569, + 0.0, + 0.0, + -2.81, + -2.81, + 0.0, + 0.0, + 3.037, + 3.037, + 0.0, + 0.0, + + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ], + "drive_pros": { + "stiffness": { + "left_arm": 10000.0, + "right_arm": 10000.0, + "left_eef": 50.0, + "right_eef": 50.0 + }, + "damping": { + "left_arm": 1000.0, + "right_arm": 1000.0, + "left_eef": 5.0, + "right_eef": 5.0 + }, + "max_effort": { + "left_arm": 10000.0, + "right_arm": 10000.0, + "left_eef": 500.0, + "right_eef": 500.0 + } + }, + "control_parts": { + "left_arm": [ + "left_fr3_joint1", + "left_fr3_joint2", + "left_fr3_joint3", + "left_fr3_joint4", + "left_fr3_joint5", + "left_fr3_joint6", + "left_fr3_joint7" + ], + "left_eef": [ + "left_finger_joint", + "left_inner_knuckle_joint", + "left_inner_finger_joint", + "left_right_outer_knuckle_joint", + "left_right_inner_knuckle_joint", + "left_right_inner_finger_joint" + ], + "right_arm": [ + "right_fr3_joint1", + "right_fr3_joint2", + "right_fr3_joint3", + "right_fr3_joint4", + "right_fr3_joint5", + "right_fr3_joint6", + "right_fr3_joint7" + ], + "right_eef": [ + "right_finger_joint", + "right_left_inner_knuckle_joint", + "right_left_inner_finger_joint", + "right_outer_knuckle_joint", + "right_inner_knuckle_joint", + "right_inner_finger_joint" + ], + "dual_arm": [ + "left_fr3_joint1", + "left_fr3_joint2", + "left_fr3_joint3", + "left_fr3_joint4", + "left_fr3_joint5", + "left_fr3_joint6", + "left_fr3_joint7", + "right_fr3_joint1", + "right_fr3_joint2", + "right_fr3_joint3", + "right_fr3_joint4", + "right_fr3_joint5", + "right_fr3_joint6", + "right_fr3_joint7" + ] + }, + "observation_joint_parts": ["left_eef", "right_eef"], + "qpos_control_part_order": ["dual_arm", "left_eef", "right_eef"], + "solver_cfg": { + "left_arm": { + "class_type": "PytorchSolver", + "urdf_path": null, + "end_link_name": "left_fr3_link8", + "root_link_name": "left_base", + "tcp": [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.2], + [0.0, 0.0, 0.0, 1.0] + ], + "num_samples": 15 + }, + "right_arm": { + "class_type": "PytorchSolver", + "urdf_path": null, + "end_link_name": "right_fr3_link8", + "root_link_name": "right_base", + "tcp": [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.2], + [0.0, 0.0, 0.0, 1.0] + ], + "num_samples": 15 + } + } +} diff --git a/embodichain/gen_sim/action_engine/generation/templates/dual_ur_robot.json b/embodichain/gen_sim/action_engine/generation/templates/dual_ur_robot.json new file mode 100644 index 000000000..8a7496547 --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/templates/dual_ur_robot.json @@ -0,0 +1,126 @@ +{ + "uid": "DualUR5", + "urdf_cfg": { + "fname": "dual_ur5_robotiq_arg2f_140_basket", + "name_case": {"joint": "lower", "link": "lower"}, + "components": [ + { + "component_type": "left_arm", + "urdf_path": "UniversalRobots/UR5/UR5.urdf", + "transform": [ + [1.0, 0.0, 0.0, -1.45], + [0.0, 1.0, 0.0, -0.3], + [0.0, 0.0, 1.0, 0.4], + [0.0, 0.0, 0.0, 1.0] + ] + }, + { + "component_type": "left_hand", + "urdf_path": "Robotiq/robotiq_arg2f_140/robotiq_arg2f_140.urdf", + "transform": [ + [0.0, 1.0, 0.0, 0.0], + [-1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0] + ] + }, + { + "component_type": "right_arm", + "urdf_path": "UniversalRobots/UR5/UR5.urdf", + "transform": [ + [1.0, 0.0, 0.0, -1.45], + [0.0, 1.0, 0.0, 0.3], + [0.0, 0.0, 1.0, 0.4], + [0.0, 0.0, 0.0, 1.0] + ] + }, + { + "component_type": "right_hand", + "urdf_path": "Robotiq/robotiq_arg2f_140/robotiq_arg2f_140.urdf", + "transform": [ + [0.0, 1.0, 0.0, 0.0], + [-1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0] + ] + } + ] + }, + "init_pos": [2.0, 0.0, 0.0], + "init_rot": [0.0, 0.0, 0.0], + "init_qpos": [ + 0, 0, -1.57, -1.57, 1.57, 1.57, -1.57, -1.57, + -1.57, -1.57, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 + ], + "drive_pros": { + "stiffness": { + "left_arm": 10000.0, + "right_arm": 10000.0, + "left_eef": 50.0, + "right_eef": 50.0 + }, + "damping": { + "left_arm": 1000.0, + "right_arm": 1000.0, + "left_eef": 5.0, + "right_eef": 5.0 + }, + "max_effort": { + "left_arm": 10000.0, + "right_arm": 10000.0, + "left_eef": 500.0, + "right_eef": 500.0 + } + }, + "control_parts": { + "left_arm": [ + "left_joint1", "left_joint2", "left_joint3", + "left_joint4", "left_joint5", "left_joint6" + ], + "left_eef": [ + "left_finger_joint", "left_inner_knuckle_joint", + "left_inner_finger_joint", "left_right_outer_knuckle_joint", + "left_right_inner_knuckle_joint", "left_right_inner_finger_joint" + ], + "right_arm": [ + "right_joint1", "right_joint2", "right_joint3", + "right_joint4", "right_joint5", "right_joint6" + ], + "right_eef": [ + "right_finger_joint", "right_left_inner_knuckle_joint", + "right_left_inner_finger_joint", "right_outer_knuckle_joint", + "right_inner_knuckle_joint", "right_inner_finger_joint" + ] + }, + "solver_cfg": { + "left_arm": { + "class_type": "URSolver", + "ur_type": "ur5", + "urdf_path": null, + "end_link_name": "left_ee_link", + "root_link_name": "left_base_link", + "ik_nearest_weight": [1.0, 4.0, 1.0, 1.0, 1.0, 1.0], + "tcp": [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.2], + [0.0, 0.0, 0.0, 1.0] + ] + }, + "right_arm": { + "class_type": "URSolver", + "ur_type": "ur5", + "urdf_path": null, + "end_link_name": "right_ee_link", + "root_link_name": "right_base_link", + "ik_nearest_weight": [1.0, 4.0, 1.0, 1.0, 1.0, 1.0], + "tcp": [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.2], + [0.0, 0.0, 0.0, 1.0] + ] + } + } +} diff --git a/embodichain/gen_sim/action_engine/generation/templates/robot_profiles.json b/embodichain/gen_sim/action_engine/generation/templates/robot_profiles.json new file mode 100644 index 000000000..31084a497 --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/templates/robot_profiles.json @@ -0,0 +1,45 @@ +{ + "dual_franka": { + "aliases": ["franka", "panda", "dual_panda", "dual_franka_panda"], + "template": "dual_franka_robot.json", + "robot_family": "franka", + "tabletop_clearance": 0.05, + "arm_component_z": 0.4, + "arm_base_x": -1.45, + "gripper_open_state": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + "gripper_close_state": [0.7, -0.7, 0.7, -0.7, -0.7, 0.7] + }, + "dual_ur3": { + "aliases": ["ur3", "dual_ur3_dh_pgi", "dual_ur3_robotiq", "dual_ur3_robotiq_arg2f_140"], + "template": "dual_ur_robot.json", + "robot_family": "ur3", + "tabletop_clearance": 0.05, + "arm_component_z": 0.4, + "arm_base_x": -1.45, + "max_effort": 56.0, + "gripper_open_state": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + "gripper_close_state": [0.7, -0.7, 0.7, -0.7, -0.7, 0.7] + }, + "dual_ur5": { + "aliases": ["ur5", "dual_ur5_dh_pgi", "dual_ur5_robotiq", "dual_ur5_robotiq_arg2f_140"], + "template": "dual_ur_robot.json", + "robot_family": "ur5", + "tabletop_clearance": 0.05, + "arm_component_z": 0.4, + "arm_base_x": -1.45, + "max_effort": 10000.0, + "gripper_open_state": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + "gripper_close_state": [0.7, -0.7, 0.7, -0.7, -0.7, 0.7] + }, + "dual_ur10": { + "aliases": ["ur10", "dual_ur10_dh_pgi", "dual_ur10_robotiq", "dual_ur10_robotiq_arg2f_140"], + "template": "dual_ur_robot.json", + "robot_family": "ur10", + "tabletop_clearance": 0.05, + "arm_component_z": 0.3, + "arm_base_x": -1.1, + "max_effort": 330.0, + "gripper_open_state": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + "gripper_close_state": [0.7, -0.7, 0.7, -0.7, -0.7, 0.7] + } +} diff --git a/embodichain/gen_sim/action_engine/generation/templates/vlm_sensors.json b/embodichain/gen_sim/action_engine/generation/templates/vlm_sensors.json new file mode 100644 index 000000000..34e964569 --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/templates/vlm_sensors.json @@ -0,0 +1,58 @@ +[ + { + "uid": "vlm_front", + "sensor_type": "Camera", + "width": 640, + "height": 480, + "intrinsics": [520, 520, 320, 240], + "enable_color": true, + "enable_depth": true, + "extrinsics": { + "eye": [-1.2, 0.0, 1.65], + "target": [0.0, 0.0, 0.75], + "up": [0.0, 0.0, 1.0] + } + }, + { + "uid": "vlm_left", + "sensor_type": "Camera", + "width": 640, + "height": 480, + "intrinsics": [520, 520, 320, 240], + "enable_color": true, + "enable_depth": true, + "extrinsics": { + "eye": [0.0, 1.2, 1.65], + "target": [0.0, 0.0, 0.75], + "up": [0.0, 0.0, 1.0] + } + }, + { + "uid": "vlm_rear", + "sensor_type": "Camera", + "width": 640, + "height": 480, + "intrinsics": [520, 520, 320, 240], + "enable_color": true, + "enable_depth": true, + "extrinsics": { + "eye": [1.2, 0.0, 1.65], + "target": [0.0, 0.0, 0.75], + "up": [0.0, 0.0, 1.0] + } + }, + { + "uid": "vlm_right", + "sensor_type": "Camera", + "width": 640, + "height": 480, + "intrinsics": [520, 520, 320, 240], + "enable_color": true, + "enable_depth": true, + "extrinsics": { + "eye": [0.0, -1.2, 1.65], + "target": [0.0, 0.0, 0.75], + "up": [0.0, 0.0, 1.0] + } + } +] diff --git a/embodichain/gen_sim/action_engine/graph_visualization.py b/embodichain/gen_sim/action_engine/graph_visualization.py new file mode 100644 index 000000000..4ad262c60 --- /dev/null +++ b/embodichain/gen_sim/action_engine/graph_visualization.py @@ -0,0 +1,938 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Headless PNG rendering for direct AtomicAction SeedGraphs. + +The renderer consumes the same validated coordinate-free v3 graph as runtime, +then builds an internal display view without grounding symbolic targets. E +TaskGroups remain the semantic grouping labels over the rendered action nodes. +Single chains use a folded timeline; DAGs use stable actor swimlanes. +""" + +from __future__ import annotations + +from collections import defaultdict +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from functools import lru_cache +from io import BytesIO +from math import hypot +from typing import Any + +import matplotlib + +# Select the non-interactive backend before importing any canvas primitives. +matplotlib.use("Agg", force=True) + +from matplotlib import patheffects +from matplotlib.backends.backend_agg import FigureCanvasAgg +from matplotlib.font_manager import FontProperties, fontManager +from matplotlib.figure import Figure +from matplotlib.patches import Circle, FancyArrowPatch, FancyBboxPatch +import networkx as nx + +from embodichain.gen_sim.action_engine.domain import ( + EXECUTION_PROGRAM_SCHEMA, + validate_execution_program, +) +from embodichain.gen_sim.action_engine.protocol import SEED_GRAPH_SCHEMA + +__all__ = ["render_seed_task_graph_png", "render_task_graph_png"] + +_PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n" +_EXECUTION_KEYS = frozenset( + { + "schema_version", + "task", + "goal_description", + "start", + "goal", + "nodes", + "edges", + "semantic_steps", + "allocation_groups", + "motion_policy_version", + } +) + +_BACKGROUND = "#F8FAFB" +_INK = "#17212B" +_MUTED = "#66727D" +_BORDER = "#CBD4DC" +_LEFT = "#168A78" +_RIGHT = "#D97706" +_AUTO = "#59636D" +_COORDINATED = "#7652A5" +_DEPENDENCY = "#8A94A0" + +# The figures are designed at this display width in inches; every type size +# below is chosen to stay readable when the PNG is shown at exactly this size. +_TARGET_WIDTH = 8.0 +_DPI = 300 +_LEVEL_STEP = 1.15 +_NODE_RADIUS = 0.16 +_SPECIAL_NODE_RADIUS = 0.20 +_SUCCESS = "#25834B" +_FAILED = "#C43E3E" +_SKIPPED = "#8B949C" +_LANE_COLORS = { + "left": _LEFT, + "auto": _AUTO, + "right": _RIGHT, + "coordinated": _COORDINATED, +} +_LANE_BACKGROUNDS = { + "left": "#EAF6F3", + "auto": "#F0F3F5", + "right": "#FFF4E6", +} +_LANE_LABELS = { + "left": "LEFT ARM [L]", + "auto": "WORLD / AUTO / COORDINATED", + "right": "RIGHT ARM [R]", +} +_STATUS_COLORS = { + "success": _SUCCESS, + "executed": _SUCCESS, + "failed": _FAILED, + "aborted": _FAILED, + "skipped": _SKIPPED, +} +_STATUS_BADGES = { + "success": "OK", + "executed": "OK", + "failed": "FAIL", + "aborted": "ABORT", + "skipped": "SKIP", +} + + +@dataclass(frozen=True) +class _RuntimeOverlay: + """Execution annotations kept separate from the immutable seed program.""" + + edge_status: Mapping[str, str] + edge_arm: Mapping[str, str] + step_status: Mapping[str, str] + graph_status: str | None = None + + +@dataclass(frozen=True) +class _GraphData: + """Validated program plus indices shared by both layout strategies.""" + + program: Mapping[str, Any] + graph: nx.MultiDiGraph + node_by_id: Mapping[str, Mapping[str, Any]] + edge_by_id: Mapping[str, Mapping[str, Any]] + step_by_id: Mapping[str, Mapping[str, Any]] + lane_override: Mapping[str, str] + runtime: _RuntimeOverlay + + +def render_seed_task_graph_png(seed_graph: Mapping[str, Any]) -> bytes: + """Render a v3 SeedGraph or package-owned legacy program through Agg.""" + program = _display_program(seed_graph) + return _render(program, _RuntimeOverlay({}, {}, {})) + + +def render_task_graph_png(task_graph: Mapping[str, Any]) -> bytes: + """Render an execution program with optional runtime event annotations. + + A bare program is accepted. Runtime events may be stored in its ``runtime`` + envelope, or beside a nested ``execution_program``, ``program``, or + ``seed_task_graph``. A record alone is rejected because it omits topology. + """ + program = _extract_execution_program(task_graph) + runtime = _extract_runtime_overlay(task_graph) + return _render(program, runtime) + + +def _render( + program: Mapping[str, Any], + runtime: _RuntimeOverlay, +) -> bytes: + data = _graph_data(program, runtime) + if _is_single_chain(data): + return _render_chain(data) + return _render_dag(data) + + +def _extract_execution_program(document: Mapping[str, Any]) -> dict[str, Any]: + """Find and validate the execution program embedded in a display document.""" + if not isinstance(document, Mapping): + raise ValueError("Task graph visualization input must be a mapping.") + + if document.get("schema_version") in {EXECUTION_PROGRAM_SCHEMA, SEED_GRAPH_SCHEMA}: + # A runtime artifact may preserve the program fields and add annotations. + if document.get("schema_version") == SEED_GRAPH_SCHEMA: + candidate = dict(document) + candidate.pop("runtime", None) + candidate.pop("runtime_record", None) + return _display_program(candidate) + candidate = {key: document[key] for key in _EXECUTION_KEYS if key in document} + return validate_execution_program(candidate) + + for key in ("execution_program", "program", "seed_task_graph"): + candidate = document.get(key) + if isinstance(candidate, Mapping): + return _display_program(candidate) + + # Supporting a full program plus a runtime schema at the top level keeps + # visualization useful for simple JSON joins without weakening validation. + if {"nodes", "edges", "semantic_steps"}.issubset(document): + candidate = {key: document[key] for key in _EXECUTION_KEYS if key in document} + return validate_execution_program(candidate) + + raise ValueError( + "Runtime records do not contain graph topology. Provide the matching " + "ExecutionProgram under 'execution_program', 'program', or " + "'seed_task_graph'." + ) + + +def _display_program(value: Mapping[str, Any]) -> dict[str, Any]: + if value.get("schema_version") == SEED_GRAPH_SCHEMA: + from embodichain.gen_sim.action_engine.compiler import ( + seed_graph_to_execution_program, + ) + + return seed_graph_to_execution_program(value, require_executable=False) + return validate_execution_program(value) + + +def _extract_runtime_overlay(document: Mapping[str, Any]) -> _RuntimeOverlay: + """Reduce a runtime record to the small set of display-only annotations.""" + record = document.get("runtime") + if record is None: + record = document.get("runtime_record", document) + if not isinstance(record, Mapping): + raise ValueError("runtime_record must be a mapping.") + raw_events = record.get("events", document.get("events", [])) + if not isinstance(raw_events, Sequence) or isinstance( + raw_events, (str, bytes, bytearray) + ): + raise ValueError("Runtime events must be a list.") + + edge_status: dict[str, str] = {} + edge_arm: dict[str, str] = {} + step_status: dict[str, str] = {} + for index, event in enumerate(raw_events): + if not isinstance(event, Mapping): + raise ValueError(f"Runtime events[{index}] must be a mapping.") + event_kind = event.get("event") + status = _optional_text(event.get("status")) + if event_kind == "edge": + edge_id = _optional_text(event.get("edge_id")) + if edge_id and status: + edge_status[edge_id] = status.lower() + arm = _optional_text(event.get("arm")) + if edge_id and arm: + edge_arm[edge_id] = arm + elif event_kind == "semantic_step": + step_id = _optional_text(event.get("semantic_step_id")) + if step_id and status: + step_status[step_id] = status.lower() + + graph_status = _optional_text(record.get("status")) + return _RuntimeOverlay( + edge_status=edge_status, + edge_arm=edge_arm, + step_status=step_status, + graph_status=graph_status.lower() if graph_status else None, + ) + + +def _graph_data( + program: Mapping[str, Any], + runtime: _RuntimeOverlay, +) -> _GraphData: + node_by_id = {str(node["id"]): node for node in program["nodes"]} + edge_by_id = {str(edge["id"]): edge for edge in program["edges"]} + step_by_id = {str(step["id"]): step for step in program["semantic_steps"]} + graph = nx.MultiDiGraph() + graph.add_nodes_from(node_by_id) + for edge in program["edges"]: + source = str(edge["source"]) + target = str(edge["target"]) + graph.add_edge(source, target, edge_id=str(edge["id"])) + if not nx.is_directed_acyclic_graph(graph): + raise ValueError("ExecutionProgram node topology must be a directed DAG.") + + return _GraphData( + program=program, + graph=graph, + node_by_id=node_by_id, + edge_by_id=edge_by_id, + step_by_id=step_by_id, + lane_override=_allocation_lane_overrides(program), + runtime=runtime, + ) + + +def _allocation_lane_overrides( + program: Mapping[str, Any], +) -> dict[str, str]: + """Give auto actors stable lanes when a distinct-arm group is declared.""" + result: dict[str, str] = {} + for group in program.get("allocation_groups", []): + if group.get("arm_constraint") != "distinct_arms": + continue + members = group.get("semantic_step_ids", []) + for index, step_id in enumerate(members): + result[str(step_id)] = "left" if index % 2 == 0 else "right" + return result + + +def _is_single_chain(data: _GraphData) -> bool: + graph = data.graph + if graph.number_of_edges() != graph.number_of_nodes() - 1: + return False + if any(graph.in_degree(node) > 1 for node in graph): + return False + if any(graph.out_degree(node) > 1 for node in graph): + return False + return ( + graph.in_degree(str(data.program["start"])) == 0 + and graph.out_degree(str(data.program["goal"])) == 0 + and nx.is_weakly_connected(graph) + ) + + +def _ordered_chain_edges(data: _GraphData) -> list[Mapping[str, Any]]: + current = str(data.program["start"]) + result: list[Mapping[str, Any]] = [] + while current != str(data.program["goal"]): + outgoing = list(data.graph.out_edges(current, data=True)) + if len(outgoing) != 1: + raise ValueError("ExecutionProgram chain has an incomplete path.") + _, target, attrs = outgoing[0] + result.append(data.edge_by_id[str(attrs["edge_id"])]) + current = str(target) + if len(result) != len(data.edge_by_id): + raise ValueError("ExecutionProgram chain does not cover every edge.") + return result + + +def _render_chain(data: _GraphData) -> bytes: + """Render a long linear program as a bounded, folded state timeline.""" + edges = _ordered_chain_edges(data) + nodes = [str(data.program["start"])] + nodes.extend(str(edge["target"]) for edge in edges) + + slots_per_row = 4 + row_count = (len(nodes) + slots_per_row - 1) // slots_per_row + width = _TARGET_WIDTH + height = max(3.4, 2.0 + row_count * 1.55) + figure, axis = _new_figure(width, height) + try: + _draw_header(axis, data, width) + left, right, first_y = 0.7, width - 0.7, 2.05 + spacing = (right - left) / (slots_per_row - 1) + positions: dict[str, tuple[float, float]] = {} + for index, node_id in enumerate(nodes): + row, column = divmod(index, slots_per_row) + visual_column = column if row % 2 == 0 else slots_per_row - 1 - column + positions[node_id] = ( + left + visual_column * spacing, + first_y + row * 1.55, + ) + + for edge in edges: + source = positions[str(edge["source"])] + target = positions[str(edge["target"])] + lane = _edge_lane(edge, data) + color = _edge_color(str(edge["id"]), lane, data.runtime) + label_position, label_align = _edge_label_position(source, target, width) + _draw_labeled_edge( + axis, + source, + target, + color=color, + label=_edge_label(edge, data), + label_position=label_position, + label_align=label_align, + ) + + for index, node_id in enumerate(nodes): + _draw_state_node( + axis, + positions[node_id], + start=node_id == str(data.program["start"]), + goal=node_id == str(data.program["goal"]), + fork=False, + join=False, + index=index, + ) + + _draw_legend(axis, width, height - 0.28) + return _figure_png_bytes(figure) + finally: + figure.clear() + + +def _render_dag(data: _GraphData) -> bytes: + """Render forks and joins against persistent actor swimlanes.""" + levels = _dag_levels(data.graph) + maximum_level = max(levels.values(), default=0) + width = _TARGET_WIDTH + height = max(5.2, 2.15 + maximum_level * _LEVEL_STEP + 1.35) + figure, axis = _new_figure(width, height) + try: + _draw_header(axis, data, width) + boundaries, lane_centers = _lane_geometry(width) + _draw_swimlanes(axis, height, boundaries, lane_centers) + positions = _dag_positions(data, levels, lane_centers) + + # Dependency arrows are drawn first and stay visually subordinate to + # physical state transitions; only constraints not already implied by + # the state topology are shown. + for source_id, target_id in _visible_dependencies(data): + _draw_dependency_arrow( + axis, + positions[source_id], + positions[target_id], + ) + + pair_groups: defaultdict[tuple[str, str], list[str]] = defaultdict(list) + for edge in data.edge_by_id.values(): + pair_groups[(str(edge["source"]), str(edge["target"]))].append( + str(edge["id"]) + ) + for edge in data.edge_by_id.values(): + edge_id = str(edge["id"]) + source_id = str(edge["source"]) + target_id = str(edge["target"]) + lane = _edge_lane(edge, data) + parallel_ids = pair_groups[(source_id, target_id)] + parallel_index = parallel_ids.index(edge_id) + curvature = (parallel_index - (len(parallel_ids) - 1) / 2.0) * 0.20 + label_position, label_align = _edge_label_position( + positions[source_id], + positions[target_id], + width, + ) + _draw_labeled_edge( + axis, + positions[source_id], + positions[target_id], + color=_edge_color(edge_id, lane, data.runtime), + label=_edge_label(edge, data), + label_position=label_position, + label_align=label_align, + curvature=curvature, + ) + + for index, node_id in enumerate(nx.topological_sort(data.graph)): + _draw_state_node( + axis, + positions[str(node_id)], + start=str(node_id) == str(data.program["start"]), + goal=str(node_id) == str(data.program["goal"]), + fork=data.graph.out_degree(node_id) > 1, + join=data.graph.in_degree(node_id) > 1, + index=index, + ) + + _draw_legend(axis, width, height - 0.28) + return _figure_png_bytes(figure) + finally: + figure.clear() + + +def _lane_geometry( + width: float, +) -> tuple[dict[str, tuple[float, float]], dict[str, float]]: + """Even thirds for lane boundaries with derived actor centers.""" + margin = 0.35 + area = width - 2 * margin + first = margin + area / 3.0 + second = margin + 2 * area / 3.0 + boundaries = { + "left": (margin, first), + "auto": (first, second), + "right": (second, width - margin), + } + centers = {lane: (left + right) / 2.0 for lane, (left, right) in boundaries.items()} + return boundaries, centers + + +def _edge_label_position( + source: tuple[float, float], + target: tuple[float, float], + width: float, +) -> tuple[tuple[float, float], str]: + """Place halo labels beside arrows instead of boxing them on the edge.""" + midpoint = _midpoint(source, target) + dx = target[0] - source[0] + dy = target[1] - source[1] + if abs(dx) < 0.3: + # Keep vertical-arrow labels inside the canvas: right side on the left + # half of the figure, left side on the right half. + if midpoint[0] > width / 2.0: + return (midpoint[0] - 0.14, midpoint[1]), "right" + return (midpoint[0] + 0.14, midpoint[1]), "left" + length = hypot(dx, dy) or 1.0 + normal_x, normal_y = dy / length, -dx / length + if normal_x < 0: + normal_x, normal_y = -normal_x, -normal_y + if abs(normal_x) < 0.2 and normal_y > 0: + # Horizontal arrows keep their label above the line in both directions. + normal_x, normal_y = -normal_x, -normal_y + return ( + (midpoint[0] + normal_x * 0.16, midpoint[1] + normal_y * 0.16), + "center", + ) + + +def _visible_dependencies(data: _GraphData) -> list[tuple[str, str]]: + """Node anchors for dependencies not implied by state continuity.""" + result: list[tuple[str, str]] = [] + for prerequisite_id, dependent_id in _dependency_pairs(data): + source = str(data.edge_by_id[prerequisite_id]["target"]) + target = str(data.edge_by_id[dependent_id]["source"]) + if source == target or nx.has_path(data.graph, source, target): + continue + result.append((source, target)) + return result + + +def _dag_levels(graph: nx.MultiDiGraph) -> dict[str, int]: + """Assign the longest-path depth so dependencies always flow downward.""" + levels: dict[str, int] = {} + for node in nx.topological_sort(graph): + predecessors = list(graph.predecessors(node)) + levels[str(node)] = ( + max(levels[str(parent)] for parent in predecessors) + 1 + if predecessors + else 0 + ) + return levels + + +def _dag_positions( + data: _GraphData, + levels: Mapping[str, int], + lane_centers: Mapping[str, float], +) -> dict[str, tuple[float, float]]: + """Place branch nodes in actor lanes and structural fork/join nodes centrally.""" + base: dict[str, tuple[str, int]] = {} + for node_id in data.node_by_id: + incoming = list(data.graph.in_edges(node_id, data=True)) + outgoing = list(data.graph.out_edges(node_id, data=True)) + if ( + node_id in {str(data.program["start"]), str(data.program["goal"])} + or len(incoming) > 1 + or len(outgoing) > 1 + ): + lane = "auto" + elif incoming: + edge = data.edge_by_id[str(incoming[0][2]["edge_id"])] + lane = _edge_lane(edge, data) + elif outgoing: + edge = data.edge_by_id[str(outgoing[0][2]["edge_id"])] + lane = _edge_lane(edge, data) + else: + lane = "auto" + if lane == "coordinated": + lane = "auto" + base[node_id] = (lane, levels[node_id]) + + groups: defaultdict[tuple[str, int], list[str]] = defaultdict(list) + for node_id, lane_level in base.items(): + groups[lane_level].append(node_id) + + result: dict[str, tuple[float, float]] = {} + for (lane, level), node_ids in groups.items(): + ordered = sorted(node_ids) + center = lane_centers[lane] + # Small symmetric offsets prevent same-level nodes from hiding each + # other while keeping every node visibly inside its actor lane. + offsets = [ + (index - (len(ordered) - 1) / 2.0) * 0.55 for index in range(len(ordered)) + ] + for node_id, offset in zip(ordered, offsets, strict=True): + result[node_id] = (center + offset, 2.15 + level * _LEVEL_STEP) + return result + + +def _dependency_pairs(data: _GraphData) -> list[tuple[str, str]]: + """Return explicit edge dependencies plus missing semantic dependencies.""" + result: list[tuple[str, str]] = [] + seen: set[tuple[str, str]] = set() + for edge in data.edge_by_id.values(): + dependent_id = str(edge["id"]) + for prerequisite_id in edge.get("depends_on", []): + pair = (str(prerequisite_id), dependent_id) + if pair not in seen: + seen.add(pair) + result.append(pair) + + for step in data.step_by_id.values(): + dependent_edges = step.get("edge_ids", []) + if not dependent_edges: + continue + for prerequisite_step_id in step.get("depends_on", []): + prerequisite = data.step_by_id[str(prerequisite_step_id)] + pair = ( + str(prerequisite["edge_ids"][-1]), + str(dependent_edges[0]), + ) + if pair not in seen: + seen.add(pair) + result.append(pair) + return result + + +def _edge_lane(edge: Mapping[str, Any], data: _GraphData) -> str: + edge_id = str(edge["id"]) + observed_arm = data.runtime.edge_arm.get(edge_id) + if observed_arm: + return _arm_lane(observed_arm) + + action_lanes = { + _actor_lane(action.get("actor", {})) for action in edge.get("actions", []) + } + action_lanes.discard("auto") + if action_lanes == {"left"}: + return "left" + if action_lanes == {"right"}: + return "right" + if "coordinated" in action_lanes or action_lanes == {"left", "right"}: + return "coordinated" + return data.lane_override.get(str(edge["semantic_step_id"]), "auto") + + +def _actor_lane(actor: Any) -> str: + if not isinstance(actor, Mapping): + return "auto" + mode = str(actor.get("mode", "auto")).lower() + if mode == "required": + return _arm_lane(str(actor.get("arm", ""))) + if mode == "coordinated": + return "coordinated" + return "auto" + + +def _arm_lane(arm: str) -> str: + normalized = arm.strip().lower() + if "left" in normalized: + return "left" + if "right" in normalized: + return "right" + if normalized in {"both", "coordinated", "dual_arm", "dual"}: + return "coordinated" + return "auto" + + +def _edge_color( + edge_id: str, + lane: str, + runtime: _RuntimeOverlay, +) -> str: + status = runtime.edge_status.get(edge_id) + return _STATUS_COLORS.get(status or "", _LANE_COLORS[lane]) + + +def _edge_label(edge: Mapping[str, Any], data: _GraphData) -> str: + """One-line semantic phrase; execution details live in the JSON artifacts.""" + edge_id = str(edge["id"]) + step = data.step_by_id[str(edge["semantic_step_id"])] + status = data.runtime.edge_status.get(edge_id) or data.runtime.step_status.get( + str(step["id"]) + ) + status_badge = f" [{_STATUS_BADGES.get(status, status.upper())}]" if status else "" + return _clip(f"{step['operator']}: {step['object']}", 40) + status_badge + + +def _draw_header(axis: Any, data: _GraphData, width: float) -> None: + status = data.runtime.graph_status + status_text = f" [{status.upper()}]" if status else "" + axis.text( + 0.4, + 0.42, + _clip(f"ACTION ENGINE / {data.program['task']}{status_text}", 84), + ha="left", + va="center", + color=_INK, + fontproperties=_font(10.0, "bold"), + zorder=20, + ) + axis.text( + 0.4, + 0.80, + _clip(str(data.program["goal_description"]), 115), + ha="left", + va="top", + color=_MUTED, + fontproperties=_font(7.0), + linespacing=1.25, + zorder=20, + ) + axis.plot( + [0.4, width - 0.4], + [1.28, 1.28], + color=_BORDER, + linewidth=0.7, + zorder=19, + ) + + +def _draw_swimlanes( + axis: Any, + height: float, + boundaries: Mapping[str, tuple[float, float]], + centers: Mapping[str, float], +) -> None: + for lane in ("left", "auto", "right"): + left, right = boundaries[lane] + axis.add_patch( + FancyBboxPatch( + (left, 1.50), + right - left, + height - 2.0, + boxstyle="round,pad=0.0,rounding_size=0.05", + facecolor=_LANE_BACKGROUNDS[lane], + edgecolor=_BORDER, + linewidth=0.6, + zorder=-10, + ) + ) + axis.plot( + [left, right], + [1.50, 1.50], + color=_LANE_COLORS[lane], + linewidth=1.1, + zorder=-9, + ) + axis.text( + centers[lane], + 1.74, + _LANE_LABELS[lane], + ha="center", + va="center", + color=_LANE_COLORS[lane], + fontproperties=_font(6.8, "bold"), + zorder=10, + ) + + +def _draw_legend(axis: Any, width: float, y: float) -> None: + """Single-row edge-type legend; START/GOAL labels are self-explanatory.""" + entries = ( + ("left", "left action", False), + ("right", "right action", False), + ("coordinated", "coordinated", False), + ("auto", "auto / world", False), + ("dependency", "dependency", True), + ) + slot = 1.32 + start = (width - slot * len(entries)) / 2.0 + for index, (key, label, dashed) in enumerate(entries): + x = start + index * slot + color = _DEPENDENCY if dashed else _LANE_COLORS[key] + axis.add_patch( + FancyArrowPatch( + (x, y), + (x + 0.3, y), + arrowstyle="-|>", + mutation_scale=7, + color=color, + linewidth=1.0, + linestyle=(0, (3.0, 2.6)) if dashed else "-", + zorder=20, + ) + ) + axis.text( + x + 0.38, + y, + label, + ha="left", + va="center", + color=_MUTED, + fontproperties=_font(6.2), + zorder=20, + ) + + +def _draw_labeled_edge( + axis: Any, + source: tuple[float, float], + target: tuple[float, float], + *, + color: str, + label: str, + label_position: tuple[float, float], + label_align: str = "center", + curvature: float = 0.0, +) -> None: + """Draw one solid state transition and its halo-backed one-line label.""" + axis.add_patch( + FancyArrowPatch( + source, + target, + arrowstyle="-|>", + mutation_scale=9, + color=color, + linewidth=1.15, + shrinkA=12, + shrinkB=12, + connectionstyle=f"arc3,rad={curvature}", + zorder=3, + ) + ) + axis.text( + *label_position, + label, + ha=label_align, + va="center", + color=_INK, + fontproperties=_font(6.5), + path_effects=[patheffects.withStroke(linewidth=1.7, foreground=_BACKGROUND)], + zorder=8, + ) + + +def _draw_dependency_arrow( + axis: Any, + source: tuple[float, float], + target: tuple[float, float], +) -> None: + if source == target: + return + axis.add_patch( + FancyArrowPatch( + source, + target, + arrowstyle="-|>", + mutation_scale=7, + color=_DEPENDENCY, + linewidth=0.9, + linestyle=(0, (3.0, 2.6)), + shrinkA=12, + shrinkB=12, + connectionstyle="arc3,rad=-0.2", + alpha=0.9, + zorder=1, + ) + ) + + +def _draw_state_node( + axis: Any, + center: tuple[float, float], + *, + start: bool, + goal: bool, + fork: bool, + join: bool, + index: int, +) -> None: + fill = "#DDEFEA" if start else ("#E7F2DD" if goal else "#FFFFFF") + edge = _SUCCESS if goal else (_LEFT if start else _INK) + radius = _SPECIAL_NODE_RADIUS if (start or goal or fork or join) else _NODE_RADIUS + axis.add_patch( + Circle( + center, + radius=radius, + facecolor=fill, + edgecolor=edge, + linewidth=1.1, + zorder=12, + ) + ) + axis.text( + center[0], + center[1], + str(index), + ha="center", + va="center", + color=_INK, + fontproperties=_font(6.5, "bold"), + zorder=13, + ) + role = ( + "START" + if start + else ("GOAL" if goal else ("FORK" if fork else "JOIN" if join else "")) + ) + if role: + axis.text( + center[0], + center[1] + radius + 0.12, + role, + ha="center", + va="top", + color=edge, + fontproperties=_font(6.0, "bold"), + zorder=13, + ) + + +def _new_figure(width: float, height: float) -> tuple[Figure, Any]: + figure = Figure(figsize=(width, height), dpi=_DPI, facecolor=_BACKGROUND) + axis = figure.subplots() + axis.set_facecolor(_BACKGROUND) + axis.set_axis_off() + axis.set_xlim(0.0, width) + axis.set_ylim(height, 0.0) + return figure, axis + + +def _figure_png_bytes(figure: Figure) -> bytes: + buffer = BytesIO() + FigureCanvasAgg(figure).print_png(buffer) + payload = buffer.getvalue() + if not payload.startswith(_PNG_SIGNATURE): + raise RuntimeError("Matplotlib did not produce a valid PNG payload.") + return payload + + +@lru_cache(maxsize=1) +def _font_family() -> str: + """Prefer a CJK-capable font while retaining a portable fallback.""" + available = {font.name for font in fontManager.ttflist} + for family in ( + "Noto Sans CJK SC", + "Noto Sans CJK JP", + "Source Han Sans CN", + "WenQuanYi Micro Hei", + "Microsoft YaHei", + "Arial Unicode MS", + "DejaVu Sans", + ): + if family in available: + return family + return "sans-serif" + + +def _font(size: float, weight: str = "normal") -> FontProperties: + return FontProperties(family=_font_family(), size=size, weight=weight) + + +def _optional_text(value: Any) -> str | None: + return value.strip() if isinstance(value, str) and value.strip() else None + + +def _clip(value: str, length: int) -> str: + return value if len(value) <= length else f"{value[: max(1, length - 3)]}..." + + +def _midpoint( + first: tuple[float, float], + second: tuple[float, float], +) -> tuple[float, float]: + return ((first[0] + second[0]) / 2.0, (first[1] + second[1]) / 2.0) diff --git a/embodichain/gen_sim/action_engine/orientation.py b/embodichain/gen_sim/action_engine/orientation.py new file mode 100644 index 000000000..3cac92e27 --- /dev/null +++ b/embodichain/gen_sim/action_engine/orientation.py @@ -0,0 +1,226 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Compile task-facing orientation goals into a small runtime contract.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +import math +from typing import Any + +__all__ = [ + "AlignAxisConstraint", + "MatchRotationConstraint", + "OrientationConstraint", + "compile_orientation_constraint", +] + +_LONG_AXES = frozenset({"long", "long_axis", "longest"}) +_SCOPES = frozenset({"terminal"}) + + +@dataclass(frozen=True) +class AlignAxisConstraint: + """Require one local object axis to align with a target axis.""" + + local_axis: str + target_axis: str = "world_up" + directed: bool = True + tolerance: float | None = None + scope: str = "terminal" + + +@dataclass(frozen=True) +class MatchRotationConstraint: + """Require a complete object rotation relative to a captured reference.""" + + reference: str + equivalence: str = "none" + tolerance: float | None = None + scope: str = "terminal" + + +OrientationTerm = AlignAxisConstraint | MatchRotationConstraint + + +@dataclass(frozen=True) +class OrientationConstraint: + """Canonical hard constraints plus a separate planning preference.""" + + terms: tuple[OrientationTerm, ...] + planning_preference: str = "minimize_rotation_from_current" + + @property + def requires_reference(self) -> bool: + """Return whether execution must capture a step-start rotation.""" + return any( + isinstance(term, MatchRotationConstraint) and term.reference == "step_start" + for term in self.terms + ) + + @property + def allows_upright_yaw_search(self) -> bool: + """Return whether all hard terms leave world-up yaw unconstrained.""" + return bool(self.terms) and all( + isinstance(term, AlignAxisConstraint) and term.target_axis == "world_up" + for term in self.terms + ) + + +def compile_orientation_constraint( + goal: Mapping[str, Any], +) -> OrientationConstraint: + """Compile legacy goal enums or a composable serialized constraint. + + Existing persisted graphs continue to carry explicit ``orientation_goal`` + values. New tasks may omit the field, which intentionally means no hard + orientation constraint while retaining a minimum-rotation preference. + """ + serialized = goal.get("orientation_constraint") + if serialized is not None: + return _compile_serialized(serialized) + + orientation_goal = str(goal.get("orientation_goal", "none")) + if orientation_goal == "none": + terms: tuple[OrientationTerm, ...] = () + elif orientation_goal == "preserve": + terms = (MatchRotationConstraint(reference="step_start"),) + elif orientation_goal == "upright": + local_axis = str(goal.get("upright_local_axis", "long_axis")) + if local_axis == "auto": + local_axis = "long_axis" + directed = goal.get( + "orientation_directed", local_axis.lower() not in _LONG_AXES + ) + if not isinstance(directed, bool): + raise ValueError("orientation_directed must be a boolean.") + terms = ( + AlignAxisConstraint( + local_axis=local_axis, + target_axis="world_up", + directed=directed, + ), + ) + elif orientation_goal in {"lay_flat", "axis_align"}: + # These established modes materialize a full target rotation. Keep that + # contract until semantic face/axis metadata can express narrower terms. + terms = (MatchRotationConstraint(reference="target_pose"),) + else: + raise ValueError(f"Unsupported orientation_goal {orientation_goal!r}.") + return OrientationConstraint(terms=terms) + + +def _compile_serialized(value: Any) -> OrientationConstraint: + if not isinstance(value, Mapping): + raise ValueError("orientation_constraint must be a mapping.") + unknown = set(value) - {"terms", "planning_preference"} + if unknown: + raise ValueError( + "orientation_constraint contains unsupported fields: " f"{sorted(unknown)}." + ) + raw_terms = value.get("terms", ()) + if not isinstance(raw_terms, Sequence) or isinstance( + raw_terms, (str, bytes, bytearray) + ): + raise ValueError("orientation_constraint.terms must be a list.") + terms = tuple(_compile_term(item, index) for index, item in enumerate(raw_terms)) + preference = str(value.get("planning_preference", "minimize_rotation_from_current")) + if preference not in {"minimize_rotation_from_current", "none"}: + raise ValueError( + "orientation_constraint.planning_preference must be " + "'minimize_rotation_from_current' or 'none'." + ) + return OrientationConstraint(terms=terms, planning_preference=preference) + + +def _compile_term(value: Any, index: int) -> OrientationTerm: + context = f"orientation_constraint.terms[{index}]" + if not isinstance(value, Mapping): + raise ValueError(f"{context} must be a mapping.") + kind = str(value.get("type", "")) + scope = str(value.get("scope", "terminal")) + if scope not in _SCOPES: + raise ValueError( + f"{context}.scope {scope!r} is unsupported by the current runtime." + ) + if kind == "align_axis": + unknown = set(value) - { + "type", + "local_axis", + "target_axis", + "directed", + "tolerance", + "scope", + } + if unknown: + raise ValueError( + f"{context} contains unsupported fields: {sorted(unknown)}." + ) + local_axis = str(value.get("local_axis", "")) + if local_axis not in {"x", "y", "z", "long_axis"}: + raise ValueError(f"{context}.local_axis {local_axis!r} is unsupported.") + target_axis = str(value.get("target_axis", "world_up")) + if target_axis != "world_up": + raise ValueError(f"{context}.target_axis {target_axis!r} is unsupported.") + directed = value.get("directed", True) + if not isinstance(directed, bool): + raise ValueError(f"{context}.directed must be a boolean.") + return AlignAxisConstraint( + local_axis=local_axis, + target_axis=target_axis, + directed=directed, + tolerance=_optional_tolerance(value, context), + scope=scope, + ) + if kind == "match_rotation": + unknown = set(value) - { + "type", + "reference", + "equivalence", + "tolerance", + "scope", + } + if unknown: + raise ValueError( + f"{context} contains unsupported fields: {sorted(unknown)}." + ) + reference = str(value.get("reference", "")) + if reference not in {"step_start", "target_pose"}: + raise ValueError(f"{context}.reference {reference!r} is unsupported.") + equivalence = str(value.get("equivalence", "none")) + if equivalence != "none": + raise ValueError(f"{context}.equivalence {equivalence!r} is unsupported.") + return MatchRotationConstraint( + reference=reference, + equivalence=equivalence, + tolerance=_optional_tolerance(value, context), + scope=scope, + ) + raise ValueError(f"{context}.type {kind!r} is unsupported.") + + +def _optional_tolerance(value: Mapping[str, Any], context: str) -> float | None: + raw = value.get("tolerance") + if raw is None: + return None + if isinstance(raw, bool) or not isinstance(raw, (int, float)): + raise ValueError(f"{context}.tolerance must be a finite positive number.") + tolerance = float(raw) + if not math.isfinite(tolerance) or tolerance <= 0.0: + raise ValueError(f"{context}.tolerance must be a finite positive number.") + return tolerance diff --git a/embodichain/gen_sim/action_engine/planning/__init__.py b/embodichain/gen_sim/action_engine/planning/__init__.py new file mode 100644 index 000000000..a02f8ce3b --- /dev/null +++ b/embodichain/gen_sim/action_engine/planning/__init__.py @@ -0,0 +1,62 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Stable route-free task planning API.""" + +from __future__ import annotations + +from .online import plan_online_seed_graph +from .dual import CandidatePair, plan_candidates_parallel +from .linker import ( + CONTRACT_LINKER_VERSION, + link_seed_graph, + link_task_dependencies, + validate_persisted_contracts, +) +from .planner import plan_task +from .selection import ( + CandidateEvaluation, + evaluate_candidate, + fuse_seed_graphs, + select_seed_graph, +) +from .vision import ( + CameraObservation, + SceneObservation, + analyze_visual_scene, + collect_scene_observation, + validate_visual_facts, +) + +__all__ = [ + "CONTRACT_LINKER_VERSION", + "CameraObservation", + "CandidatePair", + "CandidateEvaluation", + "SceneObservation", + "analyze_visual_scene", + "collect_scene_observation", + "evaluate_candidate", + "fuse_seed_graphs", + "link_seed_graph", + "link_task_dependencies", + "plan_online_seed_graph", + "plan_candidates_parallel", + "plan_task", + "select_seed_graph", + "validate_visual_facts", + "validate_persisted_contracts", +] diff --git a/embodichain/gen_sim/action_engine/planning/dual.py b/embodichain/gen_sim/action_engine/planning/dual.py new file mode 100644 index 000000000..d852464fd --- /dev/null +++ b/embodichain/gen_sim/action_engine/planning/dual.py @@ -0,0 +1,218 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Parallel offline/online candidate planning with isolated task views.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from concurrent.futures import ThreadPoolExecutor +from copy import deepcopy +from dataclasses import dataclass +from time import perf_counter +from typing import Any + +from embodichain.gen_sim.action_engine.capabilities import ( + AtomicCapabilityRegistry, + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.domain import ( + public_task_spec, + seed_graph_hash, + validate_seed_graph, + validate_task_spec, +) +from embodichain.gen_sim.action_engine.runtime.motion_policy import ( + resolve_motion_policy, +) + +from .linker import validate_persisted_contracts + +__all__ = ["CandidatePair", "plan_candidates_parallel"] + +CandidatePlanner = Callable[..., Mapping[str, Any]] + + +@dataclass(frozen=True) +class CandidatePair: + """Two independently planned graphs and branch-local planning metrics.""" + + offline: dict[str, Any] + online: dict[str, Any] + planning_metrics: dict[str, dict[str, Any]] + + +def plan_candidates_parallel( + task_spec: Mapping[str, Any], + *, + offline_planner: CandidatePlanner, + online_planner: CandidatePlanner, + known_objects: set[str] | None = None, + robot_profile: str = "dual_ur10", + registry: AtomicCapabilityRegistry | None = None, + require_executable: bool = False, +) -> CandidatePair: + """Plan both routes concurrently while hiding the oracle from online. + + Both returned graphs are validated against the same capability catalog and + motion-policy table before the pair is published. ``require_executable`` + is intentionally opt-in here: product planning may retain planning-only + candidates for inspection, while strict A/B execution enables the flag in + its final preflight. + """ + task = validate_task_spec(task_spec) + online_view = public_task_spec(task) + _reject_private_or_live_fields(online_view, "PublicTaskSpec") + capabilities = registry or build_atomic_capability_registry() + + def invoke(route: str) -> tuple[dict[str, Any], float]: + planner = offline_planner if route == "offline" else online_planner + # A planner is user/LLM supplied code. Give each route a detached + # copy so accidental mutation cannot change the other route's input or + # reintroduce private oracle fields after validation. + planner_input = deepcopy(task if route == "offline" else online_view) + started = perf_counter() + try: + result = planner(task_spec=planner_input) + except Exception as exc: + raise RuntimeError(f"{route} planner failed: {exc}") from exc + elapsed = perf_counter() - started + _reject_private_or_live_fields(result, f"{route} SeedGraph") + graph = validate_seed_graph( + result, + known_objects=known_objects, + known_actions=capabilities.names(), + executable_actions=capabilities.executable_names(), + require_executable=require_executable, + ) + if graph["planner_route"] != route: + raise ValueError( + f"{route} planner returned route {graph['planner_route']!r}." + ) + if graph["task_id"] != task["task_id"]: + raise ValueError(f"{route} planner returned a graph for another task.") + if graph["level"] != task["level"]: + raise ValueError(f"{route} planner returned a graph for another level.") + if graph["reasoning_type"] != task["reasoning_type"]: + raise ValueError( + f"{route} planner returned a graph with incompatible reasoning_type." + ) + if graph["capability_catalog_hash"] != capabilities.catalog_hash(): + raise ValueError( + f"{route} SeedGraph capability catalog does not match runtime." + ) + validate_persisted_contracts(graph, capabilities) + _validate_task_group_coverage(task, graph, route=route) + for node in graph["nodes"]: + capabilities.validate_binding(node) + if capabilities.get(str(node["atomic_action"])).runtime_available: + resolve_motion_policy( + robot_profile, + node["atomic_action"], + node["motion_policy"], + ) + return graph, elapsed + + with ThreadPoolExecutor( + max_workers=2, thread_name_prefix="action-engine-plan" + ) as pool: + futures = {route: pool.submit(invoke, route) for route in ("offline", "online")} + results: dict[str, tuple[dict[str, Any], float]] = {} + for route, future in futures.items(): + try: + results[route] = future.result() + except Exception as exc: + # Do not expose a bare Future exception; callers need to know + # which route invalidated the pair before any environment is + # allowed to move. + for other_route, other in futures.items(): + if other_route != route: + other.cancel() + raise RuntimeError( + f"A/B {route} planning/preflight failed: {exc}" + ) from exc + + metrics = { + route: { + "planning_seconds": elapsed, + "vlm_call_count": int(graph.get("metadata", {}).get("vlm_call_count", 0)), + "seed_graph_hash": seed_graph_hash(graph), + "node_count": len(graph["nodes"]), + "task_group_count": len(graph["task_groups"]), + } + for route, (graph, elapsed) in results.items() + } + return CandidatePair( + offline=results["offline"][0], + online=results["online"][0], + planning_metrics=metrics, + ) + + +def _validate_task_group_coverage( + task: Mapping[str, Any], graph: Mapping[str, Any], *, route: str +) -> None: + """Ensure every explicit L1-L3 task instance has one complete group.""" + if task.get("level") == "L4": + # L4's reference instances are intentionally hidden from the online + # route; the graph validator still enforces non-empty, coherent groups. + 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", ())} + missing = expected - actual + unexpected = actual - expected + if missing or unexpected: + raise ValueError( + f"{route} SeedGraph TaskGroup coverage mismatch; " + f"missing={sorted(missing)}, unexpected={sorted(unexpected)}." + ) + + +_PRIVATE_OR_LIVE_KEYS = frozenset( + { + "absolute_position", + "coordinates", + "grasp_pose", + "joint_positions", + "live_pose", + "live_transform", + "object_pose", + "oracle", + "pose", + "positions", + "qpos", + "target_pose", + "trajectory", + "waypoints", + "xpos", + } +) + + +def _reject_private_or_live_fields(value: Any, context: str) -> None: + """Reject private-oracle and grounded state fields in online inputs/outputs.""" + 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}]") diff --git a/embodichain/gen_sim/action_engine/planning/linker.py b/embodichain/gen_sim/action_engine/planning/linker.py new file mode 100644 index 000000000..c1db21e84 --- /dev/null +++ b/embodichain/gen_sim/action_engine/planning/linker.py @@ -0,0 +1,1018 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Deterministic causal and resource linking for SeedGraph v3.""" + +from __future__ import annotations + +from collections.abc import Collection, Mapping, Sequence +from copy import deepcopy +import hashlib +import json +from typing import Any + +from embodichain.gen_sim.action_engine.capabilities import ( + AtomicCapabilityRegistry, + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.domain import ( + validate_seed_graph, + validate_task_spec, +) +from embodichain.gen_sim.action_engine.protocol import SEED_GRAPH_SCHEMA + +__all__ = [ + "CONTRACT_LINKER_VERSION", + "link_seed_graph", + "link_task_dependencies", + "validate_persisted_contracts", +] + +CONTRACT_LINKER_VERSION = "action_contract_linker_v2" +_INITIAL_PREDICATES = frozenset({"arm_free", "object_free"}) +_REFERENCE_KEYS = frozenset( + { + "anchor", + "container", + "reference", + "reference_object", + "support", + "support_object", + "target", + "target_object", + } +) + + +def link_task_dependencies( + task_spec: Mapping[str, Any], + role_bindings: Mapping[str, str], + *, + registry: AtomicCapabilityRegistry | None = None, +) -> dict[str, Any]: + """Add the minimal stable TaskGroup dependencies implied by contracts.""" + del registry # Reserved for task-level capability specialization. + task = validate_task_spec(task_spec) + bindings = {str(key): str(value) for key, value in role_bindings.items()} + bindings_hash = hashlib.sha256( + json.dumps( + bindings, sort_keys=True, separators=(",", ":"), allow_nan=False + ).encode("utf-8") + ).hexdigest() + existing_metadata = task.get("metadata", {}) + existing_linker = ( + existing_metadata.get("action_contract_task_linker", {}) + if isinstance(existing_metadata, Mapping) + else {} + ) + if ( + isinstance(existing_linker, Mapping) + and existing_linker.get("version") == CONTRACT_LINKER_VERSION + and existing_linker.get("role_bindings_hash") == bindings_hash + ): + return task + instances = task["task_instances"] + order = [str(item["id"]) for item in instances] + dependencies = { + str(item["id"]): set(str(value) for value in item["depends_on"]) + for item in instances + } + dependency_order = { + str(item["id"]): [str(value) for value in item["depends_on"]] + for item in instances + } + claims = {str(item["id"]): _task_claims(item, bindings) for item in instances} + distinct_arm_pairs = _distinct_arm_pairs(task.get("metadata", {})) + linked: list[dict[str, str]] = [] + + latest_by_object: dict[str, tuple[str, str]] = {} + for instance in instances: + instance_id = str(instance["id"]) + task_type = str(instance["task_type"]) + primary = _task_primary_object(instance, bindings) + previous = latest_by_object.get(primary) + if ( + task_type == "E4" + and previous is not None + and previous[1] == "E2" + and previous[0] not in dependencies[instance_id] + and not _reaches(dependencies, previous[0], instance_id) + ): + dependencies[instance_id].add(previous[0]) + dependency_order[instance_id].append(previous[0]) + linked.append( + { + "from": previous[0], + "to": instance_id, + "reason": "causal", + "detail": f"object_held:{primary}", + } + ) + _assert_acyclic(dependencies, "TaskSpec causal linking") + latest_by_object[primary] = (instance_id, task_type) + + for later_index, later_id in enumerate(order): + for earlier_id in order[:later_index]: + if _reaches(dependencies, later_id, earlier_id) or _reaches( + dependencies, earlier_id, later_id + ): + continue + conflicts = _claim_conflicts(claims[earlier_id], claims[later_id]) + if frozenset({earlier_id, later_id}) in distinct_arm_pairs: + conflicts = [item for item in conflicts if item != "arm:auto"] + if not conflicts: + continue + dependencies[later_id].add(earlier_id) + dependency_order[later_id].append(earlier_id) + linked.append( + { + "from": earlier_id, + "to": later_id, + "reason": "resource", + "detail": ",".join(conflicts), + } + ) + _assert_acyclic(dependencies, "TaskSpec contract linking") + + for instance in instances: + instance_id = str(instance["id"]) + instance["depends_on"] = dependency_order[instance_id] + metadata = dict(task.get("metadata", {})) + metadata["action_contract_task_linker"] = { + "version": CONTRACT_LINKER_VERSION, + "role_bindings_hash": bindings_hash, + "linked_dependencies": linked, + } + task["metadata"] = metadata + return validate_task_spec(task) + + +def link_seed_graph( + draft: Mapping[str, Any], + *, + registry: AtomicCapabilityRegistry | None = None, + task_order: Sequence[str] = (), + completed_nodes: Collection[str] = (), + known_objects: Collection[str] | None = None, +) -> dict[str, Any]: + """Resolve contracts, link a draft graph, and return validated SeedGraph v3.""" + if not isinstance(draft, Mapping): + raise TypeError("SeedGraph draft must be a mapping.") + if draft.get("schema_version") != SEED_GRAPH_SCHEMA: + raise ValueError(f"Contract linker accepts only {SEED_GRAPH_SCHEMA!r} drafts.") + capabilities = registry or build_atomic_capability_registry() + graph = deepcopy(dict(draft)) + nodes = graph.get("nodes") + groups = graph.get("task_groups") + if not isinstance(nodes, list) or not nodes: + raise ValueError("SeedGraph draft nodes must be a non-empty list.") + if not isinstance(groups, list) or not groups: + raise ValueError("SeedGraph draft task_groups must be a non-empty list.") + + already_linked = _already_linked(graph) + for index, node in enumerate(nodes): + if not isinstance(node, dict): + raise TypeError(f"SeedGraph draft node {index} must be a mapping.") + action = str(node.get("atomic_action", "")) + expected = capabilities.get(action).resolve_contract(node).as_mapping() + persisted = node.get("contract") + if persisted is not None and persisted != expected: + raise ValueError( + f"SeedGraph node {node.get('id')!r} persisted Action Contract " + "does not match the current capability resolver." + ) + node["contract"] = expected + node.pop("resources", None) + if already_linked: + linked = validate_seed_graph( + graph, + known_objects=known_objects, + known_actions=capabilities.names(), + ) + validate_persisted_contracts(linked, capabilities) + return linked + + completed = {str(item) for item in completed_nodes} + node_by_id = _unique_by_id(nodes, "SeedGraph draft nodes") + group_by_id = _unique_by_id(groups, "SeedGraph draft task_groups") + ordered_groups = _ordered_group_ids(groups, task_order) + node_reasons: list[dict[str, str]] = [] + group_reasons: list[dict[str, str]] = [] + + for group in groups: + group_id = str(group.get("id", "")) + node_ids = [str(item) for item in group.get("node_ids", ())] + if not node_ids or any(node_id not in node_by_id for node_id in node_ids): + raise ValueError( + f"SeedGraph TaskGroup {group_id!r} has missing or unknown node IDs." + ) + _link_internal_nodes( + node_ids, + node_by_id, + completed=completed, + reasons=node_reasons, + ) + _validate_internal_symbolic_state(node_ids, node_by_id) + group.pop("contract", None) + + group_dependencies = { + group_id: set(str(item) for item in group_by_id[group_id].get("depends_on", ())) + for group_id in ordered_groups + } + original_group_dependencies = { + group_id: [str(item) for item in group_by_id[group_id].get("depends_on", ())] + for group_id in ordered_groups + } + _assert_acyclic(group_dependencies, "SeedGraph TaskGroups") + summaries = { + group_id: _summarize_group(group_by_id[group_id], node_by_id) + for group_id in ordered_groups + } + distinct_arm_pairs = _distinct_arm_pairs(graph.get("metadata", {})) + + for later_index, later_id in enumerate(ordered_groups): + for earlier_id in ordered_groups[:later_index]: + if _reaches(group_dependencies, later_id, earlier_id) or _reaches( + group_dependencies, earlier_id, later_id + ): + continue + conflicts = _claim_conflicts( + summaries[earlier_id]["claims"], summaries[later_id]["claims"] + ) + if frozenset({earlier_id, later_id}) in distinct_arm_pairs: + conflicts = [item for item in conflicts if item != "arm:auto"] + if conflicts: + _add_group_dependency( + earlier_id, + later_id, + group_dependencies, + completed, + summaries, + group_reasons, + reason="resource", + detail=",".join(conflicts), + ) + + for later_index, group_id in enumerate(ordered_groups): + for requirement in summaries[group_id]["entry_requires"]: + if requirement["predicate"] in _INITIAL_PREDICATES: + continue + candidates = [ + candidate + for candidate in ordered_groups[:later_index] + if _adds_atom(summaries[candidate]["exit_effects"], requirement) + ] + if not candidates: + raise ValueError( + f"SeedGraph TaskGroup {group_id!r} has no producer for state " + f"{requirement}." + ) + maximal = [ + candidate + for candidate in candidates + if not any( + candidate != other + and _reaches(group_dependencies, other, candidate) + for other in candidates + ) + ] + if len(maximal) != 1: + raise ValueError( + f"SeedGraph TaskGroup {group_id!r} has multiple unordered " + f"producers for state {requirement}: {maximal}." + ) + producer = maximal[0] + if not _reaches(group_dependencies, group_id, producer): + _add_group_dependency( + producer, + group_id, + group_dependencies, + completed, + summaries, + group_reasons, + reason="causal", + detail=_atom_key(requirement), + ) + + _assert_acyclic(group_dependencies, "SeedGraph contract linking") + for group_id in ordered_groups: + group = group_by_id[group_id] + group["depends_on"] = original_group_dependencies[group_id] + [ + candidate + for candidate in ordered_groups + if candidate in group_dependencies[group_id] + and candidate not in original_group_dependencies[group_id] + ] + + _link_group_boundaries( + ordered_groups, + group_dependencies, + summaries, + node_by_id, + completed, + node_reasons, + ) + for group_id in ordered_groups: + summaries[group_id] = _summarize_group(group_by_id[group_id], node_by_id) + group_by_id[group_id]["contract"] = summaries[group_id] + + _validate_symbolic_state(ordered_groups, group_dependencies, summaries, group_by_id) + metadata = dict(graph.get("metadata", {})) + metadata["action_contract_linker"] = { + "version": CONTRACT_LINKER_VERSION, + "group_dependencies": _sorted_reasons(group_reasons), + "node_dependencies": _sorted_reasons(node_reasons), + } + graph["metadata"] = metadata + graph["schema_version"] = SEED_GRAPH_SCHEMA + graph["nodes"] = nodes + graph["task_groups"] = groups + return validate_seed_graph( + graph, + known_objects=known_objects, + known_actions=capabilities.names(), + ) + + +def validate_persisted_contracts( + graph: Mapping[str, Any], registry: AtomicCapabilityRegistry +) -> None: + """Reject persisted contracts that differ from the active capability catalog.""" + metadata = graph.get("metadata", {}) + linker = ( + metadata.get("action_contract_linker", {}) + if isinstance(metadata, Mapping) + else {} + ) + if ( + not isinstance(linker, Mapping) + or linker.get("version") != CONTRACT_LINKER_VERSION + ): + raise ValueError( + "SeedGraph was not produced by the current deterministic Contract Linker; " + "regenerate the configuration bundle." + ) + for node in graph.get("nodes", ()): + expected = ( + registry.get(str(node["atomic_action"])).resolve_contract(node).as_mapping() + ) + if node.get("contract") != expected: + raise ValueError( + f"SeedGraph node {node.get('id')!r} persisted Action Contract " + "does not match the current capability resolver." + ) + node_by_id = { + str(node["id"]): node + for node in graph.get("nodes", ()) + if isinstance(node, Mapping) and "id" in node + } + for group in graph.get("task_groups", ()): + expected = _summarize_group(group, node_by_id) + if group.get("contract") != expected: + raise ValueError( + f"SeedGraph TaskGroup {group.get('id')!r} persisted contract " + "does not match its linked AtomicAction topology." + ) + + +def _task_claims( + instance: Mapping[str, Any], bindings: Mapping[str, str] +) -> list[dict[str, str]]: + task_type = str(instance["task_type"]) + params = _resolve_roles(instance.get("params", {}), bindings) + primary_key = "source_role" if task_type == "E3" else "object_role" + primary = params.get(primary_key) + claims: list[dict[str, str]] = [] + if isinstance(primary, str) and primary: + claims.append(_claim(f"object:{primary}", "exclusive")) + target = params.get("target_role") + if isinstance(target, str) and target and target != primary: + claims.append(_claim(f"object:{target}", "shared_read")) + payloads = params.get("payload_roles", []) + if isinstance(payloads, Sequence) and not isinstance( + payloads, (str, bytes, bytearray) + ): + for payload in payloads: + if isinstance(payload, str) and payload and payload != primary: + claims.append(_claim(f"object:{payload}", "exclusive")) + if task_type == "E4": + transfer = str(params.get("transfer_arm", "")) + receive = str(params.get("receive_arm", "")) + if transfer not in {"left_arm", "right_arm"} or receive not in { + "left_arm", + "right_arm", + }: + raise ValueError( + "E4 contract linking requires explicit transfer/receive arms." + ) + if transfer == receive: + raise ValueError("E4 transfer_arm and receive_arm must be distinct.") + claims.extend((_claim(f"arm:{transfer}"), _claim(f"arm:{receive}"))) + elif task_type == "E5": + claims.extend((_claim("arm:left_arm"), _claim("arm:right_arm"))) + else: + required_arm = params.get("required_arm") + if required_arm in {"left_arm", "right_arm"}: + claims.append(_claim(f"arm:{required_arm}")) + else: + claims.append(_claim("arm:auto")) + return _merge_claims(claims) + + +def _task_primary_object( + instance: Mapping[str, Any], bindings: Mapping[str, str] +) -> str: + task_type = str(instance["task_type"]) + params = _resolve_roles(instance.get("params", {}), bindings) + key = "source_role" if task_type == "E3" else "object_role" + value = params.get(key) + if not isinstance(value, str) or not value: + raise ValueError( + f"TaskGroup {instance.get('id')!r} requires a resolved {key!r}." + ) + return value + + +def _link_internal_nodes( + node_ids: Sequence[str], + node_by_id: Mapping[str, dict[str, Any]], + *, + completed: set[str], + reasons: list[dict[str, str]], +) -> None: + positions = {node_id: index for index, node_id in enumerate(node_ids)} + for later_index, later_id in enumerate(node_ids): + later = node_by_id[later_id] + for requirement in later["contract"]["requires"]: + producers = [ + earlier_id + for earlier_id in node_ids[:later_index] + if node_by_id[earlier_id]["contract"]["failure_policy"] != "best_effort" + if _adds_atom( + node_by_id[earlier_id]["contract"]["effects"], requirement + ) + ] + if producers: + _add_node_dependency( + producers[-1], + later_id, + node_by_id, + completed, + reasons, + "causal", + _atom_key(requirement), + ) + for earlier_id in node_ids[:later_index]: + earlier = node_by_id[earlier_id] + if earlier.get("sync_group") is not None and earlier.get( + "sync_group" + ) == later.get("sync_group"): + continue + if _node_reaches(node_by_id, later_id, earlier_id) or _node_reaches( + node_by_id, earlier_id, later_id + ): + continue + conflicts = _claim_conflicts( + earlier["contract"]["claims"], later["contract"]["claims"] + ) + if conflicts: + _add_node_dependency( + earlier_id, + later_id, + node_by_id, + completed, + reasons, + "resource", + ",".join(conflicts), + ) + dependencies = { + node_id: { + str(parent) + for parent in node_by_id[node_id].get("depends_on", ()) + if str(parent) in positions + } + for node_id in node_ids + } + _assert_acyclic(dependencies, "AtomicAction contract linking") + + +def _summarize_group( + group: Mapping[str, Any], node_by_id: Mapping[str, Mapping[str, Any]] +) -> dict[str, Any]: + node_ids = [str(item) for item in group["node_ids"]] + node_set = set(node_ids) + entries = [ + node_id + for node_id in node_ids + if not any( + str(parent) in node_set for parent in node_by_id[node_id]["depends_on"] + ) + ] + depended = { + str(parent) + for node_id in node_ids + for parent in node_by_id[node_id]["depends_on"] + if str(parent) in node_set + } + terminals = [node_id for node_id in node_ids if node_id not in depended] + entry_requires: list[dict[str, str]] = [] + for node_id in node_ids: + node = node_by_id[node_id] + for requirement in node["contract"]["requires"]: + if any( + producer_id in node_set + and _node_reaches(node_by_id, node_id, producer_id) + and _adds_atom( + node_by_id[producer_id]["contract"]["effects"], requirement + ) + for producer_id in node_ids + ): + continue + if requirement not in entry_requires: + entry_requires.append(deepcopy(requirement)) + + last_effect: dict[str, dict[str, Any]] = {} + effect_order: list[str] = [] + for node_id in node_ids: + if node_by_id[node_id]["contract"]["failure_policy"] == "best_effort": + continue + for effect in node_by_id[node_id]["contract"]["effects"]: + key = _atom_key(effect["atom"]) + if key not in last_effect: + effect_order.append(key) + last_effect[key] = deepcopy(effect) + claims = [ + deepcopy(claim) + for node_id in node_ids + for claim in node_by_id[node_id]["contract"]["claims"] + ] + claims.extend(_goal_read_claims(group.get("goal", {}))) + merged_claims = _merge_claims(claims) + free_resources = { + ( + f"arm:{effect['atom']['arm']}" + if effect["atom"]["predicate"] == "arm_free" + else f"object:{effect['atom']['object_uid']}" + ) + for effect in last_effect.values() + if effect["op"] == "add" + and effect["atom"]["predicate"] in {"arm_free", "object_free"} + } + for claim in merged_claims: + if claim["resource"] in free_resources: + claim["lifetime"] = "action" + completion = ( + "terminal_barrier" + if terminals + and all( + node_by_id[node_id]["contract"]["completion"] == "terminal_barrier" + for node_id in terminals + ) + else "ordinary" + ) + return { + "entry_requires": entry_requires, + "exit_effects": [last_effect[key] for key in effect_order], + "claims": merged_claims, + "entry_node_ids": entries, + "terminal_node_ids": terminals, + "completion": completion, + } + + +def _validate_internal_symbolic_state( + node_ids: Sequence[str], node_by_id: Mapping[str, Mapping[str, Any]] +) -> None: + node_set = set(node_ids) + dependencies = { + node_id: { + str(parent) + for parent in node_by_id[node_id].get("depends_on", ()) + if str(parent) in node_set + } + for node_id in node_ids + } + entry_atoms = set() + for node_id in node_ids: + for requirement in node_by_id[node_id]["contract"]["requires"]: + has_prior_producer = any( + producer_id != node_id + and _node_reaches(node_by_id, node_id, producer_id) + and _adds_atom( + node_by_id[producer_id]["contract"]["effects"], requirement + ) + for producer_id in node_ids + ) + if not has_prior_producer: + entry_atoms.add(_atom_key(requirement)) + state = set(entry_atoms) + for node_id in _stable_topological(node_ids, dependencies): + contract = node_by_id[node_id]["contract"] + for requirement in contract["requires"]: + if _atom_key(requirement) not in state: + raise ValueError( + f"SeedGraph node {node_id!r} requires unavailable state " + f"{requirement}." + ) + for effect in contract["effects"]: + key = _atom_key(effect["atom"]) + if effect["op"] == "delete": + if key not in state: + raise ValueError( + f"SeedGraph node {node_id!r} deletes unavailable state " + f"{effect['atom']}." + ) + state.remove(key) + else: + state.add(key) + + +def _add_group_dependency( + parent: str, + child: str, + dependencies: dict[str, set[str]], + completed: set[str], + summaries: Mapping[str, Mapping[str, Any]], + reasons: list[dict[str, str]], + *, + reason: str, + detail: str, +) -> None: + if any(node_id in completed for node_id in summaries[child]["entry_node_ids"]): + raise ValueError( + f"Contract linking cannot add dependency into completed TaskGroup {child!r}." + ) + dependencies[child].add(parent) + _assert_acyclic(dependencies, "SeedGraph contract linking") + reasons.append({"from": parent, "to": child, "reason": reason, "detail": detail}) + + +def _link_group_boundaries( + ordered_groups: Sequence[str], + dependencies: Mapping[str, set[str]], + summaries: Mapping[str, Mapping[str, Any]], + node_by_id: Mapping[str, dict[str, Any]], + completed: set[str], + reasons: list[dict[str, str]], +) -> None: + for child in ordered_groups: + for parent in ordered_groups: + if parent not in dependencies[child]: + continue + for child_node in summaries[child]["entry_node_ids"]: + for parent_node in summaries[parent]["terminal_node_ids"]: + _add_node_dependency( + parent_node, + child_node, + node_by_id, + completed, + reasons, + "cleanup", + f"TaskGroup {parent} terminal barrier", + ) + + +def _add_node_dependency( + parent: str, + child: str, + node_by_id: Mapping[str, dict[str, Any]], + completed: set[str], + reasons: list[dict[str, str]], + reason: str, + detail: str, +) -> None: + if parent == child: + raise ValueError(f"Contract linker cannot add self-dependency {child!r}.") + dependencies = node_by_id[child].setdefault("depends_on", []) + if parent in dependencies or _node_reaches(node_by_id, child, parent): + return + if _node_reaches(node_by_id, parent, child): + raise ValueError( + f"Contract dependency {parent!r} -> {child!r} would create a cycle." + ) + if child in completed: + raise ValueError(f"Contract linker cannot modify completed node {child!r}.") + dependencies.append(parent) + reasons.append({"from": parent, "to": child, "reason": reason, "detail": detail}) + + +def _validate_symbolic_state( + ordered_groups: Sequence[str], + dependencies: Mapping[str, set[str]], + summaries: Mapping[str, Mapping[str, Any]], + groups: Mapping[str, Mapping[str, Any]], +) -> None: + atoms = [ + atom + for summary in summaries.values() + for atom in [ + *summary["entry_requires"], + *(effect["atom"] for effect in summary["exit_effects"]), + ] + ] + state = { + _atom_key({"predicate": "arm_free", "arm": str(atom["arm"])}) + for atom in atoms + if "arm" in atom + } + state.update( + _atom_key({"predicate": "object_free", "object_uid": str(atom["object_uid"])}) + for atom in atoms + if "object_uid" in atom + ) + for group_id in _stable_topological(ordered_groups, dependencies): + if groups[group_id].get("role") == "recovery": + for requirement in summaries[group_id]["entry_requires"]: + if requirement["predicate"] == "object_free": + object_uid = str(requirement["object_uid"]) + state = { + item + for item in state + if not ( + item.startswith("object_held|") + or item.startswith("object_coordinated_held|") + ) + or f"|{object_uid}|" not in f"|{item}|" + } + state.add(_atom_key(requirement)) + for requirement in summaries[group_id]["entry_requires"]: + key = _atom_key(requirement) + if key not in state: + raise ValueError( + _unavailable_group_state_message( + group_id, + requirement, + state, + groups[group_id], + ) + ) + for effect in summaries[group_id]["exit_effects"]: + key = _atom_key(effect["atom"]) + if effect["op"] == "add": + state.add(key) + else: + state.discard(key) + + +def _unavailable_group_state_message( + group_id: str, + requirement: Mapping[str, Any], + state: Collection[str], + group: Mapping[str, Any], +) -> str: + """Explain held-object conflicts without weakening symbolic validation.""" + if requirement.get("predicate") != "arm_free": + return ( + f"SeedGraph TaskGroup {group_id!r} requires unavailable state " + f"{dict(requirement)}." + ) + arm = str(requirement.get("arm", "")) + held_objects = sorted( + parts[1] + for item in state + if len(parts := item.split("|", maxsplit=2)) == 3 + and parts[0] == "object_held" + and parts[2] == arm + and parts[1] + ) + if not held_objects: + return ( + f"SeedGraph TaskGroup {group_id!r} requires unavailable state " + f"{dict(requirement)}." + ) + primary = str(group.get("object_uid", "")) + held = ", ".join(repr(item) for item in held_objects) + return ( + f"SeedGraph TaskGroup {group_id!r} requires arm {arm!r} to be free, " + f"but it currently holds {held}; the group's primary object is " + f"{primary!r}. A post-handover continuation must preserve object " + "identity and consume object_held instead of scheduling a fresh pickup." + ) + + +def _goal_read_claims(value: Any) -> list[dict[str, str]]: + claims: list[dict[str, str]] = [] + if isinstance(value, Mapping): + for key, child in value.items(): + if key in _REFERENCE_KEYS and isinstance(child, str): + if child not in {"table_center", "world"}: + claims.append(_claim(f"object:{child}", "shared_read")) + claims.extend(_goal_read_claims(child)) + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + for child in value: + claims.extend(_goal_read_claims(child)) + return claims + + +def _claim( + resource: str, access: str = "exclusive", lifetime: str = "action" +) -> dict[str, str]: + return {"resource": resource, "access": access, "lifetime": lifetime} + + +def _merge_claims(claims: Sequence[Mapping[str, Any]]) -> list[dict[str, str]]: + merged: dict[str, dict[str, str]] = {} + order: list[str] = [] + for claim in claims: + resource = str(claim["resource"]) + if resource not in merged: + order.append(resource) + merged[resource] = _claim( + resource, + str(claim.get("access", "exclusive")), + str(claim.get("lifetime", "action")), + ) + continue + current = merged[resource] + if claim.get("access") == "exclusive": + current["access"] = "exclusive" + if claim.get("lifetime") == "until_release": + current["lifetime"] = "until_release" + return [merged[resource] for resource in order] + + +def _claim_conflicts( + first: Sequence[Mapping[str, Any]], second: Sequence[Mapping[str, Any]] +) -> list[str]: + first_by_resource = {str(item["resource"]): str(item["access"]) for item in first} + second_by_resource = {str(item["resource"]): str(item["access"]) for item in second} + conflicts = { + resource + for resource in set(first_by_resource) & set(second_by_resource) + if "exclusive" in {first_by_resource[resource], second_by_resource[resource]} + } + first_arms = {item for item in first_by_resource if item.startswith("arm:")} + second_arms = {item for item in second_by_resource if item.startswith("arm:")} + if "arm:auto" in first_arms and second_arms: + conflicts.add("arm:auto") + if "arm:auto" in second_arms and first_arms: + conflicts.add("arm:auto") + return sorted(conflicts) + + +def _distinct_arm_pairs(value: Any) -> set[frozenset[str]]: + if not isinstance(value, Mapping): + return set() + groups = value.get("legacy_allocation_groups", value.get("allocation_groups", ())) + if not isinstance(groups, Sequence) or isinstance(groups, (str, bytes, bytearray)): + return set() + result: set[frozenset[str]] = set() + for group in groups: + if ( + not isinstance(group, Mapping) + or group.get("arm_constraint") != "distinct_arms" + ): + continue + members = group.get("semantic_step_ids", group.get("task_instance_ids", ())) + if not isinstance(members, Sequence) or isinstance( + members, (str, bytes, bytearray) + ): + continue + member_ids = [str(item) for item in members] + for index, first in enumerate(member_ids): + for second in member_ids[index + 1 :]: + result.add(frozenset({first, second})) + return result + + +def _resolve_roles(value: Any, bindings: Mapping[str, str]) -> Any: + if isinstance(value, Mapping): + return { + str(key): _resolve_roles(child, bindings) for key, child in value.items() + } + if isinstance(value, list): + return [_resolve_roles(child, bindings) for child in value] + if isinstance(value, tuple): + return tuple(_resolve_roles(child, bindings) for child in value) + if isinstance(value, str): + return bindings.get(value, value) + return value + + +def _adds_atom(effects: Sequence[Mapping[str, Any]], atom: Mapping[str, Any]) -> bool: + return any( + effect.get("op") == "add" and effect.get("atom") == atom for effect in effects + ) + + +def _atom_key(atom: Mapping[str, Any]) -> str: + return "|".join( + str(atom.get(key, "")) for key in ("predicate", "object_uid", "arm") + ) + + +def _unique_by_id(items: Sequence[Mapping[str, Any]], context: str) -> dict[str, Any]: + result: dict[str, Any] = {} + for item in items: + item_id = str(item.get("id", "")) + if not item_id: + raise ValueError(f"{context} require non-empty IDs.") + if item_id in result: + raise ValueError(f"{context} contain duplicate ID {item_id!r}.") + result[item_id] = item + return result + + +def _ordered_group_ids( + groups: Sequence[Mapping[str, Any]], task_order: Sequence[str] +) -> list[str]: + available = [str(group["id"]) for group in groups] + requested = [str(item) for item in task_order] + unknown = set(requested) - set(available) + if unknown: + raise ValueError( + f"task_order references unknown TaskGroups: {sorted(unknown)}." + ) + return requested + [item for item in available if item not in set(requested)] + + +def _reaches(dependencies: Mapping[str, set[str]], child: str, parent: str) -> bool: + pending = list(dependencies.get(child, ())) + visited: set[str] = set() + while pending: + current = pending.pop() + if current == parent: + return True + if current not in visited: + visited.add(current) + pending.extend(dependencies.get(current, ())) + return False + + +def _node_reaches( + node_by_id: Mapping[str, Mapping[str, Any]], child: str, parent: str +) -> bool: + pending = [str(item) for item in node_by_id[child].get("depends_on", ())] + visited: set[str] = set() + while pending: + current = pending.pop() + if current == parent: + return True + if current not in visited and current in node_by_id: + visited.add(current) + pending.extend( + str(item) for item in node_by_id[current].get("depends_on", ()) + ) + return False + + +def _assert_acyclic(dependencies: Mapping[str, set[str]], context: str) -> None: + for item_id in dependencies: + if _reaches(dependencies, item_id, item_id): + raise ValueError(f"{context} produced a dependency cycle at {item_id!r}.") + + +def _stable_topological( + order: Sequence[str], dependencies: Mapping[str, set[str]] +) -> list[str]: + remaining = set(order) + result: list[str] = [] + while remaining: + ready = [ + item + for item in order + if item in remaining and not (dependencies[item] & remaining) + ] + if not ready: + raise ValueError("SeedGraph TaskGroups contain a dependency cycle.") + result.extend(ready) + remaining.difference_update(ready) + return result + + +def _already_linked(graph: Mapping[str, Any]) -> bool: + metadata = graph.get("metadata", {}) + linker = ( + metadata.get("action_contract_linker", {}) + if isinstance(metadata, Mapping) + else {} + ) + return ( + isinstance(linker, Mapping) + and linker.get("version") == CONTRACT_LINKER_VERSION + and all("contract" in node for node in graph.get("nodes", ())) + and all("contract" in group for group in graph.get("task_groups", ())) + ) + + +def _sorted_reasons(reasons: Sequence[Mapping[str, str]]) -> list[dict[str, str]]: + unique = { + (item["from"], item["to"], item["reason"], item["detail"]) for item in reasons + } + return [ + {"from": source, "to": target, "reason": reason, "detail": detail} + for source, target, reason, detail in sorted(unique) + ] diff --git a/embodichain/gen_sim/action_engine/planning/online.py b/embodichain/gen_sim/action_engine/planning/online.py new file mode 100644 index 000000000..fc8b36080 --- /dev/null +++ b/embodichain/gen_sim/action_engine/planning/online.py @@ -0,0 +1,371 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Online planner producing a complete direct AtomicAction SeedGraph.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from copy import deepcopy +import json +from time import perf_counter +from typing import Any + +from embodichain.gen_sim.action_engine.capabilities import ( + AtomicCapabilityRegistry, + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.domain import ( + TASK_CONTRACTS, + public_task_spec, + requested_visual_task_predicates, + validate_public_task_spec, + validate_task_spec, +) +from embodichain.gen_sim.action_engine.protocol import SEED_GRAPH_SCHEMA +from embodichain.gen_sim.action_engine.runtime.motion_policy import ( + resolve_motion_policy, +) + +from .vision import ( + SceneObservation, + _reject_live_fields as _reject_visual_live_fields, + analyze_visual_scene, + validate_visual_facts, +) +from .linker import link_seed_graph + +__all__ = ["plan_online_seed_graph"] + +GraphCaller = Callable[..., Mapping[str, Any]] + +_GRAPH_OUTPUT_SCHEMA = { + "title": "ActionEngineOnlineSeedGraphBody", + "type": "object", + "additionalProperties": False, + "required": ["nodes", "task_groups", "success"], + "properties": { + "nodes": {"type": "array", "items": {"type": "object"}}, + "task_groups": {"type": "array", "items": {"type": "object"}}, + "success": {"type": "object"}, + }, +} + + +def plan_online_seed_graph( + task_spec: Mapping[str, Any], + observation: SceneObservation, + *, + visual_facts: Mapping[str, Any] | None = None, + vlm_model: str | None = None, + fact_caller: GraphCaller | None = None, + graph_caller: GraphCaller | None = None, + registry: AtomicCapabilityRegistry | None = None, + robot_profile: str = "dual_ur10", +) -> tuple[dict[str, Any], dict[str, Any]]: + """Extract visual facts and produce one validated online SeedGraph.""" + started = perf_counter() + task = ( + validate_public_task_spec(task_spec) + if "task_instances" not in task_spec and task_spec.get("level") == "L4" + else validate_task_spec(task_spec) + ) + _reject_private_or_live_fields(public_task_spec(task), "online TaskSpec") + capabilities = registry or build_atomic_capability_registry() + _reject_visual_live_fields(observation.entities, "SceneObservation.entities") + known_uids = {str(item["uid"]) for item in observation.entities} + if len(known_uids) != len(observation.entities): + raise ValueError("Online scene observation contains duplicate entity UIDs.") + if not known_uids: + raise ValueError("Online scene observation contains no simulator entities.") + visual_call_counter = [0] + allowed_task_predicates = requested_visual_task_predicates(task) + facts = ( + validate_visual_facts( + visual_facts, + known_uids=known_uids, + camera_uids={camera.uid for camera in observation.cameras}, + allowed_task_predicates=allowed_task_predicates, + ) + if visual_facts is not None + else analyze_visual_scene( + observation, + task, + model=vlm_model, + caller=fact_caller, + call_counter=visual_call_counter, + ) + ) + _validate_fact_information(facts) + prompt = _prompt(task, facts, capabilities, robot_profile=robot_profile) + if graph_caller is None: + # Facts remain the auditable planner input, but the production VLM also + # needs the same reset-time RGB/depth evidence to bind semantic TaskSpec + # roles (for example, "the purple can") to the known simulator UIDs. + # An injected graph caller keeps the compact facts-only contract used by + # deterministic tests and alternative planners. + def caller(**kwargs: Any) -> Mapping[str, Any]: + return _default_graph_caller(observation=observation, **kwargs) + + else: + caller = graph_caller + first_error: Exception | None = None + graph_call_count = 0 + for attempt in range(2): + current_prompt = prompt + if first_error is not None: + current_prompt += ( + "\n\nThe previous graph was invalid. Correct only the JSON body. " + f"Validation error: {first_error}" + ) + graph_call_count += 1 + try: + response = caller( + prompt=current_prompt, + schema=_GRAPH_OUTPUT_SCHEMA, + model=vlm_model, + ) + graph = _wrap_graph(response, task, capabilities) + _reject_private_or_live_fields(graph, "online SeedGraph") + graph = link_seed_graph( + graph, + registry=capabilities, + task_order=[str(item["id"]) for item in task.get("task_instances", ())], + known_objects=known_uids, + ) + _validate_explicit_task_group_coverage(task, graph) + for node in graph["nodes"]: + capabilities.validate_binding(node) + if capabilities.get(str(node["atomic_action"])).runtime_available: + resolve_motion_policy( + robot_profile, + node["atomic_action"], + node["motion_policy"], + ) + graph["metadata"].update( + { + "planning_latency_seconds": perf_counter() - started, + "vlm_call_count": graph_call_count + visual_call_counter[0], + "visual_fact_call_count": visual_call_counter[0], + "graph_call_count": graph_call_count, + } + ) + return graph, facts + except (TypeError, ValueError) as error: + if attempt: + raise ValueError( + "Online SeedGraph failed validation after one repair: " f"{error}" + ) from error + first_error = error + raise AssertionError("unreachable") + + +def _prompt( + task: Mapping[str, Any], + facts: Mapping[str, Any], + capabilities: AtomicCapabilityRegistry, + *, + robot_profile: str, +) -> str: + from embodichain.gen_sim.action_engine.config import default_runtime_policy + + runtime_policy = default_runtime_policy(robot_profile) + motion_modifiers: dict[str, list[dict[str, str]]] = { + action: [] for action in runtime_policy.motion_defaults + } + for modifier_type, modes in runtime_policy.motion_modifiers.items(): + for mode, action_patches in modes.items(): + for action in action_patches: + motion_modifiers[action].append({"type": modifier_type, "mode": mode}) + grouping_instruction = ( + "Infer the necessary E TaskGroups from the abstract goal; the private " + "reference task instances are intentionally hidden." + if task["level"] == "L4" + else "Every public TaskSpec task instance must correspond to exactly one TaskGroup." + ) + return ( + "Produce the body of one coordinate-free direct AtomicAction SeedGraph. " + f"{grouping_instruction} " + "Nodes may contain only symbolic target bindings and scene UIDs; never " + "emit world coordinates, poses, qpos, trajectories, or grasp poses. " + "Do not emit Action Contracts or resource claims; the deterministic " + "Contract Linker owns those fields. " + "Use the supplied reset-time multi-view image evidence only to bind the " + "public task semantics to known UIDs; use normalized visual constraints " + "only when the facts justify them. " + "Do not output reasoning. Planning-only actions may appear but must not " + "be replaced with invented primitives.\n\n" + f"Public TaskSpec:\n{json.dumps(public_task_spec(task), ensure_ascii=False, sort_keys=True)}\n\n" + f"Visual facts:\n{json.dumps(facts, ensure_ascii=False, sort_keys=True)}\n\n" + f"E1-E9 task semantics:\n{json.dumps(_task_capability_catalog(), ensure_ascii=False, sort_keys=True)}\n\n" + f"Atomic capabilities:\n{json.dumps(capabilities.catalog(), ensure_ascii=False, sort_keys=True)}\n\n" + "Every node motion_policy must be an object with a modifiers list; " + "the AtomicAction selects its base policy implicitly. Use only the " + "typed modifiers supported by that action.\n" + f"Allowed motion modifiers by AtomicAction:\n" + f"{json.dumps(motion_modifiers, sort_keys=True)}" + ) + + +def _task_capability_catalog() -> dict[str, dict[str, Any]]: + """Return the Action Engine's runtime-aware E-task view.""" + executable = set(build_atomic_capability_registry().executable_names()) + return { + task_type: { + "semantics": contract.semantics, + "core_actions": list(contract.core_actions), + "runtime_available": set(contract.core_actions) <= executable, + } + for task_type, contract in TASK_CONTRACTS.items() + } + + +def _wrap_graph( + response: Mapping[str, Any], + task: Mapping[str, Any], + capabilities: AtomicCapabilityRegistry, +) -> dict[str, Any]: + if not isinstance(response, Mapping): + raise TypeError("Online planner output must be a mapping.") + if set(response) != {"nodes", "task_groups", "success"}: + raise ValueError( + "Online planner must return nodes, task_groups, and success only." + ) + for index, node in enumerate(response.get("nodes", ())): + if not isinstance(node, Mapping): + raise TypeError(f"Online planner node {index} must be a mapping.") + forbidden = sorted({"contract", "resources"} & set(node)) + if forbidden: + raise ValueError( + f"Online planner node {index} may not author linker-owned fields: " + f"{forbidden}." + ) + for index, group in enumerate(response.get("task_groups", ())): + if not isinstance(group, Mapping): + raise TypeError(f"Online planner TaskGroup {index} must be a mapping.") + if "contract" in group: + raise ValueError( + f"Online planner TaskGroup {index} may not author its contract." + ) + return { + "schema_version": SEED_GRAPH_SCHEMA, + "task_id": task["task_id"], + "instruction": task["instruction"], + "level": task["level"], + "reasoning_type": task["reasoning_type"], + "planner_route": "online", + "nodes": deepcopy(response["nodes"]), + "task_groups": deepcopy(response["task_groups"]), + "success": deepcopy(response["success"]), + "capability_catalog_hash": capabilities.catalog_hash(), + "metadata": { + "oracle_exposed": False, + "visual_facts_used": True, + "allocation_groups": deepcopy( + task.get("metadata", {}).get("allocation_groups", []) + ), + }, + } + + +def _default_graph_caller( + *, + prompt: str, + schema: Mapping[str, Any], + model: str | None, + observation: SceneObservation | None = None, +) -> Mapping[str, Any]: + from .vision import _camera_evidence, _default_structured_caller, _vlm_model + + images: list[str] = [] + if observation is not None: + _, images = _camera_evidence(observation) + + return _default_structured_caller( + prompt=prompt, + images=images, + schema=schema, + model=_vlm_model(model), + ) + + +def _validate_fact_information(facts: Mapping[str, Any]) -> None: + """Reject low-information visual outputs before graph planning.""" + confidence = facts.get("confidence", 0.0) + if float(confidence) < 0.5: + raise ValueError( + "VLM visual facts confidence is below the required 0.5 threshold." + ) + entities = facts.get("entities", ()) + if not any( + bool(item.get("visible", True)) and float(item.get("confidence", 0.0)) >= 0.5 + for item in entities + if isinstance(item, Mapping) + ): + raise ValueError("VLM visual facts contain no reliable visible entity.") + + +def _validate_explicit_task_group_coverage( + task: Mapping[str, Any], graph: Mapping[str, Any] +) -> None: + """Reject an online graph that drops or invents an explicit L1-L3 step.""" + 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( + "Online SeedGraph TaskGroup coverage mismatch; " + f"missing={sorted(expected - actual)}, " + f"unexpected={sorted(actual - expected)}." + ) + + +_PRIVATE_OR_LIVE_KEYS = frozenset( + { + "absolute_position", + "coordinates", + "grasp_pose", + "joint_positions", + "live_pose", + "live_transform", + "object_pose", + "oracle", + "pose", + "positions", + "qpos", + "target_pose", + "trajectory", + "waypoints", + "xpos", + } +) + + +def _reject_private_or_live_fields(value: Any, context: str) -> None: + """Reject private oracle and grounded simulator fields recursively.""" + 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, Sequence) and not isinstance(value, (str, bytes, bytearray)): + for index, child in enumerate(value): + _reject_private_or_live_fields(child, f"{context}[{index}]") diff --git a/embodichain/gen_sim/action_engine/planning/planner.py b/embodichain/gen_sim/action_engine/planning/planner.py new file mode 100644 index 000000000..a3e71d532 --- /dev/null +++ b/embodichain/gen_sim/action_engine/planning/planner.py @@ -0,0 +1,821 @@ +# ---------------------------------------------------------------------------- +# 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-free LLM planning boundary for Action Engine.""" + +from __future__ import annotations + +import json +import os +import re +from collections.abc import Callable, Mapping, Sequence +from copy import deepcopy +from pathlib import Path +from string import Template +from typing import Any + +from embodichain.gen_sim.action_engine.capabilities import build_default_registry +from embodichain.gen_sim.action_engine.domain import ( + TASK_AGENT_SCHEMA, + validate_task_agent, +) +from embodichain.gen_sim.action_engine.orientation import ( + compile_orientation_constraint, +) + +from .task_planner_prompt import TASK_PLANNER_PROMPT + +__all__ = ["plan_task"] + +LLMCaller = Callable[..., Mapping[str, Any]] + +_GEN_CONFIG_PATH = ( + Path(__file__).resolve().parents[2] + / "simready_pipeline" + / "configs" + / "gen_config.json" +) +_GEN_SIM_ENV_PATH = Path(__file__).resolve().parents[2] / ".env" +_UNSAFE_ID_RE = re.compile(r"[^0-9a-z]+") +_MODEL_STEP_KEYS = frozenset( + {"id", "operator", "object", "objects", "actor", "goal", "depends_on"} +) + +_MODEL_OUTPUT_SCHEMA: dict[str, Any] = { + "title": "ActionEngineSemanticPlan", + "type": "object", + "additionalProperties": False, + "required": ["semantic_steps", "allocation_groups"], + "properties": { + "semantic_steps": { + "type": "array", + "minItems": 1, + "items": { + "type": "object", + "additionalProperties": False, + "required": ["operator"], + "properties": { + "id": {"type": "string"}, + "operator": {"type": "string"}, + "object": {"type": "string"}, + "objects": { + "type": "array", + "items": {"type": "string"}, + }, + "actor": {"type": "object"}, + "goal": {"type": "object"}, + "depends_on": { + "type": "array", + "items": {"type": "string"}, + }, + }, + }, + }, + "allocation_groups": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": False, + "required": ["id", "semantic_step_ids", "arm_constraint"], + "properties": { + "id": {"type": "string"}, + "semantic_step_ids": { + "type": "array", + "items": {"type": "string"}, + }, + "arm_constraint": {"const": "distinct_arms"}, + }, + }, + }, + }, +} + + +def plan_task( + task_description: str, + scene_objects: Sequence[Mapping[str, Any]], + *, + task_name: str = "task", + model: str | None = None, + llm_caller: LLMCaller | None = None, +) -> dict[str, Any]: + """Plan a natural-language task as route-free semantic steps. + + The model is intentionally prohibited from emitting atomic actions, graph + edges, resources, target coordinates, or motion-policy parameters. + ``compile_task_agent`` owns all of those deterministic decisions. + + Args: + task_description: User goal in natural language. + scene_objects: JSON-like scene inventory. ``runtime_uid`` is preferred + over ``uid`` and ``source_uid`` for all generated references. + task_name: Stable task identifier stored in the TaskAgent. + model: Optional model-name override for the default LLM caller. + llm_caller: Optional injected callable accepting ``prompt=`` and + ``model=`` keyword arguments. It must return a mapping whose only + top-level key is ``semantic_steps``. + Returns: + A validated ``action_engine_task_agent_v1`` mapping. + """ + task_name = _nonempty(task_name, "task_name") + task_description = _nonempty(task_description, "task_description") + scene = _normalize_scene_objects(scene_objects) + + prompt = _render_prompt( + task_name=task_name, + task_description=task_description, + scene_objects=scene, + ) + caller = llm_caller or _default_llm_caller + response = caller(prompt=prompt, model=model) + try: + return _task_agent_from_response( + response, + task_name=task_name, + task_description=task_description, + scene=scene, + ) + except (TypeError, ValueError) as first_error: + # One bounded repair gives the model the verifier's exact complaint + # without turning generation into an unbounded conversation. + repair_prompt = ( + f"{prompt}\n\n" + "Your previous JSON did not satisfy the TaskAgent contract.\n" + f"Validation error: {first_error}\n" + "Return one corrected JSON object. Do not explain the correction." + ) + repaired = caller(prompt=repair_prompt, model=model) + try: + return _task_agent_from_response( + repaired, + task_name=task_name, + task_description=task_description, + scene=scene, + ) + except (TypeError, ValueError) as second_error: + raise ValueError( + "Action Engine planner failed validation after one repair: " + f"{second_error}" + ) from second_error + + +def _task_agent_from_response( + response: Any, + *, + task_name: str, + task_description: str, + scene: Sequence[Mapping[str, Any]], +) -> dict[str, Any]: + """Normalize and validate one model response as a TaskAgent.""" + if not isinstance(response, Mapping): + raise ValueError("Action Engine planner output must be a JSON object.") + allowed_fields = {"semantic_steps", "allocation_groups"} + if not set(response) <= allowed_fields or "semantic_steps" not in response: + raise ValueError( + "Action Engine planner output may contain only 'semantic_steps' " + "and 'allocation_groups'; " + f"received fields {sorted(str(key) for key in response)}." + ) + raw_steps = response["semantic_steps"] + if not isinstance(raw_steps, Sequence) or isinstance( + raw_steps, (str, bytes, bytearray) + ): + raise ValueError("Planner semantic_steps must be a list.") + visible_operators = set(build_default_registry().operator_names()) + for index, step in enumerate(raw_steps): + operator = step.get("operator") if isinstance(step, Mapping) else None + if operator not in visible_operators: + raise ValueError( + f"Planner semantic_steps[{index}].operator must be one of " + f"{sorted(visible_operators)}; got {operator!r}." + ) + return _wrap_agent( + task_name, + task_description, + raw_steps, + scene, + allocation_groups=response.get("allocation_groups", []), + ) + + +def _wrap_agent( + task_name: str, + task_description: str, + raw_steps: Sequence[Any], + scene: Sequence[Mapping[str, Any]], + *, + allocation_groups: Any, +) -> dict[str, Any]: + steps = _normalize_semantic_steps(raw_steps, scene) + groups = deepcopy(allocation_groups) + task_agent = validate_task_agent( + { + "schema_version": TASK_AGENT_SCHEMA, + "task": task_name, + "goal": task_description, + "semantic_steps": steps, + "allocation_groups": groups, + }, + known_objects=[_scene_runtime_uid(item) for item in scene], + ) + _validate_operator_contracts(task_agent) + return task_agent + + +def _validate_operator_contracts(task_agent: Mapping[str, Any]) -> None: + """Validate capability-specific step shapes inside the planner repair loop.""" + registry = build_default_registry() + for step in task_agent["semantic_steps"]: + operator = str(step["operator"]) + try: + expanded = registry.operator(operator).expand(step) + except (TypeError, ValueError) as error: + raise ValueError( + f"Semantic step {step['id']!r} violates the {operator!r} " + f"operator contract: {error}" + ) from error + if not expanded: + raise ValueError( + f"Semantic step {step['id']!r} produced no executable " + f"{operator!r} operation." + ) + + +def _normalize_semantic_steps( + raw_steps: Sequence[Any], + scene: Sequence[Mapping[str, Any]], +) -> list[dict[str, Any]]: + if not raw_steps: + raise ValueError("Planner semantic_steps must not be empty.") + aliases = _scene_uid_aliases(scene) + normalized: list[dict[str, Any]] = [] + known_ids: set[str] = set() + previous_id: str | None = None + + for index, raw_step in enumerate(raw_steps, start=1): + if not isinstance(raw_step, Mapping): + raise ValueError(f"Planner semantic_steps[{index - 1}] must be an object.") + step = deepcopy(dict(raw_step)) + unknown = sorted(set(step) - _MODEL_STEP_KEYS) + if unknown: + raise ValueError( + f"Planner semantic_steps[{index - 1}] contains unsupported " + f"fields: {unknown}." + ) + operator = _nonempty( + step.get("operator"), + f"semantic_steps[{index - 1}].operator", + ) + configured_id = str(step.get("id", "")).strip() + step_id = configured_id or f"s{index:02d}_{_slug(operator)}" + if step_id in known_ids: + raise ValueError( + f"Planner produced duplicate semantic step ID {step_id!r}." + ) + known_ids.add(step_id) + + result: dict[str, Any] = {"id": step_id, "operator": operator} + if "object" in step: + result["object"] = _resolve_scene_uid( + step["object"], + aliases, + f"semantic step {step_id!r} object", + ) + if "objects" in step: + objects = step["objects"] + if not isinstance(objects, Sequence) or isinstance( + objects, (str, bytes, bytearray) + ): + raise ValueError(f"Semantic step {step_id!r} objects must be a list.") + result["objects"] = [ + _resolve_scene_uid( + object_uid, + aliases, + f"semantic step {step_id!r} objects", + ) + for object_uid in objects + ] + + actor = step.get("actor", {"mode": "auto"}) + if not isinstance(actor, Mapping): + raise ValueError(f"Semantic step {step_id!r} actor must be an object.") + result["actor"] = deepcopy(dict(actor)) + raw_goal = step.get("goal", {}) + if not isinstance(raw_goal, Mapping): + raise ValueError(f"Semantic step {step_id!r} goal must be an object.") + goal = deepcopy(dict(raw_goal)) + for key in ( + "anchor", + "orientation_reference_object", + "reference_object", + "support_object", + ): + if key not in goal or goal[key] in {"table_center", "self"}: + continue + goal[key] = _resolve_scene_uid( + goal[key], + aliases, + f"semantic step {step_id!r} goal.{key}", + ) + result["goal"] = goal + + if "depends_on" in step: + depends_on = step["depends_on"] + if not isinstance(depends_on, Sequence) or isinstance( + depends_on, (str, bytes, bytearray) + ): + raise ValueError( + f"Semantic step {step_id!r} depends_on must be a list." + ) + result["depends_on"] = [str(value) for value in depends_on] + else: + # Sequential is the conservative default. The LLM must explicitly + # emit an empty list when two semantic operations are independent. + result["depends_on"] = [previous_id] if previous_id is not None else [] + normalized.append(result) + previous_id = step_id + return normalized + + +def _fuse_redundant_hold_place_steps( + steps: Sequence[Mapping[str, Any]], +) -> list[dict[str, Any]]: + """Remove a preparatory hold that a complete placement would repeat. + + ``place_relative`` already owns the pickup, transport, release, retreat, + and home phases. A model may nevertheless emit ``hold_hover(object)`` + followed by ``place_relative(object)`` as if the operators were individual + motion commands. The runtime cannot safely transfer that implicit held + state between semantic steps, so normalize the unambiguous one-consumer + pattern before TaskAgent validation. + + A hold with multiple consumers is intentionally left intact because it may + reserve one arm while unrelated branches continue. Compilation rejects any + later reuse of the held object rather than guessing an implicit handover. + """ + result = [deepcopy(dict(step)) for step in steps] + by_id = {step["id"]: step for step in result} + dependents: dict[str, list[str]] = {step_id: [] for step_id in by_id} + for step in result: + for dependency in step["depends_on"]: + if dependency in dependents: + dependents[dependency].append(step["id"]) + + removable: set[str] = set() + claimed_places: set[str] = set() + for hold in result: + if hold["operator"] != "hold_hover": + continue + consumers = dependents[hold["id"]] + if len(consumers) != 1: + continue + place = by_id[consumers[0]] + if place["operator"] != "place_relative" or place.get("object") != hold.get( + "object" + ): + continue + if place["id"] in claimed_places: + raise ValueError( + f"Semantic step {place['id']!r} cannot consume more than one " + "hold_hover state." + ) + if not _is_default_hold_goal(hold): + raise ValueError( + f"Cannot fuse {hold['id']!r} into {place['id']!r}: a " + "non-default hold_hover goal would be discarded." + ) + + place["actor"] = _merge_fused_actors( + hold["actor"], + place["actor"], + hold_id=hold["id"], + place_id=place["id"], + ) + rewritten_dependencies: list[str] = [] + for dependency in place["depends_on"]: + replacements = ( + hold["depends_on"] if dependency == hold["id"] else [dependency] + ) + for replacement in replacements: + if replacement not in rewritten_dependencies: + rewritten_dependencies.append(replacement) + place["depends_on"] = rewritten_dependencies + removable.add(hold["id"]) + claimed_places.add(place["id"]) + + return [step for step in result if step["id"] not in removable] + + +def _is_default_hold_goal(hold: Mapping[str, Any]) -> bool: + """Return whether removing a preparatory hover loses no requested state.""" + goal = hold["goal"] + if set(goal) - { + "orientation_constraint", + "orientation_axis", + "orientation_directed", + "orientation_goal", + "reference_object", + "reference_state", + }: + return False + return ( + goal.get("orientation_axis", "none") == "none" + and not compile_orientation_constraint(goal).terms + and goal.get("reference_state", "initial") == "initial" + and goal.get("reference_object", "self") in ("self", hold.get("object")) + ) + + +def _merge_fused_actors( + hold_actor: Mapping[str, Any], + place_actor: Mapping[str, Any], + *, + hold_id: str, + place_id: str, +) -> dict[str, Any]: + """Preserve an explicit arm requirement while fusing semantic steps.""" + hold = deepcopy(dict(hold_actor)) + place = deepcopy(dict(place_actor)) + hold_mode = hold.get("mode") + place_mode = place.get("mode") + hold_group = hold.get("allocation_group") + place_group = place.get("allocation_group") + if hold_group is not None and place_group is not None and hold_group != place_group: + raise ValueError( + f"Cannot fuse {hold_id!r} into {place_id!r}: conflicting " + "allocation groups would lose explicit arm-allocation intent." + ) + if hold_mode == "required" and place_mode == "required": + if hold.get("arm") != place.get("arm"): + raise ValueError( + f"Cannot fuse {hold_id!r} into {place_id!r}: conflicting " + "required arms would require an unsupported handover." + ) + merged = place + elif hold_mode == "required" and place_mode == "auto": + merged = hold + else: + merged = place + allocation_group = hold_group if hold_group is not None else place_group + if allocation_group is not None: + merged["allocation_group"] = allocation_group + return merged + + +def _render_prompt( + *, + task_name: str, + task_description: str, + scene_objects: Sequence[Mapping[str, Any]], +) -> str: + capabilities = build_default_registry() + return Template(TASK_PLANNER_PROMPT).substitute( + task_name=task_name, + task_description=task_description, + scene_objects=json.dumps( + list(scene_objects), + ensure_ascii=False, + indent=2, + sort_keys=True, + ), + operator_catalog=json.dumps( + capabilities.operator_descriptions(), + ensure_ascii=False, + indent=2, + sort_keys=True, + ), + ) + + +def _default_llm_caller(*, prompt: str, model: str | None) -> Mapping[str, Any]: + """Invoke the configured OpenAI-compatible model with structured output.""" + # Heavy client imports remain lazy so validation and deterministic + # compilation work in minimal simulation test environments. + from langchain_core.messages import HumanMessage, SystemMessage + from langchain_openai import ChatOpenAI + + settings = _load_llm_settings(model=model) + kwargs: dict[str, Any] = { + "api_key": settings["api_key"], + "model": settings["model"], + "temperature": 0, + "http_socket_options": (), + } + if settings["base_url"]: + kwargs["base_url"] = settings["base_url"] + if settings["default_query"]: + kwargs["default_query"] = settings["default_query"] + if _is_mimo_compatible(settings): + # MiMo's OpenAI-compatible endpoint supports JSON mode but not the + # OpenAI ``json_schema`` response format. Disable hidden reasoning so + # the bounded semantic response is not truncated to a few fields. + kwargs.update( + { + "max_completion_tokens": _MIMO_MAX_COMPLETION_TOKENS, + "extra_body": {"thinking": {"type": "disabled"}}, + } + ) + client = ChatOpenAI(**kwargs) + structured = _structured_output_runnable( + client, _MODEL_OUTPUT_SCHEMA, settings=settings + ) + response = structured.invoke( + [ + SystemMessage( + content=( + "Return only the requested route-free semantic plan. " + "Never emit coordinates, atomic actions, or graph edges." + ) + ), + HumanMessage(content=prompt), + ] + ) + return _coerce_model_response(response) + + +_MIMO_MAX_COMPLETION_TOKENS = 4096 + + +def _is_mimo_compatible(settings: Mapping[str, Any]) -> bool: + """Identify MiMo models or regional compatible endpoints without secrets.""" + model = str(settings.get("model", "")).casefold() + base_url = str(settings.get("base_url", "")).casefold() + return "mimo" in model or "xiaomimimo.com" in base_url + + +def _structured_output_runnable( + client: Any, + schema: Mapping[str, Any], + *, + settings: Mapping[str, Any], +) -> Any: + """Bind a portable JSON contract while retaining local strict validation. + + OpenAI-compatible providers do not share the same structured-output + dialect. MiMo documents ``json_object`` JSON mode rather than + ``json_schema``; using the latter can return HTTP 200 with sparse nested + objects. The caller still validates the decoded object against its local + schema after this transport-level binding. + """ + if not hasattr(client, "with_structured_output"): + return client + method = "json_mode" if _is_mimo_compatible(settings) else "json_schema" + try: + return client.with_structured_output(schema, method=method) + except (TypeError, ValueError): + if method == "json_mode" and hasattr(client, "bind"): + # Compatibility with older LangChain adapters that do not expose + # the ``method`` keyword but do support response_format binding. + from langchain_core.output_parsers import JsonOutputParser + + return ( + client.bind(response_format={"type": "json_object"}) + | JsonOutputParser() + ) + # Preserve the historical adapter behavior for non-MiMo providers. + return client.with_structured_output(schema) + + +def _load_llm_settings(*, model: str | None) -> dict[str, Any]: + local_env = _load_env_file(_GEN_SIM_ENV_PATH) + config: dict[str, Any] = {} + if _GEN_CONFIG_PATH.exists(): + with _GEN_CONFIG_PATH.open("r", encoding="utf-8") as stream: + raw = json.load(stream) + if isinstance(raw, Mapping): + llm = raw.get("llm", {}) + if isinstance(llm, Mapping): + configured = llm.get("openai_compatible", {}) + if isinstance(configured, Mapping): + config = dict(configured) + + # A key and endpoint identify one provider transport and must not be mixed + # across process, dotenv, and JSON configuration sources. + api_key, base_url = _resolve_transport_settings(local_env, config) + selected_model = ( + (model.strip() if isinstance(model, str) else "") + or _first_env_value( + local_env, + "ACTION_ENGINE_LLM_MODEL", + "OPENAI_MODEL", + "LLM_MODEL", + ) + or str(config.get("model", "")).strip() + ) + default_query = config.get("default_query", {}) or {} + if not api_key: + raise ValueError( + "OPENAI_API_KEY is required for Action Engine planning. Set it in " + f"the process environment or {_GEN_SIM_ENV_PATH}." + ) + if not selected_model: + raise ValueError( + "An LLM model is required through model=, OPENAI_MODEL, LLM_MODEL, " + f"or {_GEN_CONFIG_PATH}." + ) + if not isinstance(default_query, Mapping): + raise ValueError("LLM default_query must be a mapping.") + return { + "api_key": api_key, + "model": selected_model, + "base_url": base_url, + "default_query": dict(default_query), + } + + +def _resolve_transport_settings( + local_env: Mapping[str, str], + config: Mapping[str, Any], +) -> tuple[str, str]: + """Resolve an API key and endpoint from one configuration source.""" + transports = ( + ( + _mapping_value(os.environ, "OPENAI_API_KEY"), + _mapping_value( + os.environ, + "OPENAI_BASE_URL", + "OPENAI_API_BASE", + "LLM_URL", + ), + ), + ( + _mapping_value(local_env, "OPENAI_API_KEY"), + _mapping_value( + local_env, + "OPENAI_BASE_URL", + "OPENAI_API_BASE", + "LLM_URL", + ), + ), + ( + _mapping_value(config, "api_key"), + _mapping_value(config, "base_url"), + ), + ) + for api_key, base_url in transports: + if api_key and base_url: + return api_key, base_url.rstrip("/") + for api_key, base_url in transports: + if api_key: + return api_key, base_url.rstrip("/") + return "", "" + + +def _mapping_value(source: Mapping[str, Any], *names: str) -> str: + for name in names: + value = source.get(name) + if isinstance(value, str) and value.strip(): + return value.strip() + return "" + + +def _load_env_file(path: Path) -> dict[str, str]: + """Read a local dotenv file without exporting credentials process-wide.""" + if not path.is_file(): + return {} + values: dict[str, str] = {} + for line_number, raw_line in enumerate( + path.read_text(encoding="utf-8").splitlines(), + start=1, + ): + line = raw_line.strip() + if not line or line.startswith("#"): + continue + if line.startswith("export "): + line = line[len("export ") :].lstrip() + if "=" not in line: + continue + key, raw_value = line.split("=", 1) + key = key.strip() + if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", key): + raise ValueError(f"Invalid dotenv key at {path}:{line_number}.") + value = raw_value.strip() + if len(value) >= 2 and value[0] == value[-1] and value[0] in {"'", '"'}: + value = value[1:-1] + elif " #" in value: + value = value.split(" #", 1)[0].rstrip() + values[key] = value + return values + + +def _first_env_value(local_env: Mapping[str, str], *names: str) -> str | None: + """Resolve aliases while keeping every shell value above local dotenv.""" + for source in (os.environ, local_env): + for name in names: + value = source.get(name) + if isinstance(value, str) and value.strip(): + return value.strip() + return None + + +def _coerce_model_response(response: Any) -> Mapping[str, Any]: + if isinstance(response, Mapping): + return dict(response) + model_dump = getattr(response, "model_dump", None) + if callable(model_dump): + dumped = model_dump() + if isinstance(dumped, Mapping): + return dict(dumped) + content = getattr(response, "content", response) + if isinstance(content, Mapping): + return dict(content) + if isinstance(content, list): + content = "\n".join( + str(item.get("text", "")) + for item in content + if isinstance(item, Mapping) and item.get("type") == "text" + ) + if not isinstance(content, str): + raise ValueError( + f"Planner model output has unsupported type {type(content).__name__}." + ) + text = content.strip() + if text.startswith("```"): + lines = text.splitlines() + lines = lines[1:] if lines else lines + lines = lines[:-1] if lines and lines[-1].startswith("```") else lines + text = "\n".join(lines).strip() + parsed = json.loads(text) + if not isinstance(parsed, Mapping): + raise ValueError("Planner model output must decode to a JSON object.") + return dict(parsed) + + +def _normalize_scene_objects( + scene_objects: Sequence[Mapping[str, Any]], +) -> list[dict[str, Any]]: + if not isinstance(scene_objects, Sequence) or isinstance( + scene_objects, (str, bytes, bytearray) + ): + raise ValueError("scene_objects must be a list of mappings.") + normalized: list[dict[str, Any]] = [] + runtime_uids: set[str] = set() + for index, raw_object in enumerate(scene_objects): + if not isinstance(raw_object, Mapping): + raise ValueError(f"scene_objects[{index}] must be a mapping.") + item = deepcopy(dict(raw_object)) + runtime_uid = _scene_runtime_uid(item) + if runtime_uid in runtime_uids: + raise ValueError(f"Duplicate scene runtime UID {runtime_uid!r}.") + runtime_uids.add(runtime_uid) + item["runtime_uid"] = runtime_uid + normalized.append(item) + if not normalized: + raise ValueError("scene_objects must not be empty.") + return normalized + + +def _scene_uid_aliases( + scene_objects: Sequence[Mapping[str, Any]], +) -> dict[str, str]: + aliases: dict[str, str] = {} + for item in scene_objects: + runtime_uid = _scene_runtime_uid(item) + for key in ("runtime_uid", "uid", "source_uid"): + alias = item.get(key) + if isinstance(alias, str) and alias: + existing = aliases.get(alias) + if existing is not None and existing != runtime_uid: + raise ValueError(f"Ambiguous scene object alias {alias!r}.") + aliases[alias] = runtime_uid + return aliases + + +def _resolve_scene_uid(value: Any, aliases: Mapping[str, str], context: str) -> str: + uid = _nonempty(value, context) + try: + return aliases[uid] + except KeyError as exc: + raise ValueError(f"{context} references unknown scene object {uid!r}.") from exc + + +def _scene_runtime_uid(item: Mapping[str, Any]) -> str: + for key in ("runtime_uid", "uid", "source_uid"): + value = item.get(key) + if isinstance(value, str) and value.strip(): + return value + raise ValueError("Every scene object requires runtime_uid, uid, or source_uid.") + + +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() + + +def _slug(value: str) -> str: + slug = _UNSAFE_ID_RE.sub("_", value.lower()).strip("_") + return slug[:48].rstrip("_") or "step" diff --git a/embodichain/gen_sim/action_engine/planning/selection.py b/embodichain/gen_sim/action_engine/planning/selection.py new file mode 100644 index 000000000..6ef92f089 --- /dev/null +++ b/embodichain/gen_sim/action_engine/planning/selection.py @@ -0,0 +1,364 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Score, select, and conservatively fuse whole TaskGroups.""" + +from __future__ import annotations + +from collections import defaultdict, deque +from collections.abc import Collection, Mapping +from copy import deepcopy +from dataclasses import dataclass +from typing import Any + +from embodichain.gen_sim.action_engine.capabilities import ( + AtomicCapabilityRegistry, + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.domain import ( + validate_seed_graph, + validate_task_spec, +) +from embodichain.gen_sim.action_engine.runtime.motion_policy import ( + resolve_motion_policy, +) + +from .linker import link_seed_graph, validate_persisted_contracts + +__all__ = [ + "CandidateEvaluation", + "evaluate_candidate", + "fuse_seed_graphs", + "select_seed_graph", +] + + +@dataclass(frozen=True) +class CandidateEvaluation: + """Auditable static candidate score before any physical execution.""" + + route: str + valid: bool + executable: bool + coverage: float + visual_confidence: float + estimated_cost: float + score: float + errors: tuple[str, ...] = () + + +def evaluate_candidate( + graph: Mapping[str, Any], + task_spec: Mapping[str, Any], + *, + known_objects: Collection[str], + visual_confidence: float = 1.0, + exact_template_match: bool = False, + registry: AtomicCapabilityRegistry | None = None, + robot_profile: str = "dual_ur10", +) -> CandidateEvaluation: + """Apply schema, capabilities, object identity, coverage, and cost scoring.""" + task = validate_task_spec(task_spec) + capabilities = registry or build_atomic_capability_registry() + errors = [] + try: + seed = validate_seed_graph( + graph, + known_objects=known_objects, + known_actions=capabilities.names(), + ) + if seed["capability_catalog_hash"] != capabilities.catalog_hash(): + raise ValueError("SeedGraph capability catalog does not match runtime.") + validate_persisted_contracts(seed, capabilities) + for node in seed["nodes"]: + capabilities.validate_binding(node) + if capabilities.get(str(node["atomic_action"])).runtime_available: + resolve_motion_policy( + robot_profile, + node["atomic_action"], + node["motion_policy"], + ) + except (TypeError, ValueError) as error: + return CandidateEvaluation( + route=str(graph.get("planner_route", "unknown")), + valid=False, + executable=False, + coverage=0.0, + visual_confidence=0.0, + estimated_cost=float("inf"), + score=float("-inf"), + errors=(str(error),), + ) + + required = {str(item["id"]) for item in task["task_instances"]} + provided = {str(group["id"]) for group in seed["task_groups"]} + if task["level"] == "L4": + coverage = 1.0 if provided and seed["success"] else 0.0 + unexpected = set() + mismatched_types = {} + else: + coverage = len(required & provided) / max(len(required), 1) + unexpected = provided - required + if unexpected: + errors.append(f"unexpected task groups: {sorted(unexpected)}") + expected_types = { + str(item["id"]): str(item["task_type"]) for item in task["task_instances"] + } + mismatched_types = { + str(group["id"]): str(group["task_type"]) + for group in seed["task_groups"] + if group["id"] in expected_types + and group["task_type"] != expected_types[group["id"]] + } + if mismatched_types: + errors.append(f"task group type mismatches: {mismatched_types}") + unavailable = sorted( + { + str(node["atomic_action"]) + for node in seed["nodes"] + if not capabilities.get(str(node["atomic_action"])).runtime_available + } + ) + executable = not unavailable + if unavailable: + errors.append(f"planning-only actions: {unavailable}") + confidence = min(max(float(visual_confidence), 0.0), 1.0) + estimated_cost = float(len(seed["nodes"])) + score = coverage * 100.0 - estimated_cost + route = str(seed["planner_route"]) + if exact_template_match and route == "offline": + score += 15.0 + if task["level"] == "L4" and route == "online": + score += 20.0 * confidence + if not executable: + score -= 30.0 + return CandidateEvaluation( + route=route, + valid=not unexpected and not mismatched_types and coverage == 1.0, + executable=executable, + coverage=coverage, + visual_confidence=confidence, + estimated_cost=estimated_cost, + score=score, + errors=tuple(errors), + ) + + +def select_seed_graph( + offline: Mapping[str, Any], + online: Mapping[str, Any], + task_spec: Mapping[str, Any], + *, + known_objects: Collection[str], + visual_confidence: float = 1.0, + exact_template_match: bool = False, + registry: AtomicCapabilityRegistry | None = None, + robot_profile: str = "dual_ur10", +) -> tuple[dict[str, Any], dict[str, CandidateEvaluation]]: + """Choose one complete candidate; ties prefer mature offline templates.""" + evaluations = { + "offline": evaluate_candidate( + offline, + task_spec, + known_objects=known_objects, + visual_confidence=1.0, + exact_template_match=exact_template_match, + registry=registry, + robot_profile=robot_profile, + ), + "online": evaluate_candidate( + online, + task_spec, + known_objects=known_objects, + visual_confidence=visual_confidence, + registry=registry, + robot_profile=robot_profile, + ), + } + valid = [item for item in evaluations.items() if item[1].valid] + if not valid: + messages = {name: evaluation.errors for name, evaluation in evaluations.items()} + raise ValueError(f"Neither SeedGraph candidate is valid: {messages}.") + valid.sort( + key=lambda item: ( + item[1].score, + item[0] == "offline", + ), + reverse=True, + ) + selected = deepcopy(dict(offline if valid[0][0] == "offline" else online)) + selected["planner_route"] = "selected" + selected.setdefault("metadata", {})["selected_from"] = valid[0][0] + return selected, evaluations + + +def fuse_seed_graphs( + offline: Mapping[str, Any], + online: Mapping[str, Any], + group_routes: Mapping[str, str], + *, + registry: AtomicCapabilityRegistry | None = None, +) -> dict[str, Any]: + """Fuse candidates only at complete TaskGroup boundaries.""" + capabilities = registry or build_atomic_capability_registry() + if offline.get("task_id") != online.get("task_id"): + raise ValueError("Cannot fuse graphs for different tasks.") + for field in ("instruction", "level", "reasoning_type", "capability_catalog_hash"): + if offline.get(field) != online.get(field): + raise ValueError(f"Cannot fuse graphs with different {field} values.") + by_route = { + "offline": validate_seed_graph(offline, known_actions=capabilities.names()), + "online": validate_seed_graph(online, known_actions=capabilities.names()), + } + for graph in by_route.values(): + validate_persisted_contracts(graph, capabilities) + groups_by_route = { + route: {str(group["id"]): group for group in graph["task_groups"]} + for route, graph in by_route.items() + } + expected = set(groups_by_route["offline"]) + if set(groups_by_route["online"]) != expected or set(group_routes) != expected: + raise ValueError( + "Fusion requires the same complete TaskGroup set in both graphs." + ) + if set(group_routes.values()) - {"offline", "online"}: + raise ValueError("Every fused TaskGroup route must be offline or online.") + + selected_groups = { + group_id: deepcopy(groups_by_route[route][group_id]) + for group_id, route in group_routes.items() + } + _reject_state_conflicts(selected_groups) + source_nodes = { + route: {str(node["id"]): node for node in graph["nodes"]} + for route, graph in by_route.items() + } + selected_nodes_by_group: dict[str, list[dict[str, Any]]] = {} + id_map: dict[tuple[str, str], str] = {} + for group_id, route in group_routes.items(): + group = selected_groups[group_id] + group.pop("contract", None) + selected_nodes_by_group[group_id] = [] + for node_id in group["node_ids"]: + node = deepcopy(source_nodes[route][node_id]) + fused_id = f"{route}_{node_id}" + id_map[(route, node_id)] = fused_id + node["id"] = fused_id + selected_nodes_by_group[group_id].append(node) + + terminals = {} + for group_id, route in group_routes.items(): + original_ids = set(selected_groups[group_id]["node_ids"]) + referenced = { + dependency + for node_id in original_ids + for dependency in source_nodes[route][node_id]["depends_on"] + if dependency in original_ids + } + terminals[group_id] = [ + id_map[(route, node_id)] + for node_id in selected_groups[group_id]["node_ids"] + if node_id not in referenced + ] + nodes = [] + groups = [] + for group_id in _topological_groups(selected_groups): + route = group_routes[group_id] + group = selected_groups[group_id] + own_original_ids = set(group["node_ids"]) + group_nodes = selected_nodes_by_group[group_id] + for node in group_nodes: + original_id = node["id"][len(route) + 1 :] + original = source_nodes[route][original_id] + internal = [ + id_map[(route, dependency)] + for dependency in original["depends_on"] + if dependency in own_original_ids + ] + external = [ + terminal + for parent in group["depends_on"] + for terminal in terminals[parent] + ] + node["depends_on"] = list(dict.fromkeys([*internal, *external])) + nodes.append(node) + group["node_ids"] = [node["id"] for node in group_nodes] + groups.append(group) + + fused = deepcopy(by_route["offline"]) + fused["planner_route"] = "fused" + fused["nodes"] = nodes + fused["task_groups"] = groups + fused["success"] = {"op": "all", "terms": [group["success"] for group in groups]} + fused["metadata"] = { + "fusion_routes": dict(sorted(group_routes.items())), + "fusion_boundary": "task_group", + } + return link_seed_graph( + fused, + registry=capabilities, + task_order=[str(group["id"]) for group in groups], + ) + + +def _reject_state_conflicts(groups: Mapping[str, Mapping[str, Any]]) -> None: + by_object: dict[str, list[str]] = defaultdict(list) + for group_id, group in groups.items(): + by_object[str(group["object_uid"])].append(group_id) + dependencies = { + group_id: set(group["depends_on"]) for group_id, group in groups.items() + } + + def reaches(child: str, parent: str) -> bool: + pending = list(dependencies[child]) + visited = set() + while pending: + current = pending.pop() + if current == parent: + return True + if current not in visited: + visited.add(current) + pending.extend(dependencies[current]) + return False + + for object_uid, group_ids in by_object.items(): + for index, first in enumerate(group_ids): + for second in group_ids[index + 1 :]: + if not reaches(first, second) and not reaches(second, first): + raise ValueError( + f"Fusion has unordered state changes for object {object_uid!r}." + ) + + +def _topological_groups(groups: Mapping[str, Mapping[str, Any]]) -> list[str]: + outgoing = {group_id: [] for group_id in groups} + indegree = {group_id: 0 for group_id in groups} + for group_id, group in groups.items(): + for parent in group["depends_on"]: + outgoing[parent].append(group_id) + indegree[group_id] += 1 + ready = deque( + sorted(group_id for group_id, degree in indegree.items() if degree == 0) + ) + result = [] + while ready: + group_id = ready.popleft() + result.append(group_id) + for child in sorted(outgoing[group_id]): + indegree[child] -= 1 + if indegree[child] == 0: + ready.append(child) + return result diff --git a/embodichain/gen_sim/action_engine/planning/task_planner_prompt.py b/embodichain/gen_sim/action_engine/planning/task_planner_prompt.py new file mode 100644 index 000000000..9dce0b612 --- /dev/null +++ b/embodichain/gen_sim/action_engine/planning/task_planner_prompt.py @@ -0,0 +1,159 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Prompt template for the Action Engine semantic planner.""" + +from __future__ import annotations + +__all__ = ["TASK_PLANNER_PROMPT"] + +TASK_PLANNER_PROMPT = """You are the semantic planner for a tabletop robot Action Engine. + +Return exactly one JSON object with exactly these two top-level fields: + +{ + "semantic_steps": [ + { + "id": "s01_short_stable_name", + "operator": "", + "object": "", + "actor": {"mode": "auto"}, + "goal": {}, + "depends_on": [] + } + ], + "allocation_groups": [] +} + +For collective operators, replace "object" with "objects": + +{ + "id": "s01_collective_goal", + "operator": "arrange_line", + "objects": ["object_a", "object_b"], + "actor": {"mode": "auto"}, + "goal": {}, + "depends_on": [] +} + +Hard rules: + +- Plan a sequence or DAG of semantic operators. Do not select a task route. +- Emit semantic_steps and allocation_groups only. Do not emit explanations, + confidence, warnings, + atomic actions, graph nodes, graph edges, resources, motion policies, poses, + coordinates, offsets, distances, joint values, trajectories, or tolerances. +- Use runtime_uid values from the scene inventory. Never invent object IDs. +- Preserve every explicit before/after/then dependency with depends_on. +- Use depends_on=[] for genuinely independent operations that may run in + parallel. Otherwise depend on the preceding required semantic step. +- actor.mode is "auto" unless the user explicitly requires one arm. +- allocation_groups expresses an explicit distinct-arm constraint across + independent semantic steps. Use + {"id":"dual_arms_1","semantic_step_ids":["s01","s02"], + "arm_constraint":"distinct_arms"} only when the user explicitly requests + different arms. Merely independent steps must not receive a group. +- An explicitly required arm uses + {"mode": "required", "arm": "left_arm"} or "right_arm". +- Coordinated operators use + {"mode": "coordinated", "arms": ["left_arm", "right_arm"]}. +- Use named symbolic relations and policies only. Runtime observes geometry. +- Every operator is a complete skill, not an individual motion command. + place_relative already picks, transports, releases, retreats, and returns + home. Never emit individual robot motions. +- When the user asks both arms to handle two independent objects, emit two + direct object-level operators with + actor={"mode":"auto"} and depends_on=[], then reference their step IDs in one + allocation_groups entry. The deterministic compiler assigns distinct arms; + do not guess left/right from object positions. +- Spatial phrases that place objects on opposite sides describe object + locations, not an arm-allocation constraint. Emit an allocation group only + when the user explicitly requests both or distinct arms. + +Built-in operator shapes: + +1. arrange_line + - objects: at least two movable objects in requested order. + - goal fields: anchor="table_center"; axis="world_x"|"world_y"| + "table_long_axis"; order_constraint="free"|"ordered"; + order_by="explicit"|"size"|"color"; order_direction="given"| + "ascending"|"descending"; orientation_goal="none"|"preserve"|"upright"| + "lay_flat"|"axis_align"; orientation_axis="none"|"x"|"y"| + "long_axis"|"short_axis". + - In the rotated robot view, world_y is the horizontal left-to-right axis + and world_x is the front-to-back depth axis. For an unspecified line or + row direction, always use axis="world_y". Use axis="world_x" only when + the user explicitly requests a front-to-back, depth-wise, column, or + x-axis layout. Use table_long_axis only when the user explicitly names the + table's long axis; never infer it from a generic line request. + - Use order_constraint="free" when the user wants a line but does not care + which object occupies each slot. + - A line layout does not imply an orientation acceptance requirement. Use + orientation_goal="none" and orientation_axis="none" unless the task + explicitly asks to preserve orientation, make objects upright, lay them + flat, or align an axis. + +2. build_stack + - objects: bottom-to-top movable object order. + - goal fields: stack_mode="on_top"|"nested"; anchor="table_center" or a + passive support runtime_uid; orientation_goal and orientation_axis. + - A vertical stack chain is exactly one build_stack step. Always use the + plural "objects" list, never singular "object", and do not include the + passive anchor in that list. + - Repeated clauses such as "put A on anchor, then put B on top" describe one + chain: objects=[A,B], anchor=anchor. Use separate place_relative steps only + when every object should independently contact the same support. + +3. place_relative + - object: one movable object. + - goal fields: reference_object; relation="inside"|"on"|"left_of"| + "right_of"|"front_of"|"behind"|"front_left_of"|"front_right_of"| + "back_left_of"|"back_right_of"; reference_state="live"|"initial"; + orientation_goal; orientation_axis; optional + orientation_reference_object. + +4. orient_object + - object: one movable object. + - goal fields: orientation_goal="upright"|"lay_flat"|"axis_align"; + orientation_axis="none"|"x"|"y"|"long_axis"|"short_axis"; + support_object=; position_anchor="initial_xy"|"live_xy"; + upright_local_axis="auto"|"long_axis"|"x"|"y"|"z". + - Use orientation_goal="upright" only when the instruction explicitly asks + to make the object upright. + - Use support_object="table" and position_anchor="initial_xy" for an + in-place tabletop orientation request. Use upright_local_axis="auto" + unless the scene inventory explicitly supplies a local semantic axis; + never infer a mesh-local axis from an object name. + +5. coordinated_transport + - object: one shared object moved by both arms. + - goal fields: direction="none"|"world_x"|"world_y"|"front"|"back"| + "left"|"right"|"front_left"|"front_right"|"back_left"|"back_right"| + "up"|"down"; terminal_behavior="hold"|"place"; optional reference_object + and relation; orientation_goal and orientation_axis. + +Available operators: +$operator_catalog + +Task name: +$task_name + +Task description: +$task_description + +Scene inventory: +$scene_objects +""" diff --git a/embodichain/gen_sim/action_engine/planning/vision.py b/embodichain/gen_sim/action_engine/planning/vision.py new file mode 100644 index 000000000..cba17b9b7 --- /dev/null +++ b/embodichain/gen_sim/action_engine/planning/vision.py @@ -0,0 +1,808 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Auditable multi-view observation and VLM fact extraction.""" + +from __future__ import annotations + +import base64 +from collections.abc import Callable, Collection, Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass +from io import BytesIO +import json +import math +import os +from typing import Any + +import torch + +from embodichain.gen_sim.action_engine.domain import ( + VISUAL_RELATION_PARTICIPANTS, + public_task_spec, + requested_visual_task_predicates, +) + +__all__ = [ + "CameraObservation", + "SceneObservation", + "analyze_visual_scene", + "collect_scene_observation", + "validate_visual_facts", +] + +StructuredCaller = Callable[..., Mapping[str, Any]] + +_VISUAL_ENTITY_KEYS = frozenset( + { + "uid", + "camera_uid", + "bbox", + "keypoints", + "visible", + "confidence", + } +) +_VISUAL_RELATION_KEYS = frozenset({"type", "uids", "confidence"}) +_VISUAL_TASK_PREDICATE_KEYS = frozenset({"type", "confidence"}) + + +@dataclass(frozen=True) +class CameraObservation: + """One live camera sample with calibration for one vectorized env row.""" + + uid: str + rgb: torch.Tensor + depth: torch.Tensor | None + intrinsics: torch.Tensor | None + extrinsics: torch.Tensor | None + + +@dataclass(frozen=True) +class SceneObservation: + """Multi-view evidence and stable simulator entity IDs for online planning.""" + + cameras: tuple[CameraObservation, ...] + entities: tuple[dict[str, Any], ...] + env_id: int = 0 + + +_VISUAL_FACTS_SCHEMA = { + "title": "ActionEngineVisualFacts", + "type": "object", + "additionalProperties": False, + "required": ["entities", "relations", "task_predicates", "confidence"], + "properties": { + "entities": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": False, + "required": ["uid", "camera_uid", "confidence"], + "properties": { + "uid": {"type": "string"}, + "camera_uid": {"type": "string"}, + "bbox": { + "type": "array", + "minItems": 4, + "maxItems": 4, + "items": {"type": "number"}, + }, + "keypoints": { + "type": "object", + "additionalProperties": { + "type": "array", + "minItems": 2, + "maxItems": 2, + "items": {"type": "number"}, + }, + }, + "visible": {"type": "boolean"}, + "confidence": {"type": "number", "minimum": 0.0, "maximum": 1.0}, + }, + }, + }, + "relations": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": False, + "required": ["type", "uids", "confidence"], + "properties": { + "type": { + "type": "string", + "enum": sorted(VISUAL_RELATION_PARTICIPANTS), + }, + "uids": { + "type": "array", + "minItems": 2, + "maxItems": 2, + "items": {"type": "string"}, + }, + "confidence": {"type": "number", "minimum": 0.0, "maximum": 1.0}, + }, + }, + }, + "task_predicates": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": False, + "required": ["type", "confidence"], + "properties": { + "type": {"type": "string"}, + "confidence": {"type": "number", "minimum": 0.0, "maximum": 1.0}, + }, + }, + }, + "confidence": {"type": "number", "minimum": 0.0, "maximum": 1.0}, + }, +} + +# Visual facts are deliberately a much smaller contract than a simulator +# snapshot. In particular, accepting arbitrary nested ``attributes`` would +# let a caller smuggle poses/qpos into the online planner while still passing +# the top-level schema. Keep the deny-list here (rather than relying only on +# the SeedGraph validator) because visual facts are persisted and may be +# consumed by an independent planner implementation. +_FORBIDDEN_LIVE_KEYS = frozenset( + { + "absolute_position", + "coordinates", + "extrinsics", + "grasp_pose", + "joint_positions", + "live_pose", + "live_transform", + "object_pose", + "oracle", + "pose", + "positions", + "qpos", + "target_pose", + "trajectory", + "transform", + "waypoints", + "xpos", + } +) + + +def collect_scene_observation( + env: Any, + *, + camera_uids: Sequence[str] | None = None, + env_id: int = 0, +) -> SceneObservation: + """Capture current RGB/depth/calibration and a simulator entity inventory.""" + if env_id < 0 or env_id >= int(env.num_envs): + raise ValueError("env_id is outside the vectorized environment range.") + sim = env.sim + uids = ( + list(camera_uids) + if camera_uids is not None + else list(sim.get_sensor_uid_list()) + ) + cameras = [] + for uid in uids: + sensor = sim.get_sensor(str(uid)) + if sensor is None: + raise ValueError(f"Unknown camera UID {uid!r}.") + update = getattr(sensor, "update", None) + if callable(update): + update() + data = sensor.get_data() + if not isinstance(data, Mapping): + raise TypeError(f"Camera {uid!r} returned non-mapping sensor data.") + rgb_data = data.get("color", data.get("rgb")) + if rgb_data is None: + raise ValueError(f"Camera {uid!r} does not provide RGB data.") + rgb = ( + _env_row( + rgb_data, + env_id, + num_envs=int(env.num_envs), + unbatched_ndim=3, + ) + .detach() + .cpu() + ) + depth = ( + _env_row( + data["depth"], + env_id, + num_envs=int(env.num_envs), + unbatched_ndim=2, + ) + .detach() + .cpu() + if data.get("depth") is not None + else None + ) + intrinsics = _optional_call( + sensor, "get_intrinsics", env_id, num_envs=int(env.num_envs) + ) + extrinsics = _optional_call( + sensor, + "get_arena_pose", + env_id, + num_envs=int(env.num_envs), + to_matrix=True, + ) + cameras.append( + CameraObservation( + uid=str(uid), + rgb=rgb, + depth=depth, + intrinsics=intrinsics, + extrinsics=extrinsics, + ) + ) + if not cameras: + raise ValueError("Online visual planning requires at least one camera.") + + entity_uids = list(sim.get_rigid_object_uid_list()) + articulation_uids = getattr(sim, "get_articulation_uid_list", lambda: [])() + entities = [] + for uid in [*entity_uids, *articulation_uids]: + item: dict[str, Any] = {"uid": str(uid)} + # Do not expose live simulator transforms to the online planner. The + # VLM receives RGB/depth evidence and stable UIDs only; JIT grounding + # resolves world-space targets inside the runtime immediately before + # each action. This also prevents an accidental pose oracle through + # the entity inventory prompt. + entities.append(item) + return SceneObservation(tuple(cameras), tuple(entities), env_id=env_id) + + +def analyze_visual_scene( + observation: SceneObservation, + task_spec: Mapping[str, Any], + *, + model: str | None = None, + caller: StructuredCaller | None = None, + call_counter: list[int] | None = None, +) -> dict[str, Any]: + """Ask a VLM for auditable facts, never hidden reasoning or an action plan.""" + _reject_live_fields(observation.entities, "SceneObservation.entities") + public = public_task_spec(task_spec) + allowed_task_predicates = requested_visual_task_predicates(public) + relation_contracts = { + name: list(participants) + for name, participants in VISUAL_RELATION_PARTICIPANTS.items() + } + _reject_live_fields(public, "PublicTaskSpec") + camera_manifest, images = _camera_evidence(observation) + prompt = ( + "Inspect every supplied camera view. Return only observable facts needed " + "for the task. Refer to simulator entities only by the supplied UID. " + "Use normalized [0,1] bbox/keypoint values, state uncertainty explicitly, " + "and do not provide reasoning or actions. The image blocks appear in the " + "camera_evidence order: each RGB image is followed by that camera's " + "normalized depth image when depth_image_index is present. Camera " + "calibration is input evidence only; never reproduce it in the facts. " + "Use only these canonical spatial relation contracts, whose values give " + "the ordered UID participants: " + f"{json.dumps(relation_contracts, sort_keys=True)}. Put task-level visual " + "judgments in task_predicates, never in relations; their allowed types " + f"are {json.dumps(sorted(allowed_task_predicates))}.\n\n" + f"TaskSpec:\n{json.dumps(public, ensure_ascii=False, sort_keys=True)}\n\n" + f"Entity inventory:\n{json.dumps(observation.entities, ensure_ascii=False, sort_keys=True)}\n\n" + f"Camera evidence:\n{json.dumps(camera_manifest, ensure_ascii=False, sort_keys=True)}" + ) + invoke = caller or _default_structured_caller + # Test/mocked callers own their transport and may intentionally receive no + # configured model. The production caller must resolve strictly through + # the visual-model priority rather than falling back to a text-only model. + selected_model = model if caller is not None else _vlm_model(model) + first_error: Exception | None = None + for attempt in range(2): + current_prompt = prompt + if first_error is not None: + current_prompt += ( + "\n\nThe previous visual-facts JSON was invalid. Return corrected " + f"JSON only. Validation error: {first_error}" + ) + try: + if call_counter is not None: + call_counter[0] += 1 + response = invoke( + prompt=current_prompt, + images=images, + schema=_visual_facts_schema(allowed_task_predicates), + model=selected_model, + ) + facts = validate_visual_facts( + response, + known_uids={str(item["uid"]) for item in observation.entities}, + camera_uids={camera.uid for camera in observation.cameras}, + allowed_task_predicates=allowed_task_predicates, + ) + if facts["confidence"] < 0.5: + raise ValueError( + "VLM visual facts confidence is below the required 0.5 threshold." + ) + if not any( + item.get("visible", True) and item["confidence"] >= 0.5 + for item in facts["entities"] + ): + raise ValueError("VLM visual facts contain no reliable visible entity.") + return facts + except (TypeError, ValueError) as error: + if attempt: + raise ValueError( + "VLM visual facts failed validation after one repair: " f"{error}" + ) from error + first_error = error + raise AssertionError("unreachable") + + +def validate_visual_facts( + value: Mapping[str, Any], + *, + known_uids: set[str], + camera_uids: set[str], + allowed_task_predicates: Collection[str] = (), +) -> dict[str, Any]: + """Validate entity identity and normalized image-space evidence.""" + if not isinstance(value, Mapping): + raise TypeError("VLM visual facts must be a mapping.") + required_fields = {"entities", "relations", "task_predicates", "confidence"} + if set(value) != required_fields: + raise ValueError( + "VLM visual facts require exactly fields " + f"{sorted(required_fields)}; received {sorted(value)}." + ) + confidence = _confidence(value.get("confidence"), "confidence") + entities = value.get("entities") + relations = value.get("relations") + task_predicates = value.get("task_predicates") + if not isinstance(entities, Sequence) or isinstance(entities, (str, bytes)): + raise ValueError("VLM visual facts entities must be a list.") + if not isinstance(relations, Sequence) or isinstance(relations, (str, bytes)): + raise ValueError("VLM visual facts relations must be a list.") + if not isinstance(task_predicates, Sequence) or isinstance( + task_predicates, (str, bytes) + ): + raise ValueError("VLM visual facts task_predicates must be a list.") + normalized_entities = [] + for index, item in enumerate(entities): + if not isinstance(item, Mapping): + raise ValueError(f"visual entities[{index}] must be a mapping.") + unsupported = set(item) - _VISUAL_ENTITY_KEYS + if unsupported: + raise ValueError( + f"visual entities[{index}] contains unsupported fields " + f"{sorted(unsupported)}." + ) + uid = item.get("uid") + camera_uid = item.get("camera_uid") + if not isinstance(uid, str) or not uid: + raise ValueError( + f"visual entities[{index}].uid must be a non-empty string." + ) + if not isinstance(camera_uid, str) or not camera_uid: + raise ValueError( + f"visual entities[{index}].camera_uid must be a non-empty string." + ) + if uid not in known_uids: + raise ValueError( + f"visual entities[{index}] references unknown UID {uid!r}." + ) + if camera_uid not in camera_uids: + raise ValueError( + f"visual entities[{index}] references unknown camera {camera_uid!r}." + ) + normalized = dict(item) + _reject_live_fields(normalized, f"visual entities[{index}]") + if "visible" in normalized and not isinstance(normalized["visible"], bool): + raise ValueError(f"visual entities[{index}].visible must be a boolean.") + if "bbox" in normalized: + normalized["bbox"] = _normalized_vector( + normalized["bbox"], 4, f"visual entities[{index}].bbox" + ) + x_min, y_min, x_max, y_max = normalized["bbox"] + if x_min >= x_max or y_min >= y_max: + raise ValueError( + f"visual entities[{index}].bbox must have non-zero ordered bounds." + ) + keypoints = normalized.get("keypoints", {}) + if not isinstance(keypoints, Mapping): + raise ValueError(f"visual entities[{index}].keypoints must be a mapping.") + normalized["keypoints"] = { + str(name): _normalized_vector(point, 2, f"keypoint {name!r}") + for name, point in keypoints.items() + } + if ( + normalized.get("visible", True) + and "bbox" not in normalized + and not normalized["keypoints"] + ): + raise ValueError( + f"visual entities[{index}] must include a bbox or keypoint evidence." + ) + normalized["confidence"] = _confidence( + normalized.get("confidence"), f"visual entities[{index}].confidence" + ) + normalized_entities.append(normalized) + normalized_relations = [] + for index, relation in enumerate(relations): + if not isinstance(relation, Mapping): + raise ValueError(f"visual relations[{index}] must be a mapping.") + unsupported = set(relation) - _VISUAL_RELATION_KEYS + if unsupported: + raise ValueError( + f"visual relations[{index}] contains unsupported fields " + f"{sorted(unsupported)}." + ) + relation_type = relation.get("type") + if ( + not isinstance(relation_type, str) + or relation_type not in VISUAL_RELATION_PARTICIPANTS + ): + raise ValueError( + f"visual relations[{index}] relation type must be one of " + f"{sorted(VISUAL_RELATION_PARTICIPANTS)}." + ) + participants = relation.get("uids", []) + if not isinstance(participants, Sequence) or isinstance( + participants, (str, bytes) + ): + raise ValueError(f"visual relations[{index}].uids must be a list.") + if any(not isinstance(uid, str) or not uid for uid in participants): + raise ValueError( + f"visual relations[{index}].uids must contain non-empty strings." + ) + expected_count = len(VISUAL_RELATION_PARTICIPANTS[relation_type]) + if len(participants) != expected_count: + raise ValueError( + f"visual relations[{index}].uids must contain exactly " + f"{expected_count} UIDs in canonical participant order." + ) + if len(set(participants)) != len(participants): + raise ValueError( + f"visual relations[{index}].uids must contain distinct UIDs." + ) + invalid = set(participants) - known_uids + if invalid: + raise ValueError( + f"visual relations[{index}] has unknown UIDs {sorted(invalid)}." + ) + normalized = dict(relation) + _reject_live_fields(normalized, f"visual relations[{index}]") + normalized["confidence"] = _confidence( + normalized.get("confidence"), f"visual relations[{index}].confidence" + ) + normalized_relations.append(normalized) + normalized_task_predicates = [] + allowed_predicates = {str(item) for item in allowed_task_predicates} + for index, predicate in enumerate(task_predicates): + if not isinstance(predicate, Mapping): + raise ValueError(f"visual task_predicates[{index}] must be a mapping.") + unsupported = set(predicate) - _VISUAL_TASK_PREDICATE_KEYS + if unsupported or set(predicate) != _VISUAL_TASK_PREDICATE_KEYS: + raise ValueError( + f"visual task_predicates[{index}] requires exactly fields " + f"{sorted(_VISUAL_TASK_PREDICATE_KEYS)}." + ) + predicate_type = predicate.get("type") + if ( + not isinstance(predicate_type, str) + or predicate_type not in allowed_predicates + ): + raise ValueError( + f"visual task_predicates[{index}].type must be one of " + f"{sorted(allowed_predicates)}." + ) + normalized = dict(predicate) + _reject_live_fields(normalized, f"visual task_predicates[{index}]") + normalized["confidence"] = _confidence( + normalized.get("confidence"), + f"visual task_predicates[{index}].confidence", + ) + normalized_task_predicates.append(normalized) + _reject_live_fields( + { + "entities": normalized_entities, + "relations": normalized_relations, + "task_predicates": normalized_task_predicates, + }, + "VLM visual facts", + ) + return { + "entities": normalized_entities, + "relations": normalized_relations, + "task_predicates": normalized_task_predicates, + "confidence": confidence, + } + + +def _visual_facts_schema( + allowed_task_predicates: Collection[str], +) -> dict[str, Any]: + """Return the visual-fact schema specialized for the current task.""" + schema = deepcopy(_VISUAL_FACTS_SCHEMA) + predicate_schema = schema["properties"]["task_predicates"] + allowed = sorted(str(item) for item in allowed_task_predicates) + if allowed: + predicate_schema["items"]["properties"]["type"]["enum"] = allowed + else: + predicate_schema["maxItems"] = 0 + return schema + + +def _default_structured_caller( + *, + prompt: str, + images: Sequence[str], + schema: Mapping[str, Any], + model: str | None, +) -> Mapping[str, Any]: + from langchain_core.messages import HumanMessage, SystemMessage + from langchain_openai import ChatOpenAI + + from .planner import ( + _coerce_model_response, + _is_mimo_compatible, + _load_llm_settings, + _structured_output_runnable, + ) + + settings = _load_llm_settings(model=model) + kwargs: dict[str, Any] = { + "api_key": settings["api_key"], + "model": settings["model"], + "temperature": 0, + "http_socket_options": (), + } + for key in ("base_url", "default_query"): + if settings[key]: + kwargs[key] = settings[key] + if _is_mimo_compatible(settings): + kwargs.update( + { + "max_completion_tokens": 4096, + "extra_body": {"thinking": {"type": "disabled"}}, + } + ) + client = ChatOpenAI(**kwargs) + structured = _structured_output_runnable( + client, + schema, + settings=settings, + ) + schema_prompt = ( + f"{prompt}\n\nReturn one JSON object conforming exactly to this JSON " + f"Schema:\n{json.dumps(schema, ensure_ascii=False, sort_keys=True)}" + ) + content: list[dict[str, Any]] = [{"type": "text", "text": prompt}] + content[0]["text"] = schema_prompt + content.extend( + {"type": "image_url", "image_url": {"url": image}} for image in images + ) + response = structured.invoke( + [ + SystemMessage( + content="Report visual facts only. Never reveal chain-of-thought." + ), + HumanMessage(content=content), + ] + ) + return _coerce_model_response(response) + + +def _vlm_model(explicit: str | None) -> str: + if isinstance(explicit, str) and explicit.strip(): + return explicit.strip() + from .planner import _GEN_SIM_ENV_PATH, _load_env_file + + local_env = _load_env_file(_GEN_SIM_ENV_PATH) + # A VLM-specific choice wins over the generic OpenAI default regardless of + # whether it comes from the shell or the project dotenv. Within each name, + # process variables retain their normal override behavior. + for key in ("ACTION_ENGINE_VLM_MODEL", "OPENAI_MODEL"): + for source in (os.environ, local_env): + value = source.get(key) + if isinstance(value, str) and value.strip(): + return value.strip() + raise ValueError( + "A VLM model is required through --vlm-model, agent_config.vlm_model, " + "ACTION_ENGINE_VLM_MODEL, or OPENAI_MODEL." + ) + + +def _rgb_data_url(value: torch.Tensor) -> str: + from PIL import Image + + image = value + if image.ndim != 3 or image.shape[-1] not in {3, 4}: + raise ValueError("Camera RGB must have shape (H, W, 3|4).") + if image.dtype != torch.uint8: + image = image.float() + if float(image.max()) <= 1.0: + image = image * 255.0 + image = image.clamp(0, 255).to(torch.uint8) + stream = BytesIO() + Image.fromarray(image.numpy()).convert("RGB").save(stream, format="PNG") + return "data:image/png;base64," + base64.b64encode(stream.getvalue()).decode( + "ascii" + ) + + +def _camera_evidence( + observation: SceneObservation, +) -> tuple[list[dict[str, Any]], list[str]]: + """Package calibrated RGB/depth evidence in a stable camera order.""" + manifest: list[dict[str, Any]] = [] + images: list[str] = [] + for camera in observation.cameras: + rgb_index = len(images) + images.append(_rgb_data_url(camera.rgb)) + item: dict[str, Any] = { + "uid": camera.uid, + "rgb_image_index": rgb_index, + "depth_available": camera.depth is not None, + "intrinsics": _calibration_list(camera.intrinsics), + "extrinsics": _calibration_list(camera.extrinsics), + } + if camera.depth is not None: + item["depth_image_index"] = len(images) + images.append(_depth_data_url(camera.depth)) + manifest.append(item) + return manifest, images + + +def _calibration_list(value: torch.Tensor | None) -> list[Any] | None: + """Serialize finite calibration tensors for the transient VLM prompt.""" + if value is None: + return None + tensor = torch.as_tensor(value).detach().cpu() + if not bool(torch.isfinite(tensor).all()): + raise ValueError("Camera calibration contains non-finite values.") + return tensor.tolist() + + +def _depth_data_url(value: torch.Tensor) -> str: + """Render one depth frame as a normalized grayscale VLM evidence image.""" + from PIL import Image + + depth = torch.as_tensor(value).detach().cpu().float() + if depth.ndim == 3 and depth.shape[-1] == 1: + depth = depth[..., 0] + elif depth.ndim == 3 and depth.shape[0] == 1: + depth = depth[0] + if depth.ndim != 2: + raise ValueError("Camera depth must have shape (H, W) or a singleton channel.") + finite = torch.isfinite(depth) + if not bool(finite.any()): + raise ValueError("Camera depth contains no finite values.") + minimum = depth[finite].min() + maximum = depth[finite].max() + normalized = torch.zeros_like(depth) + if float(maximum - minimum) > 0.0: + normalized[finite] = (depth[finite] - minimum) / (maximum - minimum) + image = (normalized.clamp(0.0, 1.0) * 255.0).to(torch.uint8).numpy() + stream = BytesIO() + Image.fromarray(image, mode="L").save(stream, format="PNG") + return "data:image/png;base64," + base64.b64encode(stream.getvalue()).decode( + "ascii" + ) + + +def _env_row( + value: Any, + env_id: int, + *, + num_envs: int | None = None, + unbatched_ndim: int | tuple[int, ...] | None = None, +) -> torch.Tensor: + """Select one vectorized environment row without slicing image dimensions. + + Sensor APIs return either ``(num_envs, ...)`` or an unbatched ``(...)`` + tensor. The old ``shape[0] > env_id`` heuristic sliced the first image row + for an unbatched ``(H, W, C)`` RGB tensor and similarly corrupted 4x4 poses. + Prefer the known environment count and only use the legacy heuristic when + no count is available. + """ + tensor = torch.as_tensor(value) + if unbatched_ndim is not None: + allowed_ndim = ( + (unbatched_ndim,) + if isinstance(unbatched_ndim, int) + else tuple(unbatched_ndim) + ) + if tensor.ndim in allowed_ndim: + return tensor + if tensor.ndim and num_envs is not None and tensor.shape[0] == int(num_envs): + if env_id >= tensor.shape[0]: + raise ValueError("env_id is outside the sensor batch dimension.") + return tensor[env_id] + if num_envs is None and tensor.ndim and tensor.shape[0] > env_id: + return tensor[env_id] + return tensor + + +def _optional_call( + sensor: Any, name: str, env_id: int, *, num_envs: int | None = None, **kwargs: Any +) -> torch.Tensor | None: + method = getattr(sensor, name, None) + if not callable(method): + return None + try: + value = method(env_id=env_id, **kwargs) + except TypeError: + try: + value = method(env_id, **kwargs) + except TypeError: + value = method(**kwargs) + value = torch.as_tensor(value) + # Calibration methods commonly return an unbatched matrix even for a + # vectorized simulator. Select a leading environment row only when the + # shape cannot itself be a canonical calibration matrix. This preserves + # 3x3/4x4 matrices while correctly handling batched compact vectors such as + # ``(num_envs, 4)``. + unbatched_matrix = value.ndim == 2 and tuple(value.shape) in { + (3, 3), + (4, 4), + } + if ( + num_envs is not None + and value.ndim >= 1 + and value.shape[0] == int(num_envs) + and not unbatched_matrix + ): + value = value[env_id] + return value.detach().cpu() + + +def _reject_live_fields(value: Any, context: str) -> None: + """Reject nested simulator state/geometry fields in VLM facts.""" + if isinstance(value, Mapping): + for key, child in value.items(): + if str(key).lower() in _FORBIDDEN_LIVE_KEYS: + raise ValueError( + f"{context} contains forbidden live-state field {key!r}." + ) + _reject_live_fields(child, f"{context}.{key}") + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + for index, child in enumerate(value): + _reject_live_fields(child, f"{context}[{index}]") + + +def _normalized_vector(value: Any, size: int, context: str) -> list[float]: + if ( + not isinstance(value, Sequence) + or isinstance(value, (str, bytes)) + or len(value) != size + ): + raise ValueError(f"{context} must contain {size} normalized values.") + if any( + not isinstance(item, (int, float)) or isinstance(item, bool) for item in value + ): + raise ValueError(f"{context} values must be numeric.") + result = [float(item) for item in value] + if any(not math.isfinite(item) or item < 0.0 or item > 1.0 for item in result): + raise ValueError(f"{context} values must lie in [0, 1].") + return result + + +def _confidence(value: Any, context: str) -> float: + if not isinstance(value, (int, float)) or isinstance(value, bool): + raise ValueError(f"{context} must be a number in [0, 1].") + result = float(value) + if not math.isfinite(result) or result < 0.0 or result > 1.0: + raise ValueError(f"{context} must lie in [0, 1].") + return result diff --git a/embodichain/gen_sim/action_engine/protocol.py b/embodichain/gen_sim/action_engine/protocol.py new file mode 100644 index 000000000..ef704952d --- /dev/null +++ b/embodichain/gen_sim/action_engine/protocol.py @@ -0,0 +1,60 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Cross-layer identifiers owned by Action Engine. + +These values are serialized into generated artifacts, so changing one is a +protocol migration rather than a local rename. +""" + +from __future__ import annotations + +from typing import Final + +__all__ = [ + "ACTION_ENGINE_CONFIG_SCHEMA", + "ACTION_ENGINE_ENV_ID", + "AGENT_CONFIG_FILENAME", + "COMPARISON_FILENAME", + "EXECUTION_PROGRAM_FILENAME", + "EXECUTION_PROGRAM_SCHEMA", + "FAST_GYM_CONFIG_FILENAME", + "SCENE_REQUIREMENTS_FILENAME", + "SCENE_REQUIREMENTS_SCHEMA", + "SEED_TASK_GRAPH_PNG_FILENAME", + "SEED_GRAPH_SCHEMA", + "TASK_SPEC_FILENAME", + "TASK_SPEC_SCHEMA", + "TASK_AGENT_FILENAME", + "TASK_AGENT_SCHEMA", +] + +ACTION_ENGINE_ENV_ID: Final = "ActionEngine-v1" +ACTION_ENGINE_CONFIG_SCHEMA: Final = "action_engine_config_v2" +TASK_AGENT_SCHEMA: Final = "action_engine_task_agent_v1" +EXECUTION_PROGRAM_SCHEMA: Final = "action_engine_execution_graph_v1" +SEED_GRAPH_SCHEMA: Final = "action_engine_seed_graph_v3" +TASK_SPEC_SCHEMA: Final = "action_engine_task_spec_v2" +SCENE_REQUIREMENTS_SCHEMA: Final = "action_engine_scene_requirements_v2" + +FAST_GYM_CONFIG_FILENAME: Final = "fast_gym_config.json" +AGENT_CONFIG_FILENAME: Final = "agent_config.json" +TASK_AGENT_FILENAME: Final = "task_agent.json" +EXECUTION_PROGRAM_FILENAME: Final = "seed_task_graph.json" +SEED_TASK_GRAPH_PNG_FILENAME: Final = "seed_task_graph.png" +TASK_SPEC_FILENAME: Final = "task_spec.json" +SCENE_REQUIREMENTS_FILENAME: Final = "scene_requirements.json" +COMPARISON_FILENAME: Final = "comparison.json" diff --git a/embodichain/gen_sim/action_engine/runtime/__init__.py b/embodichain/gen_sim/action_engine/runtime/__init__.py new file mode 100644 index 000000000..64774425f --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/__init__.py @@ -0,0 +1,66 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Runtime API for the compositional Action Engine.""" + +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/actions.py b/embodichain/gen_sim/action_engine/runtime/actions.py new file mode 100644 index 000000000..d51121f97 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/actions.py @@ -0,0 +1,1480 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Adapt Action Engine requests to the shared typed atomic-action planner.""" + +from __future__ import annotations + +from collections.abc import Mapping +from copy import deepcopy +from dataclasses import replace +import math +from typing import Any + +import torch + +from embodichain.gen_sim.action_engine.capabilities import ( + AtomicCapability, + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.config import default_runtime_policy +from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, + ActionPlan, + AntipodalAffordance, + AtomicActionEngine, + ControlPartCommandProfile, + CoordinatedPickGoal, + DynamicCollisionMode, + EndEffectorPoseGoal, + EntityState, + ExecutionSession, + MotionPolicy, + ObjectSemantics, + PlanningContext, + RecoveryPolicy, + RobotObservation, + RigidObjectSceneProvider, + SceneProvider, + SceneSnapshot, + StateDelta, +) +from embodichain.lab.sim.planners import ( + CuroboPlannerCfg, + CuroboWorldCfg, + MotionGenCfg, + MotionGenerator, + ToppraPlannerCfg, +) +from embodichain.toolkits.graspkit.pg_grasp import ( + AntipodalSamplerCfg, + GraspGeneratorCfg, + GripperCollisionCfg, +) +from embodichain.utils.logger import log_info + +from .grasp_collision_cache import ensure_vhacd_grasp_collision_cache +from .models import ActionOutcome, GroundedAction +from .state import ExecutionState + +__all__ = ["AtomicActionAdapter"] + + +_DEFAULT_PLANNER_POLICY: dict[str, Any] = { + "backend": "curobo", + "single_arm_strategy": "motion_gen", + "coordinated_strategy": "ik_interp", + "fallback_strategy": "ik_interp", + "allow_fallback": True, + "dynamic_collision": False, + "static_obstacle_uids": [], + "dynamic_obstacle_uids": [], + "curobo": { + "log_level": "error", + "obstacle_representation": "cuboid", + "multi_env": False, + "use_cuda_graph": True, + "preserve_plan_samples": False, + "max_attempts": 5, + "collision_activation_distance": 0.01, + }, +} + +# Preserve cuRobo's fixed world shape while disabling intentional-contact objects. +_COLLISION_PARKING_Z_OFFSET = -100.0 + + +def _collision_cache_for_world( + representation: str, obstacle_count: int +) -> dict[str, int]: + """Size cuRobo's fixed collision cache for the generated scene.""" + cache = {"cuboid": 8, "mesh": 2} + if representation in cache: + cache[representation] = max(cache[representation], obstacle_count) + return cache + + +def _supported_kwargs(config_type: type, values: Mapping[str, Any]) -> dict[str, Any]: + names: set[str] = set() + for cls in reversed(config_type.__mro__): + names.update(getattr(cls, "__annotations__", {})) + return {key: value for key, value in values.items() if key in names} + + +def _as_hand_qpos(value: Any, dof: int, device: Any) -> torch.Tensor: + if dof == 0: + return torch.empty(0, dtype=torch.float32, device=device) + result = torch.as_tensor(value, dtype=torch.float32, device=device).flatten() + if result.numel() == 0: + return torch.zeros(dof, dtype=torch.float32, device=device) + if result.numel() == 1: + return result.repeat(dof) + if result.numel() >= dof: + return result[:dof] + repeats = (dof + result.numel() - 1) // result.numel() + return result.repeat(repeats)[:dof] + + +def _diagonal_approach_direction( + horizontal: torch.Tensor, + *, + vertical: float = -1.0, +) -> torch.Tensor: + """Combine one normalized horizontal role direction with a vertical component.""" + horizontal = horizontal.to(dtype=torch.float32) + norm = torch.linalg.vector_norm(horizontal) + if float(norm) <= 1.0e-6: + raise ValueError("Handover role direction must be non-zero.") + horizontal = horizontal / norm + direction = torch.stack( + (horizontal[0], horizontal[1], horizontal.new_tensor(float(vertical))) + ) + return direction / torch.linalg.vector_norm(direction) + + +class AtomicActionAdapter: + """Own the shared atomic engine and preserve Action Engine runtime contracts.""" + + def __init__( + self, + env: Any, + *, + grasp_policy: Mapping[str, Any] | None = None, + planner_policy: Mapping[str, Any] | None = None, + capability_registry: Any | None = None, + scene_provider: SceneProvider | None = None, + ) -> None: + self.env = env + self.num_envs = int(env.num_envs) + self.device = env.device + if grasp_policy is None: + profile = str(getattr(env, "agent_robot_profile", "dual_ur10")) + grasp_policy = default_runtime_policy(profile).grasp + grasp_policy = { + **grasp_policy, + **(getattr(env, "agent_grasp_runtime_defaults", {}) or {}), + } + self.grasp_policy = deepcopy(dict(grasp_policy)) + self.planner_policy = deepcopy(_DEFAULT_PLANNER_POLICY) + if planner_policy is not None: + self._merge_planner_policy(self.planner_policy, planner_policy) + if not self.planner_policy.get("static_obstacle_uids"): + configured = getattr(env, "agent_static_obstacle_uids", ()) or () + if configured: + self.planner_policy["static_obstacle_uids"] = [ + str(uid) for uid in configured + ] + else: + get_rigid_object = getattr(env.sim, "get_rigid_object", None) + if callable(get_rigid_object) and get_rigid_object("table") is not None: + self.planner_policy["static_obstacle_uids"] = ["table"] + self.capabilities = capability_registry or build_atomic_capability_registry() + self._motion_generator: MotionGenerator | None = None + self._atomic_engine: AtomicActionEngine | None = None + self._semantics: dict[str, ObjectSemantics] = {} + self._scene_time = 0.0 + if scene_provider is not None and not isinstance(scene_provider, SceneProvider): + raise TypeError("scene_provider must implement SceneProvider.") + self.scene_provider = scene_provider or self._build_scene_provider() + + @staticmethod + def _merge_planner_policy( + target: dict[str, Any], + update: Mapping[str, Any], + ) -> None: + for key, value in update.items(): + if isinstance(value, Mapping) and isinstance(target.get(key), dict): + AtomicActionAdapter._merge_planner_policy(target[key], value) + else: + target[key] = deepcopy(value) + + def initial_state(self) -> ExecutionState: + """Capture the initial full-robot planning seed.""" + return ExecutionState(last_qpos=self.env.robot.get_qpos().clone()) + + def start_session( + self, + grounded: GroundedAction, + state: ExecutionState | None = None, + ) -> ExecutionSession: + """Start one closed-loop AtomicAction session from live scene state. + + ProgramExecutor may continue using its compatibility scheduler for + compound and per-arm merged trajectories. New callers can use this + boundary to adopt feedback-driven execution without constructing + private planning contexts. + """ + capability = self.capabilities.require_executable(grounded.action_class) + state = state or self.initial_state() + grounded = self._select_upright_transport_yaw(grounded, state) + context = self._planning_context(state, grounded) + invocation = self._invocation(grounded, capability) + return self._engine().start((invocation,), context) + + def _build_scene_provider(self) -> SceneProvider | None: + """Create the shared live rigid-object provider when entities are available.""" + sim = getattr(self.env, "sim", None) + if sim is None: + return None + dynamic_uids = tuple( + str(uid) for uid in self.planner_policy.get("dynamic_obstacle_uids", ()) + ) + list_uids = getattr(sim, "get_rigid_object_uid_list", None) + uids = tuple(str(uid) for uid in list_uids()) if callable(list_uids) else () + if not uids: + uids = dynamic_uids + get_rigid_object = getattr(sim, "get_rigid_object", None) + if not callable(get_rigid_object): + return None + entities = { + uid: entity for uid in uids if (entity := get_rigid_object(uid)) is not None + } + if not entities: + return None + collision_uids = ( + dynamic_uids + if bool(self.planner_policy.get("dynamic_collision", False)) + else () + ) + return RigidObjectSceneProvider( + entities, + collision_entity_ids=collision_uids, + ) + + def semantics(self, uid: str) -> ObjectSemantics: + """Build object semantics once while retaining the live entity handle.""" + cached = self._semantics.get(uid) + if cached is not None: + return cached + entity = self.env.sim.get_rigid_object(uid) + if entity is None: + raise ValueError(f"Unknown grasp target {uid!r}.") + vertices = entity.get_vertices(env_ids=[0], scale=True) + triangles = entity.get_triangles(env_ids=[0]) + if isinstance(vertices, (tuple, list)): + vertices = vertices[0] + if isinstance(triangles, (tuple, list)): + triangles = triangles[0] + vertices = torch.as_tensor(vertices, dtype=torch.float32) + triangles = torch.as_tensor(triangles, dtype=torch.int64) + if vertices.ndim == 3 and vertices.shape[0] == 1: + vertices = vertices[0] + if triangles.ndim == 3 and triangles.shape[0] == 1: + triangles = triangles[0] + if vertices.ndim != 2 or vertices.shape[-1] != 3 or vertices.numel() == 0: + raise ValueError(f"Object {uid!r} has invalid mesh vertices.") + if triangles.ndim != 2 or triangles.shape[-1] != 3 or triangles.numel() == 0: + raise ValueError(f"Object {uid!r} has invalid mesh triangles.") + + grasp_options = self.grasp_policy + sampler = AntipodalSamplerCfg( + n_sample=int(grasp_options["antipodal_n_sample"]), + max_angle=float(grasp_options["antipodal_max_angle"]), + max_length=float(grasp_options["max_open_length"]), + min_length=float(grasp_options["min_open_length"]), + ) + generator = GraspGeneratorCfg( + viser_port=int(grasp_options["viser_port"]), + antipodal_sampler_cfg=sampler, + max_deviation_angle=float(grasp_options["max_deviation_angle"]), + n_deviated_approach_directions=int( + grasp_options["n_deviated_approach_directions"] + ), + ) + max_hulls = int(grasp_options["max_decomposition_hulls"]) + collision = GripperCollisionCfg( + max_open_length=float(grasp_options["max_open_length"]), + finger_length=float(grasp_options["finger_length"]), + point_sample_dense=float(grasp_options["point_sample_dense"]), + max_decomposition_hulls=max_hulls, + ) + cache_result = ensure_vhacd_grasp_collision_cache( + mesh_vertices=vertices, + mesh_triangles=triangles, + max_decomposition_hulls=max_hulls, + ) + if cache_result.status != "hit": + log_info(f"Prepared V-HACD grasp cache for {uid!r}: {cache_result.status}.") + + semantics = ObjectSemantics( + label=uid, + entity=entity, + geometry={"mesh_vertices": vertices, "mesh_triangles": triangles}, + affordance=AntipodalAffordance( + object_label=uid, + mesh_vertices=vertices, + mesh_triangles=triangles, + generator_cfg=generator, + gripper_collision_cfg=collision, + force_reannotate=bool(grasp_options["force_grasp_reannotate"]), + ), + ) + self._semantics[uid] = semantics + return semantics + + def plan( + self, + grounded: GroundedAction, + state: ExecutionState | None = None, + ) -> ActionOutcome: + """Plan one grounded primitive through the mainline typed contract.""" + capability = self.capabilities.require_executable(grounded.action_class) + state = state or self.initial_state() + grounded = self._select_upright_transport_yaw(grounded, state) + context = self._planning_context(state, grounded) + invocation = self._invocation(grounded, capability) + plan = self._engine().plan(invocation, context) + selected_positions = self._positions_with_agent_holds( + plan, + grounded, + capability, + ) + primary_success = plan.plan_success.to(self.device) + reachability_search = None + if bool(grounded.motion_policy.get("retreat_reachability_search", False)): + ( + grounded, + selected_positions, + primary_success, + reachability_search, + ) = self._search_reachable_retreat( + grounded=grounded, + capability=capability, + state=state, + context=context, + invocation=invocation, + initial_positions=selected_positions, + initial_success=primary_success, + ) + invocation = replace(invocation, goal=grounded.target) + combined_success = primary_success.clone() + fallback_plan: ActionPlan | None = None + use_fallback = torch.zeros_like(combined_success) + fallback_attempted = torch.zeros_like(combined_success) + fallback_success = torch.zeros_like(combined_success) + + fallback_strategy = self.planner_policy.get("fallback_strategy") + collision_safety = str(grounded.motion_policy.get("collision_safety", "auto")) + fallback_allowed = bool(self.planner_policy.get("allow_fallback", True)) and ( + collision_safety != "required" + ) + if ( + fallback_allowed + and invocation.motion_policy.strategy == "motion_gen" + and fallback_strategy in {"ik_interp"} + and not bool(combined_success.all()) + ): + fallback_attempted = ~primary_success + fallback_policy = replace( + invocation.motion_policy, + strategy=str(fallback_strategy), + dynamic_collision_mode=DynamicCollisionMode.OFF, + plan_opts=None, + ) + fallback_plan = self._engine().plan( + replace(invocation, motion_policy=fallback_policy), + context, + ) + fallback_positions = self._positions_with_agent_holds( + fallback_plan, + grounded, + capability, + ) + fallback_success = fallback_plan.plan_success.to(self.device) + use_fallback = fallback_attempted & fallback_success + selected_positions = self._merge_plan_rows( + selected_positions, + fallback_positions, + use_fallback, + state.last_qpos, + ) + combined_success |= fallback_plan.plan_success.to(self.device) + + options = invocation.skill_options + if capability.config_materializer == "handover": + combined_success &= self._handover_receiver_hold_mask( + selected_positions, + grounded, + options, + tolerance=float( + grounded.motion_policy.get( + "receiver_hold_joint_tolerance", + 2.0e-3, + ) + ), + ) + + terminal_qpos = ( + selected_positions[:, -1] + if selected_positions.shape[1] + else state.last_qpos + ) + primary_rows = combined_success & primary_success + projected_task = plan.expected_effects.apply( + context.task, + primary_rows, + ) + held_keys = set(plan.expected_effects.held_object_updates) + if fallback_plan is not None: + fallback_rows = combined_success & use_fallback + projected_task = fallback_plan.expected_effects.apply( + projected_task, + fallback_rows, + ) + held_keys.update(fallback_plan.expected_effects.held_object_updates) + committed_effects = StateDelta( + held_object_updates={ + key: projected_task.held_objects.get(key) for key in held_keys + }, + ) + next_state = ExecutionState.from_task_state( + projected_task, + last_qpos=torch.where( + combined_success[:, None], terminal_qpos, state.last_qpos + ), + ) + return ActionOutcome( + trajectory=selected_positions, + success=combined_success, + next_state=next_state, + grounded=grounded, + prior_state=state, + expected_effects=committed_effects, + planner_trace={ + **self._planner_trace( + grounded=grounded, + invocation=invocation, + context=context, + state=state, + primary_success=primary_success, + fallback_allowed=fallback_allowed, + fallback_strategy=( + str(fallback_strategy) + if invocation.motion_policy.strategy == "motion_gen" + and fallback_strategy in {"ik_interp"} + else None + ), + fallback_attempted=fallback_attempted, + fallback_success=fallback_success, + fallback_used=use_fallback, + reachability_search=reachability_search, + ), + # Auditability takes precedence over compactness here: every + # selected planner route retains its complete joint path. + "planned_trajectory": selected_positions.detach().clone(), + }, + ) + + def _search_reachable_retreat( + self, + *, + grounded: GroundedAction, + capability: AtomicCapability, + state: ExecutionState, + context: PlanningContext, + invocation: ActionInvocation, + initial_positions: torch.Tensor, + initial_success: torch.Tensor, + ) -> tuple[GroundedAction, torch.Tensor, torch.Tensor, dict[str, Any]]: + """Select the highest row-local retreat accepted by the live planner.""" + candidates = self._retreat_search_targets(grounded) + target = getattr(grounded.target, "xpos", None) + if not isinstance(target, torch.Tensor) or len(candidates) <= 1: + return ( + grounded, + initial_positions, + initial_success, + { + "strategy": "bounded_motion_planner", + "attempts": [], + "selected_target_z": ( + None + if not isinstance(target, torch.Tensor) + else target[:, 2, 3] + ), + }, + ) + + selected_target = candidates[0][1].clone() + selected_positions = initial_positions + success = initial_success.clone() + attempts: list[dict[str, Any]] = [ + { + "candidate": candidates[0][0], + "target_z": candidates[0][1][:, 2, 3].detach().clone(), + "success": initial_success.detach().clone(), + } + ] + for label, candidate_target in candidates[1:]: + unresolved = ~success + if not bool(unresolved.any()): + break + row_target = torch.where( + unresolved[:, None, None], + candidate_target, + selected_target, + ) + candidate_grounded = replace( + grounded, + target=EndEffectorPoseGoal(xpos=row_target), + ) + candidate_invocation = replace( + invocation, + goal=candidate_grounded.target, + ) + candidate_plan = self._engine().plan(candidate_invocation, context) + candidate_positions = self._positions_with_agent_holds( + candidate_plan, + candidate_grounded, + capability, + ) + candidate_success = candidate_plan.plan_success.to(self.device) + selected_rows = unresolved & candidate_success + selected_positions = self._merge_plan_rows( + selected_positions, + candidate_positions, + selected_rows, + state.last_qpos, + ) + selected_target = torch.where( + selected_rows[:, None, None], + candidate_target, + selected_target, + ) + success |= candidate_success + attempts.append( + { + "candidate": label, + "target_z": candidate_target[:, 2, 3].detach().clone(), + "success": candidate_success.detach().clone(), + } + ) + + metadata = { + "retreat_selected_target_z": selected_target[:, 2, 3].detach().clone(), + "retreat_reachability_found": success.detach().clone(), + } + selected_grounded = replace( + grounded, + target=EndEffectorPoseGoal(xpos=selected_target), + cfg={**grounded.cfg, **metadata}, + motion_policy={**grounded.motion_policy, **metadata}, + ) + return ( + selected_grounded, + selected_positions, + success, + { + "strategy": "bounded_motion_planner", + "attempts": attempts, + "selected_target_z": selected_target[:, 2, 3].detach().clone(), + }, + ) + + def _retreat_search_targets( + self, + grounded: GroundedAction, + ) -> list[tuple[str, torch.Tensor]]: + """Build bounded height and baseward retreat candidates from live poses.""" + target = getattr(grounded.target, "xpos", None) + reference = grounded.motion_policy.get("retreat_reference_pose") + if not isinstance(target, torch.Tensor) or not isinstance( + reference, torch.Tensor + ): + return [] + target = target.to(device=self.device, dtype=torch.float32) + reference = reference.to(device=self.device, dtype=torch.float32) + if target.shape == (4, 4): + target = target.unsqueeze(0).repeat(self.num_envs, 1, 1) + if reference.shape == (4, 4): + reference = reference.unsqueeze(0).repeat(self.num_envs, 1, 1) + expected = (self.num_envs, 4, 4) + if target.shape != expected or reference.shape != expected: + return [] + + sample_count = int(grounded.cfg.get("retreat_search_samples", 6)) + if not 2 <= sample_count <= 16: + raise ValueError("retreat_search_samples must be in [2, 16].") + minimum_height = float(grounded.cfg.get("minimum_retreat_height", 0.05)) + if not math.isfinite(minimum_height) or minimum_height < 0.0: + raise ValueError("minimum_retreat_height must be finite and non-negative.") + desired_height = torch.clamp( + target[:, 2, 3] - reference[:, 2, 3], + min=0.0, + ) + minimum = torch.minimum( + desired_height, + torch.full_like(desired_height, minimum_height), + ) + fractions = torch.linspace( + 1.0, + 0.0, + sample_count, + dtype=target.dtype, + device=target.device, + ) + heights = ( + minimum[:, None] + (desired_height - minimum)[:, None] * fractions[None] + ) + candidates: list[tuple[str, torch.Tensor]] = [("requested", target.clone())] + for index in range(1, sample_count): + candidate = target.clone() + candidate[:, 2, 3] = reference[:, 2, 3] + heights[:, index] + candidates.append((f"height_{index}", candidate)) + + from .frames import arm_base_poses + + left_base, right_base = arm_base_poses(self.env) + base = left_base if grounded.arm == "left_arm" else right_base + direction = base[:, :2, 3] - reference[:, :2, 3] + norm = torch.linalg.vector_norm(direction, dim=1, keepdim=True) + direction = torch.where( + norm > 1.0e-6, + direction / torch.clamp(norm, min=1.0e-6), + torch.zeros_like(direction), + ) + distance = float(grounded.cfg.get("retreat_distance", 0.10)) + if not math.isfinite(distance) or distance < 0.0: + raise ValueError("retreat_distance must be finite and non-negative.") + for index in range(sample_count): + candidate = target.clone() + candidate[:, :2, 3] = reference[:, :2, 3] + direction * distance + candidate[:, 2, 3] = reference[:, 2, 3] + heights[:, index] + candidates.append((f"baseward_{index}", candidate)) + return candidates + + def _planner_trace( + self, + *, + grounded: GroundedAction, + invocation: ActionInvocation, + context: PlanningContext, + state: ExecutionState, + primary_success: torch.Tensor, + fallback_allowed: bool, + fallback_strategy: str | None, + fallback_attempted: torch.Tensor, + fallback_success: torch.Tensor, + fallback_used: torch.Tensor, + reachability_search: Mapping[str, Any] | None = None, + ) -> dict[str, Any]: + """Build compact per-row evidence for the planner route actually used.""" + exclusions = self._collision_exclusion_masks(grounded, state) + obstacle_positions = { + uid: context.scene.entities[uid].pose[:, :3, 3].detach().clone() + for uid in context.scene.collision_entity_ids + } + revisions = torch.as_tensor( + context.scene.collision_world_revisions(self.num_envs), + dtype=torch.int64, + device=self.device, + ) + trace = { + "action_class": grounded.action_class, + "arm": grounded.arm, + "planner": invocation.motion_policy.planner, + "primary_strategy": invocation.motion_policy.strategy, + "dynamic_collision_mode": invocation.motion_policy.dynamic_collision_mode.value, + "primary_success": primary_success.detach().clone(), + "fallback_allowed": fallback_allowed, + "fallback_strategy": fallback_strategy, + "fallback_attempted": fallback_attempted.detach().clone(), + "fallback_success": fallback_success.detach().clone(), + "fallback_used": fallback_used.detach().clone(), + "search_budget": { + "primary_max_attempts": int( + self.planner_policy.get("curobo", {}).get("max_attempts", 1) + ), + "fallback_enabled": bool(fallback_allowed), + }, + "collision_world_revision": revisions, + "collision_obstacle_positions": obstacle_positions, + "collision_exclusions": { + uid: mask.detach().clone() for uid, mask in exclusions.items() + }, + } + if reachability_search is not None: + trace["reachability_search"] = deepcopy(dict(reachability_search)) + options = invocation.skill_options + object_part = getattr(options, "pick_object_part", None) + approach_direction = getattr(options, "approach_direction", None) + if object_part is None: + object_part = getattr(options, "receive_pick_object_part", None) + approach_direction = getattr( + options, + "receive_approach_direction", + approach_direction, + ) + if object_part is not None: + grasp_policy: dict[str, Any] = {"object_part": str(object_part)} + if isinstance(approach_direction, torch.Tensor): + direction = approach_direction.to(dtype=torch.float32) + norm = torch.linalg.vector_norm(direction) + if bool(torch.isfinite(norm)) and float(norm) > 0.0: + grasp_policy["approach_direction"] = ( + (direction / norm).detach().cpu().tolist() + ) + trace["grasp_policy"] = grasp_policy + return trace + + def _select_upright_transport_yaw( + self, + grounded: GroundedAction, + state: ExecutionState, + ) -> GroundedAction: + """Choose the closest IK-feasible yaw for an upright object target.""" + sample_count = int(grounded.cfg.get("upright_yaw_samples", 1)) + capability = self.capabilities.get(grounded.action_class) + if ( + capability.target_materializer != "semantic_held_object" + or sample_count <= 1 + ): + return grounded + target_pose = getattr(grounded.target, "object_target_pose", None) + if not isinstance(target_pose, torch.Tensor): + return grounded + target_pose = target_pose.to(device=self.device, dtype=torch.float32) + if target_pose.shape == (4, 4): + target_pose = target_pose.unsqueeze(0).repeat(self.num_envs, 1, 1) + if target_pose.shape != (self.num_envs, 4, 4): + raise ValueError( + "Upright transport target must have shape (4, 4) or (N, 4, 4)." + ) + + arm_part, _, _ = self._parts(grounded.arm) + held = state.get_held_object(arm_part) + if held is None: + return grounded + object_to_eef = held.object_to_eef.to( + device=self.device, + dtype=target_pose.dtype, + ) + if object_to_eef.shape == (4, 4): + object_to_eef = object_to_eef.unsqueeze(0).repeat(self.num_envs, 1, 1) + variants = self._upright_yaw_variants(target_pose, sample_count) + eef_variants = torch.matmul(variants, object_to_eef[:, None]) + joint_ids = list(self.env.robot.get_joint_ids(name=arm_part)) + start_qpos = state.last_qpos[:, joint_ids] + seeds = start_qpos[:, None].expand(-1, sample_count, -1) + success, qpos = self.env.robot.compute_batch_ik( + pose=eef_variants, + name=arm_part, + joint_seed=seeds, + ) + success = torch.as_tensor( + success, + dtype=torch.bool, + device=self.device, + ).reshape(self.num_envs, sample_count) + qpos = torch.as_tensor(qpos, dtype=torch.float32, device=self.device) + success &= torch.isfinite(qpos).all(dim=-1) + distance = torch.linalg.vector_norm(qpos - seeds, dim=-1) + distance = torch.where( + success, + distance, + torch.full_like(distance, torch.inf), + ) + best = distance.argmin(dim=1) + env_ids = torch.arange(self.num_envs, device=self.device) + selected = variants[env_ids, best] + selected = torch.where( + success.any(dim=1)[:, None, None], + selected, + target_pose, + ) + return replace( + grounded, + target=replace(grounded.target, object_target_pose=selected), + target_object_pose=selected, + ) + + @staticmethod + def _upright_yaw_variants( + target_pose: torch.Tensor, + sample_count: int, + ) -> torch.Tensor: + signed_steps = [0] + for step in range(1, (sample_count + 1) // 2): + signed_steps.extend((step, -step)) + if sample_count % 2 == 0: + signed_steps.append(sample_count // 2) + angles = target_pose.new_tensor(signed_steps) * (2.0 * math.pi / sample_count) + yaw = target_pose.new_zeros((sample_count, 3, 3)) + yaw[:, 0, 0] = torch.cos(angles) + yaw[:, 0, 1] = -torch.sin(angles) + yaw[:, 1, 0] = torch.sin(angles) + yaw[:, 1, 1] = torch.cos(angles) + yaw[:, 2, 2] = 1.0 + variants = target_pose[:, None].repeat(1, sample_count, 1, 1) + variants[:, :, :3, :3] = torch.matmul(yaw[None], target_pose[:, None, :3, :3]) + return variants + + def _planning_context( + self, + state: ExecutionState, + grounded: GroundedAction, + ) -> PlanningContext: + qpos = state.last_qpos.to(device=self.device, dtype=torch.float32) + get_qvel = getattr(self.env.robot, "get_qvel", None) + qvel = get_qvel() if callable(get_qvel) else None + if not isinstance(qvel, torch.Tensor) or qvel.shape != qpos.shape: + qvel = torch.zeros_like(qpos) + else: + qvel = qvel.to(device=self.device, dtype=qpos.dtype) + return PlanningContext( + robot=RobotObservation(timestamp=self._scene_time, qpos=qpos, qvel=qvel), + task=state.to_task_state(), + scene=self._scene_snapshot(grounded, state), + env_ids=torch.arange( + self.num_envs, + dtype=torch.long, + device=self.device, + ), + ) + + def _scene_snapshot( + self, + grounded: GroundedAction, + state: ExecutionState, + ) -> SceneSnapshot: + dynamic_uids = tuple( + str(uid) for uid in self.planner_policy.get("dynamic_obstacle_uids", ()) + ) + env_ids = torch.arange( + self.num_envs, + dtype=torch.long, + device=self.device, + ) + if self.scene_provider is None: + base = SceneSnapshot(timestamp=self._scene_time, version=0) + else: + base = self.scene_provider.snapshot( + timestamp=self._scene_time, + env_ids=env_ids, + ) + if not bool(self.planner_policy.get("dynamic_collision", False)): + return base + exclusion_masks = self._collision_exclusion_masks(grounded, state) + entities = dict(base.entities) + for uid in dynamic_uids: + entity_state = entities.get(uid) + if entity_state is None: + raise ValueError( + f"SceneProvider omitted cuRobo dynamic obstacle {uid!r}." + ) + pose = entity_state.pose.to(dtype=torch.float32, device=self.device) + if pose.shape == (4, 4): + pose = pose.unsqueeze(0).repeat(self.num_envs, 1, 1) + if pose.shape != (self.num_envs, 4, 4): + raise ValueError( + f"Dynamic obstacle {uid!r} pose must have shape (4, 4) or " + f"({self.num_envs}, 4, 4), got {tuple(pose.shape)}." + ) + excluded = exclusion_masks.get(uid) + if excluded is not None and bool(excluded.any()): + pose = pose.clone() + pose[excluded, 2, 3] += _COLLISION_PARKING_Z_OFFSET + entities[uid] = EntityState( + pose=pose, + confidence=entity_state.confidence, + ) + return SceneSnapshot( + timestamp=base.timestamp, + version=base.version, + entities=entities, + collision_world_revision=base.collision_world_revision, + collision_entity_ids=dynamic_uids, + ) + + def _collision_exclusion_masks( + self, + grounded: GroundedAction, + state: ExecutionState, + ) -> dict[str, torch.Tensor]: + """Return per-environment masks for obstacles intentionally in contact.""" + dynamic_uids = { + str(uid) for uid in self.planner_policy.get("dynamic_obstacle_uids", ()) + } + masks: dict[str, torch.Tensor] = {} + + def include(uid: str | None, env_mask: torch.Tensor | None = None) -> None: + if uid is None or uid not in dynamic_uids: + return + mask = ( + torch.ones(self.num_envs, dtype=torch.bool, device=self.device) + if env_mask is None + else torch.as_tensor( + env_mask, + dtype=torch.bool, + device=self.device, + ).reshape(-1) + ) + if mask.shape != (self.num_envs,): + raise ValueError( + f"Collision exclusion mask for {uid!r} must have shape " + f"({self.num_envs},), got {tuple(mask.shape)}." + ) + masks[uid] = masks.get(uid, torch.zeros_like(mask)) | mask + + if self.capabilities.get(grounded.action_class).allows_target_contact: + target_uid = grounded.object_uid + if target_uid is None: + target_uid = getattr( + getattr(grounded.target, "semantics", None), + "label", + None, + ) + include(target_uid) + + for held in state.held_objects.values(): + include(held.semantics.label, held.env_mask) + collision_exclusion_uids = grounded.motion_policy.get( + "collision_exclusion_uids", () + ) + if isinstance(collision_exclusion_uids, str): + collision_exclusion_uids = (collision_exclusion_uids,) + for uid in collision_exclusion_uids: + include(str(uid)) + return masks + + def _invocation( + self, + grounded: GroundedAction, + capability: AtomicCapability, + ) -> ActionInvocation: + if capability.resource_mode == "coordinated_object": + strategy = str(self.planner_policy["coordinated_strategy"]) + elif grounded.control == "hand": + strategy = "ik_interp" + else: + strategy = str(self.planner_policy["single_arm_strategy"]) + sample_count = max(2, int(grounded.cfg.get("sample_interval", 50))) + control_dt = float(getattr(self.env, "step_dt", 1.0 / 60.0)) + dynamic_collision = bool(self.planner_policy.get("dynamic_collision", False)) + collision_required = ( + grounded.motion_policy.get("collision_safety") == "required" + ) + if dynamic_collision and strategy == "motion_gen": + dynamic_mode = ( + DynamicCollisionMode.REQUIRED + if collision_required + else DynamicCollisionMode.AUTO + ) + else: + dynamic_mode = DynamicCollisionMode.OFF + goal = ( + self._coordinated_pickment_goal(grounded) + if capability.config_materializer == "coordinated_pickment" + else grounded.target + ) + return ActionInvocation( + skill_id=str(capability.action_type.skill_id), + goal=goal, + binding=self._binding(grounded, capability), + motion_policy=MotionPolicy( + planner=str(self.planner_policy["backend"]), + strategy=strategy, + sample_count=sample_count, + control_dt=control_dt, + velocity_limit=grounded.cfg.get("velocity_limit"), + acceleration_limit=grounded.cfg.get("acceleration_limit"), + dynamic_collision_mode=dynamic_mode, + ), + recovery_policy=RecoveryPolicy(), + skill_options=self._build_config(grounded, capability), + ) + + @staticmethod + def _coordinated_pickment_goal(grounded: GroundedAction) -> CoordinatedPickGoal: + """Apply GenSim-only coordinated grasp filtering to an owned goal copy.""" + target = grounded.target + if not isinstance(target, CoordinatedPickGoal): + raise TypeError("CoordinatedPickment requires a CoordinatedPickGoal.") + requested = grounded.cfg.get("is_filter_ground_collision") + if requested is None: + return target + if not isinstance(requested, bool): + raise TypeError("is_filter_ground_collision must be a boolean.") + semantics = target.semantics + affordance = semantics.affordance + if not isinstance(affordance, AntipodalAffordance): + raise TypeError( + "CoordinatedPickment requires an AntipodalAffordance for GenSim " + "grasp filtering." + ) + generator_cfg = deepcopy(affordance.generator_cfg or GraspGeneratorCfg()) + generator_cfg.is_filter_ground_collision = requested + scoped_affordance = replace(affordance, generator_cfg=generator_cfg) + scoped_semantics = replace(semantics, affordance=scoped_affordance) + return replace(target, semantics=scoped_semantics) + + def _binding( + self, + action: GroundedAction, + capability: AtomicCapability, + ) -> ActionBinding: + if capability.config_materializer == "handover": + transfer_side = str(action.cfg.get("transfer_arm", "left_arm")) + receive_side = "right_arm" if transfer_side == "left_arm" else "left_arm" + transfer_arm, transfer_hand, _ = self._parts(transfer_side) + receive_arm, receive_hand, _ = self._parts(receive_side) + if transfer_hand is None or receive_hand is None: + raise ValueError("HandOver requires two configured end effectors.") + return ActionBinding( + manipulators={"source": transfer_arm, "destination": receive_arm}, + end_effectors={"source": transfer_hand, "destination": receive_hand}, + ) + if capability.config_materializer == "coordinated_pickment": + left_arm, left_hand, _ = self._parts("left_arm") + right_arm, right_hand, _ = self._parts("right_arm") + if left_hand is None or right_hand is None: + raise ValueError("Coordinated pickup requires two end effectors.") + return ActionBinding( + manipulators={"left": left_arm, "right": right_arm}, + end_effectors={"left": left_hand, "right": right_hand}, + ) + if capability.config_materializer == "coordinated_placement": + placing_arm, placing_hand, _ = self._parts("left_arm") + support_arm, support_hand, _ = self._parts("right_arm") + if placing_hand is None or support_hand is None: + raise ValueError("Coordinated placement requires two end effectors.") + return ActionBinding( + manipulators={"placing": placing_arm, "support": support_arm}, + end_effectors={"placing": placing_hand, "support": support_hand}, + ) + + arm_part, hand_part, _ = self._parts(action.arm) + control_part = hand_part if action.control == "hand" else arm_part + if control_part is None: + raise ValueError(f"{action.arm} has no configured {action.control} part.") + end_effectors: dict[str, str] = {} + if capability.action_type.end_effector_roles: + if hand_part is None: + raise ValueError(f"{capability.name} requires an end effector.") + end_effectors["primary"] = hand_part + return ActionBinding( + manipulators={"primary": control_part}, + end_effectors=end_effectors, + ) + + def _build_config( + self, + action: GroundedAction, + capability: AtomicCapability | type, + ) -> Any: + """Build the mainline immutable ``ActionOptions`` value. + + The method name is retained as a narrow compatibility hook for existing + Action Engine tests and extensions; it no longer constructs legacy + hardware-bound ``ActionCfg`` objects. + """ + if isinstance(capability, type): + registered = self.capabilities.require_executable(action.action_class) + if registered.config_type is not capability: + raise ValueError( + f"Options type {capability.__name__!r} does not match " + f"AtomicAction {action.action_class!r}." + ) + capability = registered + if capability.config_materializer_hook is not None: + return capability.config_materializer_hook( + adapter=self, + action=action, + capability=capability, + ) + builder = getattr( + self, + f"_build_{capability.config_materializer}_config", + self._build_single_arm_config, + ) + return builder(action, capability) + + def _config_policy(self, action: GroundedAction) -> dict[str, Any]: + policy = dict(action.cfg) + for key in ( + "postcondition_tolerance", + "relation_distance", + "hover_height", + "staging_lift_height", + "transport_clearance", + "surface_clearance", + "receiver_hold_joint_tolerance", + "post_hold_steps", + ): + policy.pop(key, None) + return policy + + def _build_single_arm_config( + self, + action: GroundedAction, + capability: AtomicCapability, + ) -> Any: + policy = self._config_policy(action) + config_type = capability.config_type + if capability.target_materializer == "semantic_held_object": + from .atomic_compat import ExactTargetMoveHeldObjectOptions + + config_type = ExactTargetMoveHeldObjectOptions + if int(action.cfg.get("upright_yaw_samples", 1)) > 1: + policy["allow_automatic_transport_rotation"] = False + if capability.target_materializer == "press": + press_depth = policy.pop("press_depth", None) + if press_depth is not None and "press_distance" not in policy: + policy["press_distance"] = press_depth + approach_mode = policy.pop("approach_direction_mode", None) + if approach_mode == "handover_transfer": + from .frames import robot_frame_axes + + _, lateral = robot_frame_axes(self.env) + outward = lateral[0] if action.arm == "left_arm" else -lateral[0] + policy["approach_direction"] = _diagonal_approach_direction( + -outward.to(device=self.device) + ) + elif approach_mode is not None: + raise ValueError(f"Unknown approach_direction_mode {approach_mode!r}.") + for name in ("approach_direction", "obj_upright_direction"): + if name in policy and not isinstance(policy[name], torch.Tensor): + policy[name] = torch.as_tensor( + policy[name], dtype=torch.float32, device=self.device + ) + return config_type(**_supported_kwargs(config_type, policy)) + + def _build_coordinated_pickment_config( + self, + action: GroundedAction, + capability: AtomicCapability, + ) -> Any: + from .frames import arm_base_poses + + policy = self._config_policy(action) + left_base, right_base = arm_base_poses(self.env) + direction = right_base[0, :3, 3] - left_base[0, :3, 3] + norm = torch.linalg.vector_norm(direction) + if not torch.isfinite(direction).all() or norm <= 1.0e-6: + raise ValueError( + "Coordinated pickup requires distinct finite left/right arm bases." + ) + policy["left_to_right_arm_direction"] = direction / norm + return capability.config_type( + **_supported_kwargs(capability.config_type, policy) + ) + + def _build_coordinated_placement_config( + self, + action: GroundedAction, + capability: AtomicCapability, + ) -> Any: + return self._build_single_arm_config(action, capability) + + def _build_handover_config( + self, + action: GroundedAction, + capability: AtomicCapability, + ) -> Any: + policy = self._config_policy(action) + middle = action.cfg.get("middle_object_pose") + final = action.cfg.get("final_object_pose") + if middle is None or final is None: + raise ValueError("HandOver grounding must provide middle and final poses.") + transfer_side = str(action.cfg.get("transfer_arm", "left_arm")) + receive_side = "right_arm" if transfer_side == "left_arm" else "left_arm" + from .frames import robot_frame_axes + + _, lateral = robot_frame_axes(self.env) + receiver_outward = ( + lateral[0] if receive_side == "left_arm" else -lateral[0] + ).to(device=self.device) + receiver_inward_approach = -receiver_outward + policy.update( + { + "middle_object_pose": middle, + # Delivery is represented by a following MoveHeldObject node. + # Keep the receiver fixed while the source retreats here. + "final_object_pose": middle, + "receive_approach_direction": _diagonal_approach_direction( + receiver_inward_approach + ), + } + ) + return capability.config_type( + **_supported_kwargs(capability.config_type, policy) + ) + + def _positions_with_agent_holds( + self, + plan: ActionPlan, + grounded: GroundedAction, + capability: AtomicCapability, + ) -> torch.Tensor: + positions = plan.trajectory.positions.to( + device=self.device, + dtype=torch.float32, + ) + hold_steps = int(grounded.cfg.get("post_hold_steps", 0)) + if capability.state_effect != "release" or hold_steps <= 0: + return positions + release = next((item for item in plan.segments if item.name == "release"), None) + if release is None or release.stop <= 0 or release.stop > positions.shape[1]: + return positions + hold = positions[:, release.stop - 1 : release.stop].repeat(1, hold_steps, 1) + return torch.cat( + (positions[:, : release.stop], hold, positions[:, release.stop :]), + dim=1, + ) + + @staticmethod + def _merge_plan_rows( + primary: torch.Tensor, + fallback: torch.Tensor, + use_fallback: torch.Tensor, + hold_qpos: torch.Tensor, + ) -> torch.Tensor: + steps = max(primary.shape[1], fallback.shape[1], 1) + + def padded(value: torch.Tensor) -> torch.Tensor: + if value.shape[1] == 0: + return hold_qpos[:, None].repeat(1, steps, 1) + if value.shape[1] < steps: + value = torch.cat( + (value, value[:, -1:].repeat(1, steps - value.shape[1], 1)), + dim=1, + ) + return value + + primary = padded(primary) + fallback = padded(fallback) + return torch.where(use_fallback[:, None, None], fallback, primary) + + def _handover_receiver_hold_mask( + self, + trajectory: torch.Tensor, + grounded: GroundedAction, + options: Any, + *, + tolerance: float, + ) -> torch.Tensor: + if tolerance < 0.0: + raise ValueError("receiver_hold_joint_tolerance must be non-negative.") + retreat_steps = max(2, int(options.retreat_steps)) + if trajectory.shape[1] < retreat_steps: + return torch.zeros( + self.num_envs, dtype=torch.bool, device=trajectory.device + ) + transfer_side = str(grounded.cfg.get("transfer_arm", "left_arm")) + receive_side = "right_arm" if transfer_side == "left_arm" else "left_arm" + receive_arm, _, _ = self._parts(receive_side) + receiver_ids = self.env.robot.get_joint_ids(name=receive_arm) + receiver = trajectory[:, -retreat_steps:, receiver_ids] + drift = torch.amax(torch.abs(receiver - receiver[:, :1]), dim=(1, 2)) + return torch.isfinite(drift) & (drift <= tolerance) + + def execute_trajectory( + self, + trajectory: torch.Tensor, + *, + active: torch.Tensor, + ) -> list[torch.Tensor]: + """Advance the environment while holding inactive vectorized rows.""" + if trajectory.ndim != 3 or trajectory.shape[0] != self.num_envs: + raise ValueError("Execution trajectory must have shape (N, T, robot_dof).") + active = active.to(device=trajectory.device, dtype=torch.bool) + current = self.env.robot.get_qpos().to( + device=trajectory.device, + dtype=trajectory.dtype, + ) + commands: list[torch.Tensor] = [] + for waypoint in trajectory.unbind(dim=1): + command = torch.where(active[:, None], waypoint, current) + self.env.step(command) + self._scene_time += self._scene_step_duration() + update = getattr(self.env, "update_obj_info", None) + if callable(update): + update() + commands.append(command.detach()) + current = command + sync = getattr(self.env, "sync_agent_state_from_qpos", None) + if callable(sync) and commands: + sync(commands[-1]) + return commands + + def _scene_step_duration(self) -> float: + """Return one positive logical waypoint duration for scene timestamps.""" + sim_config = getattr(getattr(self.env, "sim", None), "sim_config", None) + candidates = ( + getattr(self.env, "physics_dt", None), + getattr(sim_config, "physics_dt", None), + ) + for value in candidates: + if isinstance(value, (int, float)) and not isinstance(value, bool): + duration = float(value) + if math.isfinite(duration) and duration > 0.0: + return duration + return 1.0 + + def combine( + self, + outcomes: Mapping[str, ActionOutcome | None], + masks: Mapping[str, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor]: + """Merge independently planned arm paths into one synchronized stream.""" + present = [item for item in outcomes.values() if item is not None] + if not present: + raise ValueError("At least one arm outcome is required.") + steps = max(int(item.trajectory.shape[1]) for item in present) + current = self.env.robot.get_qpos().to(self.device, dtype=torch.float32) + merged = current[:, None, :].repeat(1, max(steps, 1), 1) + success = torch.ones( + self.num_envs, + dtype=torch.bool, + device=self.device, + ) + for arm, outcome in outcomes.items(): + if outcome is None: + continue + mask = masks[arm].to(self.device, dtype=torch.bool) + success &= ~mask | outcome.success + trajectory = outcome.trajectory + if trajectory.shape[1] == 0: + continue + if trajectory.shape[1] < steps: + padding = trajectory[:, -1:].repeat(1, steps - trajectory.shape[1], 1) + trajectory = torch.cat((trajectory, padding), dim=1) + joint_ids = self.joint_ids(arm, include_hand=True) + if not joint_ids: + continue + selected = merged[:, :, joint_ids] + merged[:, :, joint_ids] = torch.where( + mask[:, None, None], trajectory[:, :, joint_ids], selected + ) + return merged, success + + def joint_ids(self, arm: str, *, include_hand: bool) -> list[int]: + if arm == "coordinated": + return list(range(int(self.env.robot.dof))) + side = "left" if arm == "left_arm" else "right" + result = list(getattr(self.env, f"{side}_arm_joints", ())) + if include_hand: + result.extend(getattr(self.env, f"{side}_eef_joints", ())) + return result + + def _engine(self) -> AtomicActionEngine: + if self._atomic_engine is None: + from .atomic_compat import ExactTargetMoveHeldObject + + engine = AtomicActionEngine( + self._generator(), + control_profiles=self._control_profiles(), + ) + engine.register(ExactTargetMoveHeldObject(), replace=True) + self._atomic_engine = engine + return self._atomic_engine + + def _generator(self) -> MotionGenerator: + if self._motion_generator is None: + backend = str(self.planner_policy.get("backend", "curobo")) + if backend == "curobo": + options = dict(self.planner_policy.get("curobo", {})) + obstacle_uids = tuple( + dict.fromkeys( + [ + *self.planner_policy.get("static_obstacle_uids", ()), + *self.planner_policy.get("dynamic_obstacle_uids", ()), + ] + ) + ) + rigid_objects: dict[str, Any] = {} + for uid in obstacle_uids: + obstacle_uid = str(uid) + entity = self.env.sim.get_rigid_object(obstacle_uid) + if entity is None: + raise ValueError(f"Unknown cuRobo obstacle {uid!r}.") + rigid_objects[obstacle_uid] = entity + obstacle_representation = str( + options.get("obstacle_representation", "cuboid") + ) + world = CuroboWorldCfg( + rigid_objects=rigid_objects or None, + obstacle_representation=obstacle_representation, + collision_cache=_collision_cache_for_world( + obstacle_representation, + len(rigid_objects), + ), + dynamic_obstacle_names=[ + str(uid) + for uid in self.planner_policy.get("dynamic_obstacle_uids", ()) + ], + multi_env=bool(options.get("multi_env", False)), + ) + planner_cfg = CuroboPlannerCfg( + robot_uid=self.env.robot.uid, + log_level=str(options.get("log_level", "error")), + world=world, + use_cuda_graph=bool(options.get("use_cuda_graph", True)), + preserve_plan_samples=bool( + options.get("preserve_plan_samples", False) + ), + max_attempts=int(options.get("max_attempts", 5)), + collision_activation_distance=float( + options.get("collision_activation_distance", 0.01) + ), + ) + elif backend == "toppra": + planner_cfg = ToppraPlannerCfg(robot_uid=self.env.robot.uid) + else: + raise ValueError( + f"Unsupported Action Engine planner backend {backend!r}." + ) + self._motion_generator = MotionGenerator( + cfg=MotionGenCfg(planner_cfg=planner_cfg) + ) + return self._motion_generator + + def _control_profiles(self) -> dict[str, ControlPartCommandProfile]: + profiles: dict[str, ControlPartCommandProfile] = {} + for side in ("left_arm", "right_arm"): + try: + _, hand_part, hand_dof = self._parts(side) + except ValueError: + continue + if hand_part is None or hand_dof == 0 or hand_part in profiles: + continue + profiles[hand_part] = ControlPartCommandProfile.joint_positions( + open=_as_hand_qpos(self.env.open_state, hand_dof, self.device), + grasp=_as_hand_qpos(self.env.close_state, hand_dof, self.device), + ) + return profiles + + def _parts(self, arm: str) -> tuple[str, str | None, int]: + if arm not in {"left_arm", "right_arm"}: + raise ValueError(f"Expected a physical arm, got {arm!r}.") + is_left = arm == "left_arm" + if hasattr(self.env, "get_agent_arm_control_part"): + arm_part = self.env.get_agent_arm_control_part(is_left) + hand_part = self.env.get_agent_eef_control_part(is_left) + else: + arm_part = arm + hand_part = "left_eef" if is_left else "right_eef" + hand_ids = ( + [] + if hand_part is None + else list(self.env.robot.get_joint_ids(name=hand_part)) + ) + return ( + str(arm_part), + None if hand_part is None else str(hand_part), + len(hand_ids), + ) diff --git a/embodichain/gen_sim/action_engine/runtime/atomic_compat.py b/embodichain/gen_sim/action_engine/runtime/atomic_compat.py new file mode 100644 index 000000000..d076e2b12 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/atomic_compat.py @@ -0,0 +1,84 @@ +# ---------------------------------------------------------------------------- +# 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 Engine-specific adapters for mainline atomic actions.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from embodichain.lab.sim.atomic_actions import ( + ActionPlan, + HeldObjectPoseGoal, + MoveHeldObject, + MoveHeldObjectOptions, + PlanningContext, + ResolvedActionRequest, +) + +__all__ = ["ExactTargetMoveHeldObject", "ExactTargetMoveHeldObjectOptions"] + + +@dataclass(frozen=True, slots=True, eq=False) +class ExactTargetMoveHeldObjectOptions(MoveHeldObjectOptions): + """Action Engine transport options with an exact-orientation switch.""" + + allow_automatic_transport_rotation: bool = True + """Whether the mainline transport heuristic may replace target rotation.""" + + +class ExactTargetMoveHeldObject(MoveHeldObject): + """Preserve a selected semantic orientation when explicitly requested.""" + + OptionsType = ExactTargetMoveHeldObjectOptions + + def __init__( + self, + default_options: ExactTargetMoveHeldObjectOptions | None = None, + ) -> None: + super().__init__(default_options) + self._allow_automatic_transport_rotation = True + + def _plan( + self, + request: ResolvedActionRequest[ + HeldObjectPoseGoal, + ExactTargetMoveHeldObjectOptions, + ], + context: PlanningContext, + ) -> ActionPlan: + previous = self._allow_automatic_transport_rotation + self._allow_automatic_transport_rotation = ( + request.skill_options.allow_automatic_transport_rotation + ) + try: + return super()._plan(request, context) + finally: + self._allow_automatic_transport_rotation = previous + + def _apply_automatic_transport_rotation( + self, + move_eef_xpos: torch.Tensor, + end_arm_xpos: torch.Tensor, + ) -> None: + """Apply the heuristic unless semantic grounding selected exact yaw.""" + if self._allow_automatic_transport_rotation: + super()._apply_automatic_transport_rotation( + move_eef_xpos, + end_arm_xpos, + ) 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/frames.py b/embodichain/gen_sim/action_engine/runtime/frames.py new file mode 100644 index 000000000..513308484 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/frames.py @@ -0,0 +1,165 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Resolve directional relations in live robot and world frames.""" + +from __future__ import annotations + +from typing import Any + +import torch + +from .robot_parts import arm_control_part + +__all__ = [ + "DIRECTIONAL_RELATIONS", + "arm_base_poses", + "relation_axes", + "relation_offset", + "robot_frame_axes", +] + + +_RELATION_COMPONENTS = { + "left": ("left",), + "left_of": ("left",), + "right": ("right",), + "right_of": ("right",), + "front": ("front",), + "front_of": ("front",), + "in_front_of": ("front",), + "behind": ("back",), + "back": ("back",), + "front_left": ("front", "left"), + "front_left_of": ("front", "left"), + "front_right": ("front", "right"), + "front_right_of": ("front", "right"), + "back_left": ("back", "left"), + "back_left_of": ("back", "left"), + "back_right": ("back", "right"), + "back_right_of": ("back", "right"), +} +DIRECTIONAL_RELATIONS = frozenset(_RELATION_COMPONENTS) + + +def arm_base_poses(env: Any) -> tuple[torch.Tensor, torch.Tensor]: + """Return live world poses of the left and right arm bases.""" + left_part = arm_control_part(env, "left_arm") + right_part = arm_control_part(env, "right_arm") + robot = env.robot + if hasattr(robot, "get_solver") and hasattr(robot, "get_link_pose"): + left_solver = robot.get_solver(name=left_part) + right_solver = robot.get_solver(name=right_part) + left_root = getattr(left_solver, "root_link_name", None) + right_root = getattr(right_solver, "root_link_name", None) + if left_root is None or right_root is None: + raise ValueError("Directional grounding requires both arm root links.") + left = robot.get_link_pose(link_name=left_root, to_matrix=True) + right = robot.get_link_pose(link_name=right_root, to_matrix=True) + elif hasattr(robot, "get_control_part_base_pose"): + left = robot.get_control_part_base_pose(name=left_part, to_matrix=True) + right = robot.get_control_part_base_pose(name=right_part, to_matrix=True) + elif hasattr(env, "get_current_xpos_agent"): + left, right = env.get_current_xpos_agent() + else: + raise ValueError( + "Directional grounding requires live left/right arm-base or TCP poses." + ) + + left = _batched_pose(left, env) + right = _batched_pose(right, env) + return left, right + + +def robot_frame_axes(env: Any) -> tuple[torch.Tensor, torch.Tensor]: + """Return normalized world-space forward and left axes for a dual-arm robot.""" + left, right = arm_base_poses(env) + lateral = left[:, :2, 3] - right[:, :2, 3] + norm = torch.linalg.vector_norm(lateral, dim=1, keepdim=True) + if bool((norm <= 1.0e-6).any()): + raise ValueError("Left and right arm bases must have distinct XY positions.") + lateral = lateral / norm + forward = torch.stack((lateral[:, 1], -lateral[:, 0]), dim=1) + return forward, lateral + + +def relation_axes( + env: Any, + relation: str, + *, + frame: str, +) -> tuple[torch.Tensor, ...]: + """Return signed world-space axes whose projections define a relation.""" + relation = str(relation) + if relation not in DIRECTIONAL_RELATIONS: + return () + if frame == "robot": + forward, lateral = robot_frame_axes(env) + elif frame == "world": + count = int(env.num_envs) + forward = torch.tensor( + [1.0, 0.0], dtype=torch.float32, device=env.device + ).repeat(count, 1) + lateral = torch.tensor( + [0.0, 1.0], dtype=torch.float32, device=env.device + ).repeat(count, 1) + else: + raise ValueError(f"Unsupported directional relation frame {frame!r}.") + + component_axes = { + "front": forward, + "back": -forward, + "left": lateral, + "right": -lateral, + } + return tuple(component_axes[item] for item in _RELATION_COMPONENTS[relation]) + + +def relation_offset( + env: Any, + relation: str, + *, + frame: str, + forward_distance: float, + lateral_distance: float, + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor | None: + """Resolve one directional relation into a batched world-space offset.""" + axes = relation_axes(env, relation, frame=frame) + if not axes: + return None + offset = torch.zeros((int(env.num_envs), 3), dtype=dtype, device=device) + components = _RELATION_COMPONENTS[relation] + for component, axis in zip(components, axes): + axis = axis.to(dtype=dtype, device=device) + distance = ( + forward_distance if component in {"front", "back"} else lateral_distance + ) + offset[:, :2] += axis * float(distance) + return offset + + +def _batched_pose(value: Any, env: Any) -> torch.Tensor: + pose = torch.as_tensor(value, dtype=torch.float32, device=env.device) + if pose.shape == (4, 4): + pose = pose.unsqueeze(0).repeat(int(env.num_envs), 1, 1) + if pose.shape != (int(env.num_envs), 4, 4): + raise ValueError( + "Frame pose must have shape (4, 4) or " + f"({int(env.num_envs)}, 4, 4), got {tuple(pose.shape)}." + ) + return pose diff --git a/embodichain/gen_sim/action_engine/runtime/grasp_collision_cache.py b/embodichain/gen_sim/action_engine/runtime/grasp_collision_cache.py new file mode 100644 index 000000000..252c388d6 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/grasp_collision_cache.py @@ -0,0 +1,330 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Prepare checksummed V-HACD caches for the shared grasp collision checker. + +The sidecar identifies the backend without changing Main's cache key or pickle +payload, so an unlabelled CoACD cache is never silently reused as V-HACD. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import hashlib +import io +import json +import operator +import os +from pathlib import Path +import pickle +import stat +import tempfile +from typing import Literal + +import numpy as np +import torch + +__all__ = [ + "GraspCollisionCacheError", + "GraspCollisionCacheResult", + "ensure_vhacd_grasp_collision_cache", + "grasp_collision_cache_path", +] + +_CACHE_SCHEMA_VERSION = 1 +_METADATA_SUFFIX = ".action_engine.json" +_DEFAULT_CACHE_DIR = ( + Path.home() / ".cache" / "embodichain_cache" / "convex_decomposition" +) + +CacheStatus = Literal["hit", "generated", "replaced"] + + +class GraspCollisionCacheError(RuntimeError): + """Raised when a safe, Main-compatible V-HACD cache cannot be prepared.""" + + +@dataclass(frozen=True) +class GraspCollisionCacheResult: + """Describe the prepared cache files and whether decomposition ran.""" + + status: CacheStatus + cache_path: Path + metadata_path: Path + + +def grasp_collision_cache_path( + mesh_vertices: torch.Tensor | np.ndarray, + mesh_triangles: torch.Tensor | np.ndarray, + max_decomposition_hulls: int, + *, + cache_dir: str | Path | None = None, +) -> Path: + """Return Main's exact ``_.pkl`` cache path.""" + vertices, triangles = _validate_mesh(mesh_vertices, mesh_triangles) + hull_limit = _validate_hull_limit(max_decomposition_hulls) + mesh_hash = hashlib.md5(vertices.tobytes() + triangles.tobytes()).hexdigest() + return _resolve_cache_dir(cache_dir) / f"{mesh_hash}_{hull_limit}.pkl" + + +def ensure_vhacd_grasp_collision_cache( + *, + mesh_vertices: torch.Tensor | np.ndarray, + mesh_triangles: torch.Tensor | np.ndarray, + max_decomposition_hulls: int, + cache_dir: str | Path | None = None, +) -> GraspCollisionCacheResult: + """Create or validate a V-HACD cache and its checksummed backend sidecar.""" + vertices, triangles = _validate_mesh(mesh_vertices, mesh_triangles) + hull_limit = _validate_hull_limit(max_decomposition_hulls) + mesh_hash = hashlib.md5(vertices.tobytes() + triangles.tobytes()).hexdigest() + cache_path = _resolve_cache_dir(cache_dir) / f"{mesh_hash}_{hull_limit}.pkl" + metadata_path = cache_path.with_name(f"{cache_path.name}{_METADATA_SUFFIX}") + expected_metadata: dict[str, object] = { + "schema_version": _CACHE_SCHEMA_VERSION, + "backend": "vhacd", + "mesh_hash": mesh_hash, + "max_decomposition_hulls": hull_limit, + } + + _prepare_private_directory(cache_path.parent) + _refuse_symlink(cache_path) + _refuse_symlink(metadata_path) + if _cache_matches_metadata(cache_path, metadata_path, expected_metadata): + return GraspCollisionCacheResult("hit", cache_path, metadata_path) + + exists = cache_path.exists() or metadata_path.exists() + status: CacheStatus = "replaced" if exists else "generated" + try: + plane_equations = _compute_vhacd_plane_equations( + vertices, + triangles, + hull_limit, + ) + cache_bytes = _serialize_checker_payload(plane_equations) + metadata = { + **expected_metadata, + "cache_sha256": hashlib.sha256(cache_bytes).hexdigest(), + } + + # Publish the complete pickle before its sidecar. A crash between the + # two replaces leaves a cache miss on retry, never a partial pickle. + _write_bytes_atomic(cache_path, cache_bytes) + metadata_bytes = (json.dumps(metadata, sort_keys=True) + "\n").encode() + _write_bytes_atomic(metadata_path, metadata_bytes) + except GraspCollisionCacheError: + raise + except Exception as exc: + raise GraspCollisionCacheError( + f"Failed to prepare V-HACD grasp collision cache {cache_path}: {exc}" + ) from exc + + return GraspCollisionCacheResult(status, cache_path, metadata_path) + + +def _compute_vhacd_plane_equations( + vertices: np.ndarray, + triangles: np.ndarray, + max_decomposition_hulls: int, +) -> list[tuple[np.ndarray, np.ndarray]]: + """Run DexSim V-HACD and convert its hulls to checker plane equations.""" + import open3d as o3d + from dexsim.kit.meshproc import convex_decomposition_vhacd + + from embodichain.toolkits.graspkit.pg_grasp.collision_checker import ( + extract_plane_equations, + ) + + mesh = o3d.t.geometry.TriangleMesh() + mesh.vertex.positions = o3d.core.Tensor(vertices.astype(np.float32, copy=False)) + mesh.triangle.indices = o3d.core.Tensor(triangles.astype(np.int32, copy=False)) + is_success, hull_meshes = convex_decomposition_vhacd( + mesh, + max_convex_hull_num=max_decomposition_hulls, + ) + if not is_success or not hull_meshes: + raise GraspCollisionCacheError( + "V-HACD returned no convex hulls for the grasp collision mesh." + ) + + convex_parts = [ + ( + np.asarray(hull.vertex.positions.numpy()), + np.asarray(hull.triangle.indices.numpy()), + ) + for hull in hull_meshes + ] + plane_equations = extract_plane_equations(convex_parts) + if not plane_equations: + raise GraspCollisionCacheError( + "V-HACD hulls produced no grasp collision plane equations." + ) + return plane_equations + + +def _serialize_checker_payload( + plane_equations: list[tuple[np.ndarray, np.ndarray]], +) -> bytes: + """Pack plane equations in the exact tensor dictionary Main unpickles.""" + if not plane_equations: + raise ValueError("V-HACD must produce at least one convex hull.") + + normalized: list[tuple[np.ndarray, np.ndarray]] = [] + for normals_value, offsets_value in plane_equations: + normals = np.asarray(normals_value, dtype=np.float32) + offsets = np.asarray(offsets_value, dtype=np.float32) + if normals.ndim != 2 or normals.shape[1:] != (3,) or not len(normals): + raise ValueError("Each V-HACD hull must have normals shaped [K, 3].") + if offsets.shape != (len(normals),): + raise ValueError("Each hull needs one offset per plane normal.") + if not np.isfinite(normals).all() or not np.isfinite(offsets).all(): + raise ValueError("V-HACD plane equations must contain finite values.") + normalized.append((normals, offsets)) + + max_plane_count = max(normals.shape[0] for normals, _ in normalized) + equations = torch.zeros((len(normalized), max_plane_count, 4)) + counts = torch.zeros(len(normalized), dtype=torch.int32) + for index, (normals, offsets) in enumerate(normalized): + plane_count = normals.shape[0] + equations[index, :plane_count, :3] = torch.from_numpy(normals) + equations[index, :plane_count, 3] = torch.from_numpy(offsets) + counts[index] = plane_count + + stream = io.BytesIO() + payload = {"plane_equations": equations, "plane_equation_counts": counts} + pickle.dump(payload, stream, protocol=pickle.HIGHEST_PROTOCOL) + return stream.getvalue() + + +def _cache_matches_metadata( + cache_path: Path, + metadata_path: Path, + expected_metadata: dict[str, object], +) -> bool: + if not cache_path.is_file() or not metadata_path.is_file(): + return False + try: + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + checksum = metadata.get("cache_sha256") + expected_checksum = hashlib.sha256(cache_path.read_bytes()).hexdigest() + return ( + isinstance(metadata, dict) + and all( + metadata.get(key) == value for key, value in expected_metadata.items() + ) + and isinstance(checksum, str) + and checksum == expected_checksum + ) + except (AttributeError, OSError, UnicodeDecodeError, json.JSONDecodeError): + return False + + +def _validate_mesh( + mesh_vertices: torch.Tensor | np.ndarray, + mesh_triangles: torch.Tensor | np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + if isinstance(mesh_vertices, torch.Tensor): + mesh_vertices = mesh_vertices.detach().cpu().numpy() + if isinstance(mesh_triangles, torch.Tensor): + mesh_triangles = mesh_triangles.detach().cpu().numpy() + if not isinstance(mesh_vertices, np.ndarray): + raise TypeError("mesh_vertices must be a torch.Tensor or numpy.ndarray.") + if not isinstance(mesh_triangles, np.ndarray): + raise TypeError("mesh_triangles must be a torch.Tensor or numpy.ndarray.") + vertices = np.ascontiguousarray(mesh_vertices) + triangles = np.ascontiguousarray(mesh_triangles) + if vertices.ndim != 2 or vertices.shape[1:] != (3,) or len(vertices) == 0: + raise ValueError("mesh_vertices must have non-empty shape [N, 3].") + if triangles.ndim != 2 or triangles.shape[1:] != (3,) or len(triangles) == 0: + raise ValueError("mesh_triangles must have non-empty shape [M, 3].") + if not np.issubdtype(vertices.dtype, np.number): + raise TypeError("mesh_vertices must contain numeric values.") + if not np.isfinite(vertices).all(): + raise ValueError("mesh_vertices must contain only finite values.") + if not np.issubdtype(triangles.dtype, np.integer): + raise TypeError("mesh_triangles must contain integer indices.") + if triangles.min() < 0 or triangles.max() >= len(vertices): + raise ValueError("mesh_triangles contains out-of-range vertex indices.") + return vertices, triangles + + +def _validate_hull_limit(value: int) -> int: + if isinstance(value, (bool, np.bool_)): + raise TypeError("max_decomposition_hulls must be an integer.") + try: + hull_limit = operator.index(value) + except TypeError as exc: + raise TypeError("max_decomposition_hulls must be an integer.") from exc + if hull_limit <= 0: + raise ValueError("max_decomposition_hulls must be positive.") + return hull_limit + + +def _resolve_cache_dir(cache_dir: str | Path | None) -> Path: + if cache_dir is not None: + return Path(cache_dir).expanduser().resolve() + try: + from embodichain.lab.sim import CONVEX_DECOMP_DIR + except Exception: + return _DEFAULT_CACHE_DIR + return Path(CONVEX_DECOMP_DIR).expanduser().resolve() + + +def _prepare_private_directory(path: Path) -> None: + try: + path.mkdir(parents=True, exist_ok=True, mode=0o700) + path.chmod(0o700) + except OSError as exc: + raise GraspCollisionCacheError( + f"Cannot secure grasp collision cache directory: {path}" + ) from exc + if path.stat().st_mode & (stat.S_IWGRP | stat.S_IWOTH): + raise GraspCollisionCacheError(f"Refusing writable cache directory: {path}") + + +def _refuse_symlink(path: Path) -> None: + if path.is_symlink(): + raise GraspCollisionCacheError( + f"Refusing symlinked grasp collision cache path: {path}" + ) + + +def _write_bytes_atomic(path: Path, payload: bytes) -> None: + """Publish one complete file with a same-directory atomic replacement.""" + _refuse_symlink(path) + file_descriptor, temporary_name = tempfile.mkstemp( + prefix=f".{path.name}.", + suffix=".tmp", + dir=path.parent, + ) + temporary_path = Path(temporary_name) + try: + os.fchmod(file_descriptor, 0o600) + with os.fdopen(file_descriptor, "wb") as output: + file_descriptor = -1 + output.write(payload) + output.flush() + os.fsync(output.fileno()) + _refuse_symlink(path) + os.replace(temporary_path, path) + path.chmod(0o600) + finally: + if file_descriptor >= 0: + os.close(file_descriptor) + try: + temporary_path.unlink(missing_ok=True) + except OSError: + pass diff --git a/embodichain/gen_sim/action_engine/runtime/grounding.py b/embodichain/gen_sim/action_engine/runtime/grounding.py new file mode 100644 index 000000000..e7a63337c --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/grounding.py @@ -0,0 +1,2356 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Resolve symbolic bindings from live simulator state.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass, replace +from typing import Any + +import torch + +from embodichain.gen_sim.action_engine.capabilities import ( + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.config import ( + RuntimePolicyCfg, + default_runtime_policy, +) +from embodichain.gen_sim.action_engine.domain import normalize_placement_relation +from embodichain.gen_sim.action_engine.orientation import ( + AlignAxisConstraint, + MatchRotationConstraint, + OrientationConstraint, + compile_orientation_constraint, +) +from embodichain.lab.sim.atomic_actions import ( + CoordinatedPickGoal, + CoordinatedPlacementGoal, + EndEffectorPoseGoal, + GraspGoal, + HeldObjectPoseGoal, + JointPositionGoal, + ObjectSemantics, + PlaceGoal, + PressAffordance, + PressGoal, +) +from .frames import arm_base_poses, relation_offset, robot_frame_axes +from .models import ExecutionProgram, GroundedAction, SemanticStep +from .motion_policy import resolve_motion_policy, with_motion_modifiers +from .robot_parts import arm_control_part +from .state import ExecutionState + +__all__ = ["ActionGrounder", "LiveArrangementPlan", "LivePlacementPlan"] + + +def _batched_pose(value: Any, env: Any) -> torch.Tensor: + pose = torch.as_tensor(value, dtype=torch.float32, device=env.device) + if pose.shape == (4, 4): + pose = pose.unsqueeze(0).repeat(int(env.num_envs), 1, 1) + if pose.shape != (int(env.num_envs), 4, 4): + raise ValueError( + "Live pose must have shape (4, 4) or " + f"({int(env.num_envs)}, 4, 4), got {tuple(pose.shape)}." + ) + return pose + + +def _object(env: Any, uid: str) -> Any: + entity = env.sim.get_rigid_object(uid) + if entity is None: + raise ValueError(f"Unknown rigid object {uid!r}.") + return entity + + +def _live_pose(env: Any, uid: str) -> torch.Tensor: + return _batched_pose(_object(env, uid).get_local_pose(to_matrix=True), env) + + +def _local_vertices(entity: Any, env: Any, env_id: int = 0) -> torch.Tensor: + value = entity.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=env.device) + if vertices.ndim == 3 and vertices.shape[0] == 1: + vertices = vertices[0] + if vertices.ndim != 2 or vertices.shape[-1] != 3 or vertices.numel() == 0: + raise ValueError("Rigid-object mesh vertices must have shape (N, 3).") + return vertices + + +def _world_vertices(entity: Any, env: Any, env_id: int) -> torch.Tensor: + vertices = _local_vertices(entity, env, env_id) + pose = _batched_pose(entity.get_local_pose(to_matrix=True), env)[env_id] + return vertices @ pose[:3, :3].transpose(0, 1) + pose[:3, 3] + + +@dataclass(frozen=True) +class _Geometry: + radius: torch.Tensor + half_height: torch.Tensor + + +class LiveArrangementPlan: + """Materialize collision-aware line slots independently in every env.""" + + def __init__( + self, + env: Any, + steps: Sequence[SemanticStep], + *, + slot_margin: float | None = None, + minimum_spacing: float | None = None, + clearance: float | None = None, + row_search_step: float | None = None, + row_search_radius: float | None = None, + ) -> None: + if not steps: + raise ValueError("An arrangement plan requires at least one step.") + self.env = env + self.steps = tuple(steps) + self.step_by_id = {step.id: step for step in steps} + self.num_envs = int(env.num_envs) + self.device = env.device + self.slot_count = len(steps) + self.axis = str(steps[0].goal.get("axis", "world_x")) + profile = str(getattr(env, "agent_robot_profile", "dual_ur10")) + defaults = default_runtime_policy(profile).grounding["arrangement"] + slot_margin = defaults["slot_margin"] if slot_margin is None else slot_margin + minimum_spacing = ( + defaults["minimum_spacing"] if minimum_spacing is None else minimum_spacing + ) + self.clearance = float( + defaults["layout_clearance"] if clearance is None else clearance + ) + self.row_search_step = float( + defaults["row_search_step"] if row_search_step is None else row_search_step + ) + self.row_search_radius = float( + defaults["row_search_radius"] + if row_search_radius is None + else row_search_radius + ) + + table = _object(env, "table") + bounds = [] + for env_id in range(self.num_envs): + vertices = _world_vertices(table, env, env_id) + bounds.append( + torch.stack((vertices.min(dim=0).values, vertices.max(dim=0).values)) + ) + self.table_bounds = torch.stack(bounds) + self.table_center = self.table_bounds.mean(dim=1) + self.table_top = self.table_bounds[:, 1, 2] + if self.axis == "table_long_axis": + mean_extent = ( + self.table_bounds[:, 1, :2] - self.table_bounds[:, 0, :2] + ).mean(dim=0) + self.axis_index = int(torch.argmax(mean_extent).item()) + else: + self.axis_index = 0 if self.axis in {"x", "world_x"} else 1 + self.perpendicular_index = 1 - self.axis_index + self.geometry = {step.id: self._geometry(step) for step in self.steps} + diameters = torch.stack( + [self.geometry[step.id].radius * 2.0 for step in self.steps], + dim=1, + ) + self.spacing = torch.maximum( + diameters.max(dim=1).values + float(slot_margin), + torch.full( + (self.num_envs,), + float(minimum_spacing), + dtype=torch.float32, + device=self.device, + ), + ) + self.positions = self._make_slots() + self.reassignment_reason: list[str | None] = [None] * self.num_envs + self.reassignment_cost = torch.full( + (self.num_envs,), + float("nan"), + dtype=torch.float32, + device=self.device, + ) + self.assignments = self._initial_slot_assignments() + order_by = str(self.steps[0].goal.get("order_by", "explicit")) + direction = str(self.steps[0].goal.get("order_direction", "given")) + if order_by == "size" and not any( + step.goal.get("slot_constraint") == "free_reassignable" + for step in self.steps + ): + for env_id in range(self.num_envs): + ordered = sorted( + self.steps, + key=lambda step: float(self.geometry[step.id].radius[env_id]), + reverse=direction != "ascending", + ) + for slot_id, step in enumerate(ordered): + self.assignments[step.id][env_id] = slot_id + self.completed = { + step.id: torch.zeros( + self.num_envs, + dtype=torch.bool, + device=self.device, + ) + for step in self.steps + } + + def _initial_slot_assignments(self) -> dict[str, torch.Tensor]: + """Match free-order objects to slots in their current spatial order.""" + assignments = { + step.id: torch.full( + (self.num_envs,), + int(step.goal.get("nominal_slot_index", index)), + dtype=torch.long, + device=self.device, + ) + for index, step in enumerate(self.steps) + } + free_steps = [ + step + for step in self.steps + if step.goal.get("slot_constraint") == "free_reassignable" + ] + if not free_steps: + return assignments + required_slots = { + int(step.goal.get("nominal_slot_index", index)) + for index, step in enumerate(self.steps) + if step.goal.get("slot_constraint") != "free_reassignable" + } + available_slots = [ + slot_id + for slot_id in range(self.slot_count) + if slot_id not in required_slots + ] + if len(available_slots) != len(free_steps): + raise ValueError( + "Arrangement slot constraints do not define a one-to-one assignment." + ) + axis_positions = { + step.id: _live_pose(self.env, step.object_uid)[:, self.axis_index, 3] + for step in free_steps + } + for env_id in range(self.num_envs): + ordered_steps = sorted( + free_steps, + key=lambda step: ( + float(axis_positions[step.id][env_id]), + int(step.goal.get("nominal_slot_index", 0)), + step.id, + ), + ) + ordered_slots = sorted( + available_slots, + key=lambda slot_id: ( + float(self.positions[env_id, slot_id, self.axis_index]), + slot_id, + ), + ) + matching_cost = 0.0 + changed = False + for step, slot_id in zip(ordered_steps, ordered_slots): + nominal = int(step.goal.get("nominal_slot_index", 0)) + assignments[step.id][env_id] = slot_id + changed |= slot_id != nominal + matching_cost += abs( + float(axis_positions[step.id][env_id]) + - float(self.positions[env_id, slot_id, self.axis_index]) + ) + if changed: + self.reassignment_reason[env_id] = ( + "free arrangement initialized from live spatial order" + ) + self.reassignment_cost[env_id] = matching_cost + return assignments + + def _geometry(self, step: SemanticStep) -> _Geometry: + entity = _object(self.env, step.object_uid) + radii = [] + heights = [] + for env_id in range(self.num_envs): + vertices = _local_vertices(entity, self.env, env_id) + half_extent = ( + vertices.max(dim=0).values - vertices.min(dim=0).values + ) * 0.5 + if step.goal.get("orientation_goal", "none") in {"none", "preserve"}: + rotation = _live_pose(self.env, step.object_uid)[env_id, :3, :3] + rotated = vertices @ rotation.transpose(0, 1) + radii.append(torch.linalg.vector_norm(rotated[:, :2], dim=-1).max()) + else: + # A non-preserve target may rotate the longest local dimension + # into the table plane, so retain the conservative bound. + radii.append( + torch.linalg.vector_norm(torch.topk(half_extent, k=2).values) + ) + heights.append((vertices[:, 2].max() - vertices[:, 2].min()) * 0.5) + return _Geometry(torch.stack(radii), torch.stack(heights)) + + def _make_slots(self) -> torch.Tensor: + offsets = ( + torch.arange(self.slot_count, device=self.device, dtype=torch.float32) + - (self.slot_count - 1) / 2.0 + ) + slots = torch.empty( + self.num_envs, + self.slot_count, + 3, + dtype=torch.float32, + device=self.device, + ) + radii = torch.stack( + [self.geometry[step.id].radius for step in self.steps], + dim=1, + ) + # Free slot rematching allows any remaining object to occupy any slot. + # Size every slot for the largest member in that environment rather + # than accidentally baking the nominal object order into geometry. + slot_radii = radii.max(dim=1).values[:, None].repeat(1, self.slot_count) + obstacles = self._obstacle_bounds() + search_offsets = [0.0] + steps = int(self.row_search_radius / self.row_search_step) + for index in range(1, steps + 1): + offset = self.row_search_step * index + search_offsets.extend((offset, -offset)) + for env_id in range(self.num_envs): + chosen = None + for perpendicular in search_offsets: + candidate = self.table_center[env_id].repeat(self.slot_count, 1) + candidate[:, self.axis_index] += self.spacing[env_id] * offsets + candidate[:, self.perpendicular_index] += perpendicular + candidate[:, 2] = self.table_top[env_id] + if self._safe( + candidate, + slot_radii[env_id], + self.table_bounds[env_id], + obstacles[env_id], + ): + chosen = candidate + break + if chosen is None: + raise ValueError( + f"Environment {env_id} has no collision-free arrangement row." + ) + slots[env_id] = chosen + return slots + + def _obstacle_bounds( + self, + ) -> list[list[tuple[torch.Tensor, torch.Tensor]]]: + result: list[list[tuple[torch.Tensor, torch.Tensor]]] = [ + [] for _ in range(self.num_envs) + ] + getter = getattr(self.env.sim, "get_rigid_object_uid_list", None) + if not callable(getter): + return result + movable = {step.object_uid for step in self.steps} + for uid in getter(): + if uid == "table" or uid in movable: + continue + entity = self.env.sim.get_rigid_object(uid) + if entity is None: + continue + for env_id in range(self.num_envs): + vertices = _world_vertices(entity, self.env, env_id) + if float(vertices[:, 2].max()) < float( + self.table_top[env_id] - self.clearance + ): + continue + result[env_id].append( + ( + vertices[:, :2].min(dim=0).values, + vertices[:, :2].max(dim=0).values, + ) + ) + return result + + def _safe( + self, + slots: torch.Tensor, + radii: torch.Tensor, + table_bounds: torch.Tensor, + obstacles: Sequence[tuple[torch.Tensor, torch.Tensor]], + ) -> bool: + lower = table_bounds[0, :2] + radii[:, None] + self.clearance + upper = table_bounds[1, :2] - radii[:, None] - self.clearance + if bool(((slots[:, :2] < lower) | (slots[:, :2] > upper)).any()): + return False + for center, radius in zip(slots[:, :2], radii): + for obstacle_lower, obstacle_upper in obstacles: + closest = torch.maximum( + obstacle_lower, + torch.minimum(center, obstacle_upper), + ) + if float(torch.linalg.vector_norm(center - closest)) <= float( + radius + self.clearance + ): + return False + return True + + def target( + self, + step: SemanticStep, + object_pose: torch.Tensor, + *, + phase: str, + policy: Mapping[str, Any], + ) -> torch.Tensor: + """Return a live final or collision-clear staging object pose.""" + if phase not in {"staging", "final"}: + raise ValueError(f"Unsupported arrangement phase {phase!r}.") + target = object_pose.clone() + env_ids = torch.arange(self.num_envs, device=self.device) + slot_ids = self.assignments[step.id] + target[:, :2, 3] = self.positions[env_ids, slot_ids, :2] + final_z = ( + self.table_top + + self.geometry[step.id].half_height + + float(policy["surface_clearance"]) + ) + target[:, 2, 3] = final_z + if phase == "staging": + target[:, 2, 3] = final_z + float(policy["transport_clearance"]) + return target + + def mark_completed(self, step_id: str, success: torch.Tensor) -> None: + self.completed[step_id] |= success.to(self.device, dtype=torch.bool) + + def remaining(self, env_id: int) -> list[str]: + return [ + step.id for step in self.steps if not bool(self.completed[step.id][env_id]) + ] + + def available_slots(self, env_id: int) -> list[int]: + occupied = { + int(self.assignments[step.id][env_id]) + for step in self.steps + if bool(self.completed[step.id][env_id]) + } + return [index for index in range(self.slot_count) if index not in occupied] + + def assign(self, env_id: int, assignment: Mapping[str, int]) -> None: + for step_id, slot_id in assignment.items(): + self.assignments[step_id][env_id] = int(slot_id) + + def metadata(self, step: SemanticStep, env_id: int) -> dict[str, Any]: + """Describe the live slot resolution used by one environment.""" + nominal = int(step.goal.get("nominal_slot_index", 0)) + resolved = int(self.assignments[step.id][env_id]) + return { + "nominal_slot_index": nominal, + "resolved_slot_index": resolved, + "slot_constraint": str(step.goal.get("slot_constraint", "required")), + "slot_reassigned": resolved != nominal, + "reassignment_reason": self.reassignment_reason[env_id], + "matching_cost": ( + float(self.reassignment_cost[env_id]) + if torch.isfinite(self.reassignment_cost[env_id]) + else None + ), + "spacing": float(self.spacing[env_id]), + "resolved_slot_position": self.positions[env_id, resolved].tolist(), + } + + +class LivePlacementPlan: + """Allocate non-overlapping live slots for one shared container.""" + + def __init__( + self, + env: Any, + steps: Sequence[SemanticStep], + *, + clearance: float | None = None, + ) -> None: + if not steps: + raise ValueError("A placement plan requires at least one step.") + references = {step.goal.get("reference_object") for step in steps} + if len(references) != 1 or not isinstance(next(iter(references)), str): + raise ValueError("Placement-plan steps must share one reference object.") + self.env = env + self.steps = tuple(steps) + self.reference_uid = str(next(iter(references))) + self.num_envs = int(env.num_envs) + profile = str(getattr(env, "agent_robot_profile", "dual_ur10")) + default_clearance = default_runtime_policy(profile).grounding["placement"][ + "clearance" + ] + self.clearance = float(default_clearance if clearance is None else clearance) + self.positions = self._make_slots() + + def _make_slots(self) -> dict[str, torch.Tensor]: + container = _object(self.env, self.reference_uid) + positions = { + step.id: torch.empty( + self.num_envs, + 3, + dtype=torch.float32, + device=self.env.device, + ) + for step in self.steps + } + named_slots = [str(step.goal.get("slot", "auto")) for step in self.steps] + for slot in named_slots: + if slot not in {"auto", "left", "center", "right"}: + raise ValueError(f"Unsupported container slot {slot!r}.") + + for env_id in range(self.num_envs): + vertices = _world_vertices(container, self.env, env_id) + lower = vertices.min(dim=0).values + upper = vertices.max(dim=0).values + center = (lower + upper) * 0.5 + extent = upper[:2] - lower[:2] + axis = int(torch.argmax(extent).item()) + radii = [] + for step in self.steps: + moved_vertices = _local_vertices( + _object(self.env, step.object_uid), + self.env, + env_id, + ) + half = ( + moved_vertices.max(dim=0).values - moved_vertices.min(dim=0).values + )[:2] * 0.5 + radii.append(float(torch.linalg.vector_norm(half))) + radius = max(radii) + usable_span = float(extent[axis]) - 2.0 * (radius + self.clearance) + required_span = 2.0 * radius * max(len(self.steps) - 1, 0) + if usable_span + 1.0e-6 < required_span: + raise ValueError( + f"Environment {env_id} container {self.reference_uid!r} " + "has no non-overlapping slot plan." + ) + offsets = torch.linspace( + -required_span * 0.5, + required_span * 0.5, + len(self.steps), + device=self.env.device, + ) + named_offsets = { + "left": required_span * 0.5, + "center": 0.0, + "right": -required_span * 0.5, + } + used: list[float] = [] + for index, step in enumerate(self.steps): + slot = named_slots[index] + offset = ( + float(offsets[index]) if slot == "auto" else named_offsets[slot] + ) + if any(abs(offset - item) < 2.0 * radius for item in used): + raise ValueError( + f"Container slot {slot!r} overlaps another requested slot." + ) + used.append(offset) + target = center.clone() + target[axis] += offset + target[2] = lower[2] + positions[step.id][env_id] = target + return positions + + def target( + self, + step: SemanticStep, + object_pose: torch.Tensor, + rotation: torch.Tensor, + *, + surface_clearance: float, + ) -> torch.Tensor: + """Return a slot pose corrected for the rotated object mesh bottom.""" + target = object_pose.clone() + target[:, :3, :3] = rotation + target[:, :2, 3] = self.positions[step.id][:, :2] + entity = _object(self.env, step.object_uid) + for env_id in range(self.num_envs): + bottom = ( + _local_vertices(entity, self.env, env_id) + @ rotation[env_id].transpose(0, 1) + )[:, 2].min() + target[env_id, 2, 3] = ( + self.positions[step.id][env_id, 2] + surface_clearance - bottom + ) + return target + + +class ActionGrounder: + """Translate one symbolic action into a public typed atomic-action target.""" + + def __init__( + self, + program: ExecutionProgram, + env: Any, + semantics_factory: Callable[[str], ObjectSemantics], + arrangement: ( + LiveArrangementPlan | Mapping[str, LiveArrangementPlan] | None + ) = None, + placements: Mapping[str, LivePlacementPlan] | None = None, + runtime_policy: RuntimePolicyCfg | None = None, + capability_registry: Any | None = None, + ) -> None: + self.program = program + self.env = env + self.semantics_factory = semantics_factory + self.capabilities = capability_registry or build_atomic_capability_registry() + self.robot_profile = str(getattr(env, "agent_robot_profile", "dual_ur10")) + self.runtime_policy = runtime_policy or default_runtime_policy( + self.robot_profile + ) + if isinstance(arrangement, Mapping): + self.arrangements = dict(arrangement) + elif arrangement is None: + self.arrangements = {} + else: + self.arrangements = { + step.id: arrangement + for step in program.semantic_steps + if step.operator in {"arrange_line", "place_in_line"} + } + self.placements = dict(placements or {}) + + def policy( + self, + action: Mapping[str, Any], + *, + extra_modifiers: tuple[tuple[str, str], ...] = (), + ) -> dict[str, Any]: + action_class = str(action.get("atomic_action_class", "")) + capability = self.capabilities.get(action_class) + motion_base = capability.motion_base or capability.name + policy_spec = action.get("motion_policy", {"modifiers": []}) + if extra_modifiers: + policy_spec = with_motion_modifiers(policy_spec, *extra_modifiers) + inline = action.get("motion_policy_config", action.get("cfg")) + return resolve_motion_policy( + self.robot_profile, + motion_base, + policy_spec, + motion_defaults=self.runtime_policy.motion_defaults, + motion_modifiers=self.runtime_policy.motion_modifiers, + inline_overrides=inline if isinstance(inline, Mapping) else None, + ) + + def _policy_value(self, policy: Mapping[str, Any], key: str) -> Any: + defaults = self.runtime_policy.grounding["semantic_defaults"] + return policy[key] if key in policy else defaults[key] + + def ground( + self, + action: Mapping[str, Any], + step: SemanticStep, + *, + arm: str, + state: ExecutionState, + reference_eef_pose: torch.Tensor | None = None, + orientation_reference_pose: torch.Tensor | None = None, + _handover_workspace: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> GroundedAction: + action_class = str(action["atomic_action_class"]) + capability = self.capabilities.require_executable(action_class) + self.capabilities.validate_binding(action) + control = str(action.get("control", "arm")) + binding = action.get("target_binding", {}) + if not isinstance(binding, Mapping): + raise ValueError("target_binding must be a mapping.") + kind = str(binding.get("kind", "")) + orientation = compile_orientation_constraint(step.goal) + is_handover_continuation = self._is_handover_continuation(step) + uses_handover_staging = ( + kind == "handover_staging" + and capability.target_materializer == "semantic_held_object" + ) + use_upright_yaw_search = ( + is_handover_continuation or uses_handover_staging + ) and self._uses_upright_yaw_search( + step, + orientation, + ) + extra_modifiers: tuple[tuple[str, str], ...] = () + if ( + is_handover_continuation + and use_upright_yaw_search + and capability.target_materializer + in { + "semantic_held_object", + "current_held_pose", + "eef_pose", + } + ): + extra_modifiers = (("orientation", "upright"),) + policy = self.policy(action, extra_modifiers=extra_modifiers) + if kind == "joint_state": + joint_defaults = self.runtime_policy.grounding["joint_state"] + source = binding.get("source") + if source == "gripper_closed": + policy["sample_interval"] = int( + joint_defaults["hand_close_sample_interval"] + ) + elif source == "gripper_open": + policy["sample_interval"] = int( + joint_defaults["hand_open_sample_interval"] + ) + elif source == "initial" and control == "arm": + # Returning home after release is a safety motion. If the + # collision-aware planner cannot find a route, do not silently + # replace it with collision-unaware joint interpolation. + policy["collision_safety"] = "required" + if uses_handover_staging and use_upright_yaw_search: + # Handover consumes the live payload pose immediately after this + # move. Use the existing upright-yaw feasibility search instead + # of the generic transport orientation heuristic, which can tilt + # a payload while moving it to the exchange point. + policy["upright_yaw_samples"] = max( + int(policy.get("upright_yaw_samples", 1)), + 8, + ) + object_pose = _live_pose(self.env, step.object_uid) + if step.operator == "orient_object": + policy["upright_local_axis"] = self._upright_local_axis(step) + if capability.target_materializer == "object_grasp": + policy["obj_upright_direction"] = self._upright_local_direction(step) + reference_pose = self._reference_pose(step) + target_object_pose = None + + if capability.target_materializer_hook is not None: + grounded = capability.target_materializer_hook( + grounder=self, + action=action, + step=step, + arm=arm, + state=state, + binding=binding, + policy=policy, + object_pose=object_pose, + reference_pose=reference_pose, + reference_eef_pose=reference_eef_pose, + orientation_reference_pose=orientation_reference_pose, + ) + if not isinstance(grounded, GroundedAction): + raise TypeError( + f"AtomicAction {action_class!r} target materializer must " + "return GroundedAction." + ) + return grounded + + if kind == "object": + semantics = self.semantics_factory( + str(binding.get("object", step.object_uid)) + ) + if capability.target_materializer == "object_grasp": + target: Any = GraspGoal(semantics=semantics) + elif capability.target_materializer == "coordinated_pickment": + target_object_pose = self._semantic_target( + step, + object_pose, + reference_pose, + policy, + phase="final", + orientation_reference_pose=orientation_reference_pose, + ) + target = CoordinatedPickGoal( + object_target_pose=target_object_pose, + semantics=semantics, + object_initial_pose=object_pose, + ) + elif capability.target_materializer == "press": + target_object_pose = object_pose.clone() + target = self._press_goal( + step.object_uid, + object_pose, + semantics=semantics, + ) + else: + raise ValueError( + f"{action_class} does not support object target bindings." + ) + elif kind in {"semantic_goal", "coordinated_goal"}: + phase = str(binding.get("phase", "final")) + target_object_pose = self._semantic_target( + step, + object_pose, + reference_pose, + policy, + phase=phase, + orientation_reference_pose=orientation_reference_pose, + ) + if capability.target_materializer == "coordinated_pickment": + semantics = self.semantics_factory(step.object_uid) + target = CoordinatedPickGoal( + object_target_pose=target_object_pose, + semantics=semantics, + object_initial_pose=object_pose, + ) + elif capability.target_materializer == "press": + # Press moves the TCP, not the target object. Keep the object's + # live pose as the postcondition reference while grounding a + # downward contact point from its current surface geometry. + target_object_pose = object_pose.clone() + target = self._press_goal( + step.object_uid, + object_pose, + ) + elif capability.target_materializer == "semantic_held_object": + target = HeldObjectPoseGoal(object_target_pose=target_object_pose) + else: + raise ValueError( + f"Target materializer {capability.target_materializer!r} cannot " + f"resolve {kind!r}." + ) + elif kind == "coordinated_placement_goal": + support_uid = binding.get( + "support_object", + step.goal.get("support_object"), + ) + placing_uid = binding.get("placing_object", step.object_uid) + if not isinstance(placing_uid, str) or not placing_uid: + raise ValueError("coordinated_placement_goal requires placing_object.") + if not isinstance(support_uid, str) or not support_uid: + raise ValueError("coordinated_placement_goal requires support_object.") + support_pose = _live_pose(self.env, support_uid) + target_object_pose = self._semantic_target( + step, + object_pose, + support_pose, + policy, + phase="final", + orientation_reference_pose=orientation_reference_pose, + ) + target = CoordinatedPlacementGoal( + placing_object_target_pose=target_object_pose, + support_object_target_pose=support_pose, + release=bool(step.goal.get("release", True)), + ) + elif kind == "current_held_pose": + if state.get_held_object(arm_control_part(self.env, arm)) is None: + raise ValueError("Place requires a held object from a prior PickUp.") + target = PlaceGoal( + xpos=( + reference_eef_pose + if reference_eef_pose is not None + else self._current_eef_pose(arm) + ) + ) + elif kind == "policy_pose": + source = binding.get("source") + retreat_reference = self._retreat_reference_pose( + arm, + reference_eef_pose, + ) + if binding.get("operation") == "retreat": + policy["retreat_reachability_search"] = True + policy["retreat_reference_pose"] = retreat_reference.clone() + if source in {"release", "handover"}: + policy["clearance_object_uid"] = step.object_uid + policy["collision_safety"] = "required" + contact_uids = [step.object_uid] + reference_uid = step.goal.get("reference_object") + if isinstance(reference_uid, str) and reference_uid: + contact_uids.append(reference_uid) + policy["collision_exclusion_uids"] = list(dict.fromkeys(contact_uids)) + if source == "handover": + policy.update(self.runtime_policy.grounding["handover"]) + policy["transfer_arm"] = arm + policy["transfer_role_axis"] = self._handover_role_axis( + arm, + dtype=object_pose.dtype, + device=object_pose.device, + ) + target = EndEffectorPoseGoal( + xpos=self._retreat_pose( + arm, + policy, + retreat_reference, + clear_exchange=source == "handover", + ) + ) + elif kind == "visual_constraint": + visual_pose = self._visual_target(binding, arm) + if capability.target_materializer == "semantic_held_object": + target_object_pose = object_pose.clone() + target_object_pose[:, :3, 3] = visual_pose[:, :3, 3] + target = HeldObjectPoseGoal(object_target_pose=target_object_pose) + elif capability.target_materializer == "eef_pose": + target = EndEffectorPoseGoal(xpos=visual_pose) + else: + raise ValueError( + f"Target materializer {capability.target_materializer!r} " + "cannot resolve a visual_constraint." + ) + elif kind == "joint_state": + target = JointPositionGoal( + target=self._joint_target( + arm, + control, + str(binding.get("source", "initial")), + binding, + ) + ) + elif kind in {"eef_pose", "pose"}: + target = EndEffectorPoseGoal(xpos=self._explicit_pose(binding, object_pose)) + elif kind == "handover_goal": + target, target_object_pose, policy = self._handover_target( + step, + binding, + object_pose, + reference_pose, + policy, + state, + orientation_reference_pose=orientation_reference_pose, + workspace=_handover_workspace, + ) + elif kind == "handover_staging": + transfer_arm = str(binding.get("transfer_arm", "left_arm")) + receive_arm = str(binding.get("receive_arm", "right_arm")) + middle, _ = self._handover_workspace_poses( + object_pose, + transfer_arm=transfer_arm, + receive_arm=receive_arm, + policy=policy, + step=step, + orientation_reference_pose=orientation_reference_pose, + ) + middle[:, :3, :3] = self._target_rotation( + step, + object_pose, + orientation_reference_pose=orientation_reference_pose, + ) + target_object_pose = middle + target = HeldObjectPoseGoal(object_target_pose=middle) + else: + raise ValueError(f"Unsupported target binding kind {kind!r}.") + return GroundedAction( + action_class=action_class, + arm=arm, + control=control, + target=target, + cfg=policy, + object_pose=object_pose, + reference_pose=reference_pose, + target_object_pose=target_object_pose, + motion_policy=policy, + object_uid=step.object_uid, + ) + + def _handover_role_axis( + self, + transfer_arm: str, + *, + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + """Return the world-space axis from the receiver base to transfer base.""" + if transfer_arm not in {"left_arm", "right_arm"}: + raise ValueError(f"Unknown handover arm {transfer_arm!r}.") + _, lateral = robot_frame_axes(self.env) + horizontal = lateral if transfer_arm == "left_arm" else -lateral + return torch.cat( + ( + horizontal.to(dtype=dtype, device=device), + torch.zeros( + (int(self.env.num_envs), 1), + dtype=dtype, + device=device, + ), + ), + dim=1, + ) + + def ground_candidates( + self, + action: Mapping[str, Any], + step: SemanticStep, + *, + arm: str, + state: ExecutionState, + reference_eef_pose: torch.Tensor | None = None, + orientation_reference_pose: torch.Tensor | None = None, + ) -> tuple[GroundedAction, ...]: + """Return deterministic grounding candidates for an opt-in capability.""" + binding = action.get("target_binding", {}) + if not isinstance(binding, Mapping): + return ( + self.ground( + action, + step, + arm=arm, + state=state, + reference_eef_pose=reference_eef_pose, + orientation_reference_pose=orientation_reference_pose, + ), + ) + placement_support_uid = self._placement_support_uid(step) + placement_relation = ( + normalize_placement_relation(step.goal.get("relation", "on")) + if step.operator == "place_relative" + else str(step.goal.get("relation", "none")) + ) + is_on_placement = ( + binding.get("kind") == "semantic_goal" + and binding.get("phase", "final") != "staging" + and placement_relation in {"on", "on_top", "on_top_of"} + and placement_support_uid is not None + ) + if is_on_placement: + base = self.ground( + action, + step, + arm=arm, + state=state, + reference_eef_pose=reference_eef_pose, + orientation_reference_pose=orientation_reference_pose, + ) + return self._placement_grounding_candidates( + base, + step, + support_uid=placement_support_uid, + ) + if binding.get("kind") != "handover_goal": + return ( + self.ground( + action, + step, + arm=arm, + state=state, + reference_eef_pose=reference_eef_pose, + orientation_reference_pose=orientation_reference_pose, + ), + ) + policy = self.policy(action) + object_pose = _live_pose(self.env, step.object_uid) + rotation = self._target_rotation( + step, + object_pose, + orientation_reference_pose=orientation_reference_pose, + ) + workspaces = self._handover_workspace_candidates( + step, + object_pose, + transfer_arm=str(binding.get("transfer_arm", "left_arm")), + receive_arm=str(binding.get("receive_arm", "right_arm")), + policy=policy, + rotation=rotation, + ) + return tuple( + self.ground( + action, + step, + arm=arm, + state=state, + reference_eef_pose=reference_eef_pose, + orientation_reference_pose=orientation_reference_pose, + _handover_workspace=workspace, + ) + for workspace in workspaces + ) + + def _placement_grounding_candidates( + self, + base: GroundedAction, + step: SemanticStep, + *, + support_uid: str, + ) -> tuple[GroundedAction, ...]: + """Sample bounded support-relative poses from live object geometry.""" + if base.target_object_pose is None or not isinstance( + base.target, HeldObjectPoseGoal + ): + return (base,) + support = _object(self.env, support_uid) + moved = _object(self.env, step.object_uid) + placement = self.runtime_policy.grounding["placement"] + count = int(placement["candidate_count"]) + fraction = float(placement["candidate_offset_fraction"]) + margin = float(placement["support_margin"]) + patterns = ( + (0.0, 0.0), + (1.0, 0.0), + (-1.0, 0.0), + (0.0, 1.0), + (0.0, -1.0), + (1.0, 1.0), + (1.0, -1.0), + (-1.0, 1.0), + (-1.0, -1.0), + )[:count] + candidates: list[GroundedAction] = [] + seen_offsets: list[torch.Tensor] = [] + for candidate_index, pattern in enumerate(patterns): + target_pose = base.target_object_pose.clone() + offsets = target_pose.new_zeros((int(self.env.num_envs), 2)) + for env_id in range(int(self.env.num_envs)): + support_vertices = _world_vertices(support, self.env, env_id) + moved_local = _local_vertices(moved, self.env, env_id) + rotated = moved_local @ target_pose[env_id, :3, :3].transpose(0, 1) + support_lower = support_vertices[:, :2].min(dim=0).values + support_upper = support_vertices[:, :2].max(dim=0).values + moved_lower = rotated[:, :2].min(dim=0).values + moved_upper = rotated[:, :2].max(dim=0).values + allowed_lower = support_lower + margin - moved_lower + allowed_upper = support_upper - margin - moved_upper + if bool(torch.all(allowed_lower <= allowed_upper)): + base_xy = target_pose[env_id, :2, 3].clone() + center = torch.minimum( + torch.maximum(base_xy, allowed_lower), + allowed_upper, + ) + direction = target_pose.new_tensor(pattern) + room = torch.where( + direction >= 0.0, + allowed_upper - center, + center - allowed_lower, + ) + candidate_xy = center + direction * room * fraction + offsets[env_id] = candidate_xy - base_xy + target_pose[env_id, :2, 3] = candidate_xy + + footprint_lower = target_pose[env_id, :2, 3] + moved_lower + footprint_upper = target_pose[env_id, :2, 3] + moved_upper + local_mask = torch.all( + (support_vertices[:, :2] >= footprint_lower - margin) + & (support_vertices[:, :2] <= footprint_upper + margin), + dim=1, + ) + if bool(local_mask.any()): + support_height = support_vertices[local_mask, 2].max() + else: + distances = torch.linalg.vector_norm( + support_vertices[:, :2] - target_pose[env_id, :2, 3], + dim=1, + ) + nearest_count = min(8, int(support_vertices.shape[0])) + nearest = torch.topk( + distances, + nearest_count, + largest=False, + ).indices + support_height = support_vertices[nearest, 2].max() + target_pose[env_id, 2, 3] = ( + support_height + + float(self._policy_value(base.motion_policy, "surface_clearance")) + - rotated[:, 2].min() + ) + if any(torch.allclose(offsets, prior) for prior in seen_offsets): + continue + seen_offsets.append(offsets) + candidates.append( + replace( + base, + target=replace(base.target, object_target_pose=target_pose), + target_object_pose=target_pose, + motion_policy={ + **base.motion_policy, + "placement_candidate_index": candidate_index, + "placement_xy_offset": offsets, + }, + ) + ) + return tuple(candidates) or (base,) + + @staticmethod + def _placement_support_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 + + def _is_handover_continuation(self, step: SemanticStep) -> bool: + if step.operator != "place_relative": + return False + predecessors = { + candidate.id: candidate for candidate in self.program.semantic_steps + } + return any( + (predecessor := predecessors.get(dependency)) is not None + and predecessor.operator == "handover" + and predecessor.object_uid == step.object_uid + for dependency in step.depends_on + ) + + def _visual_target( + self, + binding: Mapping[str, Any], + arm: str, + ) -> torch.Tensor: + """Unproject one normalized image keypoint using live camera depth.""" + camera_uid = str(binding.get("camera_uid", "")) + sensor = self.env.sim.get_sensor(camera_uid) + if sensor is None: + raise ValueError(f"Unknown visual-constraint camera {camera_uid!r}.") + keypoint_value = binding.get("normalized_keypoint") + if keypoint_value is None: + bbox = binding.get("normalized_bbox") + if isinstance(bbox, Sequence) and len(bbox) == 4: + keypoint_value = [ + (float(bbox[0]) + float(bbox[2])) * 0.5, + (float(bbox[1]) + float(bbox[3])) * 0.5, + ] + if keypoint_value is None: + raise ValueError( + "visual_constraint requires a normalized keypoint or bbox in [0, 1]." + ) + keypoint = torch.as_tensor( + keypoint_value, + dtype=torch.float32, + device=self.env.device, + ).flatten() + if keypoint.numel() != 2 or bool( + ((~torch.isfinite(keypoint)) | (keypoint < 0.0) | (keypoint > 1.0)).any() + ): + raise ValueError( + "visual_constraint requires a normalized keypoint or bbox in [0, 1]." + ) + data = sensor.get_data() + if "depth" not in data: + raise ValueError( + f"Camera {camera_uid!r} must provide depth for visual Grounding." + ) + depth = torch.as_tensor(data["depth"], device=self.env.device).squeeze(-1) + if depth.ndim == 2: + depth = depth.unsqueeze(0).repeat(int(self.env.num_envs), 1, 1) + if depth.ndim != 3 or depth.shape[0] != int(self.env.num_envs): + raise ValueError("Camera depth must have shape (N, H, W) or (N, H, W, 1).") + height, width = depth.shape[-2:] + pixel_x = min(max(int(round(float(keypoint[0]) * (width - 1))), 0), width - 1) + pixel_y = min(max(int(round(float(keypoint[1]) * (height - 1))), 0), height - 1) + distance = depth[:, pixel_y, pixel_x].to(torch.float32) + if bool((~torch.isfinite(distance) | (distance <= 0.0)).any()): + raise ValueError("visual_constraint keypoint has no valid live depth.") + intrinsics = torch.as_tensor( + sensor.get_intrinsics(), + dtype=torch.float32, + device=self.env.device, + ) + if intrinsics.ndim == 2: + intrinsics = intrinsics.unsqueeze(0).repeat(int(self.env.num_envs), 1, 1) + camera_pose = torch.as_tensor( + sensor.get_arena_pose(to_matrix=True), + dtype=torch.float32, + device=self.env.device, + ) + if camera_pose.ndim == 2: + camera_pose = camera_pose.unsqueeze(0).repeat(int(self.env.num_envs), 1, 1) + fx = intrinsics[:, 0, 0] + fy = intrinsics[:, 1, 1] + cx = intrinsics[:, 0, 2] + cy = intrinsics[:, 1, 2] + point = torch.stack( + ( + (float(pixel_x) - cx) * distance / fx, + (float(pixel_y) - cy) * distance / fy, + distance, + torch.ones_like(distance), + ), + dim=1, + ) + world = torch.bmm(camera_pose, point.unsqueeze(-1)).squeeze(-1) + target = self._current_eef_pose(arm).clone() + target[:, :3, 3] = world[:, :3] + return target + + def _handover_target( + self, + step: SemanticStep, + binding: Mapping[str, Any], + object_pose: torch.Tensor, + reference_pose: torch.Tensor | None, + policy: Mapping[str, Any], + state: ExecutionState, + *, + orientation_reference_pose: torch.Tensor | None, + workspace: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> tuple[GraspGoal, torch.Tensor, dict[str, Any]]: + transfer_arm = str(binding.get("transfer_arm", "left_arm")) + receive_arm = str( + binding.get( + "receive_arm", + "right_arm" if transfer_arm == "left_arm" else "left_arm", + ) + ) + if transfer_arm == receive_arm or {transfer_arm, receive_arm} != { + "left_arm", + "right_arm", + }: + raise ValueError("HandOver requires distinct left_arm/right_arm roles.") + transfer_part = arm_control_part(self.env, transfer_arm) + held = state.get_held_object(transfer_part) + if held is None: + raise ValueError( + f"HandOver requires {transfer_arm} to hold {step.object_uid!r}." + ) + + del reference_pose + if workspace is None: + middle, final = self._handover_workspace_poses( + object_pose, + transfer_arm=transfer_arm, + receive_arm=receive_arm, + policy=policy, + step=step, + orientation_reference_pose=orientation_reference_pose, + ) + else: + middle, final = (item.clone() for item in workspace) + rotation = self._target_rotation( + step, + object_pose, + orientation_reference_pose=orientation_reference_pose, + ) + middle[:, :3, :3] = rotation + final[:, :3, :3] = rotation + semantics = self.semantics_factory(step.object_uid) + grounded_policy = dict(policy) + grounded_policy.update( + { + "transfer_arm": transfer_arm, + "receive_arm": receive_arm, + "middle_object_pose": middle, + "final_object_pose": final, + } + ) + return ( + GraspGoal(semantics=semantics), + middle, + grounded_policy, + ) + + def _handover_workspace_poses( + self, + object_pose: torch.Tensor, + *, + transfer_arm: str, + receive_arm: str, + policy: Mapping[str, Any], + step: SemanticStep, + orientation_reference_pose: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Choose the highest-ranked collision-aware handover workspace.""" + rotation = self._target_rotation( + step, + object_pose, + orientation_reference_pose=orientation_reference_pose, + ) + candidates = self._handover_workspace_candidates( + step, + object_pose, + transfer_arm=transfer_arm, + receive_arm=receive_arm, + policy=policy, + rotation=rotation, + ) + return candidates[0] + + def _handover_workspace_candidates( + self, + step: SemanticStep, + object_pose: torch.Tensor, + *, + transfer_arm: str, + receive_arm: str, + policy: Mapping[str, Any], + rotation: torch.Tensor, + ) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: + """Rank exchange poses inside the two arm workspaces and above obstacles.""" + if transfer_arm == receive_arm or {transfer_arm, receive_arm} != { + "left_arm", + "right_arm", + }: + raise ValueError("Handover workspace requires distinct arm roles.") + table = self.env.sim.get_rigid_object("table") + if table is not None and hasattr(table, "get_vertices"): + centers = [] + tops = [] + bounds = [] + for env_id in range(int(self.env.num_envs)): + vertices = _world_vertices(table, self.env, env_id) + lower = vertices[:, :2].min(dim=0).values + upper = vertices[:, :2].max(dim=0).values + centers.append((lower + upper) * 0.5) + tops.append(vertices[:, 2].max()) + bounds.append(torch.stack((lower, upper))) + center = torch.stack(centers) + table_top = torch.stack(tops) + table_bounds = torch.stack(bounds) + else: + left = self._current_eef_pose("left_arm") + right = self._current_eef_pose("right_arm") + center = (left[:, :2, 3] + right[:, :2, 3]) * 0.5 + table_top = object_pose[:, 2, 3] + extent = float(policy.get("exchange_candidate_offset", 0.16)) * 2.0 + table_bounds = torch.stack((center - extent, center + extent), dim=1) + + forward, lateral = robot_frame_axes(self.env) + left_base, right_base = arm_base_poses(self.env) + transfer_base = left_base if transfer_arm == "left_arm" else right_base + receive_base = right_base if receive_arm == "right_arm" else left_base + base_midpoint = (transfer_base[:, :2, 3] + receive_base[:, :2, 3]) * 0.5 + table_forward = torch.sum((center - base_midpoint) * forward, dim=1) + shared_center = base_midpoint + forward * table_forward[:, None] + offset = float(policy.get("exchange_candidate_offset", 0.16)) + obstacle_clearance = float(policy.get("exchange_obstacle_clearance", 0.04)) + tool_horizontal_envelope = float( + policy.get("exchange_gripper_horizontal_envelope", 0.035) + ) + float(policy.get("exchange_wrist_horizontal_envelope", 0.055)) + tool_vertical_envelope = float( + policy.get("exchange_gripper_vertical_envelope", 0.025) + ) + float(policy.get("exchange_wrist_vertical_envelope", 0.04)) + minimum_reach = float(policy.get("exchange_minimum_reach", 0.10)) + maximum_reach = float(policy.get("exchange_maximum_reach", 1.00)) + if not 0.0 <= minimum_reach < maximum_reach: + raise ValueError("Handover reach bounds require 0 <= minimum < maximum.") + requested_count = max(1, int(policy.get("exchange_candidate_count", 4))) + object_clearance = float(policy.get("exchange_clearance", 0.06)) + if ( + min( + obstacle_clearance, + tool_horizontal_envelope, + tool_vertical_envelope, + object_clearance, + ) + < 0.0 + ): + raise ValueError("Handover geometry clearances must be non-negative.") + xy_coefficients = ( + (0.0, 0.0), + (1.0, 0.0), + (-1.0, 0.0), + (2.0, 0.0), + (-2.0, 0.0), + (0.0, 0.5), + (0.0, -0.5), + ) + ranked_by_env: list[list[tuple[float, torch.Tensor]]] = [] + moved = _object(self.env, step.object_uid) + obstacle_uids = ( + self.env.sim.get_rigid_object_uid_list() + if hasattr(self.env.sim, "get_rigid_object_uid_list") + else [] + ) + for env_id in range(int(self.env.num_envs)): + local_vertices = _local_vertices(moved, self.env, env_id) + rotated = local_vertices @ rotation[env_id].transpose(0, 1) + half_xy = ( + rotated[:, :2].max(dim=0).values - rotated[:, :2].min(dim=0).values + ) * 0.5 + bottom = rotated[:, 2].min() + margin = half_xy + obstacle_clearance + tool_horizontal_envelope + lower_limit = table_bounds[env_id, 0] + margin + upper_limit = table_bounds[env_id, 1] - margin + options: list[tuple[float, torch.Tensor]] = [] + for forward_scale, lateral_scale in xy_coefficients: + xy = ( + shared_center[env_id] + + forward[env_id] * (offset * forward_scale) + + lateral[env_id] * (offset * lateral_scale) + ) + if bool(((xy < lower_limit) | (xy > upper_limit)).any()): + continue + transfer_distance = torch.linalg.vector_norm( + xy - transfer_base[env_id, :2, 3] + ) + receive_distance = torch.linalg.vector_norm( + xy - receive_base[env_id, :2, 3] + ) + if not ( + minimum_reach <= float(transfer_distance) <= maximum_reach + and minimum_reach <= float(receive_distance) <= maximum_reach + ): + continue + obstacle_score, nearby_obstacle_top = self._handover_obstacle_metrics( + xy, + env_id=env_id, + object_uid=step.object_uid, + obstacle_uids=obstacle_uids, + half_xy=half_xy, + clearance=obstacle_clearance + tool_horizontal_envelope, + ) + center_cost = float( + torch.linalg.vector_norm(xy - shared_center[env_id]) + ) + pose = object_pose[env_id].clone() + pose[:3, :3] = rotation[env_id] + pose[:2, 3] = xy + safety_floor = torch.maximum(table_top[env_id], nearby_obstacle_top) + safe_z = ( + safety_floor + object_clearance + tool_vertical_envelope - bottom + ) + pose[2, 3] = torch.maximum(object_pose[env_id, 2, 3], safe_z) + lift_cost = max( + 0.0, + float(pose[2, 3] - object_pose[env_id, 2, 3]), + ) + options.append( + (obstacle_score + center_cost * 0.25 + lift_cost * 0.1, pose) + ) + if not options: + raise ValueError( + "No handover exchange pose lies inside the table bounds and " + "the reachable intersection of both arm bases." + ) + options.sort(key=lambda item: item[0]) + ranked_by_env.append(options[:requested_count]) + + candidate_count = min( + requested_count, + max(len(options) for options in ranked_by_env), + ) + candidates = [] + for candidate_index in range(candidate_count): + middle = object_pose.clone() + for env_id, options in enumerate(ranked_by_env): + middle[env_id] = options[min(candidate_index, len(options) - 1)][1] + # The built-in HandOver primitive plans its final transfer/receiver + # phase concurrently. An exchange-to-exchange target makes that + # receiver path stationary; graph-level retreat/home nodes then + # clear the transfer arm before any receiver-side continuation. + final = middle.clone() + candidates.append((middle, final)) + return tuple(candidates) + + def _handover_obstacle_metrics( + self, + xy: torch.Tensor, + *, + env_id: int, + object_uid: str, + obstacle_uids: Sequence[str], + half_xy: torch.Tensor, + clearance: float, + ) -> tuple[float, torch.Tensor]: + score = 0.0 + highest_top = torch.tensor( + -torch.inf, + dtype=xy.dtype, + device=xy.device, + ) + for uid in obstacle_uids: + if uid in {"table", object_uid}: + continue + obstacle = self.env.sim.get_rigid_object(uid) + if obstacle is None or not hasattr(obstacle, "get_vertices"): + continue + vertices = _world_vertices(obstacle, self.env, env_id) + lower = vertices[:, :2].min(dim=0).values - half_xy - clearance + upper = vertices[:, :2].max(dim=0).values + half_xy + clearance + outside = torch.maximum( + torch.maximum(lower - xy, xy - upper), + torch.zeros_like(xy), + ) + if bool((outside > 0.0).any()): + distance = float(torch.linalg.vector_norm(outside)) + score += 1.0 / max(distance, 1.0e-3) + else: + score += 1.0e3 + highest_top = torch.maximum(highest_top, vertices[:, 2].max()) + return score, highest_top + + def _handover_receiver_exit( + self, + middle: torch.Tensor, + receive_arm: str, + policy: Mapping[str, Any], + ) -> torch.Tensor: + final = middle.clone() + receive_pose = self._current_eef_pose(receive_arm) + direction = receive_pose[:, :2, 3] - middle[:, :2, 3] + norm = torch.linalg.vector_norm(direction, dim=1, keepdim=True) + fallback = direction.new_zeros(direction.shape) + fallback[:, 1] = -1.0 if receive_arm == "right_arm" else 1.0 + direction = torch.where( + norm > 1.0e-6, direction / norm.clamp_min(1.0e-6), fallback + ) + final[:, :2, 3] += direction * min( + 0.12, + float(self._policy_value(policy, "relation_distance")) * 0.5, + ) + return final + + def _reference_pose(self, step: SemanticStep) -> torch.Tensor | None: + uid = step.goal.get("reference_object", step.goal.get("support_object")) + if not isinstance(uid, str) or not uid: + return None + if step.goal.get("reference_state") == "initial": + initial = getattr(self.env, "agent_initial_object_poses", {}).get(uid) + if initial is None: + raise ValueError(f"Initial pose for {uid!r} is unavailable.") + return _batched_pose(initial, self.env) + return _live_pose(self.env, uid) + + def _semantic_target( + self, + step: SemanticStep, + object_pose: torch.Tensor, + reference_pose: torch.Tensor | None, + policy: Mapping[str, Any], + *, + phase: str, + orientation_reference_pose: torch.Tensor | None = None, + ) -> torch.Tensor: + if step.operator in {"arrange_line", "place_in_line"}: + arrangement = self.arrangements.get(step.id) + if arrangement is None: + raise ValueError("arrange_line requires a live arrangement plan.") + target = arrangement.target( + step, + object_pose, + phase=phase, + policy=policy, + ) + target[:, :3, :3] = self._target_rotation( + step, + object_pose, + orientation_reference_pose=orientation_reference_pose, + ) + moved = _object(self.env, step.object_uid) + for env_id in range(int(self.env.num_envs)): + bottom = self._rotated_local_z_min( + moved, + target[env_id, :3, :3], + env_id, + ) + target[env_id, 2, 3] = ( + arrangement.table_top[env_id] + + float(policy["surface_clearance"]) + - bottom + ) + if phase == "staging": + target[:, 2, 3] += float(policy["transport_clearance"]) + return target + placement = self.placements.get(step.id) + if placement is not None: + target = placement.target( + step, + object_pose, + self._target_rotation( + step, + object_pose, + orientation_reference_pose=orientation_reference_pose, + ), + surface_clearance=float(policy["surface_clearance"]), + ) + if phase == "staging": + target[:, 2, 3] += float(policy["transport_clearance"]) + return target + if step.operator == "orient_object": + initial = None + if step.goal.get("position_anchor", "initial_xy") == "initial_xy": + initial = getattr(self.env, "agent_initial_object_poses", {}).get( + step.object_uid + ) + target = ( + _batched_pose(initial, self.env).clone() + if initial is not None + else object_pose.clone() + ) + target[:, :3, :3] = self._target_rotation( + step, + target, + orientation_reference_pose=orientation_reference_pose, + ) + support_uid = str(step.goal.get("support_object", "table")) + support = _object(self.env, support_uid) + moved = _object(self.env, step.object_uid) + for env_id in range(int(self.env.num_envs)): + support_top = _world_vertices(support, self.env, env_id)[:, 2].max() + bottom = self._rotated_local_z_min( + moved, + target[env_id, :3, :3], + env_id, + ) + target[env_id, 2, 3] = ( + support_top + float(policy["surface_clearance"]) - bottom + ) + if phase == "staging": + target[:, 2, 3] += float(policy["staging_lift_height"]) + return target + target = object_pose.clone() + if reference_pose is not None: + target[:, :3, 3] = reference_pose[:, :3, 3] + # Operators without a relational goal (for example press or a + # direction-only coordinated transport) must preserve the live origin + # instead of being silently projected onto a synthetic table support. + relation = ( + normalize_placement_relation(step.goal.get("relation", "on")) + if step.operator == "place_relative" + else str(step.goal.get("relation", "none")) + ) + distance = float(self._policy_value(policy, "relation_distance")) + relation_frame = str(step.goal.get("relation_frame", "world")) + forward_distance = distance + lateral_distance = distance + if relation_frame == "robot" and reference_pose is not None: + nominal = float(policy.get("robot_relative_distance", 0.10)) + clearance = float(policy.get("relation_clearance", 0.02)) + reference_uid = str(step.goal.get("reference_object", "")) + if reference_uid: + forward_axis, lateral_axis = robot_frame_axes(self.env) + forward_distance = self._relative_object_spacing( + step.object_uid, + reference_uid, + axis=forward_axis, + nominal=nominal, + clearance=clearance, + ) + lateral_distance = self._relative_object_spacing( + step.object_uid, + reference_uid, + axis=lateral_axis, + nominal=nominal, + clearance=clearance, + ) + directional_offset = relation_offset( + self.env, + relation, + frame=relation_frame, + forward_distance=forward_distance, + lateral_distance=lateral_distance, + dtype=target.dtype, + device=target.device, + ) + offsets = { + "above": (0.0, 0.0, float(self._policy_value(policy, "hover_height"))), + "held_above_initial": ( + 0.0, + 0.0, + float(self._policy_value(policy, "hover_height")), + ), + } + if directional_offset is not None: + target[:, :3, 3] += directional_offset + elif (offset := offsets.get(relation)) is not None: + target[:, :3, 3] += torch.tensor( + offset, + dtype=target.dtype, + device=target.device, + ) + slot = str(step.goal.get("slot", "auto")) + if relation in {"on", "on_top", "on_top_of", "inside"} and slot in { + "left", + "right", + }: + slot_offset = relation_offset( + self.env, + slot, + frame=relation_frame, + forward_distance=forward_distance, + lateral_distance=lateral_distance, + dtype=target.dtype, + device=target.device, + ) + if slot_offset is not None: + target[:, :3, 3] += slot_offset + direction = str(step.goal.get("direction", "none")) + direction_offsets = { + "world_x": (distance, 0.0, 0.0), + "world_y": (0.0, distance, 0.0), + "up": (0.0, 0.0, distance), + "down": (0.0, 0.0, -distance), + } + planar_direction_offset = relation_offset( + self.env, + direction, + frame=relation_frame, + forward_distance=distance, + lateral_distance=distance, + dtype=target.dtype, + device=target.device, + ) + if planar_direction_offset is not None: + target[:, :3, 3] += planar_direction_offset + elif direction in direction_offsets: + target[:, :3, 3] += torch.tensor( + direction_offsets[direction], + dtype=target.dtype, + device=target.device, + ) + + root_stack_layer = ( + step.operator == "build_stack" + and int(step.goal.get("layer_index", 0)) == 0 + and reference_pose is None + ) + if root_stack_layer: + table = _object(self.env, "table") + for env_id in range(int(self.env.num_envs)): + vertices = _world_vertices(table, self.env, env_id) + target[env_id, :2, 3] = ( + vertices[:, :2].min(dim=0).values + + vertices[:, :2].max(dim=0).values + ) * 0.5 + + target[:, :3, :3] = self._target_rotation( + step, + object_pose, + orientation_reference_pose=orientation_reference_pose, + ) + if ( + step.operator == "coordinated_transport" + and relation not in {"on", "on_top", "on_top_of", "inside"} + and direction not in {"up", "down"} + ): + release = str(step.goal.get("terminal_behavior", "hold")) == "place" + if not release: + target[:, 2, 3] = object_pose[:, 2, 3] + float( + self._policy_value(policy, "transport_clearance") + ) + else: + table = _object(self.env, "table") + moved = _object(self.env, step.object_uid) + clearance = float(self._policy_value(policy, "surface_clearance")) + for env_id in range(int(self.env.num_envs)): + table_top = _world_vertices(table, self.env, env_id)[:, 2].max() + bottom = self._rotated_local_z_min( + moved, + target[env_id, :3, :3], + env_id, + ) + target[env_id, 2, 3] = table_top + clearance - bottom + if relation in {"on", "on_top", "on_top_of"} or root_stack_layer: + support_uid = ( + step.goal.get("reference_object") + or step.goal.get("support_object") + or "table" + ) + support = _object(self.env, str(support_uid)) + moved = _object(self.env, step.object_uid) + for env_id in range(int(self.env.num_envs)): + support_top = _world_vertices(support, self.env, env_id)[:, 2].max() + bottom = self._rotated_local_z_min( + moved, + target[env_id, :3, :3], + env_id, + ) + target[env_id, 2, 3] = ( + support_top + + float(self._policy_value(policy, "surface_clearance")) + - bottom + ) + elif relation == "inside" and reference_pose is not None: + # Grounding the final move happens after the staging lift. Preserve + # the pre-pick supported height rather than the lifted live height. + supported_pose = orientation_reference_pose + if supported_pose is None: + supported_pose = object_pose + supported_pose = _batched_pose(supported_pose, self.env) + target[:, 2, 3] = supported_pose[:, 2, 3] + if phase == "staging": + # Staging is a runtime waypoint, not a persisted coordinate. This + # keeps in-place orientation robust to the object's live height. + target[:, 2, 3] += float(self._policy_value(policy, "transport_clearance")) + elif self._is_handover_continuation(step) and relation not in { + "on", + "on_top", + "on_top_of", + "inside", + }: + # A handover can leave the live rigid-body center a few centimetres + # below the original table-supported height. Reusing that drifted + # height for the lateral placement target makes the can intersect + # the table during release and it may tip or slide. Preserve the + # predecessor's supported height for the final held-object pose. + supported_pose = orientation_reference_pose + if supported_pose is None: + supported_pose = object_pose + supported_pose = _batched_pose(supported_pose, self.env) + target[:, 2, 3] = torch.maximum( + target[:, 2, 3], + supported_pose[:, 2, 3], + ) + return target + + def _relative_object_spacing( + self, + moved_uid: str, + reference_uid: str, + *, + axis: int | torch.Tensor, + nominal: float, + clearance: float, + ) -> float: + """Return deterministic center spacing from live object extents.""" + moved = _object(self.env, moved_uid) + reference = _object(self.env, reference_uid) + required = float(nominal) + for env_id in range(int(self.env.num_envs)): + moved_vertices = _world_vertices(moved, self.env, env_id) + reference_vertices = _world_vertices(reference, self.env, env_id) + if isinstance(axis, torch.Tensor): + direction = axis[env_id].to( + dtype=moved_vertices.dtype, + device=moved_vertices.device, + ) + moved_axis = moved_vertices[:, :2] @ direction + reference_axis = reference_vertices[:, :2] @ direction + else: + moved_axis = moved_vertices[:, axis] + reference_axis = reference_vertices[:, axis] + moved_half = (moved_axis.max() - moved_axis.min()) * 0.5 + reference_half = (reference_axis.max() - reference_axis.min()) * 0.5 + required = max( + required, + float(moved_half + reference_half) + float(clearance), + ) + return required + + def _upright_local_direction(self, step: SemanticStep) -> torch.Tensor: + axis = self._upright_local_axis(step) + entity = _object(self.env, step.object_uid) + vertices = _local_vertices(entity, self.env, 0) + extents = vertices.max(dim=0).values - vertices.min(dim=0).values + if axis == "long_axis": + axis_index = int(torch.argmax(extents).item()) + else: + axis_index = {"x": 0, "y": 1, "z": 2}[axis] + direction = torch.zeros(3, dtype=torch.float32, device=self.env.device) + direction[axis_index] = 1.0 + return direction + + def _uses_upright_yaw_search( + self, + step: SemanticStep, + constraint: OrientationConstraint, + ) -> bool: + """Preserve a live upright state as a planning preference. + + Explicit full-frame matching cannot admit yaw search. With no hard + orientation terms, yaw search is enabled only when the live object's + long axis is already upright, so a preceding upright operation remains + stable without turning that state into a sticky acceptance constraint. + """ + if constraint.allows_upright_yaw_search: + return True + if ( + constraint.terms + or constraint.planning_preference != "minimize_rotation_from_current" + ): + return False + entity = _object(self.env, step.object_uid) + vertices = _local_vertices(entity, self.env, 0) + extents = vertices.max(dim=0).values - vertices.min(dim=0).values + axis_index = int(torch.argmax(extents).item()) + pose = _live_pose(self.env, step.object_uid) + cosine = pose[:, 2, axis_index].abs().clamp(0.0, 1.0) + tolerance = float(self.runtime_policy.predicate_fallbacks["upright_max_tilt"]) + return bool(torch.all(torch.arccos(cosine) <= tolerance).item()) + + @staticmethod + def _upright_local_axis(step: SemanticStep) -> str: + align_terms = tuple( + term + for term in compile_orientation_constraint(step.goal).terms + if isinstance(term, AlignAxisConstraint) + ) + if align_terms: + return align_terms[0].local_axis + axis = str(step.goal.get("upright_local_axis", "auto")) + return "long_axis" if axis == "auto" else axis + + def _target_rotation( + self, + step: SemanticStep, + object_pose: torch.Tensor, + *, + orientation_reference_pose: torch.Tensor | None = None, + ) -> torch.Tensor: + constraint = compile_orientation_constraint(step.goal) + if not constraint.terms: + return object_pose[:, :3, :3].clone() + if ( + len(constraint.terms) == 1 + and isinstance(constraint.terms[0], MatchRotationConstraint) + and constraint.terms[0].reference == "step_start" + ): + if orientation_reference_pose is not None: + reference = _batched_pose(orientation_reference_pose, self.env) + return reference[:, :3, :3].clone() + return object_pose[:, :3, :3].clone() + goal = str(step.goal.get("orientation_goal", "none")) + align_term = next( + ( + term + for term in constraint.terms + if isinstance(term, AlignAxisConstraint) + ), + None, + ) + if align_term is not None: + goal = "upright" + if goal not in {"upright", "lay_flat", "axis_align"}: + raise ValueError(f"Unsupported orientation_goal {goal!r}.") + + entity = _object(self.env, step.object_uid) + rotations = [] + for env_id in range(int(self.env.num_envs)): + vertices = _local_vertices(entity, self.env, env_id) + extents = vertices.max(dim=0).values - vertices.min(dim=0).values + longest_to_shortest = torch.argsort( + extents, + descending=True, + ).tolist() + if goal == "upright": + upright_axis = ( + align_term.local_axis + if align_term is not None + else self._upright_local_axis(step) + ) + vertical_axis = ( + int(longest_to_shortest[0]) + if upright_axis == "long_axis" + else {"x": 0, "y": 1, "z": 2}[upright_axis] + ) + horizontal_axis = next( + int(axis) + for axis in longest_to_shortest + if int(axis) != vertical_axis + ) + elif goal == "lay_flat": + vertical_axis = int(longest_to_shortest[-1]) + horizontal_axis = int(longest_to_shortest[0]) + else: + horizontal_axis = self._aligned_local_axis( + step, + longest_to_shortest, + ) + vertical_axis = next( + int(axis) + for axis in reversed(longest_to_shortest) + if int(axis) != horizontal_axis + ) + direction = self._horizontal_orientation( + step, + object_pose, + env_id, + horizontal_axis, + ) + rotations.append( + self._world_aligned_rotation( + direction, + horizontal_axis=horizontal_axis, + vertical_axis=vertical_axis, + ) + ) + return torch.stack(rotations) + + @staticmethod + def _aligned_local_axis( + step: SemanticStep, + longest_to_shortest: Sequence[int], + ) -> int: + axis = str(step.goal.get("orientation_axis", "long_axis")) + if axis == "x": + return 0 + if axis == "y": + return 1 + if axis == "long_axis": + return int(longest_to_shortest[0]) + if axis == "short_axis": + return int(longest_to_shortest[-1]) + raise ValueError(f"Unsupported axis_align orientation_axis {axis!r}.") + + def _horizontal_orientation( + self, + step: SemanticStep, + object_pose: torch.Tensor, + env_id: int, + local_axis: int, + ) -> torch.Tensor: + align_to = step.goal.get("orientation_reference_object") + if isinstance(align_to, str) and align_to: + reference = _object(self.env, align_to) + vertices = _local_vertices(reference, self.env, env_id) + extents = vertices.max(dim=0).values - vertices.min(dim=0).values + ordered = torch.argsort(extents, descending=True) + requested = str(step.goal.get("orientation_axis", "long_axis")) + reference_axis = int( + ordered[-1] if requested == "short_axis" else ordered[0] + ) + reference_pose = _live_pose(self.env, align_to) + direction = reference_pose[env_id, :3, reference_axis].clone() + elif step.operator in {"arrange_line", "place_in_line"}: + arrangement = self.arrangements.get(step.id) + axis_index = 0 if arrangement is None else arrangement.axis_index + direction = torch.zeros( + 3, + dtype=object_pose.dtype, + device=object_pose.device, + ) + direction[axis_index] = 1.0 + elif str(step.goal.get("orientation_axis", "")) in {"y", "world_y"}: + direction = object_pose.new_tensor([0.0, 1.0, 0.0]) + elif str(step.goal.get("orientation_axis", "")) in {"x", "world_x"}: + direction = object_pose.new_tensor([1.0, 0.0, 0.0]) + else: + direction = object_pose[env_id, :3, local_axis].clone() + direction[2] = 0.0 + norm = torch.linalg.vector_norm(direction) + if float(norm) < 1.0e-6: + return object_pose.new_tensor([1.0, 0.0, 0.0]) + return direction / norm + + @staticmethod + def _world_aligned_rotation( + horizontal_direction: torch.Tensor, + *, + horizontal_axis: int, + vertical_axis: int, + ) -> torch.Tensor: + world_up = horizontal_direction.new_tensor([0.0, 0.0, 1.0]) + remaining_axis = ({0, 1, 2} - {horizontal_axis, vertical_axis}).pop() + columns = [torch.zeros_like(world_up) for _ in range(3)] + columns[horizontal_axis] = horizontal_direction + columns[vertical_axis] = world_up + columns[remaining_axis] = torch.linalg.cross( + world_up, + horizontal_direction, + ) + rotation = torch.stack(columns, dim=1) + if float(torch.linalg.det(rotation)) < 0.0: + rotation[:, remaining_axis] *= -1.0 + return rotation + + def _rotated_local_z_min( + self, + entity: Any, + rotation: torch.Tensor, + env_id: int, + ) -> torch.Tensor: + vertices = _local_vertices(entity, self.env, env_id) + return (vertices @ rotation.transpose(0, 1))[:, 2].min() + + def _current_eef_pose(self, arm: str) -> torch.Tensor: + """Return the live TCP pose for one logical Action Engine arm.""" + if arm not in {"left_arm", "right_arm"}: + raise ValueError(f"Expected a physical arm, got {arm!r}.") + if hasattr(self.env, "get_current_xpos_agent"): + left, right = self.env.get_current_xpos_agent() + value = left if arm == "left_arm" else right + if value is not None: + return _batched_pose(value, self.env) + + is_left = arm == "left_arm" + if not hasattr(self.env, "get_agent_arm_control_part"): + raise ValueError("Coordinated placement requires live TCP poses.") + part = self.env.get_agent_arm_control_part(is_left) + qpos = self._arm_qpos(arm) + return _batched_pose( + self.env.robot.compute_fk(qpos=qpos, name=part, to_matrix=True), + self.env, + ) + + def _press_goal( + self, + uid: str, + object_pose: torch.Tensor, + *, + semantics: ObjectSemantics | None = None, + ) -> PressGoal: + """Ground the live top surface into the typed press-affordance contract.""" + if semantics is None: + semantics = self.semantics_factory(uid) + entity = _object(self.env, uid) + reference_pose = object_pose[0] + world_position = reference_pose[:3, 3].clone() + world_position[2] = _world_vertices(entity, self.env, 0)[:, 2].max() + rotation = reference_pose[:3, :3] + local_position = rotation.transpose(0, 1) @ ( + world_position - reference_pose[:3, 3] + ) + local_axis = rotation.transpose(0, 1) @ torch.tensor( + [0.0, 0.0, -1.0], + dtype=rotation.dtype, + device=rotation.device, + ) + press_semantics = replace( + semantics, + affordance=PressAffordance( + press_axis=local_axis, + press_position=tuple(float(value) for value in local_position), + ), + ) + return PressGoal( + semantics=press_semantics, + target_pose=object_pose.clone(), + ) + + def _retreat_pose( + self, + arm: str, + policy: Mapping[str, Any], + reference: torch.Tensor | None, + *, + clear_exchange: bool = False, + ) -> torch.Tensor: + target = self._retreat_reference_pose(arm, reference).clone() + desired = float(self._policy_value(policy, "retreat_height")) + if clear_exchange: + _, lateral = robot_frame_axes(self.env) + direction = lateral if arm == "left_arm" else -lateral + target[:, :2, 3] += direction.to( + dtype=target.dtype, + device=target.device, + ) * float(policy.get("retreat_distance", 0.10)) + desired = max( + desired, + float(self._policy_value(policy, "minimum_retreat_height")), + ) + ceiling = float(self._policy_value(policy, "maximum_eef_height")) + height = torch.clamp(ceiling - target[:, 2, 3], min=0.0, max=desired) + target[:, 2, 3] += height + return target + + def _retreat_reference_pose( + self, + arm: str, + reference: torch.Tensor | None, + ) -> torch.Tensor: + """Resolve the live or speculative TCP pose from which retreat starts.""" + pose = reference + if pose is None and hasattr(self.env, "get_current_xpos_agent"): + left, right = self.env.get_current_xpos_agent() + pose = left if arm == "left_arm" else right + if pose is None: + raise ValueError("Retreat grounding requires a live end-effector pose.") + return _batched_pose(pose, self.env) + + def _joint_target( + self, + arm: str, + control: str, + source: str, + binding: Mapping[str, Any], + ) -> torch.Tensor: + if source in {"gripper_closed", "gripper_open"}: + value = ( + getattr(self.env, "close_state") + if source == "gripper_closed" + else getattr(self.env, "open_state") + ) + return torch.as_tensor( + value, + dtype=torch.float32, + device=self.env.device, + ) + if source == "joint_delta": + current = self._arm_qpos(arm).clone() + index = int(binding["joint_index"]) + current[:, index] += torch.deg2rad( + torch.tensor( + float(binding.get("delta_degrees", 0.0)), + device=current.device, + ) + ) + return current + initial = getattr(self.env, "init_qpos", self.env.robot.get_qpos()) + joint_ids = self._joint_ids(arm, control) + return torch.as_tensor(initial, device=self.env.device)[:, joint_ids] + + def _arm_qpos(self, arm: str) -> torch.Tensor: + if hasattr(self.env, "get_current_qpos_agent"): + left, right = self.env.get_current_qpos_agent() + return torch.as_tensor( + left if arm == "left_arm" else right, + dtype=torch.float32, + device=self.env.device, + ) + return self.env.robot.get_qpos()[:, self._joint_ids(arm, "arm")] + + def _joint_ids(self, arm: str, control: str) -> list[int]: + side = "left" if arm == "left_arm" else "right" + key = f"{side}_{'eef' if control == 'hand' else 'arm'}_joints" + return list(getattr(self.env, key, ())) + + def _explicit_pose( + self, + binding: Mapping[str, Any], + object_pose: torch.Tensor, + ) -> torch.Tensor: + reference = str(binding.get("reference", "absolute")) + target = object_pose.clone() + if reference == "absolute": + values = binding.get("position_by_env", binding.get("position")) + position = torch.as_tensor( + values, + dtype=target.dtype, + device=target.device, + ) + if position.ndim == 1: + position = position.unsqueeze(0).repeat(int(self.env.num_envs), 1) + target[:, :3, 3] = position + return target + offset = torch.as_tensor( + binding.get("offset", (0.0, 0.0, 0.0)), + dtype=target.dtype, + device=target.device, + ) + target[:, :3, 3] += offset + return target + + def _coordinated_grasps( + self, + semantics: ObjectSemantics, + object_pose: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Build a deterministic opposing pair along the object's longest XY axis.""" + vertices = semantics.geometry.get("mesh_vertices") + vertices = torch.as_tensor( + vertices, + dtype=torch.float32, + device=self.env.device, + ) + lower = vertices.min(dim=0).values + upper = vertices.max(dim=0).values + axis = int(torch.argmax(upper[:2] - lower[:2]).item()) + center = (lower + upper) * 0.5 + grasp_policy = self.runtime_policy.grounding["coordinated_grasp"] + inset = max( + float(grasp_policy["minimum_inset"]), + float((upper[axis] - lower[axis]) * grasp_policy["inset_fraction"]), + ) + left = torch.eye(4, dtype=torch.float32, device=self.env.device) + right = left.clone() + left[:3, 3] = center + right[:3, 3] = center + left[axis, 3] = lower[axis] + inset + right[axis, 3] = upper[axis] - inset + # Keep TCP z horizontal and facing the object from opposite sides. + if axis == 0: + left[:3, :3] = torch.tensor( + [[0.0, 0.0, 1.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], + device=self.env.device, + ) + right[:3, :3] = torch.tensor( + [[0.0, 0.0, -1.0], [-1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], + device=self.env.device, + ) + else: + left[:3, :3] = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, -1.0, 0.0]], + device=self.env.device, + ) + right[:3, :3] = torch.tensor( + [[-1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, -1.0, 0.0]], + device=self.env.device, + ) + batch = int(self.env.num_envs) + return left.unsqueeze(0).repeat(batch, 1, 1), right.unsqueeze(0).repeat( + batch, 1, 1 + ) diff --git a/embodichain/gen_sim/action_engine/runtime/loader.py b/embodichain/gen_sim/action_engine/runtime/loader.py new file mode 100644 index 000000000..d4c808a07 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/loader.py @@ -0,0 +1,292 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Load or compile execution programs without publishing intermediate copies.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import replace +import json +from pathlib import Path +from typing import Any + +from embodichain.gen_sim.action_engine.protocol import ( + ACTION_ENGINE_CONFIG_SCHEMA, + EXECUTION_PROGRAM_SCHEMA, + SEED_GRAPH_SCHEMA, +) + +from .models import ExecutionProgram + +__all__ = [ + "load_agent_execution_program", + "load_execution_program", +] + + +def _read_json(path: Path, *, label: str) -> dict[str, Any]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except OSError as exc: + raise ValueError(f"Unable to read {label} at {path}: {exc}") from exc + except json.JSONDecodeError as exc: + raise ValueError(f"{label} at {path} is not valid JSON: {exc}") from exc + if not isinstance(value, Mapping): + raise ValueError(f"{label} must contain a JSON object.") + return dict(value) + + +def load_execution_program( + source: Mapping[str, Any] | str | Path, + *, + known_objects: set[str] | None = None, + registry: Any | None = None, + require_executable: bool = True, +) -> ExecutionProgram: + """Load a v3 SeedGraph and reject every legacy execution schema.""" + value = ( + dict(source) + if isinstance(source, Mapping) + else _read_json(Path(source).expanduser().resolve(), label="execution program") + ) + schema = value.get("schema_version") + if schema == SEED_GRAPH_SCHEMA: + from embodichain.gen_sim.action_engine.capabilities import ( + build_atomic_capability_registry, + ) + from embodichain.gen_sim.action_engine.compiler import ( + seed_graph_to_execution_program, + ) + from embodichain.gen_sim.action_engine.domain import validate_seed_graph + from embodichain.gen_sim.action_engine.planning.linker import ( + validate_persisted_contracts, + ) + + registry = registry or build_atomic_capability_registry() + seed = validate_seed_graph( + value, + known_objects=known_objects, + known_actions=registry.names(), + executable_actions=registry.executable_names(), + require_executable=require_executable, + ) + validate_persisted_contracts(seed, registry) + internal = seed_graph_to_execution_program( + seed, + known_objects=known_objects, + registry=registry, + require_executable=require_executable, + ) + return replace(ExecutionProgram.from_mapping(internal), seed_graph=seed) + if schema == "action_engine_seed_graph_v2": + raise ValueError( + "SeedGraph v2 lacks persisted Action Contracts and cannot be loaded; " + "regenerate seed_task_graph.json and agent_config.json with the current " + "generator to produce action_engine_seed_graph_v3." + ) + if schema == EXECUTION_PROGRAM_SCHEMA: + raise ValueError( + "Action Engine v1 execution programs are no longer accepted; " + "regenerate the task to produce action_engine_seed_graph_v3." + ) + raise ValueError(f"Unsupported Action Engine graph schema {schema!r}.") + + +def _resolve_config_path( + config: Mapping[str, Any], + config_path: str | Path, + *keys: str, +) -> Path | None: + base = Path(config_path).expanduser().resolve().parent + 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 (base / path).resolve() + return None + + +def load_agent_execution_program( + agent_config: Mapping[str, Any], + *, + agent_config_path: str | Path, + regenerate: bool = False, + require_executable: bool = True, +) -> ExecutionProgram: + """Resolve an agent config and optionally rebuild its SeedGraph in memory. + + ``--regenerate`` intentionally does not write a second graph artifact. The + deterministic compiler result is validated and handed directly to runtime. + """ + if agent_config.get("schema_version") != ACTION_ENGINE_CONFIG_SCHEMA: + raise ValueError( + "This Action Engine runtime accepts only v2 bundles. Regenerate " + "task_spec.json, scene_requirements.json, seed_task_graph.json, " + "and agent_config.json " + "with the current generator." + ) + known_objects = _known_objects(agent_config) + task_path = _resolve_config_path( + agent_config, + agent_config_path, + "task_spec", + "task_spec_path", + ) + execution_path = _resolve_config_path( + agent_config, + agent_config_path, + "seed_task_graph", + "seed_task_graph_path", + "offline_seed_task_graph", + "offline_seed_task_graph_path", + ) + if regenerate: + if task_path is None: + raise ValueError("--regenerate requires agent_config.task_spec.") + task_spec = _read_json(task_path, label="task specification") + reference_graph = ( + _read_json(execution_path, label="SeedGraph") + if execution_path is not None and execution_path.is_file() + else None + ) + program = load_execution_program( + _regenerate_seed_graph(task_spec, reference_graph=reference_graph), + known_objects=known_objects, + require_executable=require_executable, + ) + elif execution_path is None: + if task_path is None: + raise ValueError("agent_config requires seed_task_graph or task_spec.") + task_spec = _read_json(task_path, label="task specification") + program = load_execution_program( + _regenerate_seed_graph(task_spec), + known_objects=known_objects, + require_executable=require_executable, + ) + else: + program = load_execution_program( + execution_path, + known_objects=known_objects, + require_executable=require_executable, + ) + _verify_agent_program(agent_config, program) + _verify_program_objects(agent_config, program) + return program + + +def _known_objects(agent_config: Mapping[str, Any]) -> set[str] | None: + source = agent_config.get("source") + if not isinstance(source, Mapping): + return None + uid_map = source.get("uid_map") + if not isinstance(uid_map, Mapping): + return None + values = {str(uid) for uid in uid_map.values() if str(uid)} + return values or None + + +def _verify_agent_program( + agent_config: Mapping[str, Any], + program: ExecutionProgram, +) -> None: + """Reject a valid program that belongs to a different generated bundle.""" + configured_task = agent_config.get("task_name") + if configured_task is not None and configured_task != program.task: + raise ValueError( + f"agent_config.task_name {configured_task!r} does not match " + f"execution program task {program.task!r}." + ) + expected_hash = agent_config.get("seed_task_graph_hash") + if expected_hash is None: + return + if not isinstance(expected_hash, str) or not expected_hash: + raise ValueError( + "agent_config.seed_task_graph_hash must be a non-empty string." + ) + if program.seed_graph is not None: + from embodichain.gen_sim.action_engine.domain import seed_graph_hash + + actual_hash = seed_graph_hash(program.seed_graph) + else: + from embodichain.gen_sim.action_engine.domain import execution_program_hash + + actual_hash = execution_program_hash(program.raw) + if actual_hash != expected_hash: + raise ValueError( + "SeedGraph hash does not match agent_config; regenerate the " + "configuration bundle before running it." + ) + + +def _regenerate_seed_graph( + task_spec: Mapping[str, Any], + *, + reference_graph: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + from embodichain.gen_sim.action_engine.domain import validate_task_spec + + task = validate_task_spec(task_spec) + oracle = task.get("oracle", {}) + reference = oracle.get("reference_seed_graph") + if isinstance(reference, Mapping): + return dict(reference) + metadata = task.get("metadata", {}) + bindings = metadata.get("role_bindings", {}) + if not isinstance(bindings, Mapping): + raise ValueError("TaskSpec.metadata.role_bindings must be a mapping.") + if not bindings and reference_graph is not None: + graph_metadata = reference_graph.get("metadata", {}) + if not isinstance(graph_metadata, Mapping): + raise ValueError("SeedGraph.metadata must be a mapping.") + bindings = graph_metadata.get("role_bindings", {}) + if not isinstance(bindings, Mapping): + raise ValueError("SeedGraph.metadata.role_bindings must be a mapping.") + from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph + + return instantiate_seed_graph(task, bindings) + + +def _verify_program_objects( + agent_config: Mapping[str, Any], + program: ExecutionProgram, +) -> None: + known = _known_objects(agent_config) + if known is None: + return + references = {step.object_uid for step in program.semantic_steps} + for step in program.semantic_steps: + for key in ( + "reference_object", + "support_object", + "orientation_reference_object", + ): + value = step.goal.get(key) + if isinstance(value, str): + references.add(value) + for payload in step.goal.get("payloads", []): + value = payload.get("object") if isinstance(payload, Mapping) else payload + if isinstance(value, str): + references.add(value) + unknown = references - known - {"self", "table", "table_center"} + if unknown: + raise ValueError( + "Execution Program references objects not present in the scene: " + f"{sorted(unknown)}. Regenerate the configuration bundle." + ) diff --git a/embodichain/gen_sim/action_engine/runtime/models.py b/embodichain/gen_sim/action_engine/runtime/models.py new file mode 100644 index 000000000..b4596a241 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/models.py @@ -0,0 +1,314 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Small typed runtime views over the serialized execution program.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass, field +from typing import Any + +import numpy as np +import torch + +from embodichain.lab.sim.atomic_actions import StateDelta + +from .state import ExecutionState + +__all__ = [ + "ActionOutcome", + "ExecutionEdge", + "ExecutionProgram", + "ExecutionReport", + "ExecutionResult", + "GroundedAction", + "SemanticStep", +] + + +@dataclass(frozen=True) +class ExecutionEdge: + """One executable DAG edge containing symbolic atomic actions.""" + + id: str + source: str + target: str + actions: tuple[dict[str, Any], ...] + depends_on: tuple[str, ...] = () + resources: tuple[str, ...] = () + + +@dataclass(frozen=True) +class SemanticStep: + """One closed-loop intent expanded into one or more execution edges.""" + + id: str + parent_step_id: str + operator: str + object_uid: str + actor: dict[str, Any] + goal: dict[str, Any] + depends_on: tuple[str, ...] + postcondition: dict[str, Any] + edge_ids: tuple[str, ...] + + +@dataclass(frozen=True) +class ExecutionProgram: + """Validated in-memory form of ``action_engine_execution_program_v1``.""" + + raw: dict[str, Any] + task: str + start: str + goal: str + nodes: tuple[dict[str, Any], ...] + edges: tuple[ExecutionEdge, ...] + semantic_steps: tuple[SemanticStep, ...] + allocation_groups: tuple[dict[str, Any], ...] + seed_graph: dict[str, Any] | None = None + + @classmethod + def from_mapping(cls, value: Mapping[str, Any]) -> "ExecutionProgram": + """Construct an immutable runtime view from a validated mapping.""" + raw = deepcopy(dict(value)) + edges = tuple( + ExecutionEdge( + id=str(edge["id"]), + source=str(edge["source"]), + target=str(edge["target"]), + actions=tuple( + deepcopy(dict(action)) + for action in edge.get("actions", edge.get("symbolic_actions", ())) + ), + depends_on=tuple(str(item) for item in edge.get("depends_on", ())), + resources=tuple(str(item) for item in edge.get("resources", ())), + ) + for edge in raw["edges"] + ) + steps = tuple( + SemanticStep( + id=str(step["id"]), + parent_step_id=str(step["parent_step_id"]), + operator=str(step["operator"]), + object_uid=str(step.get("object", step.get("object_uid", ""))), + actor=deepcopy(dict(step["actor"])), + goal=deepcopy(dict(step.get("goal", {}))), + depends_on=tuple(str(item) for item in step.get("depends_on", ())), + postcondition=deepcopy(dict(step.get("postcondition", {}))), + edge_ids=tuple(str(item) for item in step["edge_ids"]), + ) + for step in raw["semantic_steps"] + ) + return cls( + raw=raw, + task=str(raw.get("task", raw.get("task_name", "task"))), + start=str(raw["start"]), + goal=str(raw["goal"]), + nodes=tuple(deepcopy(raw["nodes"])), + edges=edges, + semantic_steps=steps, + allocation_groups=tuple( + deepcopy(dict(group)) for group in raw.get("allocation_groups", ()) + ), + seed_graph=None, + ) + + +@dataclass(frozen=True) +class GroundedAction: + """A public atomic-action target resolved from the current simulator state.""" + + action_class: str + arm: str + control: str + target: Any + cfg: dict[str, Any] + object_pose: torch.Tensor | None = None + reference_pose: torch.Tensor | None = None + target_object_pose: torch.Tensor | None = None + motion_policy: dict[str, Any] = field(default_factory=dict) + object_uid: str | None = None + """Scene UID of the object whose semantic step produced this action.""" + + +@dataclass +class ActionOutcome: + """Planning output kept in full-robot coordinates.""" + + trajectory: torch.Tensor + success: torch.Tensor + next_state: ExecutionState + grounded: GroundedAction + prior_state: ExecutionState | None = None + expected_effects: StateDelta | None = None + planner_trace: dict[str, Any] = field(default_factory=dict) + + def state_after(self, verified: torch.Tensor) -> ExecutionState: + """Commit expected effects only for physically verified rows.""" + if self.prior_state is None or self.expected_effects is None: + return self.next_state + mask = torch.as_tensor( + verified, + dtype=torch.bool, + device=self.trajectory.device, + ).reshape(-1) + if mask.numel() != self.trajectory.shape[0]: + raise ValueError("Verified mask must match the ActionOutcome batch.") + terminal_qpos = ( + self.trajectory[:, -1] + if self.trajectory.shape[1] + else self.prior_state.last_qpos + ) + qpos = torch.where( + mask[:, None], + terminal_qpos, + self.prior_state.last_qpos, + ) + task = self.expected_effects.apply( + self.prior_state.to_task_state(), + mask, + ) + return ExecutionState.from_task_state(task, last_qpos=qpos) + + @property + def cost(self) -> torch.Tensor: + """Return joint-path length for each vectorized environment.""" + if self.trajectory.shape[1] < 2: + return torch.zeros( + self.trajectory.shape[0], + dtype=torch.float32, + device=self.trajectory.device, + ) + return torch.linalg.vector_norm( + torch.diff(self.trajectory, dim=1), + dim=-1, + ).sum(dim=1) + + +@dataclass +class ExecutionResult(Sequence[torch.Tensor]): + """Result marker used by the existing demonstration-runner contract.""" + + actions: list[torch.Tensor] + success: torch.Tensor + semantic_success: dict[str, torch.Tensor] + record_dir: str | None = None + already_executed: bool = True + retry_count: int = 0 + recovery_count: int = 0 + revision_count: int = 0 + failure_events: list[dict[str, Any]] = field(default_factory=list) + runtime_revisions: list[dict[str, Any]] = field(default_factory=list) + retry_counts: list[int] = field(default_factory=list) + + @property + def runtime_success(self) -> torch.Tensor: + return self.success + + @property + def runtime_graph_output_dir(self) -> str | None: + return self.record_dir + + def __len__(self) -> int: + return len(self.actions) + + def __iter__(self): + return iter(self.actions) + + def __getitem__(self, index): + return self.actions[index] + + +@dataclass(frozen=True) +class ExecutionReport: + """JSON-safe Task Engine result built from an ``ExecutionResult``. + + The runtime result deliberately keeps tensors because the legacy demo + runner consumes them. The Task Engine boundary instead exposes only a + compact, serializable audit view and never retains the action tensors. + """ + + task_id: str + plan_hash: str + action_graph_hash: str + status: str + run_id: str + episode_id: str + provenance: dict[str, Any] + environments: tuple[dict[str, Any], ...] = () + action_count: int = 0 + retry_count: int = 0 + recovery_count: int = 0 + revision_count: int = 0 + failure_events: tuple[dict[str, Any], ...] = () + graph_revisions: tuple[dict[str, Any], ...] = () + record_dir: str | None = None + error: str | None = None + schema_version: str = "action_engine_execution_report_v2" + + def as_mapping(self) -> dict[str, Any]: + """Return a detached mapping suitable for strict JSON serialization.""" + return { + "schema_version": self.schema_version, + "task_id": self.task_id, + "plan_hash": self.plan_hash, + "action_graph_hash": self.action_graph_hash, + "status": self.status, + "run_id": self.run_id, + "episode_id": self.episode_id, + "provenance": deepcopy(self.provenance), + "environments": deepcopy(list(self.environments)), + "action_count": self.action_count, + "retry_count": self.retry_count, + "recovery_count": self.recovery_count, + "revision_count": self.revision_count, + "failure_events": deepcopy(list(self.failure_events)), + "graph_revisions": deepcopy(list(self.graph_revisions)), + "record_dir": self.record_dir, + "error": self.error, + } + + def to_dict(self) -> dict[str, Any]: + """Compatibility spelling for artifact and CLI publishers.""" + return self.as_mapping() + + +def success_mask(value: bool | torch.Tensor, count: int, device: Any) -> torch.Tensor: + """Normalize a primitive's scalar or batched success result.""" + mask = torch.as_tensor(value, dtype=torch.bool, device=device).reshape(-1) + if mask.numel() == 1: + return mask.repeat(count) + if mask.numel() != count: + raise ValueError( + f"Atomic action success has {mask.numel()} values; expected {count}." + ) + return mask + + +def trajectory_cost_numpy(value: torch.Tensor) -> np.ndarray: + """Expose trajectory costs to assignment solvers without retaining gradients.""" + if value.shape[1] < 2: + return np.zeros(value.shape[0], dtype=np.float64) + diffs = torch.diff(value.detach(), dim=1) + return ( + torch.linalg.vector_norm(diffs, dim=-1) + .sum(dim=1) + .cpu() + .numpy() + .astype(np.float64) + ) diff --git a/embodichain/gen_sim/action_engine/runtime/motion_policy.py b/embodichain/gen_sim/action_engine/runtime/motion_policy.py new file mode 100644 index 000000000..8de7cd4bf --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/motion_policy.py @@ -0,0 +1,106 @@ +# ---------------------------------------------------------------------------- +# 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 collections.abc import Mapping +from copy import deepcopy +from typing import Any + +from embodichain.gen_sim.action_engine.config import default_runtime_policy +from embodichain.gen_sim.action_engine.domain.motion import validate_motion_policy + +__all__ = ["resolve_motion_policy", "with_motion_modifiers"] + +_PROFILE_ALIASES = dict( + franka="dual_franka", ur3="dual_ur3", ur5="dual_ur5", ur10="dual_ur10" +) + + +def resolve_motion_policy( + robot_profile: str, + atomic_action: str, + policy_spec: Mapping[str, Any], + *, + motion_defaults: Mapping[str, Mapping[str, Any]] | None = None, + motion_modifiers: ( + Mapping[str, Mapping[str, Mapping[str, Mapping[str, Any]]]] | None + ) = None, + inline_overrides: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + """Resolve an action base policy plus its composable typed modifiers.""" + profile = _PROFILE_ALIASES.get(str(robot_profile), str(robot_profile)) + runtime_policy = ( + default_runtime_policy(profile) + if motion_defaults is None or motion_modifiers is None + else None + ) + defaults = ( + runtime_policy.motion_defaults + if motion_defaults is None and runtime_policy is not None + else motion_defaults + ) + modifiers = ( + runtime_policy.motion_modifiers + if motion_modifiers is None and runtime_policy is not None + else motion_modifiers + ) + if defaults is None or modifiers is None: + raise ValueError("Motion defaults and modifiers must be provided together.") + action = str(atomic_action) + if action not in defaults: + raise ValueError(f"Unknown Action Engine motion base {action!r}.") + + spec = validate_motion_policy(policy_spec) + policy = deepcopy(dict(defaults[action])) + modifier_values: dict[str, Any] = {} + modifier_sources: dict[str, tuple[str, str]] = {} + for modifier in spec["modifiers"]: + modifier_type = modifier["type"] + mode = modifier["mode"] + patch = modifiers.get(modifier_type, {}).get(mode, {}).get(action) + if not isinstance(patch, Mapping): + raise ValueError( + f"Motion modifier {(modifier_type, mode)!r} is not supported " + f"by AtomicAction {action!r}." + ) + for key, value in patch.items(): + if key in modifier_values and modifier_values[key] != value: + raise ValueError( + f"Motion modifiers {modifier_sources[key]!r} and " + f"{(modifier_type, mode)!r} conflict on parameter {key!r}." + ) + modifier_values[key] = deepcopy(value) + modifier_sources[key] = (modifier_type, mode) + policy.update(modifier_values) + if inline_overrides is not None: + policy.update(deepcopy(dict(inline_overrides))) + return policy + + +def with_motion_modifiers( + policy_spec: Mapping[str, Any], + *modifiers: tuple[str, str], +) -> dict[str, Any]: + """Return a validated policy reference with missing modifiers appended.""" + policy = validate_motion_policy(policy_spec) + existing = { + (modifier["type"], modifier["mode"]) for modifier in policy["modifiers"] + } + for modifier_type, mode in modifiers: + if (modifier_type, mode) not in existing: + policy["modifiers"].append({"type": modifier_type, "mode": mode}) + return validate_motion_policy(policy) diff --git a/embodichain/gen_sim/action_engine/runtime/predicates.py b/embodichain/gen_sim/action_engine/runtime/predicates.py new file mode 100644 index 000000000..aff6d1df7 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/predicates.py @@ -0,0 +1,746 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Evaluate canonical closed-loop predicates against live environment state.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any + +import torch + +from embodichain.gen_sim.action_engine.config import default_runtime_policy + +from .frames import relation_axes +from .robot_parts import arm_control_part + +__all__ = ["PREDICATE_TYPES", "evaluate_predicate"] + +PREDICATE_TYPES = frozenset( + { + "both_arms_at_initial_qpos", + "both_grippers_open", + "coordinated_placed", + "grippers_clear_of_object", + "held_by_both_grippers", + "object_axis_near", + "object_axis_offset_near", + "object_held", + "object_held_by_both_grippers", + "object_held_by_gripper", + "object_in_container", + "object_lifted", + "object_not_fallen", + "object_on_object", + "object_supported_by", + "object_position_near", + "object_relative_position", + "object_upright", + "object_xy_near", + "objects_collinear", + "objects_ordered", + "pressed", + } +) +_DEFAULT_PREDICATE_FALLBACKS = default_runtime_policy("dual_ur10").predicate_fallbacks + + +def _predicate_fallbacks(env: Any) -> Mapping[str, Any]: + policy = getattr(env, "runtime_policy", None) + value = getattr(policy, "predicate_fallbacks", None) + return value if isinstance(value, Mapping) else _DEFAULT_PREDICATE_FALLBACKS + + +def _constant(env: Any, value: bool) -> torch.Tensor: + return torch.full( + (int(env.num_envs),), + value, + dtype=torch.bool, + device=env.device, + ) + + +def _pose(env: Any, uid: str) -> torch.Tensor: + entity = 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=env.device, + ) + if pose.ndim == 2: + pose = pose.unsqueeze(0).repeat(int(env.num_envs), 1, 1) + return pose + + +def _position(env: Any, uid: str) -> torch.Tensor: + return _pose(env, uid)[:, :3, 3] + + +def _world_vertices(env: Any, uid: str, env_id: int) -> torch.Tensor: + entity = env.sim.get_rigid_object(uid) + if entity is None: + raise ValueError(f"Unknown rigid object {uid!r}.") + value = entity.get_vertices(env_ids=[env_id], scale=True) + if isinstance(value, (tuple, list)): + value = value[0] + vertices = torch.as_tensor(value, dtype=torch.float32, device=env.device) + if vertices.ndim == 3 and vertices.shape[0] == 1: + vertices = vertices[0] + if vertices.ndim != 2 or vertices.shape[-1] != 3 or vertices.numel() == 0: + raise ValueError(f"Rigid object {uid!r} has invalid mesh vertices.") + pose = _pose(env, uid)[env_id] + return vertices @ pose[:3, :3].transpose(0, 1) + pose[:3, 3] + + +def _projected_center_of_mass( + env: Any, + uid: str, + env_id: int, + world_vertices: torch.Tensor, +) -> torch.Tensor: + """Return the live COM projection, with a geometry-center fallback.""" + entity = env.sim.get_rigid_object(uid) + body_data = None if entity is None else getattr(entity, "body_data", None) + com_pose = None if body_data is None else getattr(body_data, "com_pose", None) + if callable(com_pose): + com_pose = com_pose() + if com_pose is not None: + local_com = torch.as_tensor( + com_pose, + dtype=torch.float32, + device=env.device, + ) + if local_com.ndim == 1: + local_com = local_com.unsqueeze(0).repeat(int(env.num_envs), 1) + if local_com.ndim == 2 and local_com.shape[0] == int(env.num_envs): + pose = _pose(env, uid)[env_id] + return (pose[:3, :3] @ local_com[env_id, :3] + pose[:3, 3])[:2] + return ( + world_vertices[:, :2].min(dim=0).values + + world_vertices[:, :2].max(dim=0).values + ) * 0.5 + + +def _object_supported_by( + env: Any, + spec: Mapping[str, Any], + defaults: Mapping[str, Any], +) -> torch.Tensor: + """Evaluate one-frame geometric support without advancing simulation.""" + object_uid = _object(spec) + support_uid = str( + spec.get( + "support", + spec.get("reference_object", spec.get("reference", "")), + ) + ) + if not support_uid: + raise ValueError("Support predicate requires a support object uid.") + margin = float(spec.get("com_margin", defaults["support_com_margin"])) + max_gap = float(spec.get("max_vertical_gap", defaults["support_max_vertical_gap"])) + max_penetration = float( + spec.get("max_penetration", defaults["support_max_penetration"]) + ) + min_overlap = float( + spec.get("min_overlap_ratio", defaults["support_min_overlap_ratio"]) + ) + result = _constant(env, False) + for env_id in range(int(env.num_envs)): + moved = _world_vertices(env, object_uid, env_id) + support = _world_vertices(env, support_uid, env_id) + moved_lower = moved[:, :2].min(dim=0).values + moved_upper = moved[:, :2].max(dim=0).values + support_lower = support[:, :2].min(dim=0).values + support_upper = support[:, :2].max(dim=0).values + overlap_extent = torch.clamp( + torch.minimum(moved_upper, support_upper) + - torch.maximum(moved_lower, support_lower), + min=0.0, + ) + moved_extent = torch.clamp(moved_upper - moved_lower, min=1e-6) + overlap_ratio = torch.prod(overlap_extent) / torch.prod(moved_extent) + projected_center = _projected_center_of_mass( + env, + object_uid, + env_id, + moved, + ) + center_supported = torch.all( + projected_center >= support_lower + margin + ) & torch.all(projected_center <= support_upper - margin) + local_mask = torch.all( + (support[:, :2] >= moved_lower - margin) + & (support[:, :2] <= moved_upper + margin), + dim=1, + ) + if bool(local_mask.any()): + local_support_height = support[local_mask, 2].max() + else: + # Sparse meshes may have no vertex exactly under a small payload. + # Nearest vertices are a local fallback; using the mesh-wide peak + # would confuse a remote protrusion with the candidate support pose. + distances = torch.linalg.vector_norm( + support[:, :2] - projected_center, + dim=1, + ) + count = min(8, int(support.shape[0])) + local_support_height = support[ + torch.topk(distances, count, largest=False).indices, 2 + ].max() + vertical_gap = moved[:, 2].min() - local_support_height + result[env_id] = bool( + center_supported + and overlap_ratio >= min_overlap + and vertical_gap >= -max_penetration + and vertical_gap <= max_gap + ) + return result + + +def _objects(spec: Mapping[str, Any]) -> list[str]: + values = spec.get("objects", spec.get("object_uids")) + if not isinstance(values, Sequence) or isinstance(values, (str, bytes)): + raise ValueError("Predicate requires a non-empty objects list.") + return [str(value) for value in values] + + +def _object(spec: Mapping[str, Any]) -> str: + value = spec.get("object", spec.get("object_uid")) + if not isinstance(value, str) or not value: + raise ValueError("Predicate requires a non-empty object uid.") + return value + + +def _local_axis_index(env: Any, uid: str, axis: Any) -> int: + name = str(axis).lower() + if name in {"x", "y", "z"}: + return {"x": 0, "y": 1, "z": 2}[name] + if name not in {"long", "long_axis", "longest"}: + raise ValueError(f"Unsupported upright local axis {axis!r}.") + entity = env.sim.get_rigid_object(uid) + if entity is None: + raise ValueError(f"Unknown rigid object {uid!r}.") + vertices = entity.get_vertices(env_ids=[0], scale=True) + if isinstance(vertices, (tuple, list)): + vertices = vertices[0] + vertices = torch.as_tensor(vertices, dtype=torch.float32, device=env.device) + if vertices.ndim == 3: + vertices = vertices[0] + if vertices.ndim != 2 or vertices.shape[-1] != 3 or vertices.numel() == 0: + raise ValueError(f"Rigid object {uid!r} has invalid mesh vertices.") + extents = vertices.max(dim=0).values - vertices.min(dim=0).values + return int(torch.argmax(extents).item()) + + +def _arm_values( + env: Any, kind: str +) -> tuple[torch.Tensor | None, torch.Tensor | None] | None: + getter = getattr(env, f"get_current_{kind}_agent", None) + if callable(getter): + left, right = getter() + values = [] + for value in (left, right): + if value is None: + values.append(None) + continue + item = torch.as_tensor(value, device=env.device) + if kind == "xpos" and item.ndim == 2: + item = item.unsqueeze(0).repeat(int(env.num_envs), 1, 1) + elif kind == "gripper_state" and item.ndim == 1: + item = item.unsqueeze(0) + values.append(item) + return values[0], values[1] + if kind != "gripper_state": + return None + qpos = env.robot.get_qpos() + values = [] + for side in ("left", "right"): + ids = list(getattr(env, f"{side}_eef_joints", ())) + if not ids: + return None + values.append(qpos[:, ids]) + return values[0], values[1] + + +def _gripper_has_closed( + env: Any, + gripper: torch.Tensor, + *, + tolerance: float, +) -> torch.Tensor: + """Check closure intent without requiring an impossible empty-gripper pose.""" + gripper = gripper.to(device=env.device, dtype=torch.float32) + open_state = getattr(env, "open_state", None) + close_state = getattr(env, "close_state", None) + reference = open_state if open_state is not None else close_state + if reference is None: + return _constant(env, False) + expected = torch.as_tensor( + reference, + dtype=torch.float32, + device=env.device, + ).flatten() + repeats = (gripper.shape[-1] + expected.numel() - 1) // expected.numel() + expected = expected.repeat(repeats)[: gripper.shape[-1]] + distance = torch.linalg.vector_norm(gripper - expected, dim=-1) + if open_state is not None: + return distance > tolerance + return distance <= tolerance + + +def _object_held( + env: Any, + uid: str, + *, + owners: Mapping[str, Sequence[str | None]] | None, + states: Mapping[tuple[str, str], Any] | None, + position_tolerance: float, + gripper_tolerance: float, + required_arm: str | None = None, +) -> torch.Tensor: + """Verify registry ownership against live object, TCP, and gripper state.""" + result = _constant(env, False) + if owners is None or states is None or uid not in owners: + return result + eef_values = _arm_values(env, "xpos") + gripper_values = _arm_values(env, "gripper_state") + if eef_values is None or gripper_values is None: + return result + + object_pose = _pose(env, uid) + for arm_index, arm in enumerate(("left_arm", "right_arm")): + if required_arm is not None and arm != required_arm: + continue + state = states.get((uid, arm)) + held = ( + None if state is None else state.get_held_object(arm_control_part(env, arm)) + ) + actual_eef = eef_values[arm_index] + gripper = gripper_values[arm_index] + if held is None or actual_eef is None or gripper is None: + continue + label = getattr(held.semantics, "label", None) + if not label and held.semantics.entity is not None: + label = getattr(held.semantics.entity, "uid", None) + if label != uid: + continue + actual_eef = actual_eef.to(device=env.device, dtype=object_pose.dtype) + expected_eef = torch.bmm( + object_pose, + held.object_to_eef.to(device=env.device, dtype=object_pose.dtype), + ) + position_ok = ( + torch.linalg.vector_norm( + actual_eef[:, :3, 3] - expected_eef[:, :3, 3], dim=-1 + ) + <= position_tolerance + ) + closed = _gripper_has_closed( + env, + gripper, + tolerance=gripper_tolerance, + ) + owned = torch.tensor( + [item == arm for item in owners[uid]], + dtype=torch.bool, + device=env.device, + ) + result |= owned & position_ok & closed + return result + + +def _coordinated_held( + env: Any, + uid: str, + state: Any, + *, + position_tolerance: float, + gripper_tolerance: float, +) -> torch.Tensor: + result = _constant(env, False) + if state is None: + return result + held_relations = tuple( + state.get_held_object(arm_control_part(env, arm)) + for arm in ("left_arm", "right_arm") + ) + if any(held is None for held in held_relations): + return result + for held in held_relations: + assert held is not None + label = getattr(held.semantics, "label", None) + if not label and getattr(held.semantics, "entity", None) is not None: + label = getattr(held.semantics.entity, "uid", None) + if label != uid: + return result + eef_values = _arm_values(env, "xpos") + gripper_values = _arm_values(env, "gripper_state") + if eef_values is None or gripper_values is None: + return result + + object_pose = _pose(env, uid) + result = _constant(env, True) + for arm_index, held in enumerate(held_relations): + assert held is not None + if held.env_mask is not None: + result &= held.env_mask.to(device=env.device) + actual_eef = eef_values[arm_index] + gripper = gripper_values[arm_index] + if actual_eef is None or gripper is None: + return _constant(env, False) + transform = held.object_to_eef.to( + device=env.device, + dtype=object_pose.dtype, + ) + expected_eef = torch.bmm(object_pose, transform) + actual_eef = actual_eef.to(device=env.device, dtype=object_pose.dtype) + position_ok = ( + torch.linalg.vector_norm( + actual_eef[:, :3, 3] - expected_eef[:, :3, 3], + dim=-1, + ) + <= position_tolerance + ) + closed = _gripper_has_closed( + env, + gripper, + tolerance=gripper_tolerance, + ) + result &= position_ok & closed + return result + + +def evaluate_predicate( + env: Any, + spec: Mapping[str, Any] | Sequence[Mapping[str, Any]] | None, + *, + held_owners: Mapping[str, Sequence[str | None]] | None = None, + held_states: Mapping[tuple[str, str], Any] | None = None, + coordinated_state: Any | None = None, +) -> torch.Tensor: + """Evaluate one typed predicate or a boolean predicate tree.""" + runtime = { + "held_owners": held_owners, + "held_states": held_states, + "coordinated_state": coordinated_state, + } + defaults = _predicate_fallbacks(env) + if spec is None: + return _constant(env, True) + if isinstance(spec, Sequence) and not isinstance(spec, (str, bytes, Mapping)): + result = _constant(env, True) + for term in spec: + result &= evaluate_predicate(env, term, **runtime) + return result + if not isinstance(spec, Mapping): + raise TypeError("Predicate must be a mapping or a sequence of mappings.") + op = str(spec.get("op", "")).lower() + if not op and "terms" in spec: + op = "all" + if op in {"all", "and"}: + return evaluate_predicate(env, list(spec.get("terms", ())), **runtime) + if op in {"any", "or"}: + result = _constant(env, False) + for term in spec.get("terms", ()): + result |= evaluate_predicate(env, term, **runtime) + return result + if op == "not": + return ~evaluate_predicate(env, spec.get("term"), **runtime) + + kind = str(spec.get("type", spec.get("kind", ""))).lower() + if kind in {"semantic_goal", "line_member_placed", "stack_layer_supported"}: + raise ValueError( + f"Predicate {kind!r} is a compiler marker and requires the " + "executor's grounded target." + ) + if kind in {"object_held", "object_held_by_gripper"}: + required_arm = spec.get("arm") + if required_arm in {"left", "right"}: + required_arm = f"{required_arm}_arm" + return _object_held( + env, + _object(spec), + owners=held_owners, + states=held_states, + position_tolerance=float( + spec.get("position_tolerance", defaults["held_position_tolerance"]) + ), + gripper_tolerance=float( + spec.get("gripper_tolerance", defaults["held_gripper_tolerance"]) + ), + required_arm=str(required_arm) if required_arm else None, + ) + if kind == "handover_complete": + required_arm = spec.get("arm", "right_arm") + return _object_held( + env, + _object(spec), + owners=held_owners, + states=held_states, + position_tolerance=float( + spec.get("position_tolerance", defaults["held_position_tolerance"]) + ), + gripper_tolerance=float( + spec.get("gripper_tolerance", defaults["held_gripper_tolerance"]) + ), + required_arm=str(required_arm), + ) + if kind in {"held_by_both_grippers", "object_held_by_both_grippers"}: + return _coordinated_held( + env, + _object(spec), + coordinated_state, + position_tolerance=float( + spec.get("position_tolerance", defaults["held_position_tolerance"]) + ), + gripper_tolerance=float( + spec.get("gripper_tolerance", defaults["held_gripper_tolerance"]) + ), + ) + if kind in {"object_position_near", "position_near"}: + position = _position(env, _object(spec)) + target = torch.as_tensor( + spec.get("target_position", spec.get("target")), + dtype=position.dtype, + device=position.device, + ) + if target.ndim == 1: + target = target.unsqueeze(0) + return torch.linalg.vector_norm(position - target, dim=-1) <= float( + spec.get("tolerance", defaults["position_tolerance"]) + ) + if kind in {"object_xy_near", "xy_near"}: + position = _position(env, _object(spec))[:, :2] + target = torch.as_tensor( + spec.get("target_xy", spec.get("target")), + dtype=position.dtype, + device=position.device, + ).reshape(-1, 2) + return torch.linalg.vector_norm(position - target, dim=-1) <= float( + spec.get("tolerance", defaults["xy_tolerance"]) + ) + if kind in {"object_relative_position", "relative_position"}: + reference_uid = spec.get("reference_object", spec.get("reference")) + if not isinstance(reference_uid, str) or not reference_uid: + raise ValueError("Relative-position predicate requires a reference object.") + relation = str(spec.get("relation", "")) + axes = relation_axes( + env, + relation, + frame=str(spec.get("relation_frame", "world")), + ) + if not axes: + raise ValueError(f"Unsupported directional relation {relation!r}.") + delta = ( + _position(env, _object(spec))[:, :2] - _position(env, reference_uid)[:, :2] + ) + minimum_distance = float(spec.get("minimum_distance", 0.0)) + result = _constant(env, True) + for axis in axes: + projection = torch.sum( + delta * axis.to(dtype=delta.dtype, device=delta.device), dim=1 + ) + result &= projection >= minimum_distance + return result + if kind in {"object_in_container", "inside"}: + position = _position(env, _object(spec)) + container = _position( + env, str(spec.get("container", spec.get("reference_object"))) + ) + xy = torch.linalg.vector_norm(position[:, :2] - container[:, :2], dim=-1) + z = position[:, 2] - container[:, 2] + return ( + (xy <= float(spec.get("xy_radius", defaults["container_xy_radius"]))) + & (z >= float(spec.get("min_z_offset", defaults["container_min_z_offset"]))) + & (z <= float(spec.get("max_z_offset", defaults["container_max_z_offset"]))) + ) + if kind in {"object_supported_by", "object_on_object", "on"}: + return _object_supported_by(env, spec, defaults) + if kind == "object_not_fallen": + axis = _pose(env, _object(spec))[:, :3, 2] + cosine = axis[:, 2].clamp(-1.0, 1.0) + return torch.arccos(cosine) <= float( + spec.get("max_tilt", defaults["not_fallen_max_tilt"]) + ) + if kind == "object_upright": + uid = _object(spec) + local_axis = spec.get("local_axis", "long_axis") + axis_index = _local_axis_index( + env, + uid, + local_axis, + ) + axis = _pose(env, uid)[:, :3, axis_index] + cosine = axis[:, 2].clamp(-1.0, 1.0) + directed = spec.get( + "directed", + str(local_axis).lower() not in {"long", "long_axis", "longest"}, + ) + if not isinstance(directed, bool): + raise ValueError("object_upright directed must be a boolean.") + if not directed: + cosine = cosine.abs() + return torch.arccos(cosine) <= float( + spec.get("max_tilt", defaults["upright_max_tilt"]) + ) + if kind in {"object_axis_offset_near", "object_axis_near"}: + object_position = _position(env, _object(spec)) + axis = _axis_index(spec.get("axis", "x")) + reference_uid = spec.get( + "reference_object", + spec.get("reference", spec.get("support")), + ) + if isinstance(reference_uid, str) and reference_uid: + values = object_position[:, axis] - _position(env, reference_uid)[:, axis] + else: + values = object_position[:, axis] + target = spec.get( + "target_offset", + spec.get("offset", spec.get("target", 0.0)), + ) + target_value = torch.as_tensor( + target, + dtype=values.dtype, + device=values.device, + ) + return torch.abs(values - target_value) <= float( + spec.get("tolerance", defaults["axis_tolerance"]) + ) + if kind in {"objects_collinear", "collinear"}: + positions = torch.stack( + [_position(env, uid) for uid in _objects(spec)], + dim=1, + ) + axis = 0 if str(spec.get("axis", "x")) in {"x", "world_x"} else 1 + values = positions[:, :, 1 - axis] + return values.max(dim=1).values - values.min(dim=1).values <= float( + spec.get("tolerance", defaults["collinearity_tolerance"]) + ) + if kind in {"objects_ordered", "ordered"}: + positions = torch.stack( + [_position(env, uid) for uid in _objects(spec)], + dim=1, + ) + axis = 0 if str(spec.get("axis", "x")) in {"x", "world_x"} else 1 + differences = torch.diff(positions[:, :, axis], dim=1) + tolerance = float(spec.get("tolerance", defaults["ordering_tolerance"])) + if str(spec.get("direction", "ascending")) == "descending": + return torch.all(differences <= tolerance, dim=1) + return torch.all(differences >= -tolerance, dim=1) + if kind == "object_lifted": + position = _position(env, _object(spec))[:, 2] + initial = spec.get("initial_height") + if initial is None: + initial_pose = getattr(env, "agent_initial_object_poses", {}).get( + _object(spec) + ) + if initial_pose is None: + raise ValueError("object_lifted requires an initial object pose.") + initial = initial_pose[:, 2, 3] + initial = torch.as_tensor(initial, device=position.device) + return position >= initial + float( + spec.get("min_height", defaults["minimum_lift_height"]) + ) + if kind in {"both_arms_at_initial_qpos", "arms_home"}: + current = env.robot.get_qpos() + initial = getattr(env, "init_qpos", current) + return torch.all( + torch.abs(current - initial) + <= float(spec.get("tolerance", defaults["arm_initial_qpos_tolerance"])), + dim=-1, + ) + if kind in {"both_grippers_open", "grippers_open"}: + if not hasattr(env, "get_current_gripper_state_agent"): + return _constant(env, False) + left, right = env.get_current_gripper_state_agent() + expected = torch.as_tensor( + env.open_state, + dtype=torch.float32, + device=env.device, + ) + results = [] + for value in (left, right): + value = torch.as_tensor(value, dtype=torch.float32, device=env.device) + if value.ndim == 1: + value = value.unsqueeze(0).repeat(int(env.num_envs), 1) + results.append( + torch.linalg.vector_norm(value - expected, dim=-1) + <= float(spec.get("tolerance", defaults["gripper_state_tolerance"])) + ) + return results[0] & results[1] + if kind == "grippers_clear_of_object": + eef_values = _arm_values(env, "xpos") + if eef_values is None: + return _constant(env, False) + object_position = _position(env, _object(spec)) + clearance = float( + spec.get( + "min_distance", + spec.get("clearance", defaults["gripper_clear_min_distance"]), + ) + ) + result = _constant(env, True) + for eef in eef_values: + if eef is None: + return _constant(env, False) + result &= ( + torch.linalg.vector_norm( + eef[:, :3, 3] - object_position, + dim=-1, + ) + >= clearance + ) + return result + if kind == "pressed": + checker = getattr(env, "is_object_pressed", None) + if callable(checker): + value = checker(_object(spec), spec.get("terminal_state", "activated")) + result = torch.as_tensor(value, dtype=torch.bool, device=env.device) + return ( + result.repeat(int(env.num_envs)) + if result.ndim == 0 + else result.reshape(-1) + ) + states = getattr(env, "action_engine_semantic_states", {}) + value = states.get((_object(spec), "pressed")) + if value is None: + return _constant(env, False) + result = torch.as_tensor(value, dtype=torch.bool, device=env.device) + return ( + result.repeat(int(env.num_envs)) if result.ndim == 0 else result.reshape(-1) + ) + if kind == "coordinated_placed": + relation = str(spec.get("relation", "on")) + reference = spec.get("support_object", spec.get("reference_object")) + translated = { + "type": ( + "object_in_container" if relation == "inside" else "object_supported_by" + ), + "object": _object(spec), + ("container" if relation == "inside" else "support"): reference, + } + return evaluate_predicate(env, translated, **runtime) + raise ValueError(f"Unsupported execution predicate {kind!r}.") + + +def _axis_index(value: Any) -> int: + axis = str(value).lower().replace("world_", "") + if axis not in {"x", "y", "z"}: + raise ValueError(f"Unsupported predicate axis {value!r}.") + return {"x": 0, "y": 1, "z": 2}[axis] 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/runtime/robot_parts.py b/embodichain/gen_sim/action_engine/runtime/robot_parts.py new file mode 100644 index 000000000..fe6f71c58 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/robot_parts.py @@ -0,0 +1,34 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Resolve semantic Action Engine arms to physical robot control parts.""" + +from __future__ import annotations + +from typing import Any + +__all__ = ["arm_control_part"] + + +def arm_control_part(env: Any, arm: str) -> str: + """Return the physical arm control part for a semantic arm name.""" + if arm not in {"left_arm", "right_arm"}: + raise ValueError(f"Expected a semantic arm, got {arm!r}.") + if hasattr(env, "get_agent_arm_control_part"): + part = env.get_agent_arm_control_part(arm == "left_arm") + if part: + return str(part) + return arm diff --git a/embodichain/gen_sim/action_engine/runtime/solver_compat.py b/embodichain/gen_sim/action_engine/runtime/solver_compat.py new file mode 100644 index 000000000..810bc8cf7 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/solver_compat.py @@ -0,0 +1,234 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Install solver compatibility corrections scoped to Action Engine.""" + +from __future__ import annotations + +from collections.abc import Mapping +import functools +import threading +from typing import Any + +import numpy as np +import torch + +from embodichain.lab.sim.solvers import PytorchSolver, URSolver, URSolverCfg + +__all__ = [ + "install_action_engine_solver_compat", + "install_pytorch_solver_tcp_compat", + "install_ur5_solver_frame_compat", + "repair_action_engine_ur5_solver_cfg", +] + +_PYTORCH_INSTALL_MARKER = "_action_engine_tcp_inverse_compat_installed" +_UR5_INSTALL_MARKER = "_action_engine_ur5_frame_compat_installed" +_UR5_ANALYTIC_TO_URDF_EE = np.eye(4, dtype=np.float32) +_UR5_ANALYTIC_TO_URDF_EE[0, 3] = -0.01 +_UR_DH_FIELDS = ("d1", "a2", "a3", "d4", "d5", "d6") + + +def repair_action_engine_ur5_solver_cfg(robot_cfg: Any) -> int: + """Repair stale UR10 DH defaults before Action Engine creates a UR5 robot. + + ``SolverCfg.from_dict`` constructs a UR10 config before assigning a + non-default ``ur_type``. Generated UR5 Action Engine configs therefore + reach the environment with UR10 DH values. Repair only that exact stale + signature so explicitly calibrated parameters remain untouched. + + Args: + robot_cfg: Robot configuration whose solver configs will be inspected. + + Returns: + Number of unique solver configs repaired by this call. + """ + configured = getattr(robot_cfg, "solver_cfg", None) + candidates = ( + configured.values() if isinstance(configured, Mapping) else (configured,) + ) + stale_defaults = URSolverCfg() + stale_dh = tuple(float(getattr(stale_defaults, name)) for name in _UR_DH_FIELDS) + + repaired = 0 + visited: set[int] = set() + for solver_cfg in candidates: + cfg_id = id(solver_cfg) + if cfg_id in visited: + continue + visited.add(cfg_id) + if not isinstance(solver_cfg, URSolverCfg): + continue + ur_type = str(getattr(solver_cfg, "ur_type", "")) + if ur_type != "ur5": + continue + current_dh = tuple(float(getattr(solver_cfg, name)) for name in _UR_DH_FIELDS) + if not np.allclose(current_dh, stale_dh, rtol=0.0, atol=1.0e-12): + continue + canonical = URSolverCfg(ur_type=ur_type) + for name in _UR_DH_FIELDS: + setattr(solver_cfg, name, getattr(canonical, name)) + repaired += 1 + return repaired + + +def install_action_engine_solver_compat(robot: Any) -> int: + """Install all solver corrections required by the Action Engine runtime.""" + return install_pytorch_solver_tcp_compat(robot) + install_ur5_solver_frame_compat( + robot + ) + + +def install_pytorch_solver_tcp_compat(robot: Any) -> int: + """Correct TCP inversion on every PytorchSolver owned by ``robot``. + + The shared solver currently transposes a rotation into an overlapping + tensor view. This wrapper transforms the requested TCP pose with a proper + matrix inverse, temporarily presents an identity TCP to the original + implementation, and otherwise preserves its sampling and ranking behavior. + + Args: + robot: Initialized robot containing its private solver registry. + + Returns: + Number of solver instances wrapped by this call. + """ + solvers = getattr(robot, "_solvers", None) + if not isinstance(solvers, Mapping): + return 0 + + installed = 0 + visited: set[int] = set() + for solver in solvers.values(): + solver_id = id(solver) + if solver_id in visited: + continue + visited.add(solver_id) + if not isinstance(solver, PytorchSolver) or bool( + getattr(solver, _PYTORCH_INSTALL_MARKER, False) + ): + continue + _wrap_solver(solver) + installed += 1 + return installed + + +def _wrap_solver(solver: PytorchSolver) -> None: + original_get_ik = solver.get_ik + call_lock = threading.RLock() + + @functools.wraps(original_get_ik) + def corrected_get_ik( + target_xpos: torch.Tensor | np.ndarray, + *args: Any, + **kwargs: Any, + ) -> Any: + target = torch.as_tensor( + target_xpos, + dtype=torch.float32, + device=solver.device, + ) + tcp = torch.as_tensor( + solver.tcp_xpos, + dtype=torch.float32, + device=solver.device, + ) + link_target = target @ torch.linalg.inv(tcp) + + # The solver instance is shared by vectorized environments. Protect the + # temporary TCP substitution in case a caller plans from another thread. + with call_lock: + active_tcp = solver.tcp_xpos + solver.tcp_xpos = np.eye(4, dtype=np.float32) + try: + return original_get_ik( + target_xpos=link_target, + *args, + **kwargs, + ) + finally: + solver.tcp_xpos = active_tcp + + solver.get_ik = corrected_get_ik + setattr(solver, _PYTORCH_INSTALL_MARKER, True) + + +def install_ur5_solver_frame_compat(robot: Any) -> int: + """Align UR5 analytic IK targets with the URDF ``ee_link`` frame. + + The UR5 asset carries a fixed ``-0.01 m`` local-x offset on ``ee_link`` + that is absent from the analytic DH model. The correction is installed + only for UR5 solvers owned by an Action Engine environment. + + Args: + robot: Initialized robot containing its private solver registry. + + Returns: + Number of solver instances wrapped by this call. + """ + solvers = getattr(robot, "_solvers", None) + if not isinstance(solvers, Mapping): + return 0 + + installed = 0 + visited: set[int] = set() + for solver in solvers.values(): + solver_id = id(solver) + if solver_id in visited: + continue + visited.add(solver_id) + if ( + not isinstance(solver, URSolver) + or str(getattr(getattr(solver, "cfg", None), "ur_type", "")) != "ur5" + or bool(getattr(solver, _UR5_INSTALL_MARKER, False)) + ): + continue + _wrap_ur5_solver(solver) + installed += 1 + return installed + + +def _wrap_ur5_solver(solver: URSolver) -> None: + original_get_ik = solver.get_ik + + @functools.wraps(original_get_ik) + def corrected_get_ik( + target_xpos: torch.Tensor | np.ndarray, + *args: Any, + **kwargs: Any, + ) -> Any: + target = torch.as_tensor( + target_xpos, + dtype=torch.float32, + device=solver.device, + ) + tcp = torch.as_tensor( + solver.tcp_xpos, + dtype=torch.float32, + device=solver.device, + ) + analytic_to_urdf = torch.as_tensor( + _UR5_ANALYTIC_TO_URDF_EE, + dtype=torch.float32, + device=solver.device, + ) + corrected_target = ( + target @ torch.linalg.inv(tcp) @ torch.linalg.inv(analytic_to_urdf) @ tcp + ) + return original_get_ik(corrected_target, *args, **kwargs) + + solver.get_ik = corrected_get_ik + setattr(solver, _UR5_INSTALL_MARKER, True) diff --git a/embodichain/gen_sim/action_engine/runtime/state.py b/embodichain/gen_sim/action_engine/runtime/state.py new file mode 100644 index 000000000..dbafc4748 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/state.py @@ -0,0 +1,84 @@ +# ---------------------------------------------------------------------------- +# 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 Engine execution state at the atomic-planning boundary.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Mapping + +import torch + +from embodichain.lab.sim.atomic_actions import ( + HeldObjectState, + TaskState, +) + +__all__ = ["ExecutionState"] + + +@dataclass(slots=True, eq=False) +class ExecutionState: + """Projected task state paired with the next full-robot planning seed. + + The simulation atomic-action package deliberately no longer exposes the + legacy ``WorldState`` compatibility object. Action Engine keeps this narrow + orchestration state locally and converts it to immutable ``TaskState`` and + ``PlanningContext`` values immediately before invoking the shared planner. + """ + + last_qpos: torch.Tensor + held_objects: dict[str, HeldObjectState] = field(default_factory=dict) + + def get_held_object(self, control_part: str) -> HeldObjectState | None: + """Return the held-object relation for one control part.""" + return self.held_objects.get(control_part) + + def with_updates( + self, + *, + last_qpos: torch.Tensor | None = None, + held_objects: Mapping[str, HeldObjectState] | None = None, + ) -> ExecutionState: + """Return a detached successor state.""" + return ExecutionState( + last_qpos=self.last_qpos if last_qpos is None else last_qpos, + held_objects=dict( + self.held_objects if held_objects is None else held_objects + ), + ) + + def to_task_state(self) -> TaskState: + """Convert this state to the shared immutable symbolic task contract.""" + return TaskState( + batch_size=int(self.last_qpos.shape[0]), + device=self.last_qpos.device, + held_objects=self.held_objects, + ) + + @classmethod + def from_task_state( + cls, + task: TaskState, + *, + last_qpos: torch.Tensor, + ) -> ExecutionState: + """Build an orchestration state from a committed or projected task state.""" + return cls( + last_qpos=last_qpos, + held_objects=dict(task.held_objects), + ) diff --git a/embodichain/gen_sim/action_engine/tasks/__init__.py b/embodichain/gen_sim/action_engine/tasks/__init__.py new file mode 100644 index 000000000..843a55efa --- /dev/null +++ b/embodichain/gen_sim/action_engine/tasks/__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. +# ---------------------------------------------------------------------------- + +"""Task-first generation and scene hand-off for Action Engine v2.""" + +from __future__ import annotations + +from .assembly import GroundedTaskSpec +from .interpretation import ( + GroundingCaller, + INSTRUCTION_INTENT_SCHEMA, + InstructionDraftResult, + InstructionCaller, + InstructionIntent, + ground_instruction_draft, + interpret_instruction_draft, + interpret_and_ground_task_spec, + validate_instruction_intent, +) +from .recipes import instantiate_seed_graph +from .scene import SceneHandoff, validate_scene_handoff + +__all__ = [ + "GroundedTaskSpec", + "GroundingCaller", + "INSTRUCTION_INTENT_SCHEMA", + "InstructionDraftResult", + "InstructionCaller", + "InstructionIntent", + "SceneHandoff", + "ground_instruction_draft", + "instantiate_seed_graph", + "interpret_instruction_draft", + "interpret_and_ground_task_spec", + "validate_instruction_intent", + "validate_scene_handoff", +] diff --git a/embodichain/gen_sim/action_engine/tasks/assembly.py b/embodichain/gen_sim/action_engine/tasks/assembly.py new file mode 100644 index 000000000..7fdfe7246 --- /dev/null +++ b/embodichain/gen_sim/action_engine/tasks/assembly.py @@ -0,0 +1,425 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Language-neutral scene inventory and grounded TaskSpec assembly.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass, field +from typing import Any + +from embodichain.gen_sim.action_engine.domain import ( + TASK_CONTRACTS, + task_contract, + task_success_type, + validate_scene_requirements, + validate_task_spec, +) +from embodichain.gen_sim.action_engine.generation.config_builder import ( + canonical_robot_profile, +) +from embodichain.gen_sim.action_engine.protocol import ( + SCENE_REQUIREMENTS_SCHEMA, + TASK_SPEC_SCHEMA, +) + +__all__ = [ + "GroundedTaskBuilder", + "GroundedTaskSpec", + "SceneEntity", + "SceneInventory", + "validate_source_compatibility", + "validate_target_compatibility", +] + + +@dataclass(frozen=True) +class GroundedTaskSpec: + """One explicit TaskSpec plus verified scene role bindings.""" + + task_spec: dict[str, Any] + scene_requirements: dict[str, Any] + role_bindings: dict[str, str] + + +@dataclass(frozen=True) +class SceneEntity: + """One scene entity with source semantics preserved verbatim.""" + + uid: str + role: str + name: str + description: str + category: str + color: str | None + position: tuple[float, float, float] + affordances: frozenset[str] = frozenset() + initial_state: Mapping[str, Any] = field(default_factory=dict) + attributes: Mapping[str, Any] = field(default_factory=dict) + source_uid: str = "" + + +class SceneInventory: + """Structural scene index without natural-language matching rules.""" + + _PASSIVE_ROLES = frozenset( + { + "background", + "camera", + "light", + "robot", + "sensor", + "support_surface", + "table", + } + ) + + def __init__( + self, + scene_objects: Sequence[Mapping[str, Any]], + *, + robot_profile: str, + ) -> None: + self.profile = canonical_robot_profile(robot_profile) + self.entities = tuple(_scene_entity(item) for item in scene_objects) + self.by_uid = {entity.uid: entity for entity in self.entities} + if len(self.by_uid) != len(self.entities): + raise ValueError("Scene inventory contains duplicate runtime UIDs.") + self.support = tuple( + entity + for entity in self.entities + if entity.uid == "table" or entity.role in {"table", "support_surface"} + ) + self.passive = tuple( + entity + for entity in self.entities + if entity in self.support or entity.role in self._PASSIVE_ROLES + ) + self.interactive = tuple( + entity for entity in self.entities if entity not in self.passive + ) + if not self.interactive: + raise ValueError("Task planning requires at least one interaction object.") + + @property + def movable(self) -> tuple[SceneEntity, ...]: + """Compatibility alias for callers that mean source candidates.""" + return self.interactive + + def left_score(self, entity: SceneEntity) -> float: + """Return robot-relative lateral score; positive values are left. + + Generated dual-arm profiles share one final world layout: the semantic + left arm is on world ``-Y`` after all robot-level transforms. + """ + return -entity.position[1] + + +class GroundedTaskBuilder: + """Assemble grounded E1-E9 instances without parsing instruction text.""" + + def __init__( + self, + task_id: str, + instruction: str, + inventory: SceneInventory, + *, + planner: str = "structured_llm_v2", + ) -> None: + self.task_id = task_id + self.instruction = instruction + self.inventory = inventory + self.planner = planner + self.instances: list[dict[str, Any]] = [] + self.role_by_uid: dict[str, str] = {} + self.requirements: dict[str, dict[str, Any]] = {} + self.previous_object_uid: str | None = None + self.previous_arm: str | None = None + self.last_task_by_object_uid: dict[str, tuple[str, str]] = {} + + def add( + self, + task_type: str, + object_entity: SceneEntity, + *, + target: SceneEntity | None = None, + params: Mapping[str, Any] | None = None, + depends_on: Sequence[str] | None = None, + ) -> str: + values = deepcopy(dict(params or {})) + relation = str(values.get("relation", "none")) + validate_source_compatibility(task_type, (object_entity,)) + validate_target_compatibility(task_type, target, relation=relation) + + instance_id = f"task_{len(self.instances) + 1:02d}" + object_role = self._role( + object_entity, + required_affordances=task_contract(task_type).required_affordances, + initial_state={"orientation": "fallen"} if task_type == "E2" else {}, + ) + values = {"object_role": object_role, **values} + if task_type == "E3": + values["source_role"] = values.pop("object_role") + if target is not None: + values["target_role"] = self._role( + target, + required_affordances=_target_affordances(task_type, relation), + ) + if depends_on is None: + dependencies = [self.instances[-1]["id"]] if self.instances else [] + else: + dependencies = list(depends_on) + previous_for_object = self.last_task_by_object_uid.get(object_entity.uid) + if ( + task_type == "E4" + and previous_for_object is not None + and previous_for_object[1] == "E2" + and previous_for_object[0] not in dependencies + ): + dependencies.append(previous_for_object[0]) + self.instances.append( + { + "id": instance_id, + "task_type": task_type, + "params": values, + "depends_on": dependencies, + "role": "primary", + } + ) + self.last_task_by_object_uid[object_entity.uid] = (instance_id, task_type) + self.previous_object_uid = object_entity.uid + if task_type == "E4": + receive_arm = str(values.get("receive_arm", "")) + self.previous_arm = ( + receive_arm if receive_arm in {"left_arm", "right_arm"} else None + ) + elif str(values.get("required_arm", "")) in {"left_arm", "right_arm"}: + self.previous_arm = str(values["required_arm"]) + return instance_id + + def build(self) -> GroundedTaskSpec: + types = {item["task_type"] for item in self.instances} + if len(self.instances) == 1: + level = "L1" + elif len(types) == 1: + level = "L2" + else: + level = "L3" + success_terms = [ + { + "type": task_success_type(item["task_type"], item.get("params")), + "task_instance_id": item["id"], + } + for item in self.instances + ] + task = validate_task_spec( + { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": self.task_id, + "level": level, + "instruction": self.instruction, + "reasoning_type": "none", + "task_instances": self.instances, + "success": {"op": "all", "terms": success_terms}, + "oracle": { + "task_order": [item["id"] for item in self.instances], + "role_bindings": dict(sorted(self.role_bindings().items())), + }, + "metadata": {"planner": self.planner}, + } + ) + requirements = validate_scene_requirements( + { + "schema_version": SCENE_REQUIREMENTS_SCHEMA, + "task_id": self.task_id, + "objects": list(self.requirements.values()), + "cameras": [], + "spatial_constraints": [{"type": "preserve_source_scene"}], + "distractor_count": max( + 0, + len(self.inventory.interactive) - len(self.role_by_uid), + ), + "metadata": {"source": "existing_gym_project"}, + } + ) + return GroundedTaskSpec(task, requirements, self.role_bindings()) + + def role_bindings(self) -> dict[str, str]: + return {role: uid for uid, role in self.role_by_uid.items()} + + def _role( + self, + entity: SceneEntity, + task_type: str | None = None, + *, + required_affordances: Sequence[str] = (), + initial_state: Mapping[str, Any] | None = None, + ) -> str: + if task_type in TASK_CONTRACTS: + required_affordances = tuple( + set(required_affordances) + | set(task_contract(str(task_type)).required_affordances) + ) + if task_type == "E2": + initial_state = {"orientation": "fallen", **dict(initial_state or {})} + existing = self.role_by_uid.get(entity.uid) + if existing is not None: + requirement = self.requirements[existing] + requirement["affordances"] = sorted( + set(requirement["affordances"]) | set(required_affordances) + ) + requirement["initial_state"].update(dict(initial_state or {})) + return existing + role = f"object_{len(self.role_by_uid) + 1:02d}" + self.role_by_uid[entity.uid] = role + attributes = deepcopy(dict(entity.attributes)) + if entity.color is not None: + attributes.setdefault("color", entity.color) + self.requirements[role] = { + "role_id": role, + "category": entity.category or entity.role, + "count": 1, + "affordances": sorted(set(required_affordances)), + "initial_state": dict(initial_state or {}), + "attributes": attributes, + } + return role + + +def validate_source_compatibility( + task_type: str, + objects: Sequence[SceneEntity], +) -> None: + """Apply structural/explicit-affordance checks without a category taxonomy.""" + contract = task_contract(task_type) + if contract.source_structure == "articulation": + invalid = [entity.uid for entity in objects if entity.role != "articulation"] + else: + invalid = [ + entity.uid + for entity in objects + if entity.role not in {"object", "rigid_object"} + ] + if invalid: + structure_label = ( + "articulation" + if contract.source_structure == "articulation" + else "movable rigid-object" + ) + raise ValueError( + f"{task_type} requires {structure_label} structure; " + f"incompatible scene objects are {invalid}." + ) + required = set(contract.required_affordances) + for entity in objects: + if entity.affordances: + missing = required - set(entity.affordances) + if missing: + raise ValueError( + f"{task_type} is incompatible with scene object {entity.uid!r}; " + f"missing affordances {sorted(missing)}." + ) + + +def validate_target_compatibility( + task_type: str, + target: SceneEntity | None, + *, + relation: str, +) -> None: + """Reject only structural or explicitly declared target contradictions.""" + if task_type == "E1" and relation == "on" and target is not None: + # Support is a relation between two concrete bodies at a candidate + # pose. A positive affordance list is not a closed-world inventory, so + # omission of ``support_surface`` cannot prove incompatibility here. + return + requires_container = task_type == "E3" or ( + task_type == "E1" and relation == "inside" + ) + if requires_container and target is None: + raise ValueError( + f"{task_type} {relation} relation requires a target container." + ) + if not requires_container or target is None: + return + if target.role in SceneInventory._PASSIVE_ROLES: + raise ValueError( + f"{task_type} target {target.uid!r} is structurally incompatible " + "with containment." + ) + if target.affordances: + compatible = {"container", "fillable", "liquid_container", "receptacle"} + if set(target.affordances).isdisjoint(compatible): + raise ValueError( + f"{task_type} target {target.uid!r} has explicit affordances but " + f"none support containment; expected one of {sorted(compatible)}." + ) + + +def _target_affordances(task_type: str, relation: str) -> tuple[str, ...]: + if task_type == "E3" or (task_type == "E1" and relation == "inside"): + return ("container",) + return () + + +def _scene_entity(raw: Mapping[str, Any]) -> SceneEntity: + uid = str(raw.get("runtime_uid", raw.get("uid", ""))).strip() + if not uid: + raise ValueError("Every scene object requires a runtime UID.") + role = str(raw.get("role", raw.get("source_role", "object"))).strip().lower() + raw_category = raw.get("category", raw.get("object_category", "")) + category = "" if raw_category is None else str(raw_category).strip() + raw_color = raw.get("color") + attributes = raw.get("attributes", {}) + if not isinstance(attributes, Mapping): + raise ValueError(f"Scene object {uid!r} attributes must be a mapping.") + if raw_color is None: + raw_color = attributes.get("color") + color = str(raw_color).strip() if raw_color not in (None, "") else None + position = raw.get("init_pos", raw.get("position", (0.0, 0.0, 0.0))) + if ( + not isinstance(position, Sequence) + or isinstance(position, (str, bytes)) + or len(position) != 3 + ): + raise ValueError(f"Scene object {uid!r} requires a three-value init_pos.") + raw_affordances = raw.get("affordances", raw.get("capabilities", ())) + affordances = ( + frozenset( + str(item).strip().lower() for item in raw_affordances if str(item).strip() + ) + if isinstance(raw_affordances, Sequence) + and not isinstance(raw_affordances, (str, bytes)) + else frozenset() + ) + initial_state = raw.get("initial_state", raw.get("state", {})) + if not isinstance(initial_state, Mapping): + raise ValueError(f"Scene object {uid!r} initial_state must be a mapping.") + return SceneEntity( + uid=uid, + role=role, + name=str(raw.get("name", "")).strip(), + description=str(raw.get("description", "")).strip(), + category=category, + color=color, + position=tuple(float(value) for value in position), + affordances=affordances, + initial_state=dict(initial_state), + attributes=dict(attributes), + source_uid=str(raw.get("source_uid", "")).strip(), + ) diff --git a/embodichain/gen_sim/action_engine/tasks/grounding.py b/embodichain/gen_sim/action_engine/tasks/grounding.py new file mode 100644 index 000000000..de412e987 --- /dev/null +++ b/embodichain/gen_sim/action_engine/tasks/grounding.py @@ -0,0 +1,513 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Task-conditioned scene-UID grounding for structured instruction intents.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass +import json +import math +from time import perf_counter +from typing import Any + +from .assembly import SceneInventory + +__all__ = ["GroundingCaller", "GroundingResult", "ground_scene_references"] + +GroundingCaller = Callable[..., Mapping[str, Any]] + +_BINDING_KEYS = frozenset({"reference_id", "status", "uids", "confidence"}) +_QUANTIFIERS = frozenset({"one", "all", "count"}) +_REDACTED_KEYS = frozenset( + { + "absolute_position", + "bbox", + "bboxes", + "bounding_box", + "camera_matrix", + "center", + "centroid", + "coordinates", + "depth", + "dimensions", + "extrinsics", + "grasp_pose", + "init_local_pose", + "init_pos", + "init_rot", + "intrinsics", + "joint_positions", + "joints", + "keypoint", + "keypoints", + "location", + "matrix", + "pose", + "position", + "position_xyz", + "qpos", + "quaternion", + "rotation", + "scale", + "target_pose", + "trajectory", + "transform", + "translation", + "waypoints", + "world_x", + "world_y", + "world_z", + "x", + "y", + "z", + } +) + +_GROUNDING_SCHEMA: dict[str, Any] = { + "title": "ActionEngineSceneGrounding", + "type": "object", + "additionalProperties": False, + "required": ["bindings"], + "properties": { + "bindings": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": False, + "required": sorted(_BINDING_KEYS), + "properties": { + "reference_id": {"type": "string"}, + "status": { + "type": "string", + "enum": ["resolved", "ambiguous", "not_found"], + }, + "uids": { + "type": "array", + "items": {"type": "string"}, + "uniqueItems": True, + }, + "confidence": { + "type": "number", + "minimum": 0.0, + "maximum": 1.0, + }, + }, + }, + } + }, +} + + +@dataclass(frozen=True) +class GroundingResult: + """Validated scene bindings and aggregate call statistics. + + Attributes: + bindings: Mapping from ``.`` to scene UIDs. + attempts: Number of grounding-model calls, including one repair call. + latency_seconds: Total elapsed wall-clock time across the grounding stage. + """ + + bindings: dict[str, tuple[str, ...]] + attempts: int + latency_seconds: float + + +def ground_scene_references( + instruction: str, + intent: Mapping[str, Any], + inventory: SceneInventory, + scene_objects: Sequence[Mapping[str, Any]], + model: str | None, + caller: GroundingCaller, +) -> GroundingResult: + """Resolve every ``scene_ref`` selector in one task-conditioned batch. + + The grounding model can only select stable UIDs from a redacted inventory. + Its output does not add affordances, physical state, coordinates, or poses. + One failed local validation is repaired with one additional model call. + + Args: + instruction: Original user instruction for task-level context. + intent: Validated structured instruction intent. + inventory: Structural scene inventory defining authoritative candidates. + scene_objects: Original semantic inventory used to retain open labels. + model: Model name forwarded unchanged to the injected caller. + caller: Structured model transport accepting ``prompt``, ``schema``, and + ``model`` keyword arguments. + + Returns: + Validated UID bindings together with call-count and latency statistics. + + Raises: + TypeError: If the intent or response has an invalid container type. + ValueError: If requests are malformed or grounding remains invalid after + one repair attempt. + """ + if not isinstance(instruction, str) or not instruction.strip(): + raise ValueError("Grounding instruction must be a non-empty string.") + if not callable(caller): + raise TypeError("Grounding caller must be callable.") + + requests = _collect_requests(intent) + prompt_inventory = _grounding_inventory(inventory, scene_objects) + prompt = _grounding_prompt(instruction.strip(), requests, prompt_inventory) + started = perf_counter() + first_error: Exception | None = None + + for attempt in range(2): + current_prompt = prompt + if first_error is not None: + current_prompt += ( + "\n\nREPAIR OVERRIDE: the previous grounding JSON failed local " + "validation. Return one corrected JSON object only. Preserve the " + "exact output fields bindings/reference_id/status/uids/confidence, " + "cover every requested reference exactly once, and select only " + "UIDs from the supplied candidate inventory. Validation error: " + f"{first_error}" + ) + try: + response = caller( + prompt=current_prompt, + schema=deepcopy(_GROUNDING_SCHEMA), + model=model, + ) + bindings = _validate_response( + response, + requests=requests, + inventory=inventory, + ) + return GroundingResult( + bindings=bindings, + attempts=attempt + 1, + latency_seconds=perf_counter() - started, + ) + except (TypeError, ValueError) as error: + if attempt: + raise ValueError( + "Scene grounding failed validation after one repair: " f"{error}" + ) from error + first_error = error + raise AssertionError("unreachable") + + +def _collect_requests(intent: Mapping[str, Any]) -> list[dict[str, Any]]: + if not isinstance(intent, Mapping): + raise TypeError("Instruction intent must be a mapping.") + steps = intent.get("steps") + if not isinstance(steps, Sequence) or isinstance(steps, (str, bytes)): + raise ValueError("Instruction intent steps must be a list.") + + requests: list[dict[str, Any]] = [] + request_ids: set[str] = set() + for step_index, step in enumerate(steps): + context = f"InstructionIntent.steps[{step_index}]" + if not isinstance(step, Mapping): + raise ValueError(f"{context} must be a mapping.") + step_id = step.get("id") + if not isinstance(step_id, str) or not step_id.strip(): + raise ValueError(f"{context}.id must be a non-empty string.") + task_type = step.get("task_type") + if not isinstance(task_type, str) or not task_type.strip(): + raise ValueError(f"{context}.task_type must be a non-empty string.") + relation = step.get("relation", "none") + if not isinstance(relation, str): + raise ValueError(f"{context}.relation must be a string.") + + for slot in ("object", "target"): + selector = step.get(slot) + if not isinstance(selector, Mapping): + raise ValueError(f"{context}.{slot} must be a mapping.") + if selector.get("kind") != "scene_ref": + continue + reference = selector.get("reference") + if not isinstance(reference, str) or not reference.strip(): + raise ValueError( + f"{context}.{slot}.reference must be a non-empty string." + ) + quantifier = selector.get("quantifier") + if quantifier not in _QUANTIFIERS: + raise ValueError( + f"{context}.{slot}.quantifier must be one of " + f"{sorted(_QUANTIFIERS)}." + ) + count = selector.get("count") + if not isinstance(count, int) or isinstance(count, bool) or count < 0: + raise ValueError(f"{context}.{slot}.count must be an integer >= 0.") + if quantifier == "count" and count < 1: + raise ValueError( + f"{context}.{slot} quantifier=count requires count>=1." + ) + if quantifier != "count" and count != 0: + raise ValueError( + f"{context}.{slot} quantifier={quantifier} requires count=0." + ) + + request_id = f"{step_id}.{slot}" + if request_id in request_ids: + raise ValueError(f"Duplicate grounding request ID {request_id!r}.") + request_ids.add(request_id) + requests.append( + { + "reference_id": request_id, + "step_id": step_id, + "slot": slot, + "task_type": task_type, + "relation": relation, + "reference": reference.strip(), + "quantifier": quantifier, + "count": count, + } + ) + return requests + + +def _grounding_inventory( + inventory: SceneInventory, + scene_objects: Sequence[Mapping[str, Any]], +) -> list[dict[str, Any]]: + raw_by_uid: dict[str, Mapping[str, Any]] = {} + for item_index, raw in enumerate(scene_objects): + if not isinstance(raw, Mapping): + raise ValueError(f"Scene inventory item {item_index} must be a mapping.") + uid = str(raw.get("runtime_uid", raw.get("uid", ""))).strip() + if uid: + raw_by_uid[uid] = raw + + ranked = sorted( + inventory.entities, + key=lambda entity: (-inventory.left_score(entity), entity.uid), + ) + rank_by_uid = {entity.uid: rank for rank, entity in enumerate(ranked, start=1)} + payload = [] + for entity in sorted(inventory.entities, key=lambda item: item.uid): + raw = raw_by_uid.get(entity.uid, {}) + score = inventory.left_score(entity) + side = "left" if score > 0.0 else "right" if score < 0.0 else "center" + raw_category = raw.get( + "category", + raw.get("object_category", entity.category), + ) + attributes = _redact_semantic_mapping(entity.attributes) + if entity.color is not None: + attributes.setdefault("color", entity.color) + payload.append( + { + "uid": entity.uid, + "role": entity.role, + "name": str(raw.get("name", entity.name)).strip(), + "category": str(raw_category).strip() or entity.category, + "description": entity.description, + "affordances": sorted(entity.affordances), + "attributes": attributes, + "initial_state": _redact_semantic_mapping(entity.initial_state), + "side": side, + "rank": rank_by_uid[entity.uid], + } + ) + return payload + + +def _grounding_prompt( + instruction: str, + requests: Sequence[Mapping[str, Any]], + inventory: Sequence[Mapping[str, Any]], +) -> str: + return ( + "Ground the requested natural-language scene references to the supplied " + "scene inventory. Resolve all requests together using the original task, " + "step type, relation, quantifier, and reference text as context. Select " + "only exact inventory UIDs. The inventory's affordances and states are " + "source evidence only: never infer, add, authorize, or return an " + "affordance, capability, physical state, coordinate, pose, orientation, " + "path, or action. The side and rank fields are discrete robot-relative " + "labels; rank 1 is leftmost. Object requests may select only movable " + "inventory entities. Target requests may also select support surfaces. " + "Use status=ambiguous or status=not_found instead of guessing when the " + "evidence is insufficient. Return exactly one binding per reference_id " + "with only reference_id, status, uids, and confidence.\n\n" + f"Instruction:\n{instruction}\n\n" + "Grounding requests:\n" + f"{json.dumps(list(requests), ensure_ascii=False, sort_keys=True)}\n\n" + "Redacted scene inventory:\n" + f"{json.dumps(list(inventory), ensure_ascii=False, sort_keys=True)}" + ) + + +def _validate_response( + value: Mapping[str, Any], + *, + requests: Sequence[Mapping[str, Any]], + inventory: SceneInventory, +) -> dict[str, tuple[str, ...]]: + if not isinstance(value, Mapping): + raise TypeError("Scene grounding output must be a mapping.") + if set(value) != {"bindings"}: + raise ValueError( + "Scene grounding output must contain exactly the 'bindings' field." + ) + raw_bindings = value["bindings"] + if not isinstance(raw_bindings, Sequence) or isinstance(raw_bindings, (str, bytes)): + raise ValueError("Scene grounding bindings must be a list.") + + request_by_id = {str(request["reference_id"]): request for request in requests} + bindings: dict[str, tuple[str, ...]] = {} + for binding_index, raw in enumerate(raw_bindings): + context = f"SceneGrounding.bindings[{binding_index}]" + if not isinstance(raw, Mapping): + raise ValueError(f"{context} must be a mapping.") + if set(raw) != _BINDING_KEYS: + missing = sorted(_BINDING_KEYS - set(raw)) + extra = sorted(set(raw) - _BINDING_KEYS) + raise ValueError( + f"{context} fields must be exactly {sorted(_BINDING_KEYS)}; " + f"missing={missing}, unsupported={extra}." + ) + reference_id = raw["reference_id"] + if not isinstance(reference_id, str) or not reference_id: + raise ValueError(f"{context}.reference_id must be a non-empty string.") + if reference_id not in request_by_id: + raise ValueError(f"{context} references unknown request {reference_id!r}.") + if reference_id in bindings: + raise ValueError(f"Duplicate grounding binding for {reference_id!r}.") + + status = raw["status"] + if status not in {"resolved", "ambiguous", "not_found"}: + raise ValueError( + f"{context}.status must be resolved, ambiguous, or not_found." + ) + if status != "resolved": + raise ValueError( + f"Grounding request {reference_id!r} was not resolved: {status}." + ) + confidence = raw["confidence"] + if ( + not isinstance(confidence, (int, float)) + or isinstance(confidence, bool) + or not math.isfinite(float(confidence)) + or not 0.0 <= float(confidence) <= 1.0 + ): + raise ValueError(f"{context}.confidence must be a number in [0, 1].") + if float(confidence) < 0.5: + raise ValueError( + f"Grounding request {reference_id!r} confidence is below 0.5." + ) + + raw_uids = raw["uids"] + if not isinstance(raw_uids, Sequence) or isinstance(raw_uids, (str, bytes)): + raise ValueError(f"{context}.uids must be a list.") + uids = tuple(raw_uids) + if any(not isinstance(uid, str) or not uid for uid in uids): + raise ValueError(f"{context}.uids must contain non-empty strings.") + if len(set(uids)) != len(uids): + raise ValueError( + f"Grounding request {reference_id!r} contains duplicate UIDs." + ) + unknown = sorted(set(uids) - set(inventory.by_uid)) + if unknown: + raise ValueError( + f"Grounding request {reference_id!r} selected unknown UIDs {unknown}." + ) + + request = request_by_id[reference_id] + allowed = ( + {entity.uid for entity in inventory.interactive} + if request["slot"] == "object" + else {entity.uid for entity in (*inventory.interactive, *inventory.support)} + ) + disallowed = sorted(set(uids) - allowed) + if disallowed: + raise ValueError( + f"Grounding request {reference_id!r} selected UIDs outside its " + f"{request['slot']} candidate range: {disallowed}." + ) + _validate_cardinality(request, uids) + bindings[reference_id] = uids + + missing = sorted(set(request_by_id) - set(bindings)) + if missing: + raise ValueError(f"Scene grounding omitted requests {missing}.") + _reject_self_references(requests, bindings) + return bindings + + +def _validate_cardinality( + request: Mapping[str, Any], + uids: Sequence[str], +) -> None: + request_id = str(request["reference_id"]) + quantifier = str(request["quantifier"]) + if quantifier == "one" and len(uids) != 1: + raise ValueError( + f"Grounding request {request_id!r} quantifier=one requires exactly one UID." + ) + if quantifier == "count" and len(uids) != int(request["count"]): + raise ValueError( + f"Grounding request {request_id!r} requires exactly " + f"{request['count']} UIDs." + ) + if quantifier == "all" and not uids: + raise ValueError( + f"Grounding request {request_id!r} quantifier=all requires at " + "least one UID." + ) + + +def _reject_self_references( + requests: Sequence[Mapping[str, Any]], + bindings: Mapping[str, tuple[str, ...]], +) -> None: + slots_by_step: dict[str, dict[str, str]] = {} + for request in requests: + slots_by_step.setdefault(str(request["step_id"]), {})[str(request["slot"])] = ( + str(request["reference_id"]) + ) + for step_id, slots in slots_by_step.items(): + object_id = slots.get("object") + target_id = slots.get("target") + if object_id is None or target_id is None: + continue + overlap = sorted(set(bindings[object_id]) & set(bindings[target_id])) + if overlap: + raise ValueError( + f"Grounding step {step_id!r} uses the same UID as object and " + f"target: {overlap}." + ) + + +def _redact_semantic_mapping(value: Mapping[str, Any]) -> dict[str, Any]: + result: dict[str, Any] = {} + for key, child in value.items(): + name = str(key) + normalized = name.strip().lower().replace("-", "_") + if normalized in _REDACTED_KEYS: + continue + if isinstance(child, Mapping): + nested = _redact_semantic_mapping(child) + if nested: + result[name] = nested + elif isinstance(child, (str, int, float, bool)) and not isinstance( + child, complex + ): + result[name] = child + elif isinstance(child, Sequence) and not isinstance(child, (str, bytes)): + semantic_values = [item for item in child if isinstance(item, (str, bool))] + if semantic_values and len(semantic_values) == len(child): + result[name] = semantic_values + return result diff --git a/embodichain/gen_sim/action_engine/tasks/interpretation.py b/embodichain/gen_sim/action_engine/tasks/interpretation.py new file mode 100644 index 000000000..f84f6d2bd --- /dev/null +++ b/embodichain/gen_sim/action_engine/tasks/interpretation.py @@ -0,0 +1,404 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Compatibility bridge from Task Engine drafts to Action Engine TaskSpec v2.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any + +from embodichain.gen_sim.task_engine.interpretation import ( + INSTRUCTION_INTENT_SCHEMA, + InstructionCaller, + InstructionDraftResult, + InstructionIntent, + _default_instruction_caller, + _instruction_prompt, + _instruction_selector_rules, + interpret_instruction_draft, + validate_instruction_intent, +) + +from .assembly import ( + GroundedTaskBuilder, + GroundedTaskSpec, + SceneEntity, + SceneInventory, + validate_source_compatibility, + validate_target_compatibility, +) +from .grounding import GroundingCaller, ground_scene_references + +__all__ = [ + "GroundingCaller", + "INSTRUCTION_INTENT_SCHEMA", + "InstructionCaller", + "InstructionDraftResult", + "InstructionIntent", + "ground_instruction_draft", + "interpret_and_ground_task_spec", + "interpret_instruction_draft", + "validate_instruction_intent", +] + + +def interpret_and_ground_task_spec( + task_name: str, + task_description: str, + scene_objects: Sequence[Mapping[str, Any]], + *, + robot_profile: str, + model: str | None = None, + caller: InstructionCaller | None = None, + grounding_caller: GroundingCaller | None = None, +) -> GroundedTaskSpec: + """Interpret through Task Engine, then ground through Action Engine.""" + task_id = str(task_name).strip() + instruction = str(task_description).strip() + if not task_id or not instruction: + raise ValueError("task_name and task_description must be non-empty.") + inventory = SceneInventory(scene_objects, robot_profile=robot_profile) + draft = interpret_instruction_draft(instruction, model=model, caller=caller) + invoke = caller or _default_instruction_caller + selected_model = None if draft.model == "injected_caller" else draft.model + grounding = ground_scene_references( + instruction=instruction, + intent=draft.intent, + inventory=inventory, + scene_objects=scene_objects, + model=selected_model, + caller=grounding_caller or invoke, + ) + grounded = _ground_intent( + task_id, + instruction, + draft.intent, + inventory, + grounding.bindings, + ) + grounded.task_spec["metadata"].update( + { + "instruction_interpreter": "structured_llm_v2", + "instruction_model": draft.model, + "instruction_call_count": draft.attempts, + "instruction_latency_seconds": draft.latency_seconds, + "scene_grounding_model": selected_model or "injected_caller", + "scene_grounding_call_count": grounding.attempts, + "scene_grounding_latency_seconds": grounding.latency_seconds, + } + ) + if draft.normalizations: + grounded.task_spec["metadata"]["instruction_intent_normalizations"] = list( + draft.normalizations + ) + return grounded + + +def ground_instruction_draft( + task_id: str, + instruction: str, + intent: Mapping[str, Any], + scene_objects: Sequence[Mapping[str, Any]], + *, + robot_profile: str, + reference_bindings: Mapping[str, Sequence[str]], +) -> GroundedTaskSpec: + """Lower a Task Engine draft using verified scene bindings.""" + normalized_task_id = str(task_id).strip() + normalized_instruction = str(instruction).strip() + if not normalized_task_id or not normalized_instruction: + raise ValueError("task_id and instruction must be non-empty.") + inventory = SceneInventory(scene_objects, robot_profile=robot_profile) + return _ground_intent( + normalized_task_id, + normalized_instruction, + validate_instruction_intent(intent), + inventory, + reference_bindings, + ) + + +def _ground_intent( + task_id: str, + instruction: str, + intent: Mapping[str, Any], + inventory: SceneInventory, + scene_bindings: Mapping[str, Sequence[str]], +) -> GroundedTaskSpec: + builder = GroundedTaskBuilder( + task_id, + instruction, + inventory, + planner="structured_llm_v2", + ) + objects_by_step: dict[str, list[SceneEntity]] = {} + task_ids_by_step: dict[str, list[str]] = {} + for step in _topological_steps(intent["steps"]): + step_id = str(step["id"]) + objects = _resolve_reference( + step["object"], + inventory, + objects_by_step, + context=f"instruction step {step_id!r} object", + reference_id=f"{step_id}.object", + scene_bindings=scene_bindings, + ) + validate_source_compatibility(str(step["task_type"]), objects) + target_objects = _resolve_reference( + step["target"], + inventory, + objects_by_step, + context=f"instruction step {step_id!r} target", + reference_id=f"{step_id}.target", + scene_bindings=scene_bindings, + allow_none=True, + exclude={item.uid for item in objects}, + allow_support=True, + ) + if len(target_objects) > 1: + raise ValueError(f"Instruction step {step_id!r} target is ambiguous.") + validate_target_compatibility( + str(step["task_type"]), + target_objects[0] if target_objects else None, + relation=str(step["relation"]), + ) + dependencies_by_step = list(step["depends_on"]) + for selector in (step["object"], step["target"]): + if selector["kind"] == "step_result": + reference = str(selector["step_id"]) + if reference not in dependencies_by_step: + dependencies_by_step.append(reference) + dependencies = [ + emitted_id + for dependency in dependencies_by_step + for emitted_id in task_ids_by_step[str(dependency)] + ] + emitted = _emit_step( + builder, + step, + objects, + target_objects[0] if target_objects else None, + dependencies, + ) + objects_by_step[step_id] = objects + task_ids_by_step[step_id] = emitted + return builder.build() + + +def _emit_step( + builder: GroundedTaskBuilder, + step: Mapping[str, Any], + objects: Sequence[SceneEntity], + target: SceneEntity | None, + dependencies: Sequence[str], +) -> list[str]: + task_type = str(step["task_type"]) + if step["layout"] == "line": + roles = [builder._role(entity, "E1") for entity in objects] + parent = str(step["id"]) + return [ + builder.add( + "E1", + entity, + params={ + "target_role": "table", + "relation": "on", + "layout": "line", + "objects_roles": roles, + "axis": "world_y" if step["axis"] == "none" else step["axis"], + "order_by": "explicit", + "order_direction": "given", + "order_constraint": "free", + "orientation_goal": step["orientation_goal"], + "orientation_axis": "none", + "nominal_slot_index": slot, + "slot_constraint": "free_reassignable", + "parent_task_instance_id": parent, + }, + depends_on=dependencies, + ) + for slot, entity in enumerate(objects) + ] + + emitted = [] + for entity in objects: + params: dict[str, Any] = {} + required_arm = str(step["required_arm"]) + if required_arm in {"left_arm", "right_arm"}: + params["required_arm"] = required_arm + if task_type == "E1": + relation = str(step["relation"]) + if relation == "none": + if target is None or target not in builder.inventory.support: + raise ValueError( + "E1 omitted relation is only valid for a unique table " + "support target." + ) + relation = "on" + params.update( + { + "relation": relation, + "relation_frame": "robot", + "orientation_goal": step["orientation_goal"], + "orientation_axis": "none", + } + ) + elif task_type == "E2": + params.update( + { + "orientation_goal": "upright", + "support_role": "table", + "upright_local_axis": "long_axis", + } + ) + elif task_type == "E3": + params.update({"relation": "above", "relation_frame": "robot"}) + elif task_type == "E4": + params.update( + { + "transfer_arm": step["transfer_arm"], + "receive_arm": step["receive_arm"], + "orientation_goal": step["orientation_goal"], + } + ) + elif task_type == "E5": + params.update( + { + "direction": step["direction"], + "terminal_behavior": step["terminal_behavior"], + "relation": step["relation"], + "relation_frame": "robot", + } + ) + elif task_type in {"E6", "E7"}: + params["target_state"] = step["target_state"] + elif task_type == "E8": + params["target_setting"] = int(step["target_setting"]) + elif task_type == "E9": + params["terminal_state"] = step["target_state"] + emitted.append( + builder.add( + task_type, + entity, + target=target, + params=params, + depends_on=dependencies, + ) + ) + return emitted + + +def _resolve_reference( + selector: Mapping[str, Any], + inventory: SceneInventory, + objects_by_step: Mapping[str, Sequence[SceneEntity]], + *, + context: str, + reference_id: str, + scene_bindings: Mapping[str, Sequence[str]], + allow_none: bool = False, + exclude: set[str] | None = None, + allow_support: bool = False, +) -> list[SceneEntity]: + kind = str(selector["kind"]) + if kind == "none": + if allow_none: + return [] + raise ValueError(f"{context} is required.") + if kind == "step_result": + step_id = str(selector["step_id"]) + if step_id not in objects_by_step: + raise ValueError(f"{context} references unavailable step {step_id!r}.") + objects = list(objects_by_step[step_id]) + if len(objects) != 1: + raise ValueError( + f"{context} references step {step_id!r}, which has {len(objects)} objects." + ) + if exclude and objects[0].uid in exclude: + raise ValueError( + f"{context} references the same object as its source; " + "self-referential placement is not allowed." + ) + return objects + + if reference_id not in scene_bindings: + raise ValueError(f"{context} has no verified scene-grounding binding.") + excluded = exclude or set() + source_uids = ( + {entity.uid for entity in inventory.entities} + if allow_support + else {entity.uid for entity in inventory.interactive} + ) + resolved_uids = tuple(str(uid) for uid in scene_bindings[reference_id]) + pool = [ + inventory.by_uid[uid] + for uid in resolved_uids + if uid in source_uids and uid not in excluded + ] + pool = sorted(pool, key=lambda item: item.uid) + if not pool: + raise ValueError(f"{context} did not bind an eligible scene object.") + quantifier = str(selector["quantifier"]) + count = int(selector["count"]) + if quantifier == "one" and len(pool) != 1: + raise ValueError( + f"{context} is ambiguous; matched scene UIDs {[item.uid for item in pool]}." + ) + if quantifier == "count" and (count < 1 or len(pool) != count): + raise ValueError( + f"{context} requested exactly {count} objects but matched {len(pool)}." + ) + if quantifier == "all" and count not in {0, len(pool)}: + raise ValueError( + f"{context} quantifier=all cannot carry count={count}; use count for an exact quantity." + ) + return pool + + +def _topological_steps( + steps: Sequence[Mapping[str, Any]], +) -> list[dict[str, Any]]: + by_id = {str(step["id"]): dict(step) for step in steps} + effective_dependencies: dict[str, tuple[str, ...]] = {} + for step_id, step in by_id.items(): + deps = [str(dep) for dep in step["depends_on"]] + for selector in (step["object"], step["target"]): + if selector["kind"] == "step_result": + reference = str(selector["step_id"]) + if reference not in deps: + deps.append(reference) + effective_dependencies[step_id] = tuple(deps) + pending = set(by_id) + ordered: list[dict[str, Any]] = [] + original = [str(step["id"]) for step in steps] + while pending: + ready = [ + step_id + for step_id in original + if step_id in pending + and all(str(dep) not in pending for dep in effective_dependencies[step_id]) + ] + if not ready: + raise ValueError("Instruction intent dependencies contain a cycle.") + # Select one earliest-ready step at a time. Emitting the whole ready + # frontier lets a later independent step leapfrog an earlier step that + # becomes ready after its predecessor, changing the instruction's + # resource-order tie break without any causal reason. + step_id = ready[0] + ordered.append(by_id[step_id]) + pending.remove(step_id) + return ordered diff --git a/embodichain/gen_sim/action_engine/tasks/recipes.py b/embodichain/gen_sim/action_engine/tasks/recipes.py new file mode 100644 index 000000000..73cabc44a --- /dev/null +++ b/embodichain/gen_sim/action_engine/tasks/recipes.py @@ -0,0 +1,997 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Direct AtomicAction recipes for E1-E9 TaskSpec instances.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +from typing import Any + +from embodichain.gen_sim.action_engine.capabilities import ( + AtomicCapabilityRegistry, + build_atomic_capability_registry, + capability_precondition, +) +from embodichain.gen_sim.action_engine.domain import ( + TERMINAL_BEHAVIORS, + TRANSPORT_DIRECTIONS, + motion_policy, + task_success_type, + validate_task_spec, +) +from embodichain.gen_sim.action_engine.planning.linker import ( + link_seed_graph, + link_task_dependencies, +) +from embodichain.gen_sim.action_engine.protocol import SEED_GRAPH_SCHEMA + +__all__ = ["instantiate_seed_graph"] + + +def instantiate_seed_graph( + task_spec: Mapping[str, Any], + role_bindings: Mapping[str, str], + *, + planner_route: str = "offline", + registry: AtomicCapabilityRegistry | None = None, +) -> dict[str, Any]: + """Instantiate a coordinate-free SeedGraph after Scene Engine hand-off.""" + task = validate_task_spec(task_spec) + bindings = _validate_bindings(task, role_bindings) + capabilities = registry or build_atomic_capability_registry() + task, payload_links = _propagate_direct_payloads(task, bindings) + task = link_task_dependencies(task, bindings, registry=capabilities) + instances = _topological_instances(task["task_instances"]) + nodes: list[dict[str, Any]] = [] + groups = [] + terminal_by_group: dict[str, list[str]] = {} + held_after_group: dict[str, tuple[str, str] | None] = {} + for instance in instances: + group_id = str(instance["id"]) + task_type = str(instance["task_type"]) + params = _resolve_params(instance["params"], bindings) + object_uid = _primary_object(task_type, params) + incoming_held_arm = _incoming_held_arm( + task_type, + object_uid, + instance["depends_on"], + held_after_group, + ) + actor = _actor(task_type, params, incoming_held_arm=incoming_held_arm) + dependency_nodes = [ + node_id + for dependency in instance["depends_on"] + for node_id in terminal_by_group[str(dependency)] + ] + recipe_nodes, operator, goal, success = _recipe( + group_id, + task_type, + object_uid, + actor, + params, + dependency_nodes, + role=str(instance["role"]), + incoming_held_arm=incoming_held_arm, + ) + for node in recipe_nodes: + node["precondition"] = capability_precondition( + capabilities.get(str(node["atomic_action"])), + object_uid=str(node["object_uid"]), + actor=node["actor"], + target_binding=node["target_binding"], + ) + nodes.extend(recipe_nodes) + terminal_by_group[group_id] = _terminal_nodes(recipe_nodes) + held_after_group[group_id] = _terminal_hold( + task_type, + object_uid, + params, + ) + groups.append( + { + "id": group_id, + "task_type": task_type, + "role": str(instance["role"]), + "operator": operator, + "object_uid": object_uid, + "actor": actor, + "goal": goal, + "depends_on": list(instance["depends_on"]), + "parent_task_instance_id": str( + params.get("parent_task_instance_id", group_id) + ), + "node_ids": [node["id"] for node in recipe_nodes], + "success": success, + } + ) + + graph_metadata = { + "task_spec_id": task["task_id"], + "role_bindings": dict(sorted(bindings.items())), + "allocation_groups": deepcopy( + task.get("metadata", {}).get("allocation_groups", []) + ), + "direct_payload_links": payload_links, + "oracle_exposed": False, + "planning_latency_seconds": 0.0, + "vlm_call_count": 0, + } + task_linker = task.get("metadata", {}).get("action_contract_task_linker") + if isinstance(task_linker, Mapping): + graph_metadata["action_contract_task_linker"] = deepcopy(dict(task_linker)) + + graph = { + "schema_version": SEED_GRAPH_SCHEMA, + "task_id": task["task_id"], + "instruction": task["instruction"], + "level": task["level"], + "reasoning_type": task["reasoning_type"], + "planner_route": planner_route, + "nodes": nodes, + "task_groups": groups, + "success": { + "op": "all", + "terms": [deepcopy(group["success"]) for group in groups], + }, + "capability_catalog_hash": capabilities.catalog_hash(), + "metadata": graph_metadata, + } + known_objects = set(bindings.values()) | {"table"} + graph = link_seed_graph( + graph, + registry=capabilities, + task_order=[str(instance["id"]) for instance in instances], + known_objects=known_objects, + ) + for node in graph["nodes"]: + capabilities.validate_binding(node) + return graph + + +def _topological_instances( + instances: Sequence[Mapping[str, Any]], +) -> list[dict[str, Any]]: + """Emit task groups in dependency order even for externally authored specs.""" + by_id = {str(instance["id"]): dict(instance) for instance in instances} + original = [str(instance["id"]) for instance in instances] + pending = set(by_id) + ordered: list[dict[str, Any]] = [] + while pending: + ready = [ + instance_id + for instance_id in original + if instance_id in pending + and all( + str(dependency) not in pending + for dependency in by_id[instance_id]["depends_on"] + ) + ] + if not ready: + raise ValueError("TaskSpec task instances contain a dependency cycle.") + for instance_id in ready: + ordered.append(by_id[instance_id]) + pending.remove(instance_id) + return ordered + + +def _propagate_direct_payloads( + task: Mapping[str, Any], + bindings: Mapping[str, str], +) -> tuple[dict[str, Any], list[dict[str, str]]]: + """Carry direct E1 support relations into a later single-arm E1 move. + + This is intentionally a one-hop physical relation rather than a general + scene-state planner: an object placed on or inside a carrier becomes that + carrier's direct payload until the object itself is manipulated again. + """ + result = deepcopy(dict(task)) + role_by_uid = {uid: role for role, uid in bindings.items()} + direct_by_carrier: dict[str, list[tuple[str, str, str]]] = {} + carrier_by_payload: dict[str, str] = {} + links: list[dict[str, str]] = [] + changed = False + + for instance in _topological_instances(result["task_instances"]): + task_type = str(instance["task_type"]) + params = instance["params"] + primary_key = "source_role" if task_type == "E3" else "object_role" + primary_role = params.get(primary_key) + if not isinstance(primary_role, str) or not primary_role: + continue + primary_uid = bindings.get(primary_role, primary_role) + direct_payloads = list(direct_by_carrier.get(primary_uid, ())) + if direct_payloads: + if task_type != "E1": + raise ValueError( + f"TaskGroup {instance['id']!r} moves carrier {primary_uid!r} " + "with direct payloads, but payload propagation currently " + "supports only single-arm E1 placement." + ) + payload_roles = [payload_role for _, payload_role, _ in direct_payloads] + if params.get("payload_roles") != payload_roles: + params["payload_roles"] = payload_roles + changed = True + for payload_uid, _payload_role, producer_id in direct_payloads: + links.append( + { + "producer": producer_id, + "consumer": str(instance["id"]), + "carrier": primary_uid, + "payload": payload_uid, + "relation": "direct_support", + } + ) + + if task_type in {"E1", "E2", "E3", "E4", "E5"}: + old_carrier = carrier_by_payload.pop(primary_uid, None) + if old_carrier is not None: + direct_by_carrier[old_carrier] = [ + item + for item in direct_by_carrier.get(old_carrier, ()) + if item[0] != primary_uid + ] + + if task_type != "E1" or str(params.get("relation")) not in {"on", "inside"}: + continue + target_role = params.get("target_role") + if not isinstance(target_role, str) or not target_role: + continue + target_uid = bindings.get(target_role, target_role) + if target_uid in {"table", "table_center"} or target_uid == primary_uid: + continue + payload_role = role_by_uid.get(primary_uid, primary_role) + direct_by_carrier.setdefault(target_uid, []).append( + (primary_uid, payload_role, str(instance["id"])) + ) + carrier_by_payload[primary_uid] = target_uid + + if changed: + metadata = dict(result.get("metadata", {})) + metadata.pop("action_contract_task_linker", None) + result["metadata"] = metadata + return validate_task_spec(result), links + + +def _payload_goal(params: Mapping[str, Any], object_uid: str) -> list[dict[str, str]]: + raw_payloads = params.get("payload_roles", []) + if not isinstance(raw_payloads, Sequence) or isinstance( + raw_payloads, (str, bytes, bytearray) + ): + raise ValueError("E1 payload_roles must be a list.") + payloads = [str(value) for value in raw_payloads] + if any(not value for value in payloads): + raise ValueError("E1 payload_roles must contain non-empty object IDs.") + if object_uid in payloads: + raise ValueError("An E1 carrier cannot be its own payload.") + if len(payloads) != len(set(payloads)): + raise ValueError("E1 direct payload objects must be unique.") + return [{"object": value, "slot": "center"} for value in payloads] + + +def _orientation_extensions(params: Mapping[str, Any]) -> dict[str, Any]: + """Copy optional compiled-orientation fields from one task instance.""" + return { + key: deepcopy(params[key]) + for key in ("orientation_constraint", "orientation_directed") + if key in params + } + + +def _recipe( + group_id: str, + task_type: str, + object_uid: str, + actor: Mapping[str, Any], + params: Mapping[str, Any], + dependencies: list[str], + *, + role: str, + incoming_held_arm: str | None = None, +) -> tuple[list[dict[str, Any]], str, dict[str, Any], dict[str, Any]]: + if task_type == "E1": + target = str(params.get("target_role", "table")) + relation = str(params.get("relation", "on")) + layout = str(params.get("layout", "")) + if layout == "line": + goal = { + "layout": "line", + "objects": list(params["objects_roles"]), + "axis": str(params.get("axis", "world_y")), + "anchor": "table_center", + "order_by": str(params.get("order_by", "explicit")), + "order_direction": str(params.get("order_direction", "given")), + "order_constraint": str(params.get("order_constraint", "free")), + "participation": str(params.get("participation", "auto")), + "orientation_goal": str(params.get("orientation_goal", "none")), + "orientation_axis": str(params.get("orientation_axis", "none")), + **_orientation_extensions(params), + "nominal_slot_index": int(params["nominal_slot_index"]), + "slot_constraint": str( + params.get("slot_constraint", "free_reassignable") + ), + } + payloads = _payload_goal(params, object_uid) + if payloads: + goal["payloads"] = payloads + success = { + "type": "line_member_placed", + "nominal_slot_index": goal["nominal_slot_index"], + "slot_constraint": goal["slot_constraint"], + "order_constraint": goal["order_constraint"], + } + return ( + _single_arm_manipulation( + group_id, + task_type, + object_uid, + actor, + dependencies, + role=role, + already_held=incoming_held_arm is not None, + payloads=payloads, + ), + "arrange_line", + goal, + success, + ) + goal = { + "reference_object": target, + "reference_state": "live", + "relation": relation, + "relation_frame": str(params.get("relation_frame", "world")), + "orientation_goal": str(params.get("orientation_goal", "none")), + "orientation_axis": str(params.get("orientation_axis", "none")), + **_orientation_extensions(params), + "slot": str(params.get("slot", "auto")), + } + if "visual_constraint" in params: + goal["visual_constraint"] = deepcopy(params["visual_constraint"]) + payloads = _payload_goal(params, object_uid) + if payloads: + goal["payloads"] = payloads + success = { + "type": "semantic_goal", + "relation": relation, + "reference_object": target, + } + return ( + _single_arm_manipulation( + group_id, + task_type, + object_uid, + actor, + dependencies, + role=role, + already_held=incoming_held_arm is not None, + payloads=payloads, + ), + "place_relative", + goal, + success, + ) + if task_type == "E2": + terminal_behavior = str(params.get("terminal_behavior", "place")) + if terminal_behavior == "hold" and role != "recovery": + raise ValueError( + "Ordinary E2 groups must release their supported object at the " + "TaskGroup boundary." + ) + goal = { + "relation": "none", + "reference_state": "live", + "orientation_goal": str(params.get("orientation_goal", "upright")), + "orientation_axis": str(params.get("orientation_axis", "none")), + "position_anchor": "initial_xy", + "support_object": str(params.get("support_role", "table")), + "upright_local_axis": str(params.get("upright_local_axis", "long_axis")), + **_orientation_extensions(params), + } + if terminal_behavior == "hold": + goal["terminal_behavior"] = "hold" + success = { + "type": task_success_type(task_type, params), + "object": object_uid, + "local_axis": goal["upright_local_axis"], + } + return ( + _single_arm_manipulation( + group_id, + task_type, + object_uid, + actor, + dependencies, + role=role, + already_held=incoming_held_arm is not None, + leave_held=str(params.get("terminal_behavior", "place")) == "hold", + ), + "orient_object", + goal, + success, + ) + if task_type == "E3": + target = str(params["target_role"]) + goal = { + "reference_object": target, + "relation": "above", + "amount": "task_defined", + } + return ( + [ + _node( + group_id, + 1, + "Pour", + task_type, + object_uid, + actor, + "arm", + { + "kind": "pour_goal", + "object": object_uid, + "reference_object": target, + }, + dependencies, + role, + { + "type": "poured", + "object": object_uid, + "reference_object": target, + }, + motion_policy(), + ) + ], + "pour", + goal, + {"type": "poured", "object": object_uid, "reference_object": target}, + ) + if task_type == "E4": + transfer = str(params.get("transfer_arm", "left_arm")) + receive = str(params.get("receive_arm", "right_arm")) + if incoming_held_arm == "coordinated": + raise ValueError( + "E4 cannot consume a coordinated hold; an explicit single-arm " + "handover state is required." + ) + if incoming_held_arm is not None and transfer != incoming_held_arm: + raise ValueError( + f"E4 transfer_arm {transfer!r} conflicts with the predecessor " + f"holder {incoming_held_arm!r}." + ) + pickup_actor = {"mode": "required", "arm": transfer} + pickup = None + if incoming_held_arm is None: + pickup = _node( + group_id, + 1, + "PickUp", + task_type, + object_uid, + pickup_actor, + "arm", + {"kind": "object", "object": object_uid}, + dependencies, + role, + {"type": "object_held", "object": object_uid, "arm": transfer}, + motion_policy(("handover_role", "transfer")), + ) + staging = _node( + group_id, + 1 if pickup is None else 2, + "MoveHeldObject", + task_type, + object_uid, + pickup_actor, + "arm", + { + "kind": "handover_staging", + "object": object_uid, + "transfer_arm": transfer, + "receive_arm": receive, + }, + dependencies if pickup is None else [pickup["id"]], + role, + {"type": "object_held", "object": object_uid, "arm": transfer}, + motion_policy(), + ) + handover = _node( + group_id, + 2 if pickup is None else 3, + "HandOver", + task_type, + object_uid, + {"mode": "coordinated", "arms": ["left_arm", "right_arm"]}, + "coordinated", + { + "kind": "handover_goal", + "object": object_uid, + "transfer_arm": transfer, + "receive_arm": receive, + }, + [staging["id"]], + role, + {"type": "handover_complete", "object": object_uid, "arm": receive}, + motion_policy(), + ) + # Grounding configures HandOver as exchange-to-exchange, so its receiver + # stays at the grasp while the transfer arm performs the built-in lift. + # This ordered retreat/home suffix then verifies and completes clearance + # before any receiver-side continuation may carry the object away. + retreat = _node( + group_id, + 3 if pickup is None else 4, + "MoveEndEffector", + task_type, + object_uid, + pickup_actor, + "arm", + { + "kind": "policy_pose", + "source": "handover", + "operation": "retreat", + }, + [handover["id"]], + "cleanup", + {}, + motion_policy(), + ) + home = _node( + group_id, + 4 if pickup is None else 5, + "MoveJoints", + task_type, + object_uid, + pickup_actor, + "arm", + { + "kind": "joint_state", + "source": "initial", + "operation": "handover_home", + }, + [retreat["id"]], + "cleanup", + {}, + motion_policy(), + ) + return ( + [ + item + for item in (pickup, staging, handover, retreat, home) + if item is not None + ], + "handover", + { + "relation": "handover", + "orientation_goal": str(params.get("orientation_goal", "none")), + "orientation_axis": "none", + **_orientation_extensions(params), + "transfer_arm": transfer, + "receive_arm": receive, + }, + {"type": "handover_complete", "object": object_uid, "arm": receive}, + ) + if task_type == "E5": + terminal_behavior = str(params.get("terminal_behavior", "hold")) + if terminal_behavior not in TERMINAL_BEHAVIORS - {"none"}: + raise ValueError("E5 terminal_behavior must be 'hold' or 'place'.") + direction = str(params.get("direction", "up")) + if direction not in TRANSPORT_DIRECTIONS: + raise ValueError(f"E5 direction {direction!r} is unsupported.") + goal = { + "direction": direction, + "terminal_behavior": terminal_behavior, + "orientation_goal": str(params.get("orientation_goal", "none")), + "orientation_axis": "none", + **_orientation_extensions(params), + "relation_frame": str(params.get("relation_frame", "robot")), + } + target = params.get("target_role") + relation = str(params.get("relation", "none")) + if isinstance(target, str) and target: + if relation == "none": + raise ValueError("E5 target_role requires a symbolic relation.") + goal.update( + { + "reference_object": target, + "reference_state": "live", + "relation": relation, + "direction": "none", + } + ) + elif direction == "none" and terminal_behavior != "place": + raise ValueError("E5 requires a direction or target_role relation.") + pick = _node( + group_id, + 1, + "CoordinatedPickment", + task_type, + object_uid, + actor, + "coordinated", + {"kind": "coordinated_goal", "object": object_uid}, + dependencies, + role, + {"type": "held_by_both_grippers", "object": object_uid}, + motion_policy(), + ) + nodes = [pick] + if terminal_behavior == "place": + release_sync_group = f"{group_id}__dual_release" + for index, arm, release_role in ( + (2, "left_arm", "participant"), + (3, "right_arm", "commit"), + ): + release = _node( + group_id, + index, + "MoveJoints", + task_type, + object_uid, + {"mode": "required", "arm": arm}, + "hand", + { + "kind": "joint_state", + "source": "gripper_open", + "coordinated_release_role": release_role, + }, + [pick["id"]], + role, + {}, + motion_policy(), + ) + release["sync_group"] = release_sync_group + nodes.append(release) + success_type = task_success_type(task_type, params) + success = ( + {"type": success_type, "object": object_uid} + if success_type == "held_by_both_grippers" + else { + "type": success_type, + "relation": relation, + **( + {"reference_object": target} + if isinstance(target, str) and target + else {} + ), + } + ) + return ( + nodes, + "coordinated_transport", + goal, + success, + ) + planning = { + "E6": ("PullArticulatedPart", "pull_articulated_part"), + "E7": ("PushArticulatedPart", "push_articulated_part"), + "E8": ("TurnKnob", "turn_knob"), + } + if task_type in planning: + action_name, operator = planning[task_type] + success = { + "type": "articulation_joint_near", + "object": object_uid, + "target_state": params.get("target_state", params.get("target_setting")), + } + return ( + [ + _node( + group_id, + 1, + action_name, + task_type, + object_uid, + actor, + "arm", + {"kind": "articulation_goal", "object": object_uid}, + dependencies, + role, + success, + motion_policy(), + ) + ], + operator, + { + key: deepcopy(value) + for key, value in params.items() + if not key.endswith("_role") + }, + success, + ) + if task_type == "E9": + success = { + "type": "pressed", + "object": object_uid, + "terminal_state": str(params.get("terminal_state", "activated")), + } + return ( + [ + _node( + group_id, + 1, + "Press", + task_type, + object_uid, + actor, + "arm", + {"kind": "object", "object": object_uid}, + dependencies, + role, + success, + motion_policy(), + ) + ], + "press", + {"terminal_state": success["terminal_state"]}, + success, + ) + raise ValueError(f"Unsupported task type {task_type!r}.") + + +def _single_arm_manipulation( + group_id: str, + task_type: str, + object_uid: str, + actor: Mapping[str, Any], + dependencies: list[str], + *, + role: str, + already_held: bool = False, + leave_held: bool = False, + payloads: Sequence[Mapping[str, Any]] = (), +) -> list[dict[str, Any]]: + orientation_modifiers: tuple[tuple[str, str], ...] = ( + (("orientation", "upright"),) if task_type == "E2" else () + ) + payload_binding = deepcopy(list(payloads)) + specs = ( + ( + "PickUp", + { + "kind": "object", + "object": object_uid, + **({"payloads": payload_binding} if payload_binding else {}), + }, + motion_policy(*orientation_modifiers), + ), + ( + "MoveHeldObject", + { + "kind": "semantic_goal", + "semantic_step": group_id, + "phase": "staging", + **({"payloads": payload_binding} if payload_binding else {}), + }, + motion_policy(*orientation_modifiers), + ), + ( + "MoveHeldObject", + { + "kind": "semantic_goal", + "semantic_step": group_id, + "phase": "final", + **({"payloads": payload_binding} if payload_binding else {}), + }, + motion_policy(*orientation_modifiers), + ), + ( + "Place", + { + "kind": "current_held_pose", + **({"payloads": payload_binding} if payload_binding else {}), + }, + motion_policy(*orientation_modifiers), + ), + ( + "MoveEndEffector", + {"kind": "policy_pose", "source": "release", "operation": "retreat"}, + motion_policy(*orientation_modifiers), + ), + ( + "MoveJoints", + {"kind": "joint_state", "source": "initial"}, + motion_policy(), + ), + ) + if already_held: + specs = specs[1:] + if leave_held: + # A held continuation must not retreat or home the arm after the final + # semantic move: those cleanup phases would move away from the + # handover staging state while still owning the object. + place_index = next( + (index for index, spec in enumerate(specs) if spec[0] == "Place"), + len(specs), + ) + specs = specs[:place_index] + nodes = [] + previous = list(dependencies) + for index, (action, binding, policy) in enumerate(specs, start=1): + node_role = "cleanup" if action in {"MoveEndEffector", "MoveJoints"} else role + node = _node( + group_id, + index, + action, + task_type, + object_uid, + actor, + "arm", + binding, + previous, + node_role, + {}, + policy, + ) + nodes.append(node) + previous = [node["id"]] + return nodes + + +def _node( + group_id: str, + index: int, + action: str, + task_type: str, + object_uid: str, + actor: Mapping[str, Any], + control: str, + binding: Mapping[str, Any], + dependencies: list[str], + role: str, + postcondition: Mapping[str, Any], + motion_policy: Mapping[str, Any], +) -> dict[str, Any]: + return { + "id": f"{group_id}__a{index:02d}", + "atomic_action": action, + "object_uid": object_uid, + "actor": deepcopy(dict(actor)), + "control": control, + "target_binding": deepcopy(dict(binding)), + "depends_on": list(dependencies), + "task_instance_id": group_id, + "task_type": task_type, + "role": role, + "precondition": {}, + "postcondition": deepcopy(dict(postcondition)), + "motion_policy": deepcopy(dict(motion_policy)), + } + + +def _terminal_nodes(nodes: list[Mapping[str, Any]]) -> list[str]: + depended = {dependency for node in nodes for dependency in node["depends_on"]} + return [str(node["id"]) for node in nodes if node["id"] not in depended] + + +def _primary_object(task_type: str, params: Mapping[str, Any]) -> str: + key = "source_role" if task_type == "E3" else "object_role" + value = params.get(key) + if not isinstance(value, str) or not value: + raise ValueError(f"{task_type} requires resolved parameter {key!r}.") + return value + + +def _actor( + task_type: str, + params: Mapping[str, Any], + *, + incoming_held_arm: str | None = None, +) -> dict[str, Any]: + required_arm = params.get("required_arm") + if ( + incoming_held_arm is not None + and required_arm in {"left_arm", "right_arm"} + and str(required_arm) != incoming_held_arm + ): + raise ValueError( + f"Continuation requires {incoming_held_arm!r}, but the task " + f"requested {required_arm!r}." + ) + if incoming_held_arm is not None: + return {"mode": "required", "arm": incoming_held_arm} + if required_arm in {"left_arm", "right_arm"}: + return {"mode": "required", "arm": str(required_arm)} + if task_type == "E5": + return {"mode": "coordinated", "arms": ["left_arm", "right_arm"]} + if task_type == "E4": + return {"mode": "required", "arm": str(params.get("transfer_arm", "left_arm"))} + return {"mode": "auto"} + + +def _incoming_held_arm( + task_type: str, + object_uid: str, + dependencies: list[str], + held_after_group: Mapping[str, tuple[str, str] | None], +) -> str | None: + """Resolve a predecessor-provided hold for a continuation recipe.""" + if task_type not in {"E1", "E2", "E4"}: + return None + candidates = { + held[1] + for dependency in dependencies + if (held := held_after_group.get(str(dependency))) is not None + and held[0] == object_uid + } + if len(candidates) > 1: + raise ValueError( + f"Task instance has conflicting predecessor holders for {object_uid!r}." + ) + return next(iter(candidates), None) + + +def _terminal_hold( + task_type: str, + object_uid: str, + params: Mapping[str, Any], +) -> tuple[str, str] | None: + if task_type == "E4": + return object_uid, str(params.get("receive_arm", "right_arm")) + if task_type == "E2" and str(params.get("terminal_behavior", "place")) == "hold": + arm = str(params.get("required_arm", "")) + if arm in {"left_arm", "right_arm"}: + return object_uid, arm + if task_type == "E5" and str(params.get("terminal_behavior", "hold")) == "hold": + return object_uid, "coordinated" + return None + + +def _validate_bindings( + task: Mapping[str, Any], + role_bindings: Mapping[str, str], +) -> dict[str, str]: + bindings = dict(role_bindings) + for role, uid in bindings.items(): + if not isinstance(role, str) or not role or not isinstance(uid, str) or not uid: + raise ValueError("role_bindings must map non-empty role IDs to scene UIDs.") + required = set() + for instance in task["task_instances"]: + required.update(_role_references(instance["params"])) + required.discard("table") + missing = sorted(required - set(bindings)) + if missing: + raise ValueError(f"Scene hand-off is missing role bindings: {missing}.") + if len(bindings.values()) != len(set(bindings.values())): + raise ValueError("Scene role bindings must resolve to unique object UIDs.") + return bindings + + +def _role_references(value: Any, key: str = "") -> set[str]: + if isinstance(value, Mapping): + return { + role + for child_key, child in value.items() + for role in _role_references(child, str(child_key)) + } + if isinstance(value, list): + return {role for child in value for role in _role_references(child, key)} + if isinstance(value, str) and (key.endswith("_role") or key.endswith("_roles")): + return {value} + return set() + + +def _resolve_params(value: Any, bindings: Mapping[str, str], key: str = "") -> Any: + if isinstance(value, Mapping): + return { + child_key: _resolve_params(child, bindings, str(child_key)) + for child_key, child in value.items() + } + if isinstance(value, list): + return [_resolve_params(child, bindings, key) for child in value] + if isinstance(value, str) and (key.endswith("_role") or key.endswith("_roles")): + return bindings.get(value, value) + return deepcopy(value) diff --git a/embodichain/gen_sim/action_engine/tasks/scene.py b/embodichain/gen_sim/action_engine/tasks/scene.py new file mode 100644 index 000000000..242b15dd4 --- /dev/null +++ b/embodichain/gen_sim/action_engine/tasks/scene.py @@ -0,0 +1,177 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Validate a Scene Engine result against task-first requirements.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +from embodichain.gen_sim.action_engine.domain import validate_scene_requirements + +__all__ = ["SceneHandoff", "validate_scene_handoff"] + + +@dataclass(frozen=True) +class SceneHandoff: + """Validated role-to-UID resolution returned by an external Scene Engine.""" + + task_id: str + role_bindings: dict[str, Any] + object_uids: tuple[str, ...] + camera_uids: tuple[str, ...] + + +def validate_scene_handoff( + requirements: Mapping[str, Any], + scene: Mapping[str, Any], + role_bindings: Mapping[str, str], +) -> SceneHandoff: + """Reject scenes that do not satisfy roles, affordances, state, or cameras.""" + required = validate_scene_requirements(requirements) + objects = scene.get("objects") + if not isinstance(objects, Sequence) or isinstance(objects, (str, bytes)): + raise ValueError("Scene hand-off requires an objects list.") + object_by_uid: dict[str, Mapping[str, Any]] = {} + for index, item in enumerate(objects): + if not isinstance(item, Mapping): + raise ValueError(f"Scene objects[{index}] must be a mapping.") + uid = item.get("uid") + if not isinstance(uid, str) or not uid: + raise ValueError(f"Scene objects[{index}] requires a UID.") + if uid in object_by_uid: + raise ValueError(f"Scene contains duplicate object UID {uid!r}.") + object_by_uid[uid] = item + + bindings = dict(role_bindings) + required_roles = {item["role_id"] for item in required["objects"]} + if set(bindings) != required_roles: + missing = sorted(required_roles - set(bindings)) + extra = sorted(set(bindings) - required_roles) + raise ValueError( + f"Scene role bindings mismatch; missing={missing}, extra={extra}." + ) + normalized_bindings: dict[str, str | tuple[str, ...]] = {} + assigned_uids: list[str] = [] + for requirement in required["objects"]: + role = requirement["role_id"] + count = int(requirement["count"]) + binding = bindings[role] + if isinstance(binding, str): + uids = [binding] + elif isinstance(binding, Sequence) and not isinstance(binding, (str, bytes)): + uids = [str(uid) for uid in binding] + else: + raise ValueError(f"Scene role {role!r} has an invalid UID binding.") + if len(uids) != count or any(not uid for uid in uids): + raise ValueError( + f"Scene role {role!r} requires exactly {count} UID binding(s)." + ) + normalized_bindings[role] = uids[0] if count == 1 else tuple(uids) + assigned_uids.extend(uids) + for uid in uids: + _validate_bound_object(object_by_uid, uid, role, requirement) + if len(assigned_uids) != len(set(assigned_uids)): + raise ValueError("Each scene requirement role must resolve to unique UIDs.") + + cameras = scene.get("cameras", []) + if not isinstance(cameras, Sequence) or isinstance(cameras, (str, bytes)): + raise ValueError("Scene cameras must be a list.") + camera_uids = [] + normalized_cameras = [] + for camera in cameras: + if not isinstance(camera, Mapping) or not isinstance(camera.get("uid"), str): + raise ValueError("Every scene camera requires a UID.") + camera_uids.append(str(camera["uid"])) + normalized_cameras.append(camera) + for camera_requirement in required["cameras"]: + modalities = set(camera_requirement.get("modalities", ())) + coverage = camera_requirement.get("coverage") + if not any( + modalities <= set(camera.get("modalities", ())) + and (coverage is None or camera.get("coverage") == coverage) + for camera in normalized_cameras + ): + raise ValueError( + "Scene cameras do not satisfy requirement " + f"{dict(camera_requirement)!r}." + ) + reported_constraints = scene.get("satisfied_spatial_constraints", []) + if not isinstance(reported_constraints, Sequence) or isinstance( + reported_constraints, (str, bytes) + ): + raise ValueError("Scene satisfied_spatial_constraints must be a list.") + reported = {_canonical(item) for item in reported_constraints} + missing_constraints = [ + constraint + for constraint in required["spatial_constraints"] + if _canonical(constraint) not in reported + ] + if missing_constraints: + raise ValueError( + "Scene does not satisfy spatial constraints: " f"{missing_constraints}." + ) + return SceneHandoff( + task_id=required["task_id"], + role_bindings=normalized_bindings, + object_uids=tuple(sorted(object_by_uid)), + camera_uids=tuple(sorted(camera_uids)), + ) + + +def _validate_bound_object( + object_by_uid: Mapping[str, Mapping[str, Any]], + uid: str, + role: str, + requirement: Mapping[str, Any], +) -> None: + if uid not in object_by_uid: + raise ValueError(f"Scene role {role!r} references unknown UID {uid!r}.") + actual = object_by_uid[uid] + if actual.get("category") != requirement["category"]: + raise ValueError( + f"Scene object {uid!r} category does not satisfy role {role!r}." + ) + missing_affordances = set(requirement["affordances"]) - set( + actual.get("affordances", ()) + ) + if missing_affordances: + raise ValueError( + f"Scene object {uid!r} lacks affordances {sorted(missing_affordances)}." + ) + for field in ("initial_state", "attributes"): + actual_values = actual.get(field, {}) + if not isinstance(actual_values, Mapping): + raise ValueError(f"Scene object {uid!r} {field} must be a mapping.") + mismatched = { + key: expected + for key, expected in requirement[field].items() + if actual_values.get(key) != expected + } + if mismatched: + raise ValueError( + f"Scene object {uid!r} does not satisfy {field} {mismatched}." + ) + + +def _canonical(value: Any) -> str: + import json + + if not isinstance(value, Mapping): + raise ValueError("Every satisfied spatial constraint must be a mapping.") + return json.dumps(dict(value), sort_keys=True, separators=(",", ":")) 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/embodichain/gen_sim/scene_engine/cli/start.py b/embodichain/gen_sim/scene_engine/cli/start.py index 59454b09f..6d46f32dc 100644 --- a/embodichain/gen_sim/scene_engine/cli/start.py +++ b/embodichain/gen_sim/scene_engine/cli/start.py @@ -21,15 +21,31 @@ from pathlib import Path from embodichain.gen_sim.scene_engine.pipeline.generate import generate_scene_from_image +from embodichain.gen_sim.scene_engine.pipeline.edit import edit_scene _SUPPORTED_IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png"} def cli_scene_engine( - image: str | Path, + image: str | Path | None, output_root: str | Path, + *, + edit_prompt: str | None = None, ) -> None: - """Generate one scene using the required ``gen_sim/.env`` settings.""" + """Generate a scene from an image, edit an export, or do both in sequence.""" + resolved_output_root = Path(output_root).expanduser().resolve() + if edit_prompt is not None: + edit_prompt = edit_prompt.strip() + if not edit_prompt: + raise ValueError("Edit prompt must not be empty.") + + if image is None: + if edit_prompt is None: + raise ValueError("Provide --image, --edit_prompt, or both.") + edit_scene(output_root=resolved_output_root, edit_prompt=edit_prompt) + print("Successfully completed!") + return + resolved_image_path = Path(image).expanduser().resolve() if not resolved_image_path.exists(): raise FileNotFoundError(f"Image input not found: {resolved_image_path}") @@ -40,27 +56,30 @@ def cli_scene_engine( "Image input must have one of these extensions: .jpg, .jpeg, .png" ) - resolved_output_root = Path(output_root).expanduser().resolve() resolved_output_root.mkdir(parents=True, exist_ok=True) - generate_scene_from_image( image_path=resolved_image_path, output_root=resolved_output_root, ) + if edit_prompt is not None: + edit_scene( + output_root=resolved_output_root, + edit_prompt=edit_prompt, + ) print("Successfully completed!") def main(argv: Sequence[str] | None = None) -> None: parser = argparse.ArgumentParser( prog="embodichain scene-engine", - description="Generate a Scene Engine export from one input image.", + description="Generate a Scene Engine export, edit one, or do both.", epilog="Service settings are read from embodichain/gen_sim/.env.", ) parser.add_argument( "--image", type=str, - required=True, - help="Path to the required input image file (.jpg, .jpeg, or .png)", + required=False, + help="Optional input image file (.jpg, .jpeg, or .png)", ) parser.add_argument( "--output_root", @@ -68,9 +87,15 @@ def main(argv: Sequence[str] | None = None) -> None: required=True, help="Path to the output directory", ) + parser.add_argument( + "--edit_prompt", + type=str, + default=None, + help="Text instruction for editing an existing or newly generated output root", + ) args = parser.parse_args(argv) - cli_scene_engine(args.image, args.output_root) + cli_scene_engine(args.image, args.output_root, edit_prompt=args.edit_prompt) if __name__ == "__main__": diff --git a/embodichain/gen_sim/scene_engine/clients/geometry_generation.py b/embodichain/gen_sim/scene_engine/clients/geometry_generation.py index c84fe4c26..8f4d13a13 100644 --- a/embodichain/gen_sim/scene_engine/clients/geometry_generation.py +++ b/embodichain/gen_sim/scene_engine/clients/geometry_generation.py @@ -25,6 +25,8 @@ import requests +from embodichain.gen_sim.scene_engine.errors import SceneServiceError + from embodichain.gen_sim.scene_engine.configs.environment import ( read_scene_engine_env_values, ) @@ -77,7 +79,7 @@ def check_health(self) -> None: last_error = exc assert last_error is not None - raise RuntimeError( + raise SceneServiceError( "Geometry Generation Server health check failed after " f"{self._max_attempts} attempts." ) from last_error @@ -91,6 +93,7 @@ def generate_objects( image_path: str | Path, object_masks: list[tuple[str, Path]], output_root: str | Path, + seed: int | None = None, ) -> tuple[dict[str, Any], list[dict[str, Any]]]: """Generate objects through the geometry server's mask-list endpoint. @@ -125,6 +128,7 @@ def generate_objects( response_data, response_objects = self._request_objects( image_path=resolved_image_path, object_masks=resolved_object_masks, + seed=seed, ) resolved_output_root = Path(output_root).expanduser().resolve() @@ -158,6 +162,7 @@ def _request_objects( *, image_path: Path, object_masks: list[tuple[str, Path]], + seed: int | None, ) -> tuple[dict[str, Any], list[dict[str, Any]]]: last_error: Exception | None = None for _ in range(self._max_attempts): @@ -171,6 +176,7 @@ def _request_objects( ] response = self._session.post( self._url(self._generate_objects_path), + data=(None if seed is None else {"seed": str(int(seed))}), files=[ ( "image", @@ -201,6 +207,11 @@ def _request_objects( "Geometry Generation Server response is not valid JSON." ) from exc response_data = self._wait_for_task_if_needed(response_data) + if seed is not None and _response_seed(response_data) != int(seed): + raise RuntimeError( + "Geometry Generation Server did not acknowledge the " + f"requested seed {int(seed)}." + ) response_objects = _parse_objects_response( response_data, object_ids=[object_id for object_id, _ in object_masks], @@ -210,7 +221,7 @@ def _request_objects( last_error = exc assert last_error is not None - raise RuntimeError( + raise SceneServiceError( "Geometry Generation Server request failed after " f"{self._max_attempts} attempts." ) from last_error @@ -294,7 +305,7 @@ def _download_glb(self, mesh_path: str, output_path: Path) -> None: last_error = exc assert last_error is not None - raise RuntimeError( + raise SceneServiceError( "Geometry Generation Server GLB download failed after " f"{self._max_attempts} attempts: {mesh_path}" ) from last_error @@ -374,6 +385,25 @@ def _parse_objects_response( return parsed_objects +def _response_seed(value: object) -> int | None: + """Return a seed acknowledged by a geometry response envelope.""" + if not isinstance(value, dict): + return None + candidates = [value.get("seed")] + for key in ("result", "metadata"): + nested = value.get(key) + if isinstance(nested, dict): + candidates.append(nested.get("seed")) + for candidate in candidates: + if candidate is None or isinstance(candidate, bool): + continue + try: + return int(candidate) + except (TypeError, ValueError): + continue + return None + + def _parse_numeric_list( value: object, *, diff --git a/embodichain/gen_sim/scene_engine/clients/image_generation.py b/embodichain/gen_sim/scene_engine/clients/image_generation.py new file mode 100644 index 000000000..4f59fff11 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/clients/image_generation.py @@ -0,0 +1,192 @@ +# ---------------------------------------------------------------------------- +# 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 pathlib import Path +from typing import Any + +import requests + +from embodichain.gen_sim.scene_engine.errors import SceneServiceError + +from embodichain.gen_sim.scene_engine.configs.environment import ( + read_scene_engine_env_values, +) + + +class ImageGenerationClient: + """Manage the Image Generation Server connection.""" + + def __init__( + self, + *, + base_url: str, + timeout_s: int, + max_attempts: int, + health_path: str, + generate_image_by_prompt_path: str, + session: requests.Session | None = None, + ) -> None: + self._base_url = base_url.rstrip("/") + self._timeout_s = timeout_s + self._max_attempts = max_attempts + self._health_path = health_path + self._generate_image_by_prompt_path = generate_image_by_prompt_path + self._session = session or requests.Session() + + @classmethod + def from_dotenv(cls) -> "ImageGenerationClient": + """Create a client from its required ``gen_sim/.env`` settings.""" + return cls(**_load_dotenv_config()) + + def check_health(self) -> None: + last_error: requests.RequestException | RuntimeError | None = None + for _ in range(self._max_attempts): + try: + response = self._session.get( + self._url(self._health_path), + timeout=10, # Use a shorter timeout for avoiding long waits. + ) + response.raise_for_status() + response_data = response.json() + if ( + not isinstance(response_data, dict) + or response_data.get("ok") is not True + ): + raise RuntimeError( + "Image Generation Server health response does not contain ok=true." + ) + return + except (requests.RequestException, ValueError, RuntimeError) as exc: + last_error = exc + + assert last_error is not None + raise SceneServiceError( + "Image Generation Server health check failed after " + f"{self._max_attempts} attempts." + ) from last_error + + def close(self) -> None: + self._session.close() + + def generate_image_by_prompt( + self, + *, + prompt: str, + output_path: str | Path, + seed: int | None = None, + ) -> Path: + """Generate one PNG image from ``prompt`` and save it to ``output_path``.""" + prompt = prompt.strip() + if not prompt: + raise ValueError("Image generation prompt must not be empty.") + + resolved_output_path = Path(output_path).expanduser().resolve() + resolved_output_path.parent.mkdir(parents=True, exist_ok=True) + + last_error: Exception | None = None + for _ in range(self._max_attempts): + try: + response = self._session.post( + self._url(self._generate_image_by_prompt_path), + json={ + "prompt": prompt, + **({} if seed is None else {"seed": int(seed)}), + }, + timeout=self._timeout_s, + ) + response.raise_for_status() + if seed is not None: + acknowledged = response.headers.get( + "x-generation-seed", + response.headers.get("x-seed"), + ) + try: + acknowledged_seed = int(acknowledged) + except (TypeError, ValueError): + acknowledged_seed = None + if acknowledged_seed != int(seed): + raise RuntimeError( + "Image Generation Server did not acknowledge the " + f"requested seed {int(seed)}." + ) + content_type = response.headers.get("content-type", "").split(";")[0] + if content_type != "image/png": + raise RuntimeError( + "Image Generation Server response is not a PNG image." + ) + resolved_output_path.write_bytes(response.content) + return resolved_output_path + except (requests.RequestException, RuntimeError) as exc: + last_error = exc + + assert last_error is not None + raise SceneServiceError( + "Image Generation Server request failed after " + f"{self._max_attempts} attempts." + ) from last_error + + def _url(self, path: str) -> str: + return f"{self._base_url}/{path.lstrip('/')}" + + +def _load_dotenv_config() -> dict[str, Any]: + values = read_scene_engine_env_values( + "SCENE_ENGINE_IMAGE_GENERATION_BASE_URL", + "SCENE_ENGINE_IMAGE_GENERATION_TIMEOUT_S", + "SCENE_ENGINE_IMAGE_GENERATION_MAX_ATTEMPTS", + "SCENE_ENGINE_IMAGE_GENERATION_HEALTH_PATH", + "SCENE_ENGINE_IMAGE_GENERATION_BY_PROMPT_PATH", + ) + try: + timeout_s = int(values["SCENE_ENGINE_IMAGE_GENERATION_TIMEOUT_S"]) + except (TypeError, ValueError) as exc: + raise ValueError( + "SCENE_ENGINE_IMAGE_GENERATION_TIMEOUT_S must be an integer." + ) from exc + if timeout_s < 1: + raise ValueError("SCENE_ENGINE_IMAGE_GENERATION_TIMEOUT_S must be at least 1.") + + try: + max_attempts = int(values["SCENE_ENGINE_IMAGE_GENERATION_MAX_ATTEMPTS"]) + except (TypeError, ValueError) as exc: + raise ValueError( + "SCENE_ENGINE_IMAGE_GENERATION_MAX_ATTEMPTS must be an integer." + ) from exc + if max_attempts < 1: + raise ValueError( + "SCENE_ENGINE_IMAGE_GENERATION_MAX_ATTEMPTS must be at least 1." + ) + + string_keys = ( + "SCENE_ENGINE_IMAGE_GENERATION_BASE_URL", + "SCENE_ENGINE_IMAGE_GENERATION_HEALTH_PATH", + "SCENE_ENGINE_IMAGE_GENERATION_BY_PROMPT_PATH", + ) + for key in string_keys: + if not values[key].strip(): + raise ValueError(f"Scene Engine .env key {key} must be a non-empty string.") + + return { + "base_url": values["SCENE_ENGINE_IMAGE_GENERATION_BASE_URL"].strip(), + "timeout_s": timeout_s, + "max_attempts": max_attempts, + "health_path": values["SCENE_ENGINE_IMAGE_GENERATION_HEALTH_PATH"].strip(), + "generate_image_by_prompt_path": values[ + "SCENE_ENGINE_IMAGE_GENERATION_BY_PROMPT_PATH" + ].strip(), + } diff --git a/embodichain/gen_sim/scene_engine/clients/image_segmentation.py b/embodichain/gen_sim/scene_engine/clients/image_segmentation.py index 2c56af91a..8414b9627 100644 --- a/embodichain/gen_sim/scene_engine/clients/image_segmentation.py +++ b/embodichain/gen_sim/scene_engine/clients/image_segmentation.py @@ -36,14 +36,14 @@ def __init__( timeout_s: int, max_attempts: int, health_path: str, - segment_single_object_path: str, + segment_by_prompt_path: str, session: requests.Session | None = None, ) -> None: self._base_url = base_url.rstrip("/") self._timeout_s = timeout_s self._max_attempts = max_attempts self._health_path = health_path - self._segment_single_object_path = segment_single_object_path + self._segment_by_prompt_path = segment_by_prompt_path self._session = session or requests.Session() @classmethod @@ -96,7 +96,7 @@ def segment_single_object( try: with resolved_image_path.open("rb") as image_file: response = self._session.post( - self._url(self._segment_single_object_path), + self._url(self._segment_by_prompt_path), data={"prompt": prompt}, files={"image": (resolved_image_path.name, image_file)}, timeout=self._timeout_s, @@ -138,7 +138,7 @@ def _load_dotenv_config() -> dict[str, Any]: "SCENE_ENGINE_IMAGE_SEGMENTATION_TIMEOUT_S", "SCENE_ENGINE_IMAGE_SEGMENTATION_MAX_ATTEMPTS", "SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH", - "SCENE_ENGINE_IMAGE_SEGMENTATION_SINGLE_OBJECT_PATH", + "SCENE_ENGINE_IMAGE_SEGMENTATION_BY_PROMPT_PATH", ) try: timeout_s = int(values["SCENE_ENGINE_IMAGE_SEGMENTATION_TIMEOUT_S"]) @@ -165,7 +165,7 @@ def _load_dotenv_config() -> dict[str, Any]: string_keys = ( "SCENE_ENGINE_IMAGE_SEGMENTATION_BASE_URL", "SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH", - "SCENE_ENGINE_IMAGE_SEGMENTATION_SINGLE_OBJECT_PATH", + "SCENE_ENGINE_IMAGE_SEGMENTATION_BY_PROMPT_PATH", ) for key in string_keys: if not values[key].strip(): @@ -176,8 +176,8 @@ def _load_dotenv_config() -> dict[str, Any]: "timeout_s": timeout_s, "max_attempts": max_attempts, "health_path": values["SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH"].strip(), - "segment_single_object_path": values[ - "SCENE_ENGINE_IMAGE_SEGMENTATION_SINGLE_OBJECT_PATH" + "segment_by_prompt_path": values[ + "SCENE_ENGINE_IMAGE_SEGMENTATION_BY_PROMPT_PATH" ].strip(), } diff --git a/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py new file mode 100644 index 000000000..0f78b9699 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py @@ -0,0 +1,280 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Literal + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + OrientationState, + SceneConstraintType, + SceneGraph, + TableRegion, + TABLE_OBJECT_ID, +) + +__all__ = ["SceneEditOperation", "SceneEditPlan"] + +SceneEditOperationType = Literal["add", "move", "delete"] + + +@dataclass(frozen=True) +class SceneEditOperation: + """One normalized edit operation produced from an LLM edit draft.""" + + op: SceneEditOperationType + object_id: str | None = None + target_id: str | None = None + relation: SceneConstraintType | None = None + table_region: TableRegion | None = None + category: str | None = None + name: str | None = None + description: str | None = None + orientation_state: OrientationState | None = None + + def to_dict(self) -> dict[str, object]: + """Serialize one normalized edit operation.""" + return { + "op": self.op, + "object_id": self.object_id, + "target_id": self.target_id, + "relation": self.relation, + "table_region": self.table_region, + "category": self.category, + "name": self.name, + "description": self.description, + "orientation_state": self.orientation_state, + } + + +@dataclass +class SceneEditPlan: + """Validated operations against one immutable pre-edit scene state.""" + + scene: Scene + scene_graph: SceneGraph + operations: list[SceneEditOperation] = field(default_factory=list) + + def __post_init__(self) -> None: + """Validate the plan before later stages prepare assets or edit layouts.""" + self.validate() + + def to_dict(self) -> dict[str, object]: + """Serialize the input scene state and normalized edit operations.""" + return { + "scene": self.scene.to_dict(), + "scene_graph": self.scene_graph.to_dict(), + "operations": [operation.to_dict() for operation in self.operations], + } + + def validate(self) -> None: + """Validate object references and edit conflicts against the input scene.""" + # Edit-plan rules: + # - move and delete identify one existing non-table object with object_id. + # - add carries generated object_id plus non-empty category, name, and description. + # - add may preserve an explicit standing or lying user placement intent. + # - move always supplies target_id and relation; add may omit both. + # - table_region is only valid with target_id=table and relation=on. + # - target_id and relation are otherwise supplied together or both absent. + # - every target is from the pre-edit scene; new and deleted objects are invalid targets. + # - an existing object has at most one move or delete operation in one plan. + # - delete carries no placement or new-object metadata and must delete every descendant. + # - these checks validate intent only; they do not mutate Scene or SceneGraph. + # Scene object IDs must remain a one-to-one lookup key for edit operations. + scene_object_ids = {scene_object.id for scene_object in self.scene.objects} + if len(scene_object_ids) != len(self.scene.objects): + raise ValueError("Scene edit input must contain unique object ids.") + # Parent and child checks require the graph to describe this exact scene. + if set(self.scene_graph.node_by_id()) != scene_object_ids: + raise ValueError("Scene edit plan graph nodes must match scene object ids.") + + existing_object_ids = set(scene_object_ids) + added_object_ids: set[str] = set() + # Collect deletion intents first so move and add cannot target them regardless of order. + deleted_object_ids = { + operation.object_id + for operation in self.operations + if operation.op == "delete" + } + if None in deleted_object_ids: + raise ValueError("Delete operations must identify an existing object.") + + edited_object_ids: set[str] = set() + # Validate each operation against the unchanged input scene and graph. + for operation in self.operations: + self._validate_operation( + operation=operation, + existing_object_ids=existing_object_ids, + deleted_object_ids=deleted_object_ids, + edited_object_ids=edited_object_ids, + added_object_ids=added_object_ids, + ) + # A removed support object must not leave any child objects orphaned. + self._validate_deleted_subtrees(deleted_object_ids) + + def _validate_operation( + self, + *, + operation: SceneEditOperation, + existing_object_ids: set[str], + deleted_object_ids: set[str], + edited_object_ids: set[str], + added_object_ids: set[str], + ) -> None: + if operation.op == "add": + self._validate_add_operation( + operation, + existing_object_ids, + deleted_object_ids, + added_object_ids, + ) + return + if operation.op not in {"move", "delete"}: + raise ValueError(f"Unsupported scene edit operation: {operation.op!r}") + if operation.object_id not in existing_object_ids: + raise ValueError( + "Move and delete operations must reference existing objects." + ) + if operation.object_id == TABLE_OBJECT_ID: + raise ValueError("The table cannot be moved or deleted.") + # Existing objects accept only one move or delete instruction per plan. + if operation.object_id in edited_object_ids: + raise ValueError("An existing object may have only one edit operation.") + edited_object_ids.add(operation.object_id) + + if operation.op == "delete": + # Delete carries no new metadata or spatial placement. + if any( + value is not None + for value in ( + operation.target_id, + operation.relation, + operation.table_region, + operation.category, + operation.name, + operation.description, + operation.orientation_state, + ) + ): + raise ValueError("Delete operations may only specify object_id.") + return + + if operation.target_id is None or operation.relation is None: + raise ValueError("Move operations must specify target_id and relation.") + existing_orientation_state = self.scene_graph.node_by_id()[ + operation.object_id + ].orientation_state + if operation.orientation_state not in {None, existing_orientation_state}: + raise ValueError( + "Move operations may only preserve the existing orientation_state." + ) + self._validate_position_reference( + operation=operation, + existing_object_ids=existing_object_ids, + deleted_object_ids=deleted_object_ids, + ) + if any( + value is not None + for value in ( + operation.category, + operation.name, + operation.description, + ) + ): + raise ValueError("Move operations must not declare a new object.") + + @staticmethod + def _validate_add_operation( + operation: SceneEditOperation, + existing_object_ids: set[str], + deleted_object_ids: set[str], + added_object_ids: set[str], + ) -> None: + if not operation.object_id: + raise ValueError("Add operations must have a generated object_id.") + # Generated IDs must not collide with the input scene or this add batch. + if ( + operation.object_id in existing_object_ids + or operation.object_id in added_object_ids + ): + raise ValueError("Add operations must use unique new object ids.") + added_object_ids.add(operation.object_id) + if not all( + isinstance(value, str) and value.strip() + for value in (operation.category, operation.name, operation.description) + ): + raise ValueError("Add operations require category, name, and description.") + if operation.orientation_state not in {None, "standing", "lying"}: + raise ValueError("Add operation orientation_state is invalid.") + SceneEditPlan._validate_position_reference( + operation=operation, + existing_object_ids=existing_object_ids, + deleted_object_ids=deleted_object_ids, + ) + + @staticmethod + def _validate_position_reference( + *, + operation: SceneEditOperation, + existing_object_ids: set[str], + deleted_object_ids: set[str], + ) -> None: + if (operation.target_id is None) != (operation.relation is None): + raise ValueError("target_id and relation must be specified together.") + if operation.table_region is not None and ( + operation.target_id != TABLE_OBJECT_ID or operation.relation != "on" + ): + raise ValueError( + "table_region requires target_id='table' and relation='on'." + ) + if operation.target_id is None: + return + # Targets come only from the pre-edit scene, so new objects cannot be targets. + if operation.target_id not in existing_object_ids: + raise ValueError("Edit targets must reference existing scene objects.") + # One edit may not position an object relative to a deleted target. + if operation.target_id in deleted_object_ids: + raise ValueError("Edit targets must not reference deleted objects.") + + def _validate_deleted_subtrees(self, deleted_object_ids: set[str]) -> None: + # Index the support graph once before checking every deleted parent. + children_by_parent: dict[str, list[str]] = {} + for node in self.scene_graph.nodes: + if node.parent_id is not None: + children_by_parent.setdefault(node.parent_id, []).append(node.object_id) + + for object_id in deleted_object_ids: + descendants = self._descendant_ids(object_id, children_by_parent) + if not descendants.issubset(deleted_object_ids): + raise ValueError( + "Deleting a parent requires deleting all of its children." + ) + + @staticmethod + def _descendant_ids( + object_id: str, + children_by_parent: dict[str, list[str]], + ) -> set[str]: + descendants: set[str] = set() + # Traverse every support descendant, not only direct children. + pending = list(children_by_parent.get(object_id, [])) + while pending: + child_id = pending.pop() + descendants.add(child_id) + pending.extend(children_by_parent.get(child_id, [])) + return descendants diff --git a/embodichain/gen_sim/scene_engine/core/scene_graph.py b/embodichain/gen_sim/scene_engine/core/scene_graph.py new file mode 100644 index 000000000..500980b53 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/core/scene_graph.py @@ -0,0 +1,546 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Literal + +TABLE_OBJECT_ID = "table" + +# Static type constraint for the nine regions of the tabletop 3x3 grid. +TableRegion = Literal[ + "left_back", + "back_center", + "right_back", + "left_center", + "center", + "right_center", + "left_front", + "front_center", + "right_front", +] +# Runtime membership set for validating serialized and user-provided regions. +TABLE_REGIONS = frozenset( + { + "left_back", + "back_center", + "right_back", + "left_center", + "center", + "right_center", + "left_front", + "front_center", + "right_front", + } +) + +# A on B, then B is the parent node of A. +SupportRelationType = Literal["on"] + +# A PlanarRelation with B, then A and B must have the same parent node. +PlanarRelationType = Literal["left_of", "right_of", "in_front_of", "behind"] +SceneConstraintType = SupportRelationType | PlanarRelationType +OrientationState = Literal["standing", "lying"] + + +@dataclass +class SceneGraphNode: + """One object node in the edit-time scene hierarchy. + + ``orientation_state`` is an image-derived placement semantic, rather than + an edge to the node itself or an exact three-dimensional transform. + """ + + object_id: str + parent_id: str | None + parent_relation: SupportRelationType | None = None + table_region: TableRegion | None = None + # Preserves image-observed placement semantics for later pose refinement. + orientation_state: OrientationState | None = None + + def __post_init__(self) -> None: + """Validate local node fields before graph-level checks.""" + if not self.object_id: + raise ValueError("object_id must be non-empty.") + if self.table_region not in {None, *TABLE_REGIONS}: + raise ValueError("table_region is invalid.") + if self.orientation_state not in {None, "standing", "lying"}: + raise ValueError("orientation_state is invalid.") + # If the node is the table. + if self.object_id == TABLE_OBJECT_ID: + if self.parent_id is not None: + raise ValueError("table must not have a parent.") + if self.parent_relation is not None: + raise ValueError("table must not have a parent relation.") + if self.orientation_state is not None: + raise ValueError("table must not have an orientation state.") + # If the node is not the table. + elif self.parent_id is None: + raise ValueError("non-table nodes must have a parent.") + elif self.parent_relation not in {None, "on"}: + raise ValueError("non-table nodes must be on their parent.") + + def to_dict(self) -> dict[str, object]: + """Serialize this node for scene graph debugging artifacts.""" + return { + "object_id": self.object_id, + "parent_id": self.parent_id, + "parent_relation": self.parent_relation, + "table_region": self.table_region, + "orientation_state": self.orientation_state, + } + + +@dataclass +class SceneGraphRelation: + """One edit-time spatial relation between two non-table objects.""" + + source_id: str + relation: PlanarRelationType + target_id: str + + def __post_init__(self) -> None: + """Validate local relation fields before graph-level checks.""" + if not self.source_id or not self.target_id: + raise ValueError("relation endpoints must be non-empty.") + if self.source_id == self.target_id: + raise ValueError("relation endpoints must be different.") + + def to_dict(self) -> dict[str, object]: + """Serialize this planar relation for scene graph debugging artifacts.""" + return { + "source_id": self.source_id, + "relation": self.relation, + "target_id": self.target_id, + } + + +@dataclass +class SceneGraph: + """Layered support graph plus planar relations for scene editing.""" + + nodes: list[SceneGraphNode] = field(default_factory=list) + relations: list[SceneGraphRelation] = field(default_factory=list) + validate_on_refresh: bool = True # Validate after each automatic refresh. + + def __post_init__(self) -> None: + """Normalize new graphs so downstream stages see canonical constraints.""" + self.refresh() + + def refresh(self) -> None: + """Normalize the graph and optionally validate semantic constraints.""" + # First normalize then validate (if applicable). + self.normalize() + if self.validate_on_refresh: + self.validate() + + def node_by_id(self) -> dict[str, SceneGraphNode]: + """Return nodes keyed by object id, raising on duplicate ids.""" + nodes_by_id: dict[str, SceneGraphNode] = {} + for node in self.nodes: + if node.object_id in nodes_by_id: + raise ValueError(f"Duplicate scene graph node: {node.object_id}") + nodes_by_id[node.object_id] = node + return nodes_by_id + + def remove_nodes(self, object_ids: set[str]) -> None: + """Remove nodes and their incident planar relations, then validate.""" + # If no node to be removed, return directly. + if not object_ids: + return + if TABLE_OBJECT_ID in object_ids: + raise ValueError("The table cannot be removed from a scene graph.") + unknown_object_ids = object_ids - set(self.node_by_id()) + if unknown_object_ids: + raise ValueError( + f"Cannot remove unknown scene graph nodes: {sorted(unknown_object_ids)}" + ) + + # Removing every incident relation prevents dangling planar endpoints. + self.nodes = [node for node in self.nodes if node.object_id not in object_ids] + self.relations = [ + relation + for relation in self.relations + if relation.source_id not in object_ids + and relation.target_id not in object_ids + ] + # Refresh. + self.refresh() + + def add_node(self, node: SceneGraphNode) -> None: + """Add one node and validate the resulting graph.""" + if node.object_id in self.node_by_id(): + raise ValueError(f"Duplicate scene graph node: {node.object_id}") + self.nodes.append(node) + self.refresh() + + def apply_updates( + self, + *, + deleted_object_ids: set[str], + added_object_ids: list[str], + added_orientation_states_by_id: dict[str, OrientationState | None], + on_parent_updates: list[tuple[str, str, TableRegion | None]], + planar_relation_updates: list[tuple[str, PlanarRelationType, str]], + ) -> None: + """Apply one atomic batch of node and relationship updates.""" + if TABLE_OBJECT_ID in deleted_object_ids: + raise ValueError("The table cannot be removed from a scene graph.") + + existing_object_ids = set(self.node_by_id()) + unknown_object_ids = deleted_object_ids - existing_object_ids + if unknown_object_ids: + raise ValueError( + f"Cannot remove unknown scene graph nodes: {sorted(unknown_object_ids)}" + ) + + # Delete all requested nodes before resolving new parents and relations. + self.nodes = [ + node for node in self.nodes if node.object_id not in deleted_object_ids + ] + self.relations = [ + relation + for relation in self.relations + if relation.source_id not in deleted_object_ids + and relation.target_id not in deleted_object_ids + ] + + remaining_object_ids = set(self.node_by_id()) + if len(added_object_ids) != len(set(added_object_ids)): + raise ValueError("Added scene graph node ids must be unique.") + duplicate_object_ids = set(added_object_ids) & remaining_object_ids + if duplicate_object_ids: + raise ValueError( + f"Duplicate scene graph nodes: {sorted(duplicate_object_ids)}" + ) + if set(added_orientation_states_by_id) != set(added_object_ids): + raise ValueError("Added orientation states must match added node ids.") + + # New nodes default to the table; later updates replace that parent when needed. + self.nodes.extend( + SceneGraphNode( + object_id=object_id, + parent_id=TABLE_OBJECT_ID, + parent_relation="on", + orientation_state=added_orientation_states_by_id[object_id], + ) + for object_id in added_object_ids + ) + + # Apply support-parent changes before planar updates need the final parent. + for object_id, parent_id, table_region in on_parent_updates: + self._set_on_parent( + object_id=object_id, + parent_id=parent_id, + table_region=table_region, + ) + + # Resolve chained planar parent inheritance before adding final relations. + self._resolve_planar_parent_updates(planar_relation_updates) + for source_id, relation, target_id in planar_relation_updates: + self._clear_incident_planar_relations(source_id) + self.relations.append( + SceneGraphRelation( + source_id=source_id, + relation=relation, + target_id=target_id, + ) + ) + + # Normalize inverse relations and reject invalid final graph constraints. + self.refresh() + + def _set_on_parent( + self, + *, + object_id: str, + parent_id: str, + table_region: TableRegion | None = None, + ) -> None: + """Replace one node's support parent and stale planar constraints.""" + nodes_by_id = self.node_by_id() + if object_id == TABLE_OBJECT_ID: + raise ValueError("The table cannot be moved onto another object.") + if object_id not in nodes_by_id or parent_id not in nodes_by_id: + raise ValueError( + "Parent updates must reference existing scene graph nodes." + ) + if object_id == parent_id: + raise ValueError("A scene graph node cannot be its own parent.") + + node = nodes_by_id[object_id] + node.parent_id = parent_id + node.parent_relation = "on" + node.table_region = table_region + self._clear_incident_planar_relations(object_id) + + def _resolve_planar_parent_updates( + self, + planar_relation_updates: list[tuple[str, PlanarRelationType, str]], + ) -> None: + """Make every planar source share its target's final support parent.""" + for _ in range(len(planar_relation_updates)): + changed = False + for source_id, _, target_id in planar_relation_updates: + nodes_by_id = self.node_by_id() + if source_id == TABLE_OBJECT_ID: + raise ValueError("The table cannot have a planar relation.") + if source_id not in nodes_by_id or target_id not in nodes_by_id: + raise ValueError( + "Planar updates must reference existing scene graph nodes." + ) + target_parent_id = nodes_by_id[target_id].parent_id + if target_parent_id is None: + raise ValueError("Planar relation targets must have a parent.") + source = nodes_by_id[source_id] + if source.parent_id != target_parent_id: + source.parent_id = target_parent_id + source.parent_relation = "on" + source.table_region = None + changed = True + if not changed: + return + + def _clear_incident_planar_relations(self, object_id: str) -> None: + """Remove planar constraints invalidated when one node changes parent.""" + self.relations = [ + relation + for relation in self.relations + if relation.source_id != object_id and relation.target_id != object_id + ] + + def normalize(self) -> None: + """Materialize inverse planar relations and remove duplicates.""" + self._materialize_inverse_planar_relations() + self._deduplicate_relations() + + def layer_by_id(self) -> dict[str, int]: + """Return the layer depth of each node inferred from parent links.""" + # Build fast lookup tables before walking the table-rooted tree. + nodes_by_id = self.node_by_id() + children_by_parent = self._children_by_parent() + layers: dict[str, int] = {} + visiting: set[str] = set() + visited: set[str] = set() + + def visit(node_id: str, layer: int) -> None: + # A node already on the recursion path means the parent chain loops. + if node_id in visiting: + raise ValueError(f"Parent cycle detected at node: {node_id}") + if node_id in visited: + return + # Missing parent nodes cannot contribute a valid table-rooted layer. + if node_id not in nodes_by_id: + raise ValueError(f"Parent node does not exist: {node_id}") + + visiting.add(node_id) + layers[node_id] = layer + # Children are exactly one support level above their parent. + for child in children_by_parent.get(node_id, []): + visit(child.object_id, layer + 1) + visiting.remove(node_id) + visited.add(node_id) + + if TABLE_OBJECT_ID not in nodes_by_id: + raise ValueError("Scene graph must contain a table node.") + visit(TABLE_OBJECT_ID, 0) + return layers + + def derive_constraints(self) -> list[dict[str, str]]: + """Return support constraints plus materialized planar relations.""" + self.refresh() + constraints: list[dict[str, str]] = [] + for node in self.nodes: + if node.parent_id is None: + continue + # Parent links become direct support constraints. + constraints.append( + self._constraint_dict( + source_id=node.object_id, + relation=node.parent_relation, + target_id=node.parent_id, + ), + ) + for relation in self.relations: + # Inverse planar relations are already stored during normalization. + constraints.append( + self._constraint_dict( + source_id=relation.source_id, + relation=relation.relation, + target_id=relation.target_id, + ), + ) + return self._deduplicate_constraints(constraints) + + def validate(self) -> None: + """Validate hierarchy, table regions, and planar relation constraints.""" + # id -> Node mapping. + nodes_by_id = self.node_by_id() + # Table node must exist. + if TABLE_OBJECT_ID not in nodes_by_id: + raise ValueError("Scene graph must contain a table node.") + + # Validate the table-rooted support tree before checking sibling relations. + for node in self.nodes: + if node.object_id == TABLE_OBJECT_ID: + if node.table_region is not None: + raise ValueError("table must not have a table_region.") + continue + parent = nodes_by_id.get(node.parent_id) + # Parent must exist, except for the table. (root node) + if parent is None: + raise ValueError(f"Parent node does not exist: {node.parent_id}") + if node.parent_relation is None: + raise ValueError( + f"Node {node.object_id} must define its parent relation." + ) + # Table regions are only valid for objects directly on the table. + if node.table_region is not None and node.parent_id != TABLE_OBJECT_ID: + raise ValueError("table_region is only valid for objects on the table.") + # Get id -> layer mapping. + layers = self.layer_by_id() + if len(layers) != len(nodes_by_id): + raise ValueError("All scene graph nodes must be reachable from the table.") + # Validate planar relations between nodes with the same parent. + for relation in self.relations: + source = nodes_by_id.get(relation.source_id) + target = nodes_by_id.get(relation.target_id) + if source is None or target is None: + raise ValueError("Planar relation endpoint does not exist.") + if source.parent_id != target.parent_id: + raise ValueError("Planar relation endpoints must share one parent.") + if source.parent_relation != "on" or target.parent_relation != "on": + raise ValueError("Planar relation endpoints must be on their parent.") + # Validate and imply planar relation. + self._validate_planar_relation_conflicts() + + def to_dict(self) -> dict[str, object]: + """Serialize the normalized graph state.""" + self.refresh() + return { + "nodes": [node.to_dict() for node in self.nodes], + "relations": [relation.to_dict() for relation in self.relations], + } + + def _children_by_parent(self) -> dict[str, list[SceneGraphNode]]: + children_by_parent: dict[str, list[SceneGraphNode]] = {} + for node in self.nodes: + if node.parent_id is not None: + children_by_parent.setdefault(node.parent_id, []).append(node) + return children_by_parent + + def _deduplicate_relations(self) -> None: + """Remove duplicate planar relations while preserving the first occurrence.""" + deduplicated: list[SceneGraphRelation] = [] + seen: set[tuple[str, PlanarRelationType, str]] = set() + for relation in self.relations: + key = (relation.source_id, relation.relation, relation.target_id) + # Only identical triples are duplicates; inverse relations are both retained. + if key in seen: + continue + seen.add(key) + deduplicated.append(relation) + self.relations = deduplicated + + def _materialize_inverse_planar_relations(self) -> None: + """Add the inverse of every planar relation to the graph.""" + inverse_relations = [ + SceneGraphRelation( + source_id=relation.target_id, + relation=self._inverse_planar_relation(relation.relation), + target_id=relation.source_id, + ) + for relation in self.relations + ] + self.relations.extend(inverse_relations) + + def _deduplicate_constraints( + self, + constraints: list[dict[str, str]], + ) -> list[dict[str, str]]: + deduplicated: list[dict[str, str]] = [] + seen: set[tuple[str, SceneConstraintType, str]] = set() + for constraint in constraints: + key = ( + constraint["source_id"], + constraint["relation"], + constraint["target_id"], + ) + if key in seen: + continue + seen.add(key) + deduplicated.append(constraint) + return deduplicated + + def _constraint_dict( + self, + *, + source_id: str, + relation: SceneConstraintType, + target_id: str, + ) -> dict[str, str]: + return { + "source_id": source_id, + "relation": relation, + "target_id": target_id, + } + + def _validate_planar_relation_conflicts(self) -> None: + implied_relations: dict[tuple[str, str], PlanarRelationType] = {} + for relation in self.relations: + self._add_implied_planar_relation( + implied_relations, + source_id=relation.source_id, + relation=relation.relation, + target_id=relation.target_id, + ) + self._add_implied_planar_relation( + implied_relations, + source_id=relation.target_id, + relation=self._inverse_planar_relation(relation.relation), + target_id=relation.source_id, + ) + + def _add_implied_planar_relation( + self, + implied_relations: dict[tuple[str, str], PlanarRelationType], + *, + source_id: str, + relation: PlanarRelationType, + target_id: str, + ) -> None: + key = (source_id, target_id) + existing_relation = implied_relations.get(key) + if existing_relation is not None and existing_relation != relation: + raise ValueError( + f"Conflicting planar relations: {source_id} " + f"{existing_relation} and {relation} {target_id}" + ) + implied_relations[key] = relation + + @classmethod + def _inverse_planar_relation( + cls, + relation: PlanarRelationType, + ) -> PlanarRelationType: + if relation == "left_of": + return "right_of" + if relation == "right_of": + return "left_of" + if relation == "in_front_of": + return "behind" + return "in_front_of" diff --git a/embodichain/gen_sim/scene_engine/core/scene_object.py b/embodichain/gen_sim/scene_engine/core/scene_object.py index d9a837e87..2b868e3c3 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_object.py +++ b/embodichain/gen_sim/scene_engine/core/scene_object.py @@ -66,6 +66,10 @@ class SceneObject: rot: list[float] | None = None # Final y-up Euler XYZ rotation in degrees. pos: list[float] | None = None # Final y-up world position in metres. scale: list[float] | None = None # Final y-up object scale. + center_xy: list[float] | None = None # Z-up table-frame XY AABB center. + support_surface_z: float | None = None # Detected tabletop height in z-up. + support_contour_xy: list[list[float]] | None = None # Outer support contour. + support_optimization_rect_xy: list[list[float]] | None = None # Safe XY rectangle. physics: ObjectPhysics | None = None # Assigned when SimReady processing succeeds. def to_dict(self) -> dict[str, object]: @@ -81,5 +85,9 @@ def to_dict(self) -> dict[str, object]: "rot": self.rot, "pos": self.pos, "scale": self.scale, + "center_xy": self.center_xy, + "support_surface_z": self.support_surface_z, + "support_contour_xy": self.support_contour_xy, + "support_optimization_rect_xy": self.support_optimization_rect_xy, "physics": self.physics.to_dict() if self.physics is not None else None, } diff --git a/embodichain/gen_sim/scene_engine/errors.py b/embodichain/gen_sim/scene_engine/errors.py new file mode 100644 index 000000000..dd23d2aa9 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/errors.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. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +__all__ = ["SceneServiceError"] + + +class SceneServiceError(RuntimeError): + """A transient or remote Scene Engine service failure.""" diff --git a/embodichain/gen_sim/scene_engine/pipeline/__init__.py b/embodichain/gen_sim/scene_engine/pipeline/__init__.py index 015c41510..ecf448d22 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/__init__.py +++ b/embodichain/gen_sim/scene_engine/pipeline/__init__.py @@ -16,4 +16,26 @@ from __future__ import annotations -__all__: list[str] = [] +from .api import ( + SCENE_BLUEPRINT_SCHEMA, + SCENE_EDIT_BLUEPRINT_SCHEMA, + SceneBlueprintPackage, + SceneEditBlueprintPackage, + SceneMaterialization, + analyze_edit, + analyze_image, + materialize_blueprint, + materialize_edit, +) + +__all__ = [ + "SCENE_BLUEPRINT_SCHEMA", + "SCENE_EDIT_BLUEPRINT_SCHEMA", + "SceneBlueprintPackage", + "SceneEditBlueprintPackage", + "SceneMaterialization", + "analyze_edit", + "analyze_image", + "materialize_blueprint", + "materialize_edit", +] diff --git a/embodichain/gen_sim/scene_engine/pipeline/api.py b/embodichain/gen_sim/scene_engine/pipeline/api.py new file mode 100644 index 000000000..6f219615c --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/api.py @@ -0,0 +1,351 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Auditable stage boundaries for Scene Engine generation and editing.""" + +from __future__ import annotations + +from copy import deepcopy +from dataclasses import dataclass +import hashlib +import json +from pathlib import Path +from typing import Any, Final + +from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( + GeometryGenerationClient, +) +from embodichain.gen_sim.scene_engine.clients.image_generation import ( + ImageGenerationClient, +) +from embodichain.gen_sim.scene_engine.clients.image_segmentation import ( + ImageSegmentationClient, +) +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_edit_plan import SceneEditPlan +from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraph +from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) +from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_asset_preparation import ( + prepare_scene_edit_assets, +) +from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_layout_generation import ( + edit_layout, +) +from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_understanding import ( + understand_scene_edit, +) +from embodichain.gen_sim.scene_engine.pipeline.generation.scene_generation import ( + generate_scene_and_refine, +) +from embodichain.gen_sim.scene_engine.pipeline.generation.scene_understanding import ( + understand_scene, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_exporter import SceneExporter +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_importer import ( + SceneExportImporter, +) +from embodichain.utils.logger import log_info + +__all__ = [ + "SCENE_BLUEPRINT_SCHEMA", + "SCENE_EDIT_BLUEPRINT_SCHEMA", + "SceneBlueprintPackage", + "SceneEditBlueprintPackage", + "SceneMaterialization", + "analyze_edit", + "analyze_image", + "materialize_blueprint", + "materialize_edit", +] + +SCENE_BLUEPRINT_SCHEMA: Final = "embodichain.scene-blueprint/v1" +SCENE_EDIT_BLUEPRINT_SCHEMA: Final = "embodichain.scene-edit-blueprint/v1" + + +@dataclass(frozen=True) +class SceneBlueprintPackage: + """In-process scene semantics plus their persisted audit document.""" + + blueprint_id: str + image_path: Path + output_root: Path + manifest_path: Path + scene: Scene + scene_graph: SceneGraph + + +@dataclass(frozen=True) +class SceneEditBlueprintPackage: + """Validated edit intent before added assets and layout are materialized.""" + + blueprint_id: str + edit_prompt: str + output_root: Path + manifest_path: Path + scene_edit_plan: SceneEditPlan + updated_scene_graph: SceneGraph + + +@dataclass(frozen=True) +class SceneMaterialization: + """One exported materialized scene revision.""" + + scene: Scene + scene_graph: SceneGraph + output_root: Path + scene_config_path: Path + + +def analyze_image( + image_path: str | Path, + output_root: str | Path, + *, + vlm_client: OpenAICompatibleVLM | None = None, + image_segmentation_client: ImageSegmentationClient | None = None, +) -> SceneBlueprintPackage: + """Understand an image and persist the pre-generation semantic blueprint.""" + resolved_image = Path(image_path).expanduser().resolve() + resolved_output = Path(output_root).expanduser().resolve() + resolved_output.mkdir(parents=True, exist_ok=True) + effective_vlm = vlm_client or OpenAICompatibleVLM.from_dotenv() + segmentation = image_segmentation_client or ImageSegmentationClient.from_dotenv() + owns_segmentation = image_segmentation_client is None + log_info("Starting Scene Understanding") + try: + segmentation.check_health() + scene, scene_graph = understand_scene( + scene=Scene(), + image_path=resolved_image, + output_root=resolved_output, + vlm_client=effective_vlm, + image_segmentation_client=segmentation, + ) + finally: + if owns_segmentation: + segmentation.close() + log_info("Completed Scene Understanding") + + payload = { + "schema_version": SCENE_BLUEPRINT_SCHEMA, + "image_path": resolved_image.as_posix(), + "scene": scene.to_dict(), + "scene_graph": scene_graph.to_dict(), + "artifacts": _artifact_records(resolved_output / "scene_understanding"), + } + blueprint_id = _canonical_hash(payload) + document = {**payload, "blueprint_id": blueprint_id} + manifest_path = resolved_output / "scene_blueprint.json" + _write_json(manifest_path, document) + return SceneBlueprintPackage( + blueprint_id=blueprint_id, + image_path=resolved_image, + output_root=resolved_output, + manifest_path=manifest_path, + scene=scene, + scene_graph=scene_graph, + ) + + +def materialize_blueprint( + blueprint: SceneBlueprintPackage, + *, + vlm_client: OpenAICompatibleVLM | None = None, + geometry_generation_client: GeometryGenerationClient | None = None, + seed: int | None = None, +) -> SceneMaterialization: + """Generate assets and layout for one image-derived blueprint.""" + scene = deepcopy(blueprint.scene) + scene_graph = deepcopy(blueprint.scene_graph) + effective_vlm = vlm_client or OpenAICompatibleVLM.from_dotenv() + geometry = geometry_generation_client or GeometryGenerationClient.from_dotenv() + owns_geometry = geometry_generation_client is None + log_info("Starting Objects + Coarse Layout Generation") + try: + geometry.check_health() + scene = generate_scene_and_refine( + image_path=blueprint.image_path, + output_root=blueprint.output_root, + scene=scene, + scene_graph=scene_graph, + geometry_generation_client=geometry, + vlm_client=effective_vlm, + seed=seed, + ) + finally: + if owns_geometry: + geometry.close() + log_info("Completed Objects + Coarse Layout Generation") + return _export_materialization( + scene=scene, + scene_graph=scene_graph, + output_root=blueprint.output_root, + ) + + +def analyze_edit( + *, + output_root: str | Path, + edit_prompt: str, + vlm_client: OpenAICompatibleVLM | None = None, +) -> SceneEditBlueprintPackage: + """Interpret and persist one edit against an already generated scene.""" + resolved_output = Path(output_root).expanduser().resolve() + normalized_prompt = str(edit_prompt).strip() + if not normalized_prompt: + raise ValueError("Edit prompt must not be empty.") + scene, scene_graph = SceneExportImporter( + output_root=resolved_output + ).import_scene_and_graph() + effective_vlm = vlm_client or OpenAICompatibleVLM.from_dotenv() + log_info("Starting Edit Understanding") + scene_edit_plan, updated_scene_graph = understand_scene_edit( + scene=scene, + scene_graph=scene_graph, + edit_prompt=normalized_prompt, + vlm_client=effective_vlm, + ) + log_info("Completed Edit Understanding") + payload = { + "schema_version": SCENE_EDIT_BLUEPRINT_SCHEMA, + "edit_prompt": normalized_prompt, + "scene_edit_plan": scene_edit_plan.to_dict(), + "updated_scene_graph": updated_scene_graph.to_dict(), + } + blueprint_id = _canonical_hash(payload) + manifest_path = resolved_output / "scene_edit" / "scene_edit_blueprint.json" + _write_json(manifest_path, {**payload, "blueprint_id": blueprint_id}) + return SceneEditBlueprintPackage( + blueprint_id=blueprint_id, + edit_prompt=normalized_prompt, + output_root=resolved_output, + manifest_path=manifest_path, + scene_edit_plan=scene_edit_plan, + updated_scene_graph=updated_scene_graph, + ) + + +def materialize_edit( + blueprint: SceneEditBlueprintPackage, + *, + vlm_client: OpenAICompatibleVLM | None = None, + image_generation_client: ImageGenerationClient | None = None, + geometry_generation_client: GeometryGenerationClient | None = None, + image_segmentation_client: ImageSegmentationClient | None = None, + seed: int | None = None, +) -> SceneMaterialization: + """Generate added assets, apply layout edits, and export the new revision.""" + scene_edit_plan = deepcopy(blueprint.scene_edit_plan) + updated_scene_graph = deepcopy(blueprint.updated_scene_graph) + effective_vlm = vlm_client or OpenAICompatibleVLM.from_dotenv() + image_generation = image_generation_client or ImageGenerationClient.from_dotenv() + geometry = geometry_generation_client or GeometryGenerationClient.from_dotenv() + segmentation = image_segmentation_client or ImageSegmentationClient.from_dotenv() + owned_clients = ( + (image_generation, image_generation_client is None), + (geometry, geometry_generation_client is None), + (segmentation, image_segmentation_client is None), + ) + log_info("Starting Objects Preparation") + try: + for client, _ in owned_clients: + client.check_health() + added_assets = prepare_scene_edit_assets( + scene_edit_plan=scene_edit_plan, + output_root=blueprint.output_root, + image_generation_client=image_generation, + geometry_generation_client=geometry, + image_segmentation_client=segmentation, + vlm_client=effective_vlm, + seed=seed, + ) + finally: + for client, owned in owned_clients: + if owned: + client.close() + log_info("Completed Objects Preparation") + log_info("Starting Layout Generation") + scene = edit_layout( + scene=scene_edit_plan.scene, + scene_edit_plan=scene_edit_plan, + updated_scene_graph=updated_scene_graph, + added_assets=added_assets, + output_root=blueprint.output_root, + ) + log_info("Completed Layout Generation") + return _export_materialization( + scene=scene, + scene_graph=updated_scene_graph, + output_root=blueprint.output_root, + ) + + +def _export_materialization( + *, + scene: Scene, + scene_graph: SceneGraph, + output_root: Path, +) -> SceneMaterialization: + log_info("Starting Scene Export") + scene_config_path = SceneExporter( + scene=scene, + scene_graph=scene_graph, + output_root=output_root, + ).export() + log_info("Completed Scene Export") + return SceneMaterialization( + scene=scene, + scene_graph=scene_graph, + output_root=output_root, + scene_config_path=scene_config_path, + ) + + +def _artifact_records(root: Path) -> list[dict[str, Any]]: + if not root.is_dir(): + return [] + records = [] + for path in sorted(item for item in root.rglob("*") if item.is_file()): + records.append( + { + "path": path.resolve().as_posix(), + "sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + "size": path.stat().st_size, + } + ) + return records + + +def _canonical_hash(value: Any) -> str: + payload = json.dumps( + value, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def _write_json(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text( + json.dumps(value, indent=2, ensure_ascii=False, allow_nan=False) + "\n", + encoding="utf-8", + ) + temporary.replace(path) diff --git a/embodichain/gen_sim/scene_engine/pipeline/edit.py b/embodichain/gen_sim/scene_engine/pipeline/edit.py new file mode 100644 index 000000000..ec993c8dc --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/edit.py @@ -0,0 +1,44 @@ +# ---------------------------------------------------------------------------- +# 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 pathlib import Path + +from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) +from embodichain.gen_sim.scene_engine.pipeline.api import ( + analyze_edit, + materialize_edit, +) + + +def edit_scene( + *, + output_root: str | Path, + edit_prompt: str, +) -> None: + """Apply one text edit instruction to an existing Scene Engine output.""" + resolved_output_root = Path(output_root).expanduser().resolve() + vlm_client = OpenAICompatibleVLM.from_dotenv() + blueprint = analyze_edit( + output_root=resolved_output_root, + edit_prompt=edit_prompt, + vlm_client=vlm_client, + ) + materialize_edit(blueprint, vlm_client=vlm_client) + return None diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/__init__.py b/embodichain/gen_sim/scene_engine/pipeline/editing/__init__.py new file mode 100644 index 000000000..015c41510 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/__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 + +__all__: list[str] = [] diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py new file mode 100644 index 000000000..c1e3516d2 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py @@ -0,0 +1,341 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +import shutil + +from PIL import Image + +from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( + GeometryGenerationClient, +) +from embodichain.gen_sim.scene_engine.clients.image_generation import ( + ImageGenerationClient, +) +from embodichain.gen_sim.scene_engine.clients.image_segmentation import ( + ImageSegmentationClient, +) +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_edit_plan import SceneEditPlan +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.image_segmentation_utils import ( + MaskCandidate, + build_mask_candidates, + invert_mask_if_foreground_is_off_center, + save_binary_mask, + union_overlapping_mask_candidates, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor import ( + SimReadyProcessor, + SimReadyProcessorConfig, +) + +__all__ = ["prepare_scene_edit_assets"] + + +@dataclass(frozen=True) +class _AddedAssetInfo: + """Semantic information needed while preparing one newly added asset.""" + + object_id: str + category: str + name: str + description: str + + +def prepare_scene_edit_assets( + *, + scene_edit_plan: SceneEditPlan, + output_root: str | Path, + image_generation_client: ImageGenerationClient, + geometry_generation_client: GeometryGenerationClient, + image_segmentation_client: ImageSegmentationClient, + vlm_client: OpenAICompatibleVLM | None = None, + seed: int | None = None, +) -> list[SceneObject]: + """Prepare and return SimReady assets required by add operations.""" + # Prepare descriptions for all newly added objects. + added_asset_descriptions = _collect_added_asset_descriptions(scene_edit_plan) + # Skip asset generation when the edit plan only moves or deletes existing objects. + if not added_asset_descriptions: + return [] + + # Recreate this stage only when new assets need image, segmentation, and geometry outputs. + stage_output_root = ( + Path(output_root).expanduser().resolve() / "scene_editing" / "asset_preparation" + ) + if stage_output_root.exists(): + shutil.rmtree(stage_output_root) + stage_output_root.mkdir(parents=True, exist_ok=True) + + generated_asset_images = _generate_added_asset_images( + added_asset_descriptions=added_asset_descriptions, + stage_output_root=stage_output_root, + image_generation_client=image_generation_client, + seed=seed, + ) + generated_asset_masks = _segment_generated_added_asset_images( + added_asset_descriptions=added_asset_descriptions, + generated_asset_images=generated_asset_images, + stage_output_root=stage_output_root, + image_segmentation_client=image_segmentation_client, + ) + + generated_asset_glbs = _generate_added_assets_coarse_geometry( + generated_asset_images=generated_asset_images, + generated_asset_masks=generated_asset_masks, + stage_output_root=stage_output_root, + geometry_generation_client=geometry_generation_client, + seed=seed, + ) + # Build a list of added SceneObjects. + added_assets = _build_added_scene_objects( + added_asset_descriptions=added_asset_descriptions, + generated_asset_glbs=generated_asset_glbs, + ) + # The temporary scene contains only new assets because the existing table is reused. + tmp_scene = Scene(objects=added_assets) + simready_processor = SimReadyProcessor( + scene=tmp_scene, + coarse_layout_by_id=_coarse_layouts_by_id(generated_asset_glbs), + coarse_geometry_root=stage_output_root / "coarse_geometry", + simready_geometry_root=stage_output_root / "simready_geometry", + # Every added asset uses the VLM's pose and post-pose XY footprint scale. + config=SimReadyProcessorConfig( + use_vlm_scale=vlm_client is not None, + use_vlm_rotation=vlm_client is not None, + # An explicit edit state overrides the default stable tabletop pose. + orientation_states_by_id={ + operation.object_id: operation.orientation_state + for operation in scene_edit_plan.operations + if operation.op == "add" + and operation.object_id is not None + and operation.orientation_state is not None + }, + ), + vlm_client=vlm_client, + ) + # process_assets() validates and processes assets only; it does not require a table. + simready_processor.process_assets() + # Canonical GLBs use identity edit-time poses; layout editing sets them later. + _reset_added_asset_layouts(added_assets) + return added_assets + + +def _build_added_scene_objects( + *, + added_asset_descriptions: list[_AddedAssetInfo], + generated_asset_glbs: list[tuple[str, Path]], +) -> list[SceneObject]: + """Build temporary SceneObjects from generated coarse GLBs.""" + glbs_by_id = dict(generated_asset_glbs) + if len(glbs_by_id) != len(generated_asset_glbs): + raise ValueError("Generated asset GLBs must use unique object ids.") + assets: list[SceneObject] = [] + for asset_info in added_asset_descriptions: + glb_path = glbs_by_id.get(asset_info.object_id) + if glb_path is None: + raise ValueError(f"Generated asset {asset_info.object_id!r} has no GLB.") + assets.append( + SceneObject( + id=asset_info.object_id, + kind="asset", + category=asset_info.category, + name=asset_info.name, + description=asset_info.description, + simready_glb_path=str(glb_path), + ) + ) + return assets + + +def _reset_added_asset_layouts(added_assets: list[SceneObject]) -> None: + """Reset added asset poses after SimReady canonicalization.""" + for asset in added_assets: + asset.rot = [0.0, 0.0, 0.0] + asset.pos = [0.0, 0.0, 0.0] + asset.scale = [1.0, 1.0, 1.0] + + +def _coarse_layouts_by_id( + generated_asset_glbs: list[tuple[str, Path]], +) -> dict[str, dict[str, object]]: + """Build edit-time layouts with fixed identity poses and scale.""" + return { + object_id: { + "rot": [0.0, 0.0, 0.0], + "pos": [0.0, 0.0, 0.0], + "scale": [1.0, 1.0, 1.0], + } + for object_id, _ in generated_asset_glbs + } + + +def _collect_added_asset_descriptions( + scene_edit_plan: SceneEditPlan, +) -> list[_AddedAssetInfo]: + """Return complete semantic information for add operations in plan order.""" + # Existing assets already have SimReady GLBs, so only add operations need assets. + added_asset_descriptions: list[_AddedAssetInfo] = [] + for operation in scene_edit_plan.operations: + if operation.op != "add": + continue + if ( + operation.object_id is None + or operation.category is None + or operation.name is None + or operation.description is None + ): + raise ValueError( + "Add operations must have an object_id, category, name, and description." + ) + added_asset_descriptions.append( + _AddedAssetInfo( + object_id=operation.object_id, + category=operation.category, + name=operation.name, + description=operation.description, + ) + ) + return added_asset_descriptions + + +def _generate_added_asset_images( + *, + added_asset_descriptions: list[_AddedAssetInfo], + stage_output_root: Path, + image_generation_client: ImageGenerationClient, + seed: int | None, +) -> list[tuple[str, Path]]: + """Generate one stable PNG for each new object description.""" + # Prepare a list. + generated_asset_images: list[tuple[str, Path]] = [] + # Create a subdir. + image_output_root = stage_output_root / "generated_images" + image_output_root.mkdir(parents=True, exist_ok=True) + + for index, asset_info in enumerate(added_asset_descriptions): + object_id = asset_info.object_id + # Stable object IDs preserve the image-to-asset mapping across later stages. + generation_kwargs = { + "prompt": asset_info.description, + "output_path": image_output_root / f"{object_id}.png", + } + if seed is not None: + generation_kwargs["seed"] = int(seed) + index + image_path = image_generation_client.generate_image_by_prompt( + **generation_kwargs + ) + generated_asset_images.append((object_id, image_path)) + return generated_asset_images + + +def _segment_generated_added_asset_images( + *, + added_asset_descriptions: list[_AddedAssetInfo], + generated_asset_images: list[tuple[str, Path]], + stage_output_root: Path, + image_segmentation_client: ImageSegmentationClient, +) -> list[tuple[str, Path]]: + """Segment each generated image with its description and return binary masks.""" + asset_info_by_id = { + asset_info.object_id: asset_info for asset_info in added_asset_descriptions + } + if len(asset_info_by_id) != len(added_asset_descriptions): + raise ValueError("Added asset descriptions must use unique object ids.") + + masks_output_root = stage_output_root / "generated_masks" + masks_output_root.mkdir(parents=True, exist_ok=True) + generated_asset_masks: list[tuple[str, Path]] = [] + for object_id, image_path in generated_asset_images: + asset_info = asset_info_by_id.get(object_id) + if asset_info is None: + raise ValueError(f"Generated image {object_id!r} has no description.") + + candidates: list[MaskCandidate] = [] + # Retry with simpler semantic prompts when the detailed description is not found. + for prompt in (asset_info.description, asset_info.name, asset_info.category): + candidates = union_overlapping_mask_candidates( + build_mask_candidates( + image_segmentation_client.segment_single_object( + image_path=image_path, + prompt=prompt, + ) + ), + min_iou=0.8, + ) + if candidates: + break + # A single generated object may still yield multiple SAM3 candidates; use the first one. + if not candidates: + raise ValueError( + f"Generated asset {object_id!r} produced no segmentation candidates." + ) + with Image.open(image_path) as image: + image_size = image.size + mask_path = save_binary_mask( + invert_mask_if_foreground_is_off_center(candidates[0]), + image_size=image_size, + output_path=masks_output_root / f"{object_id}_mask.png", + ) + generated_asset_masks.append((object_id, mask_path)) + return generated_asset_masks + + +def _generate_added_assets_coarse_geometry( + *, + generated_asset_images: list[tuple[str, Path]], + generated_asset_masks: list[tuple[str, Path]], + stage_output_root: Path, + geometry_generation_client: GeometryGenerationClient, + seed: int | None, +) -> list[tuple[str, Path]]: + """Generate one coarse GLB for each generated image and binary mask.""" + masks_by_id = dict(generated_asset_masks) + if len(masks_by_id) != len(generated_asset_masks): + raise ValueError("Generated asset masks must use unique object ids.") + if set(masks_by_id) != {object_id for object_id, _ in generated_asset_images}: + raise ValueError( + "Generated asset images and masks must have matching object ids." + ) + + geometry_output_root = stage_output_root / "coarse_geometry" + geometry_output_root.mkdir(parents=True, exist_ok=True) + generated_asset_glbs: list[tuple[str, Path]] = [] + for index, (object_id, image_path) in enumerate(generated_asset_images): + # Each generated object has its own color image, so it needs an individual request. + generation_kwargs = { + "image_path": image_path, + "object_masks": [(object_id, masks_by_id[object_id])], + "output_root": geometry_output_root, + } + if seed is not None: + generation_kwargs["seed"] = int(seed) + index + geometry_generation_client.generate_objects(**generation_kwargs) + glb_path = geometry_output_root / f"{object_id}.glb" + if not glb_path.is_file(): + raise FileNotFoundError( + f"Geometry generation did not produce a GLB for {object_id!r}: {glb_path}" + ) + generated_asset_glbs.append((object_id, glb_path)) + return generated_asset_glbs diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_layout_generation.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_layout_generation.py new file mode 100644 index 000000000..9c6b06410 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_layout_generation.py @@ -0,0 +1,77 @@ +# ---------------------------------------------------------------------------- +# 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 shutil +from pathlib import Path + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_edit_plan import SceneEditPlan +from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraph +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_constructor import ( + SceneLayoutConstructor, +) + + +def edit_layout( + *, + scene: Scene, + scene_edit_plan: SceneEditPlan, + updated_scene_graph: SceneGraph, + added_assets: list[SceneObject], + output_root: str | Path, +) -> Scene: + """Dispatch one edit-layout optimization from the goal scene graph.""" + formal_scene = scene + goal_scene_graph = updated_scene_graph + generated_scene_objects = added_assets + # Recreate this stage only when new assets need image, segmentation, and geometry outputs. + stage_output_root = ( + Path(output_root).expanduser().resolve() + / "scene_editing" + / "layout_optimization" + ) + if stage_output_root.exists(): + shutil.rmtree(stage_output_root) + stage_output_root.mkdir(parents=True, exist_ok=True) + + # Add and move operations are the only layout variables in this edit pass. + layout_variable_ids = { + operation.object_id + for operation in scene_edit_plan.operations + if operation.op in {"add", "move"} + } + if None in layout_variable_ids: + raise ValueError("Add and move operations must identify an object.") + layout_variable_ids = { + object_id for object_id in layout_variable_ids if object_id is not None + } + + # Optimize the layout constrained by the goal scene graph. + layout_constructor = SceneLayoutConstructor( + formal_scene=formal_scene, + goal_scene_graph=goal_scene_graph, + layout_variable_ids=layout_variable_ids, + generated_scene_objects=generated_scene_objects, + output_root=stage_output_root, + ) + # Optimize. + post_edit_scene = layout_constructor.construct() + + return post_edit_scene diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py new file mode 100644 index 000000000..92720c7ec --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py @@ -0,0 +1,492 @@ +# ---------------------------------------------------------------------------- +# 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 embodichain.gen_sim.scene_engine.core.scene_edit_plan import ( + SceneEditOperation, + SceneEditPlan, +) +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + OrientationState, + PlanarRelationType, + SceneGraph, + SceneGraphNode, + SceneGraphRelation, + TABLE_REGIONS, + TableRegion, +) +from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) + +_EDIT_SYSTEM_PROMPT = """You convert one user instruction into edits for an existing tabletop scene. + +Use an existing object ID only when it appears in the supplied Existing object +IDs list. IDs identify existing objects exactly; never invent, correct, or +renumber them. The table ID is "table" and cannot be moved or deleted. + +Each operation is one of: +1. move: move one existing object. object_id, target_id, and relation must all + be provided. +2. delete: delete one existing object. Only object_id is provided. +3. add: create one new object. object_id must be null. Provide a lower-case + singular snake_case category, name, and description. Multiple add operations + may have the same category and name; their final IDs are assigned by the + program in operation order. target_id and relation are either both provided + or both null. Set orientation_state to standing or lying only when the user + explicitly asks for that placement; otherwise set it to null so the object + uses its natural, physically stable tabletop pose. + +For every move and every positioned add, target_id must be an Existing object +ID and relation must be one of on, left_of, right_of, in_front_of, or behind. +When the target is the tabletop, use target_id "table", relation "on", and set +table_region to one of left_back, back_center, right_back, left_center, center, +right_center, left_front, front_center, or right_front. Do not use a planar +relation with the table. For non-table placement, table_region must be null. +If an add operation has no target_id and relation, it is placed on the table by +default and table_region must be null. +In the tabletop 9-grid, smaller x means left, larger x means right, smaller y +means back, and larger y means front: left_back is the upper-left/back cell, +back_center is the upper-center/back cell, right_back is the upper-right/back +cell, left_center/center/right_center are the middle row, and +left_front/front_center/right_front are the lower/front row. +Do not position a new object relative to another newly added object. + +Each existing object's center_xy is its center position [x, y] in the +table-frame Z-up world coordinate system. Smaller x is left, larger x is right, +larger y is in front, and smaller y is behind. Use center_xy only to disambiguate +references such as "the bottle on the left"; do not output coordinates. Express +the requested position using target_id and one allowed relation instead. + +For every newly added object, category is its lower-case singular snake_case +class. name contains only color, material, texture, shape, and object details. +description contains only visible category, material, color, texture, shape, +and structural details. name and description must not mention position, the +table, relations to any object, or orientation. orientation_state must be null +unless the user explicitly requests standing/upright/vertical or lying/flat/ +horizontal placement. Follow that explicit user intent even if it is not the +object's natural stable pose. + +Return JSON only: no Markdown, comments, or prose. Every operation must contain +exactly these fields: op, object_id, target_id, relation, table_region, category, +name, description, and orientation_state. Use null for every field that does not apply: +{ + "operations": [ + { + "op": "move", + "object_id": "bottle_001", + "target_id": "book_001", + "relation": "right_of", + "table_region": null, + "category": null, + "name": null, + "description": null, + "orientation_state": null + }, + { + "op": "delete", + "object_id": "cup_001", + "target_id": null, + "relation": null, + "table_region": null, + "category": null, + "name": null, + "description": null, + "orientation_state": null + }, + { + "op": "add", + "object_id": null, + "target_id": "table", + "relation": "on", + "table_region": "back_center", + "category": "orange", + "name": "small orange", + "description": "small round orange with a textured peel", + "orientation_state": null + }, + { + "op": "add", + "object_id": null, + "target_id": "book_001", + "relation": "right_of", + "table_region": null, + "category": "orange", + "name": "small orange", + "description": "small round orange with a textured peel", + "orientation_state": null + }, + { + "op": "add", + "object_id": null, + "target_id": null, + "relation": null, + "table_region": null, + "category": "bottle", + "name": "blue glass bottle", + "description": "tall transparent blue glass bottle with a narrow neck", + "orientation_state": "standing" + }, + { + "op": "add", + "object_id": null, + "target_id": null, + "relation": null, + "table_region": null, + "category": "fork", + "name": "silver metal fork", + "description": "four-tined silver stainless-steel fork with a plain handle", + "orientation_state": "lying" + } + ] +} +The two orange additions intentionally share category and name. The bottle +example represents an explicit user request to stand it upright, and the fork +example represents an explicit user request to lay it flat. Only add operations +may introduce a new non-null orientation_state. A move may use null or repeat +its existing orientation_state from the supplied scene metadata, but it must not +change that state. Delete operations must use null. Do not add fields beyond the +required schema.""" + + +def understand_scene_edit( + *, + scene: Scene, + scene_graph: SceneGraph, + edit_prompt: str, + vlm_client: OpenAICompatibleVLM, + json_max_attempts: int = 3, +) -> tuple[SceneEditPlan, SceneGraph]: + """Understand one text edit instruction for an existing scene.""" + edit_prompt = edit_prompt.strip() + if not edit_prompt: + raise ValueError("Edit prompt must not be empty.") + if json_max_attempts < 1: + raise ValueError("json_max_attempts must be at least 1.") + # Give the VLM only the scene metadata needed to identify existing objects. + simplified_scene_info = _simplify_scene_info( + scene=scene, + scene_graph=scene_graph, + ) + operations = _vlm_understand_scene_edit( + scene=scene, + edit_prompt=edit_prompt, + simplified_scene_info=simplified_scene_info, + vlm_client=vlm_client, + json_max_attempts=json_max_attempts, + ) + + # SceneEditPlan validates all references against the immutable input scene graph. + scene_edit_plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=operations, + ) + updated_scene_graph = _build_updated_scene_graph( + scene_graph=scene_graph, + scene_edit_plan=scene_edit_plan, + ) + return scene_edit_plan, updated_scene_graph + + +def _build_updated_scene_graph( + *, + scene_graph: SceneGraph, + scene_edit_plan: SceneEditPlan, +) -> SceneGraph: + """Build and validate the target graph implied by one edit plan.""" + # Copy every mutable graph value so the pre-edit graph remains unchanged. + updated_scene_graph = SceneGraph( + nodes=[ + SceneGraphNode( + object_id=node.object_id, + parent_id=node.parent_id, + parent_relation=node.parent_relation, + table_region=node.table_region, + orientation_state=node.orientation_state, + ) + for node in scene_graph.nodes + ], + relations=[ + SceneGraphRelation( + source_id=relation.source_id, + relation=relation.relation, + target_id=relation.target_id, + ) + for relation in scene_graph.relations + ], + validate_on_refresh=scene_graph.validate_on_refresh, + ) + _apply_scene_edit_plan_to_scene_graph( + scene_graph=updated_scene_graph, + scene_edit_plan=scene_edit_plan, + ) + return updated_scene_graph + + +def _apply_scene_edit_plan_to_scene_graph( + *, + scene_graph: SceneGraph, + scene_edit_plan: SceneEditPlan, +) -> None: + """Apply the target graph updates implied by add and move operations.""" + deleted_object_ids: set[str] = set() + added_object_ids: list[str] = [] + added_orientation_states_by_id: dict[str, OrientationState | None] = {} + on_parent_updates: list[tuple[str, str, TableRegion | None]] = [] + planar_relation_updates: list[tuple[str, PlanarRelationType, str]] = [] + for operation in scene_edit_plan.operations: + if operation.op == "delete": + if operation.object_id is not None: + deleted_object_ids.add(operation.object_id) + continue + if operation.object_id is None: + raise ValueError("Add and move operations must have an object_id.") + if operation.op == "add": + added_object_ids.append(operation.object_id) + added_orientation_states_by_id[operation.object_id] = ( + operation.orientation_state + ) + if operation.target_id is None or operation.relation is None: + continue + if operation.relation == "on": + on_parent_updates.append( + ( + operation.object_id, + operation.target_id, + operation.table_region, + ) + ) + continue + planar_relation_updates.append( + (operation.object_id, operation.relation, operation.target_id) + ) + + # Apply all graph changes atomically so intermediate edit states need not be valid. + scene_graph.apply_updates( + deleted_object_ids=deleted_object_ids, + added_object_ids=added_object_ids, + added_orientation_states_by_id=added_orientation_states_by_id, + on_parent_updates=on_parent_updates, + planar_relation_updates=planar_relation_updates, + ) + + +def _simplify_scene_info( + *, + scene: Scene, + scene_graph: SceneGraph, +) -> dict[str, object]: + """Return the object metadata needed for edit instruction resolution.""" + table_regions_by_id = { + node.object_id: node.table_region for node in scene_graph.nodes + } + orientation_states_by_id = { + node.object_id: node.orientation_state for node in scene_graph.nodes + } + return { + "existing_object_ids": [scene_object.id for scene_object in scene.objects], + "objects": [ + { + "id": scene_object.id, + "category": scene_object.category, + "name": scene_object.name, + "description": scene_object.description, + "center_xy": scene_object.center_xy, + "table_region": table_regions_by_id.get(scene_object.id), + "orientation_state": orientation_states_by_id.get(scene_object.id), + } + for scene_object in scene.objects + ], + } + + +def _vlm_understand_scene_edit( + *, + scene: Scene, + edit_prompt: str, + simplified_scene_info: dict[str, object], + vlm_client: OpenAICompatibleVLM, + json_max_attempts: int, +) -> list[SceneEditOperation]: + """Return parsed edit operations from the VLM with assigned add IDs.""" + # Construct user prompt. + user_prompt = ( + f"User edit instruction:\n{edit_prompt}\n\n" + "Existing scene metadata:\n" + f"{json.dumps(simplified_scene_info, indent=2, ensure_ascii=False)}" + ) + last_error: ValueError | None = None + for _ in range(json_max_attempts): + response_text = vlm_client.complete( + system_prompt=_EDIT_SYSTEM_PROMPT, + user_prompt=user_prompt, + ) + try: + value = json.loads(_strip_json_code_fence(response_text)) + return _parse_scene_edit_operations(value, scene=scene) + except (json.JSONDecodeError, ValueError) as exc: + last_error = ValueError(f"VLM returned invalid scene edit JSON: {exc}") + continue + + assert last_error is not None + raise ValueError( + "VLM returned invalid scene edit JSON after " + f"{json_max_attempts} attempts: {last_error}" + ) from last_error + + +def _strip_json_code_fence(response_text: str) -> str: + """Remove one optional Markdown JSON fence from a VLM response.""" + stripped = response_text.strip() + if not stripped.startswith("```"): + return stripped + lines = stripped.splitlines() + if len(lines) < 3 or not lines[-1].strip().startswith("```"): + raise ValueError("VLM response contains an incomplete JSON code fence.") + return "\n".join(lines[1:-1]).strip() + + +def _parse_scene_edit_operations( + value: object, + *, + scene: Scene, +) -> list[SceneEditOperation]: + """Parse the strict VLM edit-draft schema into typed operations with add IDs.""" + if not isinstance(value, dict) or set(value) != {"operations"}: + raise ValueError("Scene edit draft must contain exactly operations.") + # Get and validate a list operation value. + operations_value = value["operations"] + if not isinstance(operations_value, list): + raise ValueError("Scene edit draft operations must be a list.") + + expected_keys = { + "op", + "object_id", + "target_id", + "relation", + "table_region", + "category", + "name", + "description", + "orientation_state", + } + # Get ids and counts of existing objects to assign new add IDs. + assigned_object_ids = {scene_object.id for scene_object in scene.objects} + category_counts = { + category: sum( + scene_object.category == category for scene_object in scene.objects + ) + for category in {scene_object.category for scene_object in scene.objects} + } + operations: list[SceneEditOperation] = [] + for value in operations_value: + if not isinstance(value, dict) or not isinstance(value.get("op"), str): + raise ValueError("Scene edit operations must contain a string op.") + op = value["op"] + if op not in {"add", "move", "delete"}: + raise ValueError("Scene edit operation op is invalid.") + if set(value) != expected_keys: + raise ValueError("Scene edit operations must use the required schema.") + object_id = _optional_string(value.get("object_id"), field_name="object_id") + category = _optional_string(value.get("category"), field_name="category") + orientation_state = _optional_orientation_state(value.get("orientation_state")) + if op == "add": + if object_id is not None: + raise ValueError("VLM add operations must set object_id to null.") + if category is None: + raise ValueError("VLM add operations must provide a category.") + # Add operation should generate new id here. + # Never believe the LLM could always generate a valid id. + object_id = _next_add_object_id( + category=category, + category_counts=category_counts, + assigned_object_ids=assigned_object_ids, + ) + operations.append( + SceneEditOperation( + op=op, + object_id=object_id, + target_id=_optional_string( + value.get("target_id"), field_name="target_id" + ), + relation=_optional_relation(value.get("relation")), + table_region=_optional_table_region(value.get("table_region")), + category=category, + name=_optional_string(value.get("name"), field_name="name"), + description=_optional_string( + value.get("description"), field_name="description" + ), + orientation_state=orientation_state, + ) + ) + return operations + + +def _next_add_object_id( + *, + category: str, + category_counts: dict[str, int], + assigned_object_ids: set[str], +) -> str: + """Assign the next available ID for one new object category.""" + index = category_counts.get(category, 0) + 1 + object_id = f"{category}_{index:03d}" + # In case the scene have orange_001 and orange_003. + while object_id in assigned_object_ids: + index += 1 + object_id = f"{category}_{index:03d}" + category_counts[category] = index + assigned_object_ids.add(object_id) + return object_id + + +def _optional_string(value: object, *, field_name: str) -> str | None: + if value is None: + return None + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"Scene edit operation {field_name} must be a string or null.") + return value.strip() + + +def _optional_relation(value: object) -> str | None: + if value is None: + return None + if value not in {"on", "left_of", "right_of", "in_front_of", "behind"}: + raise ValueError("Scene edit operation relation is invalid.") + return value + + +def _optional_table_region(value: object) -> TableRegion | None: + if value is None: + return None + if value not in TABLE_REGIONS: + raise ValueError("Scene edit operation table_region is invalid.") + return value + + +def _optional_orientation_state(value: object) -> OrientationState | None: + """Validate an optional explicit upright or lying edit intent.""" + if value is None: + return None + if value not in {"standing", "lying"}: + raise ValueError("Scene edit operation orientation_state is invalid.") + return value diff --git a/embodichain/gen_sim/scene_engine/pipeline/generate.py b/embodichain/gen_sim/scene_engine/pipeline/generate.py index 773a897bf..603591684 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generate.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generate.py @@ -22,20 +22,11 @@ from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( OpenAICompatibleVLM, ) -from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( - GeometryGenerationClient, +from embodichain.gen_sim.scene_engine.pipeline.api import ( + analyze_image, + materialize_blueprint, ) -from embodichain.gen_sim.scene_engine.pipeline.scene_understanding import ( - understand_scene, -) -from embodichain.utils.logger import log_info - -from embodichain.gen_sim.scene_engine.pipeline.scene_generation import ( - generate_scene_and_refine, -) -from embodichain.gen_sim.scene_engine.pipeline.utils.scene_exporter import SceneExporter - def generate_scene_from_image( image_path: str | Path, @@ -43,45 +34,10 @@ def generate_scene_from_image( ) -> Scene: """Generate the initial core scene state from an input image.""" resolved_output_root = Path(output_root).expanduser().resolve() - resolved_output_root.mkdir(parents=True, exist_ok=True) - - # Initialize the VLM client and the Scene data structure. vlm_client = OpenAICompatibleVLM.from_dotenv() - scene = Scene() - - # 1. Scene Understanding - log_info("Starting Scene Understanding") - scene = understand_scene( - scene=scene, - image_path=image_path, - output_root=resolved_output_root, + blueprint = analyze_image( + image_path, + resolved_output_root, vlm_client=vlm_client, ) - log_info("Completed Scene Understanding") - - # 2. Objects + Coarse Layout Generation - log_info("Starting Objects + Coarse Layout Generation") - # Load .env settings and fail if the Geometry Generation Server is unavailable. - geometry_generation_client = GeometryGenerationClient.from_dotenv() - try: - geometry_generation_client.check_health() # Error raising will happen internally. - scene = generate_scene_and_refine( - image_path=image_path, - output_root=resolved_output_root, - scene=scene, - geometry_generation_client=geometry_generation_client, - ) - finally: - geometry_generation_client.close() # Kill the session to avoid resource leaks. - log_info("Completed Objects + Coarse Layout Generation") - - # 3. Scene Export - log_info("Starting Scene Export") - scene_exporter = SceneExporter( - scene=scene, - output_root=resolved_output_root, - ) - scene_exporter.export() - log_info("Completed Scene Export") - - return scene + return materialize_blueprint(blueprint, vlm_client=vlm_client).scene diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/__init__.py b/embodichain/gen_sim/scene_engine/pipeline/generation/__init__.py new file mode 100644 index 000000000..015c41510 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/__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 + +__all__: list[str] = [] diff --git a/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py similarity index 64% rename from embodichain/gen_sim/scene_engine/pipeline/scene_generation.py rename to embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index d14ce364f..1884ba409 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -22,13 +22,19 @@ import shutil import numpy as np +from scipy.spatial.transform import Rotation import trimesh +from shapely.geometry import Polygon from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( GeometryGenerationClient, ) from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraph from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_support_clamp import ( AssetsGroupSupportClamp, ) @@ -38,8 +44,9 @@ from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_layout_optimizer import ( AssetsSupportLayoutOptimizer, ) -from embodichain.gen_sim.scene_engine.pipeline.utils.assets_gravity_settler import ( - AssetsGravitySettler, +from embodichain.gen_sim.scene_engine.pipeline.utils.gravity_settler import ( + GravitySettleBody, + GravitySettler, ) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( layout_object_to_transform_matrix, @@ -47,11 +54,9 @@ quaternion_wxyz_to_euler_xyz_degrees, transform_matrix_to_layout_object, ) -from embodichain.gen_sim.scene_engine.pipeline.utils.simready_scene_processor import ( - SimReadySceneProcessor, -) -from embodichain.gen_sim.scene_engine.pipeline.utils.table_support_surface import ( - TableSupportSurfaceDetector, +from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor import ( + SimReadyProcessor, + SimReadyProcessorConfig, ) from embodichain.utils.logger import log_info @@ -62,11 +67,16 @@ def generate_scene_and_refine( image_path: str | Path, output_root: str | Path, scene: Scene, + scene_graph: SceneGraph, *, geometry_generation_client: GeometryGenerationClient, + vlm_client: OpenAICompatibleVLM, + seed: int | None = None, ) -> Scene: resolved_image_path = _validate_image_path(image_path) + # Validate the scene graph before layout refinement consumes it. + scene_graph.validate() # Create stage output directory. stage_output_root = Path(output_root).expanduser().resolve() / "scene_generation" if stage_output_root.exists(): @@ -93,6 +103,7 @@ def generate_scene_and_refine( coarse_geometry_output_root=coarse_geometry_output_root, scene=scene, # Use the masks which are kept in the scene data structure. geometry_generation_client=geometry_generation_client, + seed=seed, ) # Simready all the assets(includes table). @@ -101,11 +112,26 @@ def generate_scene_and_refine( coarse_layout_by_id = { layout_object["id"]: layout_object for layout_object in coarse_layout } - simready_processor = SimReadySceneProcessor( + # Coarse poses already preserve lying and unconstrained assets; only standing + # assets need a VLM semantic-axis correction before later z-up calibration. + standing_orientation_states_by_id = { + node.object_id: node.orientation_state + for node in scene_graph.nodes + if node.orientation_state == "standing" + } + simready_processor = SimReadyProcessor( scene=scene, coarse_layout_by_id=coarse_layout_by_id, coarse_geometry_root=coarse_geometry_output_root, simready_geometry_root=simready_geometry_output_root, + debug_output_root=debug_output_root, + # Keep the geometry-server scale and only correct unstable standing poses. + config=SimReadyProcessorConfig( + use_vlm_scale=False, + use_vlm_rotation=False, + orientation_states_by_id=standing_orientation_states_by_id, + ), + vlm_client=vlm_client, ) simready_assets_layout = simready_processor.process_assets() simready_table_layout = simready_processor.process_table() @@ -119,6 +145,7 @@ def generate_scene_and_refine( # Layout refinement will start with the table. refined_table_layout, refined_assets_layout = _layout_refinement( scene=scene, # Update this data structure internally. + scene_graph=scene_graph, simready_geometry_output_root=simready_geometry_output_root, # Contains simready assets and their current coarse layout JSON. debug_output_root=debug_output_root, # Keep the table support surface info + optimized layout info (render with matplotlib) for debugging. ) @@ -138,6 +165,7 @@ def _generate_coarse_results_from_masks( scene: Scene, *, geometry_generation_client: GeometryGenerationClient, + seed: int | None = None, ) -> None: # Parse whether the scene has each assets' binary masks. @@ -164,10 +192,15 @@ def _generate_coarse_results_from_masks( ) # id + mask, for avoiding the download glbs order confusion. # Sent the request, wait, then save the intermediate results. + generation_kwargs = { + "image_path": image_path, + "object_masks": object_masks, + "output_root": coarse_geometry_output_root, + } + if seed is not None: + generation_kwargs["seed"] = seed response_data, response_objects = geometry_generation_client.generate_objects( - image_path=image_path, - object_masks=object_masks, - output_root=coarse_geometry_output_root, # Keep the coarse geometries + **generation_kwargs ) # Write the response JSON which contains all the layout info the server gave us. # Keep original response for getting the sam3d coarse layout matrix. @@ -195,16 +228,18 @@ def _generate_coarse_results_from_masks( return None -def _update_scene_final_y_up_layout( +def _update_scene_final_y_up_layout_and_z_up_centers( *, scene: Scene, table_layout: dict[str, object], assets_layout: list[dict[str, object]], + geometry_root: str | Path, ) -> None: - """Copy final y-up layout values into the matching table and asset objects.""" + """Write final y-up layouts and z-up XY centers into the scene.""" if scene.table is None: raise ValueError("Cannot update a final layout without a table.") + # Keep final poses in the y-up layout convention used by exported GLBs. _copy_y_up_layout_to_scene_object(scene.table, table_layout) assets_by_id = {asset.id: asset for asset in scene.assets} layout_ids = set() @@ -223,6 +258,29 @@ def _update_scene_final_y_up_layout( f"Final layout is missing scene assets: {sorted(missing_assets)}." ) + # Measure final geometry in z-up so scene edits can compare tabletop XY positions. + table_mesh, assets_aabb_corners_by_id = _measure_table_and_assets_in_z_up_world( + table_layout=table_layout, + assets_layout=assets_layout, + geometry_root=geometry_root, + ) + # Persist AABB centers for future scene-edit object disambiguation. + scene.table.center_xy = table_mesh.bounds[:, :2].mean(axis=0).tolist() + if scene.table.support_contour_xy is not None: + # Move SimReady-local support geometry into the final table-frame position. + table_center_xy = np.asarray(scene.table.center_xy, dtype=float) + scene.table.support_contour_xy = [ + (np.asarray(point, dtype=float) + table_center_xy).tolist() + for point in scene.table.support_contour_xy + ] + if scene.table.support_optimization_rect_xy is not None: + scene.table.support_optimization_rect_xy = [ + (np.asarray(point, dtype=float) + table_center_xy).tolist() + for point in scene.table.support_optimization_rect_xy + ] + for asset in scene.assets: + asset.center_xy = assets_aabb_corners_by_id[asset.id].mean(axis=0).tolist() + def _copy_y_up_layout_to_scene_object( scene_object: SceneObject, @@ -252,6 +310,7 @@ def _copy_y_up_layout_to_scene_object( def _layout_refinement( *, scene: Scene, + scene_graph: SceneGraph, simready_geometry_output_root: str | Path, debug_output_root: str | Path, ) -> tuple[dict[str, object], list[dict[str, object]]]: @@ -318,7 +377,14 @@ def _layout_refinement( ) ) - # 3. Move all assets as one rigid group so its lowest AABB point is 2cm above + # 3. Correct image-observed standing containers before every geometry-based + # layout stage measures their footprint. + refined_assets_layout = _scene_graph_based_calibration( + scene_graph=scene_graph, + assets_layout=refined_assets_layout, + ) + + # 4. Move all assets as one rigid group so its lowest AABB point is 2cm above # the table. This preserves the initial relative poses for the later # gravity simulation, which can settle individual assets physically. @@ -332,31 +398,36 @@ def _layout_refinement( log_info("Scene has no movable assets; skipping support-region clamping.") return refined_table_layout, [] - # 4. Detect the actual upward support triangles instead of projecting the - # entire table mesh to one convex hull. The result retains concavities - # (for example, an L-shaped tabletop) and is the only boundary used for - # placement below. - ( - table_world_mesh_z_up, - assets_aabb_2d_z_up_world_corners_by_id, - ) = _measure_table_and_assets_in_z_up_world( - table_layout=refined_table_layout, - assets_layout=refined_assets_layout, - geometry_root=simready_geometry_output_root, - ) - support_detector = TableSupportSurfaceDetector( - table_world_mesh=table_world_mesh_z_up, - debug_output_root=debug_output_root, + # 5. Reuse support geometry detected during SimReady processing. + if ( + scene.table is None + or scene.table.support_contour_xy is None + or scene.table.support_optimization_rect_xy is None + ): + raise ValueError("Scene table has no persisted support geometry.") + table_support_polygon = Polygon(scene.table.support_contour_xy) + table_optimization_rectangle = Polygon(scene.table.support_optimization_rect_xy) + if not table_support_polygon.is_valid or table_support_polygon.is_empty: + raise ValueError("Scene table support contour is not a valid polygon.") + if ( + not table_optimization_rectangle.is_valid + or table_optimization_rectangle.is_empty + ): + raise ValueError("Scene table optimization rectangle is not valid.") + _, assets_aabb_2d_z_up_world_corners_by_id = ( + _measure_table_and_assets_in_z_up_world( + table_layout=refined_table_layout, + assets_layout=refined_assets_layout, + geometry_root=simready_geometry_output_root, + ) ) - table_support_region = support_detector.detect() - support_detector.save_support_surface_debug_images() - # 5. Keep the complete clutter rigid in the table plane. A successful + # 6. Keep the complete clutter rigid in the table plane. A successful # result applies one shared z-up XY delta to every AABB, so it preserves # all existing asset-to-asset relations. It is *not* an asset packing # pass: pre-existing overlap is deliberately left to a later optimizer. group_clamp = AssetsGroupSupportClamp( - support_region=table_support_region.support_polygon, + support_region=table_support_polygon, assets_aabb_2d_z_up_world_corners_by_id=( assets_aabb_2d_z_up_world_corners_by_id ), @@ -377,13 +448,10 @@ def _layout_refinement( ) ) - # 6. Restore the previous pairwise AABB separation stage, but constrain - # every candidate with the actual support polygon rather than the legacy - # largest internal rectangle. Assets may now move independently only as - # much as needed to remove overlap; every resulting AABB remains on the - # L-shaped, circular, or otherwise non-convex support region. + # 7. Optimize independent asset positions inside the conservative rectangle. + # The clamp above already used the exact outer contour for the shared shift. overlap_optimizer = AssetsSupportLayoutOptimizer( - support_region=table_support_region.support_polygon, + support_region=table_optimization_rectangle, assets_aabb_2d_z_up_world_corners_by_id=( clamped_assets_aabb_2d_z_up_world_corners_by_id ), @@ -396,24 +464,138 @@ def _layout_refinement( refined_assets_layout = overlap_optimizer.optimize() overlap_optimizer.save_overlap_optimization_debug_images() - # 7. Gravity simulation, to let all the assets to be stable and placed well on the table's support surface. - # Notice that: we do not consider the assets like a bottle, which should be standing on the table but laid down - # after the simulation. - gravity_settler = AssetsGravitySettler( + # 8. The initial image graph has one on-table level, so every asset settles + # dynamically against the table in this first generic gravity pass. + assets_by_id = {asset.id: asset for asset in scene.assets} + # All the assets are dynamic; the table is static. + settled_pose_by_id = GravitySettler( + table_body=GravitySettleBody( + scene_object=scene.table, + y_up_layout=refined_table_layout, + ), + participant_bodies=[ + GravitySettleBody( + scene_object=assets_by_id[str(asset_layout["id"])], + y_up_layout=asset_layout, + ) + for asset_layout in refined_assets_layout + ], + dynamic_asset_ids=set(assets_by_id), + static_asset_ids=set(), + ).settle() + # Update. + for asset_layout in refined_assets_layout: + asset_id = str(asset_layout["id"]) + settled_pose = settled_pose_by_id[asset_id] + asset_layout["pos"] = settled_pose["pos"] + asset_layout["rot"] = settled_pose["rot"] + + # Update the scene data structure with the final layout and spatial metadata. + _update_scene_final_y_up_layout_and_z_up_centers( scene=scene, table_layout=refined_table_layout, assets_layout=refined_assets_layout, geometry_root=simready_geometry_output_root, ) - refined_assets_layout = gravity_settler.settle() + return refined_table_layout, refined_assets_layout - # Update the scene data structure with the final y-up layout values. - _update_scene_final_y_up_layout( - scene=scene, - table_layout=refined_table_layout, - assets_layout=refined_assets_layout, + +def _scene_graph_based_calibration( + *, + scene_graph: SceneGraph, + assets_layout: list[dict[str, object]], +) -> list[dict[str, object]]: + """Minimally align graph-marked standing assets with the z-up table frame.""" + # This is the extension point for future image-conditioned scene generation + # calibration. The scene graph may later provide richer image-grounded + # constraints, but the current implementation deliberately consumes only + # ``orientation_state`` to correct standing container axes before layout. + y_up_to_z_up_matrix = np.eye(4) + y_up_to_z_up_matrix[:3, :3] = np.array( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] ) - return refined_table_layout, refined_assets_layout + z_up_to_y_up_matrix = np.linalg.inv(y_up_to_z_up_matrix) + nodes_by_id = scene_graph.node_by_id() + calibrated_assets_layout: list[dict[str, object]] = [] + + for asset_layout in assets_layout: + asset_id = asset_layout.get("id") + if not isinstance(asset_id, str) or not asset_id: + raise ValueError("Each asset layout must contain a non-empty string id.") + node = nodes_by_id.get(asset_id) + if node is None: + raise ValueError(f"Scene graph does not contain asset {asset_id!r}.") + # Only correct the standing assets. + if node.orientation_state != "standing": + calibrated_assets_layout.append(asset_layout) + continue + + # Conjugate the y-up pose so the SimReady container axis is local z. + z_up_asset_to_table_matrix = ( + y_up_to_z_up_matrix + @ layout_object_to_transform_matrix(asset_layout) + @ z_up_to_y_up_matrix + ) + linear_matrix = z_up_asset_to_table_matrix[:3, :3] + # Layout transforms store rotation and per-axis scale in the same matrix. + scale = np.linalg.norm(linear_matrix, axis=0) + if np.any(scale <= 1e-8): + raise ValueError(f"Asset {asset_id!r} has a zero scale axis.") + rotation_matrix = linear_matrix / scale + if not np.allclose(rotation_matrix.T @ rotation_matrix, np.eye(3), atol=1e-6): + raise ValueError(f"Asset {asset_id!r} layout contains shear.") + + local_z_axis_in_table = rotation_matrix[:, 2] + # Treat the long axis as unsigned to avoid an unnecessary 180-degree flip. + target_z_axis = np.array( + [0.0, 0.0, 1.0 if local_z_axis_in_table[2] >= 0.0 else -1.0] + ) + # Left multiplication applies the correction in the table/world frame. + z_up_asset_to_table_matrix[:3, :3] = ( + _minimum_axis_alignment_rotation( + source_axis=local_z_axis_in_table, + target_axis=target_z_axis, + ) + @ rotation_matrix + @ np.diag(scale) + ) + calibrated_assets_layout.append( + transform_matrix_to_layout_object( + asset_id, + z_up_to_y_up_matrix @ z_up_asset_to_table_matrix @ y_up_to_z_up_matrix, + ) + ) + return calibrated_assets_layout + + +def _minimum_axis_alignment_rotation( + *, + source_axis: np.ndarray, + target_axis: np.ndarray, +) -> np.ndarray: + """Return the smallest proper rotation mapping one nonzero axis to another.""" + source = np.asarray(source_axis, dtype=float) + target = np.asarray(target_axis, dtype=float) + source_norm = np.linalg.norm(source) + target_norm = np.linalg.norm(target) + if source_norm <= 1e-8 or target_norm <= 1e-8: + raise ValueError("Axis alignment requires nonzero axes.") + source /= source_norm + target /= target_norm + + cross_product = np.cross(source, target) + sine = np.linalg.norm(cross_product) + cosine = float(np.clip(np.dot(source, target), -1.0, 1.0)) + if sine <= 1e-8: + if cosine > 0.0: + return np.eye(3) + basis_axis = np.eye(3)[np.argmin(np.abs(source))] + rotation_axis = np.cross(source, basis_axis) + rotation_axis /= np.linalg.norm(rotation_axis) + return Rotation.from_rotvec(np.pi * rotation_axis).as_matrix() + + rotation_axis = cross_product / sine + return Rotation.from_rotvec(np.arctan2(sine, cosine) * rotation_axis).as_matrix() def _measure_table_and_assets_in_z_up_world( diff --git a/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py similarity index 74% rename from embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py rename to embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py index f91c665e6..d93ada019 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py @@ -29,6 +29,11 @@ ImageSegmentationClient, ) from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, + TABLE_OBJECT_ID, +) from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( OpenAICompatibleVLM, @@ -36,6 +41,7 @@ from embodichain.gen_sim.scene_engine.pipeline.utils.image_segmentation_utils import ( MaskCandidate, build_mask_candidates, + render_asset_mask_id_overlay, render_image_without_masks, render_numbered_mask_candidates, save_binary_mask, @@ -44,12 +50,6 @@ _SUPPORTED_IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png"} _CATEGORY_PATTERN = re.compile(r"^[a-z][a-z0-9_]*$") -_LOCATION_WORD_PATTERN = re.compile( - r"\b(?:left|right|front|back|center|middle|top|bottom|upper|lower|" - r"foreground|background|near|next|beside|behind|between|on|in|inside|" - r"under|above|below|against)\b", - flags=re.IGNORECASE, -) _SYSTEM_PROMPT = """You inspect one tabletop-scene image. Identify the main table and every visible, physically distinct object that should be segmented and later generated as an independent 3D asset. @@ -62,19 +62,23 @@ 3. Do not merge objects merely resting on another object. A mug on a table and the table are separate entries. 4. List every visible physical instance separately. If two objects look alike, - keep the same category and name, but distinguish them in description using - location. Do not add location to name. + keep the same category and name. Do not encode location or spatial context + in any semantic field. 5. category is a lower-case singular snake_case class, such as mug, book, potted_plant, or coffee_table. It must not contain color or material. -6. name contains only color, material, texture, shape, and object description. - It must not contain position or relations, such as left, right, on, in, or - near. +6. name is a concise human-readable phrase containing only color, material, + texture, shape, and object details. It may contain spaces, but must not + contain position or relations, such as left, right, on, in, or near. 7. For table, description contains only its category, material, color, texture, shape, and visible structural details. Do not mention image coverage, image position, camera framing, or viewpoint. For example, do not write "occupying most of the image" or "at the center of the image". -8. For assets, description may include all visible details, including location - and spatial context. +8. For assets, description contains only visible category, material, color, + texture, shape, and structural details. Do not mention location, the table, + or any relationship to another object. + Structural direction words are allowed when they describe the object itself: + "bottle with a black cap on top" is valid, while "bottle on the left of the + table" is not. Return JSON only: no Markdown, comments, or prose outside this exact schema: { @@ -87,14 +91,14 @@ { "category": "mug", "name": "blue ceramic mug", - "description": "small blue ceramic mug on the left side of the table" + "description": "small blue ceramic mug with a curved handle" } ] } -For two identical blue mugs, output two asset entries with the same category and -name, and use their descriptions to state left/right or front/back. Do not -infer objects that are not visible. Use an empty assets array when no objects -are visible. Every field must be a non-empty string.""" +For two identical blue mugs, output two asset entries with the same category, +name, and description. Do not infer objects that are not visible. Use an empty +assets array when no objects are visible. Every field must be a non-empty +string.""" _USER_PROMPT = "Analyze the provided image and return only the required JSON object." @@ -141,6 +145,23 @@ Return JSON only, with exactly one key: assignments. It must be null or an array of asset_id and mask_index objects. Do not include Markdown or any other text.""" +_ORIENTATION_STATE_SYSTEM_PROMPT = """You inspect an outlined tabletop-scene image. +Each visible asset has an outline and an ID label. Determine whether each listed +asset is standing, lying, or unknown in the image. + +Use a non-null state only for an elongated object with a clear primary long axis. +Use "standing" when its primary axis is approximately vertical to the tabletop. +Use "lying" when its primary axis is approximately parallel to the tabletop. +Use null for every object without a clear primary long axis or when uncertain. + +Return JSON only, with exactly this schema. Include every supplied asset ID +exactly once. Never include the table or any ID that was not supplied: +{ + "orientation_states": [ + {"object_id": "bottle_001", "orientation_state": "standing"}, + {"object_id": "book_001", "orientation_state": null} + ] +}""" def understand_scene( @@ -149,8 +170,9 @@ def understand_scene( output_root: str | Path, *, vlm_client: OpenAICompatibleVLM, + image_segmentation_client: ImageSegmentationClient, json_max_attempts: int = 3, -) -> Scene: +) -> tuple[Scene, SceneGraph]: resolved_image_path = _validate_image_path(image_path) # The output in this stage will keep a JSON which contains @@ -167,26 +189,165 @@ def understand_scene( json_max_attempts=json_max_attempts, ) - # Load .env settings and fail if the Image Segmentation Server is unavailable. - image_segmentation_client = ImageSegmentationClient.from_dotenv() - try: - image_segmentation_client.check_health() # Error raising will happen internally. - _segment_scene( - image_path=resolved_image_path, - stage_output_root=stage_output_root, - scene=scene, - vlm_client=vlm_client, - image_segmentation_client=image_segmentation_client, - ) - finally: - image_segmentation_client.close() # Kill the session to avoid resource leaks. + # Receive the validated whole-scene mask, for VLM output the scene graph. + asset_mask_id_overlay_path = _segment_scene( + image_path=resolved_image_path, + stage_output_root=stage_output_root, + scene=scene, + vlm_client=vlm_client, + image_segmentation_client=image_segmentation_client, + ) + + # Use the segmented image to initialize the scene graph + # with the help of the VLM client. + # But at here, we do with the simplest way (hard code). + scene_graph = _initialize_scene_graph_from_segmented_scene( + scene, + asset_mask_id_overlay_path=asset_mask_id_overlay_path, + vlm_client=vlm_client, + json_max_attempts=json_max_attempts, + ) # Write the Updated scene JSON for debugging. (stage_output_root / "scene.json").write_text( json.dumps(scene.to_dict(), indent=2, ensure_ascii=False) + "\n", encoding="utf-8", ) - return scene + (stage_output_root / "scene_graph.json").write_text( + json.dumps(scene_graph.to_dict(), indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + return scene, scene_graph + + +def _initialize_scene_graph_from_segmented_scene( + scene: Scene, + *, + asset_mask_id_overlay_path: str | Path, + vlm_client: OpenAICompatibleVLM, + json_max_attempts: int = 3, +) -> SceneGraph: + """Build the initial graph assuming every segmented asset rests on the table.""" + # Get simplified scene info for VLM. + scene_info = _simplify_scene_info_for_graph_initialization(scene=scene) + resolved_asset_mask_id_overlay_path = _validate_image_path( + asset_mask_id_overlay_path + ) + if scene.table is None: + raise ValueError("Cannot initialize a scene graph without a table.") + orientation_states_by_id = _query_orientation_states( + scene_info=scene_info, + asset_mask_id_overlay_path=resolved_asset_mask_id_overlay_path, + vlm_client=vlm_client, + json_max_attempts=json_max_attempts, + ) + return SceneGraph( + nodes=[ + SceneGraphNode(object_id=TABLE_OBJECT_ID, parent_id=None), + *[ + SceneGraphNode( + object_id=asset.id, + parent_id=TABLE_OBJECT_ID, + parent_relation="on", # semi-hard-code. + orientation_state=orientation_states_by_id[asset.id], + ) + for asset in scene.assets + ], + ], + ) + + +def _simplify_scene_info_for_graph_initialization( + *, + scene: Scene, +) -> dict[str, object]: + """Return the object metadata needed to initialize an image-based graph.""" + return { + "asset_ids": [asset.id for asset in scene.assets], + } + + +def _query_orientation_states( + *, + scene_info: dict[str, object], + asset_mask_id_overlay_path: Path, + vlm_client: OpenAICompatibleVLM, + json_max_attempts: int, +) -> dict[str, str | None]: + """Return validated image-observed orientation states keyed by asset ID.""" + if json_max_attempts < 1: + raise ValueError("json_max_attempts must be at least 1.") + last_validation_error: ValueError | None = None + for _ in range(json_max_attempts): + response_text = vlm_client.complete( + image_path=asset_mask_id_overlay_path, + system_prompt=_ORIENTATION_STATE_SYSTEM_PROMPT, + user_prompt=json.dumps(scene_info, ensure_ascii=False), + ) + try: + return _parse_orientation_states_response( + response_text=response_text, + asset_ids=scene_info["asset_ids"], + ) + except ValueError as exc: + last_validation_error = exc + assert last_validation_error is not None + raise ValueError( + "VLM returned invalid orientation-state JSON after " + f"{json_max_attempts} attempts: {last_validation_error}" + ) from last_validation_error + + +def _parse_orientation_states_response( + *, + response_text: str, + asset_ids: object, +) -> dict[str, str | None]: + """Parse a complete VLM orientation-state response for known asset IDs.""" + if not isinstance(asset_ids, list) or not all( + isinstance(object_id, str) for object_id in asset_ids + ): + raise ValueError("Scene graph initialization requires string asset IDs.") + json_text = _strip_json_code_fence(response_text) + try: + payload = json.loads(json_text) + except json.JSONDecodeError as exc: + raise ValueError(f"VLM response is not valid JSON: {exc.msg}") from exc + if not isinstance(payload, dict) or set(payload) != {"orientation_states"}: + raise ValueError("VLM JSON must contain exactly the key: orientation_states.") + states_value = payload["orientation_states"] + if not isinstance(states_value, list): + raise ValueError("VLM JSON key orientation_states must be an array.") + + orientation_states_by_id: dict[str, str | None] = {} + for index, state_value in enumerate(states_value): + if not isinstance(state_value, dict) or set(state_value) != { + "object_id", + "orientation_state", + }: + raise ValueError( + "VLM JSON orientation_states[" + f"{index}] must contain exactly object_id and orientation_state." + ) + object_id = state_value["object_id"] + orientation_state = state_value["orientation_state"] + if not isinstance(object_id, str) or not object_id: + raise ValueError( + f"VLM JSON orientation_states[{index}].object_id is invalid." + ) + if orientation_state not in {None, "standing", "lying"}: + raise ValueError( + f"VLM JSON orientation_states[{index}].orientation_state is invalid." + ) + if object_id in orientation_states_by_id: + raise ValueError(f"VLM JSON repeats orientation state for {object_id!r}.") + orientation_states_by_id[object_id] = orientation_state + + if set(orientation_states_by_id) != set(asset_ids): + raise ValueError( + "VLM JSON orientation states must match all supplied asset IDs." + ) + return orientation_states_by_id def _analyze_image_objects( @@ -330,13 +491,6 @@ def _parse_scene_object_fields( f"VLM JSON key {field_name}.category must be a lower-case snake_case " "class name." ) - if _LOCATION_WORD_PATTERN.search( - fields["name"] - ): # Check whether the name contains location. - raise ValueError( - f"VLM JSON key {field_name}.name must not contain location or " - "relationship words." - ) return fields @@ -353,8 +507,8 @@ def _segment_scene( scene: Scene, vlm_client: OpenAICompatibleVLM, image_segmentation_client: ImageSegmentationClient, -) -> None: - """Add validated table and asset mask paths to a semantic scene.""" +) -> Path: + """Add validated masks and return an asset-only ID overlay image.""" debug_output_root = ( Path(stage_output_root) / "debug" ) # Keeps the mask debug images. @@ -396,6 +550,16 @@ def _segment_scene( vlm_client=vlm_client, image_segmentation_client=image_segmentation_client, ) + asset_masks: list[tuple[str, str]] = [] + for asset in scene.assets: + if asset.mask_path is None: + raise ValueError(f"Asset {asset.id!r} has no validated mask path.") + asset_masks.append((asset.id, asset.mask_path)) + return render_asset_mask_id_overlay( + image_path=image_path, + asset_masks=asset_masks, + output_path=Path(masks_output_root) / "asset_masks_with_ids.png", + ) def _segment_table( diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/assets_gravity_settler.py b/embodichain/gen_sim/scene_engine/pipeline/utils/assets_gravity_settler.py deleted file mode 100644 index 31e9e8443..000000000 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/assets_gravity_settler.py +++ /dev/null @@ -1,349 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -from __future__ import annotations - -from dataclasses import dataclass -from pathlib import Path -from typing import Sequence - -import numpy as np -from scipy.spatial.transform import Rotation -import trimesh - -from embodichain.gen_sim.scene_engine.core.scene import Scene -from embodichain.gen_sim.scene_engine.core.scene_object import ( - ObjectPhysics, - SceneObject, -) -from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( - layout_object_to_transform_matrix, - load_glb_mesh, - transform_matrix_to_layout_object, -) -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg -from embodichain.lab.sim.shapes import MeshCfg -from embodichain.utils.logger import log_info - - -@dataclass(frozen=True) -class AssetsGravitySettlerConfig: - """Physics controls for table-top asset settling.""" - - clearance_m: float = 0.02 # Initial gap between each asset and the table top. - settle_steps: int = 300 # Fixed number of simulator steps to execute. - physics_dt: float = 1.0 / 100.0 # Physics timestep in seconds. - sim_device: str = "cpu" # Simulation device requested from EmbodiChain Lab. - - -class AssetsGravitySettler: - """Settle all assets together on one kinematic table in a z-up simulation.""" - - def __init__( - self, - *, - scene: Scene, - table_layout: dict[str, object], - assets_layout: list[dict[str, object]], - geometry_root: str | Path, - config: AssetsGravitySettlerConfig | None = None, - ) -> None: - self.scene = scene - self.table_layout = table_layout - self.assets_layout = assets_layout - self.geometry_root = Path(geometry_root).expanduser().resolve() - self.settled_assets_layout: list[dict[str, object]] | None = None - self.config = config if config is not None else AssetsGravitySettlerConfig() - # Check. - if self.config.clearance_m < 0.0: - raise ValueError("Gravity-settle clearance_m must be non-negative.") - if self.config.settle_steps <= 0: - raise ValueError("Gravity-settle settle_steps must be positive.") - if self.config.physics_dt <= 0.0: - raise ValueError("Gravity-settle physics_dt must be positive.") - - def settle(self) -> list[dict[str, object]]: - """Run gravity settling and return the resulting y-up asset layouts.""" - self.settled_assets_layout = None - if not self.assets_layout: - self.settled_assets_layout = [] - log_info("Scene has no movable assets; skipping gravity settling.") - return self.settled_assets_layout - - table_id = self._require_layout_id(self.table_layout, name="Table") - table_object = self._require_scene_object(table_id, kind="table") - asset_ids: set[str] = set() - asset_objects_by_id: dict[str, SceneObject] = {} - for asset_layout in self.assets_layout: - asset_id = self._require_layout_id(asset_layout, name="Asset") - if asset_id in asset_ids: - raise ValueError(f"Asset layouts contain duplicate id {asset_id!r}.") - asset_ids.add(asset_id) - asset_objects_by_id[asset_id] = self._require_scene_object( - asset_id, kind="asset" - ) - expected_asset_ids = {asset.id for asset in self.scene.assets} - if asset_ids != expected_asset_ids: - raise ValueError( - "Gravity-settle layouts must contain exactly the scene asset ids." - ) - - y_up_to_z_up_matrix = np.eye(4) - y_up_to_z_up_matrix[:3, :3] = np.array( - [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] - ) - z_up_to_y_up_matrix = np.linalg.inv(y_up_to_z_up_matrix) - table_info = self._prepare_sim_body( - layout_object=self.table_layout, - y_up_to_z_up_matrix=y_up_to_z_up_matrix, - ) - table_world_mesh = self._mesh_to_z_up_world_for_aabb( - y_up_mesh=table_info["mesh"], - z_up_rigid_layout=table_info["rigid_layout"], - z_up_scale=table_info["z_up_scale"], - y_up_to_z_up_matrix=y_up_to_z_up_matrix, - ) - table_top_z = float(table_world_mesh.bounds[1, 2]) - - prepared_assets: dict[str, dict[str, object]] = {} - for asset_layout in self.assets_layout: - asset_id = str(asset_layout["id"]) - asset_info = self._prepare_sim_body( - layout_object=asset_layout, - y_up_to_z_up_matrix=y_up_to_z_up_matrix, - ) - asset_world_mesh = self._mesh_to_z_up_world_for_aabb( - y_up_mesh=asset_info["mesh"], - z_up_rigid_layout=asset_info["rigid_layout"], - z_up_scale=asset_info["z_up_scale"], - y_up_to_z_up_matrix=y_up_to_z_up_matrix, - ) - asset_bottom_z = float(asset_world_mesh.bounds[0, 2]) - asset_info["rigid_layout"]["pos"][2] += ( - table_top_z + self.config.clearance_m - asset_bottom_z - ) - prepared_assets[asset_id] = asset_info - - log_info( - "Gravity settling started: " - f"assets={len(prepared_assets)}, steps={self.config.settle_steps}, " - f"physics_dt={self.config.physics_dt:.4f} s." - ) - sim = SimulationManager( - SimulationManagerCfg( - headless=True, - physics_dt=self.config.physics_dt, - sim_device=self.config.sim_device, - ) - ) - try: - # Add table. - sim.add_rigid_object( - RigidObjectCfg( - uid=table_id, - shape=MeshCfg(fpath=str(table_info["mesh_path"])), - init_pos=tuple(table_info["rigid_layout"]["pos"]), - init_rot=tuple( - self._simulation_euler_xyz_degrees(table_info["rigid_layout"]) - ), - body_scale=tuple(table_info["y_up_scale"]), - attrs=self._rigid_body_attrs(table_object.physics), - body_type=table_object.physics.body_type, - max_convex_hull_num=table_object.physics.max_convex_hull_num, - acd_method="vhacd", - ) - ) - # Add assets. - simulated_assets: dict[str, object] = {} - for asset_id, asset_info in prepared_assets.items(): - rigid_layout = asset_info["rigid_layout"] - simulated_assets[asset_id] = sim.add_rigid_object( - RigidObjectCfg( - uid=asset_id, - shape=MeshCfg(fpath=str(asset_info["mesh_path"])), - init_pos=tuple(rigid_layout["pos"]), - init_rot=tuple( - self._simulation_euler_xyz_degrees(rigid_layout) - ), - body_scale=tuple(asset_info["y_up_scale"]), - attrs=self._rigid_body_attrs( - asset_objects_by_id[asset_id].physics - ), - body_type=asset_objects_by_id[asset_id].physics.body_type, - max_convex_hull_num=( - asset_objects_by_id[asset_id].physics.max_convex_hull_num - ), - acd_method="vhacd", - ) - ) - # Run simulation to settle all assets. - sim.update(step=self.config.settle_steps) - - # Update the final layouts. - settled_layout_by_id: dict[str, dict[str, object]] = {} - for asset_id, simulated_asset in simulated_assets.items(): - final_rigid_pose_z_up = np.asarray( - simulated_asset.get_local_pose(to_matrix=True)[0] - .detach() - .cpu() - .numpy(), - dtype=float, - ) - scale_matrix = np.eye(4) - scale_matrix[:3, :3] = np.diag(prepared_assets[asset_id]["z_up_scale"]) - final_z_up_layout_matrix = final_rigid_pose_z_up @ scale_matrix - settled_layout_by_id[asset_id] = transform_matrix_to_layout_object( - asset_id, - z_up_to_y_up_matrix - @ final_z_up_layout_matrix - @ y_up_to_z_up_matrix, - ) - finally: - # Release resources. - sim.destroy(exit_process=False) - SimulationManager.flush_cleanup_queue() - - self.settled_assets_layout = [ - settled_layout_by_id[str(asset_layout["id"])] - for asset_layout in self.assets_layout - ] - log_info("Gravity settling completed for all assets.") - return self.settled_assets_layout - - def _prepare_sim_body( - self, - *, - layout_object: dict[str, object], - y_up_to_z_up_matrix: np.ndarray, - ) -> dict[str, object]: - """Load one y-up GLB and prepare its z-up simulation pose.""" - object_id = self._require_layout_id(layout_object, name="Layout object") - source_mesh_path = self.geometry_root / f"{object_id}.glb" - source_mesh = load_glb_mesh(source_mesh_path) - z_up_layout = self._convert_layout_coordinate_system( - layout_object, - source_to_target_matrix=y_up_to_z_up_matrix, - ) - return { - "mesh_path": source_mesh_path, - "mesh": source_mesh, - "rigid_layout": { - "id": object_id, - "rot": self._three_floats(z_up_layout.get("rot"), field_name="rot"), - "pos": self._three_floats(z_up_layout.get("pos"), field_name="pos"), - "scale": [1.0, 1.0, 1.0], - }, - "y_up_scale": self._three_floats( - layout_object.get("scale"), field_name="scale" - ), - "z_up_scale": self._three_floats( - z_up_layout.get("scale"), field_name="scale" - ), - } - - def _require_scene_object(self, object_id: str, *, kind: str) -> SceneObject: - """Return one physics-ready scene object with the expected semantic kind.""" - matching_objects = [ - scene_object - for scene_object in self.scene.objects - if scene_object.id == object_id - ] - if len(matching_objects) != 1: - raise ValueError( - f"Gravity settling requires exactly one scene object {object_id!r}." - ) - scene_object = matching_objects[0] - if scene_object.kind != kind: - raise ValueError( - f"Scene object {object_id!r} must have kind {kind!r} before " - "gravity settling." - ) - if scene_object.physics is None: - raise ValueError( - f"Scene object {object_id!r} has no SimReady physics settings." - ) - return scene_object - - @staticmethod - def _rigid_body_attrs(physics: ObjectPhysics | None) -> RigidBodyAttributesCfg: - """Convert persisted SceneObject physics attributes into Lab config.""" - if physics is None: - raise ValueError("Gravity settling requires SimReady physics settings.") - return RigidBodyAttributesCfg(**physics.attrs) - - @staticmethod - def _mesh_to_z_up_world_for_aabb( - *, - y_up_mesh: trimesh.Trimesh, - z_up_rigid_layout: dict[str, object], - z_up_scale: Sequence[float], - y_up_to_z_up_matrix: np.ndarray, - ) -> trimesh.Trimesh: - """Transform a y-up mesh into its z-up world pose for AABB measurement.""" - mesh = y_up_mesh.copy() - mesh.apply_transform(y_up_to_z_up_matrix) - scale_matrix = np.eye(4) - scale_matrix[:3, :3] = np.diag(z_up_scale) - mesh.apply_transform(scale_matrix) - mesh.apply_transform(layout_object_to_transform_matrix(z_up_rigid_layout)) - return mesh - - @staticmethod - def _simulation_euler_xyz_degrees(layout_object: dict[str, object]) -> list[float]: - """Convert lowercase-xyz layout rotation to SimulationManager's XYZ order.""" - layout_rotation = Rotation.from_euler( - "xyz", - AssetsGravitySettler._three_floats( - layout_object.get("rot"), field_name="rot" - ), - degrees=True, - ) - return layout_rotation.as_euler("XYZ", degrees=True).tolist() - - @staticmethod - def _convert_layout_coordinate_system( - layout_object: dict[str, object], - *, - source_to_target_matrix: np.ndarray, - ) -> dict[str, object]: - """Convert one layout object between coordinate frames through its matrix.""" - return transform_matrix_to_layout_object( - str(layout_object["id"]), - source_to_target_matrix - @ layout_object_to_transform_matrix(layout_object) - @ np.linalg.inv(source_to_target_matrix), - ) - - @staticmethod - def _require_layout_id(layout_object: dict[str, object], *, name: str) -> str: - """Check id.""" - object_id = layout_object.get("id") - if not isinstance(object_id, str) or not object_id: - raise ValueError(f"{name} layout must contain a non-empty string id.") - return object_id - - @staticmethod - def _three_floats(value: object, *, field_name: str) -> list[float]: - """Check three values.""" - if not isinstance(value, list) or len(value) != 3: - raise ValueError(f"Layout field {field_name} must contain three values.") - try: - return [float(item) for item in value] - except (TypeError, ValueError) as exc: - raise ValueError( - f"Layout field {field_name} must contain numeric values." - ) from exc diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_layout_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_layout_optimizer.py index d4021fede..8b3933b09 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_layout_optimizer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_layout_optimizer.py @@ -120,24 +120,20 @@ def optimize(self) -> list[dict[str, object]]: base_aabbs = np.stack([aabbs_by_id[asset_id] for asset_id in asset_ids]) offsets = np.zeros((len(asset_ids), 2), dtype=float) if not self._all_contained(safe_support, base_aabbs, offsets): - log_warning( - "AABB overlap optimization requires all input AABBs to be inside " - "the support region." - ) - raise ValueError( - "Overlap optimization requires AABBs already inside support; " - "run AssetsGroupSupportClamp first." - ) + # Independently project each AABB into the rectangular optimization region. + offsets = self._project_aabbs_inside_rectangle(safe_support, base_aabbs) initial_overlaps = self._overlaps(base_aabbs, offsets) + projected_asset_count = int(np.count_nonzero(np.any(offsets != 0.0, axis=1))) log_info( "Support-constrained AABB overlap optimization started: " f"assets={len(asset_ids)}, initial_overlaps={len(initial_overlaps)}, " + f"initial_projections={projected_asset_count}, " f"boundary_margin={self.config.margin_m:.4f} m, " f"aabb_clearance={self.config.aabb_clearance_m:.4f} m, " f"max_rounds={self.config.max_rounds}." ) if not initial_overlaps: # Return directly if there are no overlaps to resolve. - log_info("AABB overlap optimization succeeded without movement.") + log_info("AABB overlap optimization succeeded without pair separation.") self.refined_assets_layout = self._apply_offsets_to_y_up_layouts( asset_ids=asset_ids, offsets=offsets, @@ -213,6 +209,50 @@ def optimize(self) -> list[dict[str, object]]: "inside the detected table support region." ) + @staticmethod + def _project_aabbs_inside_rectangle( + support: Polygon | MultiPolygon, + base_aabbs: np.ndarray, + ) -> np.ndarray: + """Return minimum per-AABB offsets that place AABBs in a rectangle.""" + if not isinstance(support, Polygon) or support.interiors: + raise ValueError( + "Initial AABB projection requires an axis-aligned rectangular " + "support region." + ) + minimum_x, minimum_y, maximum_x, maximum_y = support.bounds + rectangle = Polygon( + [ + (minimum_x, minimum_y), + (maximum_x, minimum_y), + (maximum_x, maximum_y), + (minimum_x, maximum_y), + ] + ) + if not support.equals(rectangle): + raise ValueError( + "Initial AABB projection requires an axis-aligned rectangular " + "support region." + ) + + aabb_minimums, aabb_maximums = base_aabbs.min(axis=1), base_aabbs.max(axis=1) + half_extents = (aabb_maximums - aabb_minimums) / 2.0 + support_minimum = np.array([minimum_x, minimum_y], dtype=float) + support_maximum = np.array([maximum_x, maximum_y], dtype=float) + valid_center_minimums = support_minimum + half_extents + valid_center_maximums = support_maximum - half_extents + if np.any(valid_center_minimums > valid_center_maximums + 1e-9): + raise ValueError( + "An asset AABB is larger than the rectangular support region." + ) + + centers = (aabb_minimums + aabb_maximums) / 2.0 + # A center must stay inset from each boundary by its AABB half extent. + projected_centers = np.clip( + centers, valid_center_minimums, valid_center_maximums + ) + return projected_centers - centers + def _apply_offsets_to_y_up_layouts( self, *, asset_ids: list[str], offsets: np.ndarray ) -> list[dict[str, object]]: diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/gravity_settler.py b/embodichain/gen_sim/scene_engine/pipeline/utils/gravity_settler.py new file mode 100644 index 000000000..b5078c75f --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/gravity_settler.py @@ -0,0 +1,356 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +from scipy.spatial.transform import Rotation + +from embodichain.gen_sim.scene_engine.core.scene_object import ( + ObjectPhysics, + SceneObject, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( + layout_object_to_transform_matrix, + transform_matrix_to_layout_object, +) +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg +from embodichain.lab.sim.shapes import MeshCfg +from embodichain.utils.logger import log_info + + +@dataclass(frozen=True) +class GravitySettlerConfig: + """Physics controls for one caller-defined gravity-settlement pass.""" + + settle_steps: int = 300 + physics_dt: float = 1.0 / 100.0 + sim_device: str = "cpu" + + def __post_init__(self) -> None: + """Reject invalid numerical controls before starting a simulation.""" + if self.settle_steps <= 0: + raise ValueError("Gravity-settle settle_steps must be positive.") + if self.physics_dt <= 0.0: + raise ValueError("Gravity-settle physics_dt must be positive.") + + +@dataclass(frozen=True) +class GravitySettleBody: + """One scene object and its latest complete y-up pipeline layout.""" + + scene_object: SceneObject + y_up_layout: dict[str, object] + + +class GravitySettler: + """Settle caller-selected dynamic assets against a mandatory table body. + + All supplied layouts use Scene Engine's y-up pipeline convention. The + settler converts them to z-up only at the simulation boundary. The caller + explicitly classifies every participant as dynamic or static; static + participants and the table are kinematic collision bodies. + """ + + def __init__( + self, + *, + table_body: GravitySettleBody, + participant_bodies: list[GravitySettleBody], + dynamic_asset_ids: set[str], + static_asset_ids: set[str], + config: GravitySettlerConfig | None = None, + ) -> None: + self.table_body = table_body + self.participant_bodies = participant_bodies + self.dynamic_asset_ids = set(dynamic_asset_ids) + self.static_asset_ids = set(static_asset_ids) + self.config = config if config is not None else GravitySettlerConfig() + + def settle(self) -> dict[str, dict[str, list[float]]]: + """Return final y-up poses for dynamic participants only. + + Input layouts are used as-is. Placement clearance and support-surface + alignment remain the responsibility of the calling layout optimizer. + Static participants and every object's scale are unchanged, so they are + deliberately omitted from the result. + """ + # Check table. + table = self.table_body.scene_object + if table.kind != "table": + raise ValueError("Gravity settling requires a table body.") + table_id = self._require_body_layout_id(self.table_body, name="Table") + + participant_bodies_by_id: dict[str, GravitySettleBody] = {} + for participant_body in self.participant_bodies: + asset_id = self._require_body_layout_id( + participant_body, name="Participant asset" + ) + if asset_id == table_id: + raise ValueError( + "Gravity-settle participants cannot include the table." + ) + if participant_body.scene_object.kind != "asset": + raise ValueError( + f"Gravity-settle participant {asset_id!r} must be an asset body." + ) + if asset_id in participant_bodies_by_id: + raise ValueError( + f"Gravity-settle participant assets repeat id {asset_id!r}." + ) + participant_bodies_by_id[asset_id] = participant_body + + participant_ids = set(participant_bodies_by_id) + classified_ids = self.dynamic_asset_ids | self.static_asset_ids + if self.dynamic_asset_ids & self.static_asset_ids: + raise ValueError( + "Gravity-settle dynamic and static asset IDs must not overlap." + ) + if classified_ids != participant_ids: + raise ValueError( + "Gravity-settle dynamic and static asset IDs must exactly match " + f"participants; participants={sorted(participant_ids)}, " + f"classified={sorted(classified_ids)}." + ) + if not self.dynamic_asset_ids: + log_info("Gravity settle has no dynamic participants; skipping simulation.") + return {} + + y_up_to_z_up_matrix = self._y_up_to_z_up_matrix() + z_up_to_y_up_matrix = np.linalg.inv(y_up_to_z_up_matrix) + table_info = self._prepare_sim_body( + body=self.table_body, + y_up_to_z_up_matrix=y_up_to_z_up_matrix, + ) + participant_infos_by_id = { + asset_id: self._prepare_sim_body( + body=participant_body, + y_up_to_z_up_matrix=y_up_to_z_up_matrix, + ) + for asset_id, participant_body in participant_bodies_by_id.items() + } + + log_info( + "Gravity settling started: " + f"dynamic_assets={len(self.dynamic_asset_ids)}, " + f"kinematic_assets={len(participant_infos_by_id) - len(self.dynamic_asset_ids)}, " + f"steps={self.config.settle_steps}, " + f"physics_dt={self.config.physics_dt:.4f} s." + ) + sim = SimulationManager( + SimulationManagerCfg( + headless=True, + physics_dt=self.config.physics_dt, + sim_device=self.config.sim_device, + ) + ) + try: + self._add_sim_body( + sim=sim, + object_id=table_id, + body_info=table_info, + physics=table.physics, + body_type="kinematic", + ) + simulated_assets: dict[str, object] = {} + for asset_id, asset_info in participant_infos_by_id.items(): + simulated_assets[asset_id] = self._add_sim_body( + sim=sim, + object_id=asset_id, + body_info=asset_info, + physics=participant_bodies_by_id[asset_id].scene_object.physics, + body_type=( + "dynamic" if asset_id in self.dynamic_asset_ids else "kinematic" + ), + ) + sim.update(step=self.config.settle_steps) + + settled_pose_by_id: dict[str, dict[str, list[float]]] = {} + for asset_id in self.dynamic_asset_ids: + simulated_asset = simulated_assets[asset_id] + final_rigid_pose_z_up = np.asarray( + simulated_asset.get_local_pose(to_matrix=True)[0] + .detach() + .cpu() + .numpy(), + dtype=float, + ) + scale_matrix = np.eye(4) + scale_matrix[:3, :3] = np.diag( + participant_infos_by_id[asset_id]["z_up_scale"] + ) + final_y_up_layout = transform_matrix_to_layout_object( + asset_id, + z_up_to_y_up_matrix + @ final_rigid_pose_z_up + @ scale_matrix + @ y_up_to_z_up_matrix, + ) + settled_pose_by_id[asset_id] = { + "pos": self._three_floats( + final_y_up_layout.get("pos"), field_name="pos" + ), + "rot": self._three_floats( + final_y_up_layout.get("rot"), field_name="rot" + ), + } + finally: + sim.destroy(exit_process=False) + SimulationManager.flush_cleanup_queue() + + log_info("Gravity settling completed for participating assets.") + return settled_pose_by_id + + def _prepare_sim_body( + self, + *, + body: GravitySettleBody, + y_up_to_z_up_matrix: np.ndarray, + ) -> dict[str, object]: + """Convert one supplied y-up layout into a simulator body pose.""" + scene_object = body.scene_object + if scene_object.simready_glb_path is None: + raise ValueError( + f"Gravity-settle object {scene_object.id!r} has no SimReady GLB path." + ) + mesh_path = Path(scene_object.simready_glb_path).expanduser().resolve() + if not mesh_path.is_file(): + raise FileNotFoundError( + f"Gravity-settle GLB for {scene_object.id!r} not found: {mesh_path}" + ) + y_up_layout = body.y_up_layout + z_up_layout = self._convert_layout_coordinate_system( + y_up_layout, + source_to_target_matrix=y_up_to_z_up_matrix, + ) + return { + "mesh_path": mesh_path, + "rigid_layout": { + "id": scene_object.id, + "rot": self._three_floats(z_up_layout.get("rot"), field_name="rot"), + "pos": self._three_floats(z_up_layout.get("pos"), field_name="pos"), + "scale": [1.0, 1.0, 1.0], + }, + "y_up_scale": self._three_floats( + y_up_layout.get("scale"), field_name="scale" + ), + "z_up_scale": self._three_floats( + z_up_layout.get("scale"), field_name="scale" + ), + } + + def _add_sim_body( + self, + *, + sim: SimulationManager, + object_id: str, + body_info: dict[str, object], + physics: ObjectPhysics | None, + body_type: str, + ) -> object: + """Add one supplied body with a pass-specific dynamic or kinematic type.""" + rigid_layout = body_info["rigid_layout"] + if not isinstance(rigid_layout, dict): + raise ValueError("Gravity-settle body has invalid rigid layout.") + return sim.add_rigid_object( + RigidObjectCfg( + uid=object_id, + shape=MeshCfg(fpath=str(body_info["mesh_path"])), + init_pos=tuple( + self._three_floats(rigid_layout.get("pos"), field_name="pos") + ), + init_rot=tuple(self._simulation_euler_xyz_degrees(rigid_layout)), + body_scale=tuple( + self._three_floats(body_info["y_up_scale"], field_name="scale") + ), + attrs=self._rigid_body_attrs(physics), + body_type=body_type, + max_convex_hull_num=self._max_convex_hull_num(physics), + acd_method="vhacd", + ) + ) + + @staticmethod + def _rigid_body_attrs(physics: ObjectPhysics | None) -> RigidBodyAttributesCfg: + """Convert persisted collision material data into one Lab config.""" + if physics is None: + raise ValueError("Gravity settling requires SimReady physics settings.") + return RigidBodyAttributesCfg(**physics.attrs) + + @staticmethod + def _max_convex_hull_num(physics: ObjectPhysics | None) -> int: + """Read the persisted collision-hull budget after validating physics.""" + if physics is None: + raise ValueError("Gravity settling requires SimReady physics settings.") + return physics.max_convex_hull_num + + @staticmethod + def _require_body_layout_id(body: GravitySettleBody, *, name: str) -> str: + """Validate that a body layout belongs to its scene object.""" + object_id = body.y_up_layout.get("id") + if not isinstance(object_id, str) or not object_id: + raise ValueError(f"{name} layout requires a non-empty string id.") + if object_id != body.scene_object.id: + raise ValueError( + f"{name} layout id {object_id!r} does not match its scene object." + ) + return object_id + + @staticmethod + def _three_floats(value: object, *, field_name: str) -> list[float]: + """Return three finite layout values as Python floats.""" + if not isinstance(value, (list, tuple)) or len(value) != 3: + raise ValueError(f"Gravity-settle {field_name} must contain three values.") + result = [float(component) for component in value] + if not np.all(np.isfinite(result)): + raise ValueError(f"Gravity-settle {field_name} must contain finite values.") + return result + + @staticmethod + def _simulation_euler_xyz_degrees(layout_object: dict[str, object]) -> list[float]: + """Convert lowercase-xyz layout rotation to SimulationManager's XYZ order.""" + layout_rotation = Rotation.from_euler( + "xyz", + GravitySettler._three_floats(layout_object.get("rot"), field_name="rot"), + degrees=True, + ) + return layout_rotation.as_euler("XYZ", degrees=True).tolist() + + @staticmethod + def _convert_layout_coordinate_system( + layout_object: dict[str, object], + *, + source_to_target_matrix: np.ndarray, + ) -> dict[str, object]: + """Convert one complete layout through the y-up/z-up basis change.""" + return transform_matrix_to_layout_object( + str(layout_object["id"]), + source_to_target_matrix + @ layout_object_to_transform_matrix(layout_object) + @ np.linalg.inv(source_to_target_matrix), + ) + + @staticmethod + def _y_up_to_z_up_matrix() -> np.ndarray: + """Return the coordinate conversion used by Scene Engine layouts.""" + matrix = np.eye(4) + matrix[:3, :3] = np.array([[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]]) + return matrix diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py index e0c3f772a..302ec2cbf 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py @@ -94,6 +94,35 @@ def decode_rle_mask(mask_rle: dict[str, Any]) -> Image.Image: return Image.frombytes("L", (width, height), bytes(pixels)) +def invert_mask_if_foreground_is_off_center(candidate: MaskCandidate) -> MaskCandidate: + """Invert a mask when its foreground is less concentrated at image center. + + This heuristic is intended for generated single-object images, where the + object is expected near the center and SAM3 may return its background. + """ + mask = decode_rle_mask(candidate.mask_rle) + width, height = mask.size + left, top = width // 6, height // 6 + right, bottom = width - left, height - top + center_mask = mask.crop((left, top, right, bottom)) + + center_foreground_ratio = _foreground_ratio(center_mask) + total_foreground_pixels = _foreground_pixel_count(mask) + outside_foreground_pixels = total_foreground_pixels - _foreground_pixel_count( + center_mask + ) + outside_pixel_count = width * height - center_mask.width * center_mask.height + outside_foreground_ratio = outside_foreground_pixels / outside_pixel_count + if center_foreground_ratio >= outside_foreground_ratio: + return candidate + + inverted_mask = mask.point(lambda value: 0 if value else 255) + return MaskCandidate( + index=candidate.index, + mask_rle=_encode_binary_mask_rle(inverted_mask), + ) + + def union_overlapping_mask_candidates( candidates: list[MaskCandidate], *, @@ -277,6 +306,72 @@ def render_numbered_mask_candidates( return resolved_output_path +def render_asset_mask_id_overlay( + *, + image_path: str | Path, + asset_masks: list[tuple[str, str | Path]], + output_path: str | Path, +) -> Path: + """Overlay outlined asset masks and stable asset IDs on a scene image. + + The table mask is intentionally omitted so its large contour does not + obscure the asset labels or their visual context in the source image. + """ + asset_ids = [asset_id for asset_id, _ in asset_masks] + if any(not asset_id for asset_id in asset_ids): + raise ValueError("Every asset mask must have a non-empty asset id.") + if len(asset_ids) != len(set(asset_ids)): + raise ValueError("Asset mask ids must be unique.") + + image = Image.open(image_path).convert("RGBA") + overlay = Image.new("RGBA", image.size, (0, 0, 0, 0)) + colors = ( + (239, 83, 80, 255), + (66, 165, 245, 255), + (102, 187, 106, 255), + (255, 202, 40, 255), + (171, 71, 188, 255), + (38, 198, 218, 255), + ) + decoded_masks: list[tuple[str, Image.Image]] = [] + for index, (asset_id, mask_path) in enumerate(asset_masks): + mask = Image.open(mask_path).convert("L") + _require_image_size(mask, image.size) + decoded_masks.append((asset_id, mask)) + color_layer = Image.new("RGBA", image.size, colors[index % len(colors)]) + transparent_layer = Image.new("RGBA", image.size, (0, 0, 0, 0)) + overlay.alpha_composite( + Image.composite( + color_layer, + transparent_layer, + _mask_outer_outline(mask, image.size), + ) + ) + + draw = ImageDraw.Draw(overlay) + for asset_id, mask in decoded_masks: + bbox = mask.getbbox() + if bbox is None: + raise ValueError(f"Asset mask {asset_id!r} is empty.") + font = _load_asset_id_label_font( + image_size=image.size, + mask_bbox=bbox, + label=asset_id, + ) + _draw_number_label( + draw=draw, + label=asset_id, + center=((bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2), + font=font, + minimum_padding=2, + ) + + resolved_output_path = Path(output_path).expanduser().resolve() + resolved_output_path.parent.mkdir(parents=True, exist_ok=True) + Image.alpha_composite(image, overlay).convert("RGB").save(resolved_output_path) + return resolved_output_path + + def _require_image_size(mask: Image.Image, image_size: tuple[int, int]) -> None: if mask.size != image_size: raise ValueError( @@ -361,6 +456,19 @@ def _mask_iou(first_mask: Image.Image, second_mask: Image.Image) -> float: return intersection.histogram()[255] / union_pixels +def _foreground_pixel_count(mask: Image.Image) -> int: + """Return the number of white pixels in one binary mask.""" + return mask.convert("L").histogram()[255] + + +def _foreground_ratio(mask: Image.Image) -> float: + """Return the white-pixel ratio in one non-empty image region.""" + pixel_count = mask.width * mask.height + if pixel_count == 0: + raise ValueError("Mask region must contain at least one pixel.") + return _foreground_pixel_count(mask) / pixel_count + + def _encode_binary_mask_rle(mask: Image.Image) -> dict[str, Any]: binary_mask = mask.convert("L").point( lambda value: 255 if value else 0 @@ -401,6 +509,47 @@ def _union_parent(parents: list[int], first_index: int, second_index: int) -> No def _load_label_font(image_size: tuple[int, int]) -> ImageFont.ImageFont: font_size = max(16, round(min(image_size) / 32)) + return _load_label_font_at_size(font_size) + + +def _load_asset_id_label_font( + *, + image_size: tuple[int, int], + mask_bbox: tuple[int, int, int, int], + label: str, +) -> ImageFont.ImageFont: + """Choose an ID-label font constrained by both image and mask dimensions.""" + # The image sets the readable upper bound; the individual mask then caps it. + image_font_size = min(32, max(8, round(min(image_size) / 48))) + mask_width = mask_bbox[2] - mask_bbox[0] + mask_height = mask_bbox[3] - mask_bbox[1] + maximum_label_width = max(24, round(mask_width * 0.9)) + maximum_label_height = max(16, round(mask_height * 0.75)) + # Measure the complete text-and-background rectangle, not glyphs alone. + probe_draw = ImageDraw.Draw(Image.new("RGBA", image_size)) + smallest_font = _load_label_font_at_size(6) + for font_size in range(image_font_size, 5, -1): + font = _load_label_font_at_size(font_size) + label_bounds = _number_label_bounds( + draw=probe_draw, + label=label, + center=(0.0, 0.0), + font=font, + minimum_padding=2, + ) + if ( + label_bounds[2] - label_bounds[0] <= maximum_label_width + and label_bounds[3] - label_bounds[1] <= maximum_label_height + ): + return font + smallest_font = font + return smallest_font + + +def _load_label_font_at_size(font_size: int) -> ImageFont.ImageFont: + """Load the shared bold label font at one validated pixel size.""" + if font_size < 1: + raise ValueError("Label font size must be positive.") try: return ImageFont.truetype("DejaVuSans-Bold.ttf", font_size) except OSError: @@ -413,10 +562,15 @@ def _draw_number_label( label: str, center: tuple[float, float], font: ImageFont.ImageFont, + minimum_padding: int = 4, ) -> None: """Draw a numbered label with red background and white text at the given center position.""" label_bounds = _number_label_bounds( - draw=draw, label=label, center=center, font=font + draw=draw, + label=label, + center=center, + font=font, + minimum_padding=minimum_padding, ) label_box = draw.textbbox((0, 0), label, font=font) label_width = label_box[2] - label_box[0] @@ -438,12 +592,15 @@ def _number_label_bounds( label: str, center: tuple[float, float], font: ImageFont.ImageFont, + minimum_padding: int = 4, ) -> tuple[int, int, int, int]: """Return the red label rectangle bounds for a label centre.""" label_box = draw.textbbox((0, 0), label, font=font) label_width = label_box[2] - label_box[0] label_height = label_box[3] - label_box[1] - padding = max(4, round(max(label_width, label_height) / 4)) + if minimum_padding < 0: + raise ValueError("Label minimum padding must be non-negative.") + padding = max(minimum_padding, round(max(label_width, label_height) / 4)) x = center[0] - label_width / 2 y = center[1] - label_height / 2 return ( diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/parent_surface_layout_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/parent_surface_layout_optimizer.py new file mode 100644 index 000000000..6d2a3b0a5 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/parent_surface_layout_optimizer.py @@ -0,0 +1,546 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import numpy as np +from scipy.optimize import minimize + +from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraphRelation +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_utils import ( + load_scene_object_z_up_mesh, + measure_scene_object_z_up_world_aabb, +) + +if TYPE_CHECKING: + from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_constructor import ( + SceneLayoutGroup, + SceneLayoutProblem, + ) + + +@dataclass +class ParentSurfaceLayoutProblem: + """All geometry and layout-state inputs for one parent-surface solve.""" + + assets_by_id: dict[str, SceneObject] + child_ids: list[str] + child_seed_xy_by_id: dict[str, list[float]] + imported_child_ids: set[str] + fixed_child_xy_by_id: dict[str, list[float] | None] + parent_aabb_xy: list[list[float]] + parent_top_z: float + child_relations: list[SceneGraphRelation] + + @classmethod + def from_layout_problem( + cls, + *, + layout_problem: SceneLayoutProblem, + group: SceneLayoutGroup, + current_xy_by_id: dict[str, list[float] | None], + ) -> ParentSurfaceLayoutProblem: + """Build one parent-surface problem without mutating layout state.""" + assets_by_id = { + asset.id: asset for asset in layout_problem.post_edit_scene.assets + } + child_ids = set(group.child_ids) + parent = assets_by_id.get(group.parent_id) + if parent is None: + raise ValueError(f"Parent {group.parent_id!r} is not an asset.") + parent_aabb = measure_scene_object_z_up_world_aabb(scene_object=parent) + parent_aabb_xy = [ + [parent_aabb[0][0], parent_aabb[0][1]], + [parent_aabb[1][0], parent_aabb[1][1]], + ] + parent_center_xy = [ + (parent_aabb[0][0] + parent_aabb[1][0]) / 2.0, + (parent_aabb[0][1] + parent_aabb[1][1]) / 2.0, + ] + child_seed_xy_by_id = {} + for child_id in group.child_ids: + inherited_xy = current_xy_by_id[child_id] + # New children begin from the solved parent's AABB center. + child_seed_xy_by_id[child_id] = ( + parent_center_xy if inherited_xy is None else list(inherited_xy) + ) + return cls( + assets_by_id=assets_by_id, + child_ids=group.child_ids, + child_seed_xy_by_id=child_seed_xy_by_id, + imported_child_ids={ + child_id + for child_id in group.child_ids + if layout_problem.initial_xy_by_id[child_id] is not None + }, + fixed_child_xy_by_id={ + child_id: ( + None + if child_id in layout_problem.layout_variable_ids + else current_xy_by_id[child_id] + ) + for child_id in group.child_ids + }, + parent_aabb_xy=parent_aabb_xy, + parent_top_z=parent_aabb[1][2], + child_relations=[ + relation + for relation in layout_problem.goal_scene_graph.relations + if relation.source_id in child_ids and relation.target_id in child_ids + ], + ) + + +@dataclass(frozen=True) +class ParentSurfaceLayoutOptimizerConfig: + """Numerical controls for one non-table parent-surface sibling solve.""" + + relation_clearance_m: float = 0.03 + collision_margin_m: float = 0.02 + max_slsqp_iterations: int = 500 + slsqp_ftol: float = 1e-6 + max_collision_rounds: int = 8 + max_added_collision_pairs: int = 64 + imported_seed_weight: float = 5.0 + min_center_distance_m: float = 0.01 + min_center_distance_weight: float = 0.05 + + def __post_init__(self) -> None: + """Reject invalid controls before assembling parent-surface constraints.""" + if self.relation_clearance_m < 0.0: + raise ValueError("relation_clearance_m must be non-negative.") + if self.collision_margin_m < 0.0: + raise ValueError("collision_margin_m must be non-negative.") + if self.max_slsqp_iterations <= 0: + raise ValueError("max_slsqp_iterations must be positive.") + if self.slsqp_ftol <= 0.0: + raise ValueError("slsqp_ftol must be positive.") + if self.max_collision_rounds <= 0: + raise ValueError("max_collision_rounds must be positive.") + if self.max_added_collision_pairs <= 0: + raise ValueError("max_added_collision_pairs must be positive.") + + +class ParentSurfaceLayoutOptimizer: + """Solve direct ``on`` children inside one parent's current XY footprint.""" + + def __init__( + self, + *, + config: ParentSurfaceLayoutOptimizerConfig | None = None, + ) -> None: + self.config = ( + config if config is not None else ParentSurfaceLayoutOptimizerConfig() + ) + + def optimize( + self, + problem: ParentSurfaceLayoutProblem, + ) -> dict[str, list[float]]: + """Return sibling XY centers inside the parent AABB without overlap.""" + child_half_extents_xy = _asset_half_extents_xy( + assets_by_id=problem.assets_by_id, + object_ids=problem.child_ids, + ) + inequality_constraints, equality_constraints = _build_constraints( + problem=problem, + child_half_extents_xy=child_half_extents_xy, + config=self.config, + ) + solved_child_xy_by_id = _solve_root_xy( + root_ids=problem.child_ids, + root_seed_xy_by_id=problem.child_seed_xy_by_id, + imported_root_ids=problem.imported_child_ids, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + config=self.config, + ) + return _refine_root_collisions( + root_ids=problem.child_ids, + root_seed_xy_by_id=problem.child_seed_xy_by_id, + imported_root_ids=problem.imported_child_ids, + root_half_extents_xy=child_half_extents_xy, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + fixed_root_xy_by_id=problem.fixed_child_xy_by_id, + solved_root_xy_by_id=solved_child_xy_by_id, + config=self.config, + ) + + +class _LayoutInfeasibleError(ValueError): + """Internal marker for an SLSQP failure while testing one collision direction.""" + + +def _build_constraints( + *, + problem: ParentSurfaceLayoutProblem, + child_half_extents_xy: dict[str, np.ndarray], + config: ParentSurfaceLayoutOptimizerConfig, +) -> tuple[list[tuple[np.ndarray, float]], list[tuple[np.ndarray, float]]]: + """Build hard parent-AABB, planar-relation, and fixed-child constraints.""" + root_index = {child_id: index for index, child_id in enumerate(problem.child_ids)} + parent_bounds = _bounds_from_points(problem.parent_aabb_xy) + inequality_constraints: list[tuple[np.ndarray, float]] = [] + equality_constraints: list[tuple[np.ndarray, float]] = [] + for child_id in problem.child_ids: + # Keep each child's 2D AABB inside the parent AABB support proxy. + _append_aabb_center_bounds( + constraints=inequality_constraints, + root_index=root_index, + root_id=child_id, + bounds=parent_bounds, + half_extents_xy=child_half_extents_xy[child_id], + ) + fixed_xy = problem.fixed_child_xy_by_id[child_id] + if fixed_xy is not None: + _append_fixed_root_constraints( + constraints=equality_constraints, + root_index=root_index, + root_id=child_id, + fixed_xy=fixed_xy, + ) + for relation in problem.child_relations: + # Apply planar relations between direct on-children of this parent. + _append_planar_relation_constraint( + constraints=inequality_constraints, + root_index=root_index, + source_id=relation.source_id, + relation=relation.relation, + target_id=relation.target_id, + source_half_extents_xy=child_half_extents_xy[relation.source_id], + target_half_extents_xy=child_half_extents_xy[relation.target_id], + relation_clearance_m=config.relation_clearance_m, + ) + return inequality_constraints, equality_constraints + + +def _asset_half_extents_xy( + *, assets_by_id: dict[str, SceneObject], object_ids: list[str] +) -> dict[str, np.ndarray]: + """Measure each optimized child asset's z-up XY half-extents.""" + result = {} + for object_id in object_ids: + asset = assets_by_id.get(object_id) + if asset is None: + raise ValueError(f"Parent child {object_id!r} is not an asset.") + mesh = load_scene_object_z_up_mesh(scene_object=asset) + result[object_id] = (mesh.bounds[1, :2] - mesh.bounds[0, :2]) / 2.0 + return result + + +def _bounds_from_points(points: list[list[float]]) -> np.ndarray: + """Return finite XY minimum and maximum bounds from polygon points.""" + coordinates = np.asarray(points, dtype=float) + if ( + coordinates.ndim != 2 + or coordinates.shape[1] != 2 + or len(coordinates) < 2 + or not np.all(np.isfinite(coordinates)) + ): + raise ValueError("XY bounds must contain at least two finite points.") + return np.stack([coordinates.min(axis=0), coordinates.max(axis=0)]) + + +def _append_aabb_center_bounds( + *, + constraints: list[tuple[np.ndarray, float]], + root_index: dict[str, int], + root_id: str, + bounds: np.ndarray, + half_extents_xy: np.ndarray, +) -> None: + """Constrain one child AABB center to lie completely inside XY bounds.""" + minimum, maximum = bounds[0] + half_extents_xy, bounds[1] - half_extents_xy + if np.any(minimum > maximum): + raise ValueError(f"Asset {root_id!r} cannot fit inside its parent AABB.") + # root_id is the child whose center is constrained in this AABB bound. + offset, count = 2 * root_index[root_id], 2 * len(root_index) + # offset selects this child's XY pair; count is the full flattened XY vector size. + for axis in range(2): + upper, lower = np.zeros(count), np.zeros(count) + upper[offset + axis], lower[offset + axis] = 1.0, -1.0 + constraints.extend( + [(upper, float(maximum[axis])), (lower, -float(minimum[axis]))] + ) + + +def _append_fixed_root_constraints( + *, + constraints: list[tuple[np.ndarray, float]], + root_index: dict[str, int], + root_id: str, + fixed_xy: list[float], +) -> None: + """Lock one fixed child's center to its imported XY coordinates.""" + offset, count = 2 * root_index[root_id], 2 * len(root_index) + for axis, coordinate in enumerate(fixed_xy): + row = np.zeros(count) + row[offset + axis] = 1.0 + constraints.append((row, float(coordinate))) + + +def _append_planar_relation_constraint( + *, + constraints: list[tuple[np.ndarray, float]], + root_index: dict[str, int], + source_id: str, + relation: str, + target_id: str, + source_half_extents_xy: np.ndarray, + target_half_extents_xy: np.ndarray, + relation_clearance_m: float, +) -> None: + """Append one world-XY separation constraint for sibling planar semantics.""" + axis, sign = { + "left_of": (0, 1.0), + "right_of": (0, -1.0), + "behind": (1, 1.0), + "in_front_of": (1, -1.0), + }.get(relation, (None, None)) + if axis is None or source_id not in root_index or target_id not in root_index: + raise ValueError(f"Unsupported parent-child planar relation {relation!r}.") + row = np.zeros(2 * len(root_index)) + row[2 * root_index[source_id] + axis] = sign + row[2 * root_index[target_id] + axis] = -sign + constraints.append( + ( + row, + -float( + source_half_extents_xy[axis] + + target_half_extents_xy[axis] + + relation_clearance_m + ), + ) + ) + + +def _solve_root_xy( + *, + root_ids: list[str], + root_seed_xy_by_id: dict[str, list[float]], + imported_root_ids: set[str], + inequality_constraints: list[tuple[np.ndarray, float]], + equality_constraints: list[tuple[np.ndarray, float]], + config: ParentSurfaceLayoutOptimizerConfig, +) -> dict[str, list[float]]: + """Solve one parent child-group's XY positions with SLSQP.""" + initial = np.asarray( + [root_seed_xy_by_id[root_id] for root_id in root_ids], dtype=float + ) + + def objective(values: np.ndarray) -> float: + xy = values.reshape(-1, 2) + loss = 0.0 + for index, root_id in enumerate(root_ids): + if root_id in imported_root_ids: + delta = xy[index] - initial[index] + loss += config.imported_seed_weight * float(delta @ delta) + return loss + + constraints = [ + { + "type": "ineq", + "fun": lambda values, row=row, bound=bound: bound - float(row @ values), + } + for row, bound in inequality_constraints + ] + [ + { + "type": "eq", + "fun": lambda values, row=row, bound=bound: float(row @ values) - bound, + } + for row, bound in equality_constraints + ] + result = minimize( + objective, + initial.reshape(-1), + method="SLSQP", + constraints=constraints, + options={ + "maxiter": config.max_slsqp_iterations, + "ftol": config.slsqp_ftol, + "disp": False, + }, + ) + if not result.success: + raise _LayoutInfeasibleError( + f"Parent layout optimization failed: {result.message}" + ) + return { + root_id: [float(result.x[2 * index]), float(result.x[2 * index + 1])] + for index, root_id in enumerate(root_ids) + } + + +def _refine_root_collisions( + *, + root_ids: list[str], + root_seed_xy_by_id: dict[str, list[float]], + imported_root_ids: set[str], + root_half_extents_xy: dict[str, np.ndarray], + inequality_constraints: list[tuple[np.ndarray, float]], + equality_constraints: list[tuple[np.ndarray, float]], + fixed_root_xy_by_id: dict[str, list[float] | None], + solved_root_xy_by_id: dict[str, list[float]], + config: ParentSurfaceLayoutOptimizerConfig, +) -> dict[str, list[float]]: + """Iteratively add separation constraints for overlapping child AABBs.""" + current = solved_root_xy_by_id + seen: set[tuple[str, str]] = set() + for _ in range(config.max_collision_rounds): + overlaps = [ + pair + for pair in _root_aabb_overlaps( + root_ids=root_ids, half_extents=root_half_extents_xy, xy_by_id=current + ) + if fixed_root_xy_by_id[pair[1]] is None + or fixed_root_xy_by_id[pair[2]] is None + ] + if not overlaps: + return current + added = 0 + for _, first_id, second_id in overlaps[: config.max_added_collision_pairs]: + key = tuple(sorted((first_id, second_id))) + if key in seen: + continue + # Earlier pair updates may already have separated this stale overlap. + if key not in { + tuple(sorted((first, second))) + for _, first, second in _root_aabb_overlaps( + root_ids=root_ids, + half_extents=root_half_extents_xy, + xy_by_id=current, + ) + }: + continue + for separation_constraint in _aabb_separation_constraints( + root_ids=root_ids, + first_id=first_id, + second_id=second_id, + half_extents=root_half_extents_xy, + xy_by_id=current, + margin=config.collision_margin_m, + ): + # Keep a candidate only when it is compatible with all hard constraints. + try: + solved_xy_by_id = _solve_root_xy( + root_ids=root_ids, + root_seed_xy_by_id=current, + imported_root_ids=imported_root_ids, + inequality_constraints=[ + *inequality_constraints, + separation_constraint, + ], + equality_constraints=equality_constraints, + config=config, + ) + except _LayoutInfeasibleError: + continue + inequality_constraints.append(separation_constraint) + current = solved_xy_by_id + seen.add(key) + added += 1 + break + else: + raise ValueError( + "Parent-child AABB pair has no feasible separation direction: " + f"{first_id!r}, {second_id!r}." + ) + if not added: + break + raise ValueError("Parent-child AABB collisions remain after layout refinement.") + + +def _root_aabb_overlaps( + *, + root_ids: list[str], + half_extents: dict[str, np.ndarray], + xy_by_id: dict[str, list[float]], +) -> list[tuple[float, str, str]]: + """Return overlapping child pairs with their minimum XY overlap distance.""" + result = [] + for index, first_id in enumerate(root_ids): + for second_id in root_ids[index + 1 :]: + overlap = np.minimum( + np.asarray(xy_by_id[first_id]) + half_extents[first_id], + np.asarray(xy_by_id[second_id]) + half_extents[second_id], + ) - np.maximum( + np.asarray(xy_by_id[first_id]) - half_extents[first_id], + np.asarray(xy_by_id[second_id]) - half_extents[second_id], + ) + if np.all(overlap > 1e-9): + result.append((float(np.min(overlap)), first_id, second_id)) + return sorted(result, reverse=True) + + +def _aabb_separation_constraints( + *, + root_ids: list[str], + first_id: str, + second_id: str, + half_extents: dict[str, np.ndarray], + xy_by_id: dict[str, list[float]], + margin: float, +) -> list[tuple[np.ndarray, float]]: + """Return ordered feasible-direction candidates for one overlapping AABB pair.""" + first, second = np.asarray(xy_by_id[first_id]), np.asarray(xy_by_id[second_id]) + overlap = np.minimum( + first + half_extents[first_id], second + half_extents[second_id] + ) - np.maximum(first - half_extents[first_id], second - half_extents[second_id]) + axes = np.argsort(overlap) + constraints = [] + for axis in axes: + current_order = first[axis] < second[axis] or ( + first[axis] == second[axis] and first_id < second_id + ) + for first_is_lower in (current_order, not current_order): + constraints.append( + _aabb_separation_constraint_for_direction( + root_ids=root_ids, + first_id=first_id, + second_id=second_id, + half_extents=half_extents, + axis=int(axis), + first_is_lower=first_is_lower, + margin=margin, + ) + ) + return constraints + + +def _aabb_separation_constraint_for_direction( + *, + root_ids: list[str], + first_id: str, + second_id: str, + half_extents: dict[str, np.ndarray], + axis: int, + first_is_lower: bool, + margin: float, +) -> tuple[np.ndarray, float]: + """Return one directed AABB separation inequality on a selected axis.""" + index = {root_id: i for i, root_id in enumerate(root_ids)} + row = np.zeros(2 * len(root_ids)) + sign = 1.0 if first_is_lower else -1.0 + row[2 * index[first_id] + axis], row[2 * index[second_id] + axis] = sign, -sign + return row, -float( + half_extents[first_id][axis] + half_extents[second_id][axis] + margin + ) diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py index fb66c30c4..1ad106ac0 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py @@ -26,6 +26,7 @@ from scipy.spatial.transform import Rotation from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraph from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject from embodichain.utils.logger import log_info @@ -40,18 +41,28 @@ class SceneExporter: - """Write one generated scene and its SimReady meshes as a scene export.""" + """Write one generated scene and its SimReady meshes as a scene export. + + By default, the complete z-up scene is rotated 180 degrees around the + table center before serialization. The input ``Scene`` is not mutated. + """ def __init__( self, *, scene: Scene, + scene_graph: SceneGraph, output_root: str | Path, + rotate_z_up_180: bool = True, ) -> None: self.scene = scene + self.scene_graph = scene_graph self.output_root = Path(output_root).expanduser().resolve() + self.rotate_z_up_180 = rotate_z_up_180 # Keep the legacy frame on request. self.export_root = self.output_root / "scene_export" self.scene_config_path: Path | None = None + self.scene_graph_path: Path | None = None + self.scene_json_path: Path | None = None def export(self) -> Path: """Write a scene-only config and copy SimReady GLBs into ``mesh_assets``. @@ -59,7 +70,9 @@ def export(self) -> Path: Scene layouts are y-up. The simulator automatically converts each y-up GLB to z-up, so this exporter copies each GLB unchanged and converts only its world position and rotation for ``init_pos`` and ``init_rot``. - ``body_scale`` remains the original y-up scale associated with the GLB. + The default applies one additional 180-degree global z-up rotation about + the table center to every object and XY layout metadata. ``body_scale`` + remains the original y-up scale associated with the GLB. This is not a complete ``EmbodiedEnv``/``run-env`` configuration because a generated scene does not determine a robot, its placement, or control. """ @@ -72,6 +85,11 @@ def export(self) -> Path: object_ids = [scene_object.id for scene_object in scene_objects] if len(set(object_ids)) != len(object_ids): raise ValueError("Scene export requires unique table and asset ids.") + self.scene_graph.validate() + if set(self.scene_graph.node_by_id()) != set(object_ids): + raise ValueError("Scene graph nodes must match exported scene object ids.") + + z_up_rotation, z_up_pivot_xy = self._z_up_export_transform() exported_entries = { scene_object.id: self._copy_scene_object_to_assets( @@ -90,12 +108,16 @@ def export(self) -> Path: self._scene_object_config( scene_object=self.scene.table, asset_relative_path=exported_entries[self.scene.table.id], + z_up_rotation=z_up_rotation, + z_up_pivot_xy=z_up_pivot_xy, ) ], "rigid_object": [ self._scene_object_config( scene_object=asset, asset_relative_path=exported_entries[asset.id], + z_up_rotation=z_up_rotation, + z_up_pivot_xy=z_up_pivot_xy, ) for asset in self.scene.assets ], @@ -106,6 +128,37 @@ def export(self) -> Path: encoding="utf-8", ) log_info(f"Exported scene config: {self.scene_config_path}") + self.scene_graph_path = self.export_root / "scene_graph.json" + self.scene_graph_path.write_text( + json.dumps(self.scene_graph.to_dict(), indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + log_info(f"Exported scene graph: {self.scene_graph_path}") + self.scene_json_path = self.export_root / "scene.json" + self.scene_json_path.write_text( + json.dumps( + { + "objects": [ + self._scene_object_y_up_dict( + scene_object=scene_object, + z_up_rotation=z_up_rotation, + z_up_pivot_xy=z_up_pivot_xy, + ) + for scene_object in scene_objects + ] + }, + indent=2, + ensure_ascii=False, + ) + + "\n", + encoding="utf-8", + ) + log_info(f"Exported scene JSON: {self.scene_json_path}") + # Remove only assets absent from the completed scene export. + self._remove_stale_mesh_assets( + mesh_assets_root=mesh_assets_root, + object_ids=set(object_ids), + ) return self.scene_config_path @staticmethod @@ -135,28 +188,64 @@ def _copy_scene_object_to_assets( ) destination_glb_path = mesh_assets_root / object_id / f"{object_id}.glb" destination_glb_path.parent.mkdir(parents=True, exist_ok=True) - shutil.copy2(source_glb_path, destination_glb_path) + # Imported assets already live at their export destination. + if not destination_glb_path.is_file() or not source_glb_path.samefile( + destination_glb_path + ): + shutil.copy2(source_glb_path, destination_glb_path) return destination_glb_path.relative_to(mesh_assets_root.parent).as_posix() + @staticmethod + def _remove_stale_mesh_assets( + *, + mesh_assets_root: Path, + object_ids: set[str], + ) -> None: + """Remove copied asset directories that no longer belong to the scene.""" + for asset_root in mesh_assets_root.iterdir(): + if asset_root.name in object_ids: + continue + if asset_root.is_dir(): + shutil.rmtree(asset_root) + else: + asset_root.unlink() + + def _z_up_export_transform(self) -> tuple[np.ndarray, np.ndarray]: + """Return the optional global z-up rotation and its table-center pivot.""" + table = self.scene.table + if table is None: + raise ValueError("Cannot transform a scene export without a table.") + table_pos_z_up, _ = self._final_z_up_pose( + scene_object=table, + z_up_rotation=np.eye(3), + z_up_pivot_xy=np.zeros(2), + ) + z_up_rotation = ( + Rotation.from_euler("z", 180.0, degrees=True).as_matrix() + if self.rotate_z_up_180 + else np.eye(3) + ) + return z_up_rotation, table_pos_z_up[:2] + @staticmethod def _scene_object_config( *, scene_object: SceneObject, asset_relative_path: str, + z_up_rotation: np.ndarray, + z_up_pivot_xy: np.ndarray, ) -> dict[str, object]: """Build one z-up scene-only object config from a final y-up object.""" - pos_y_up = SceneExporter._scene_vector(scene_object, "pos") - rot_y_up = SceneExporter._scene_vector(scene_object, "rot") scale_y_up = SceneExporter._scene_vector(scene_object, "scale") if scene_object.physics is None: raise ValueError( f"Scene object {scene_object.id!r} has no SimReady physics settings." ) - pos_z_up = _Y_UP_TO_Z_UP_ROTATION @ np.asarray(pos_y_up, dtype=float) - rotation_y_up = Rotation.from_euler("xyz", rot_y_up, degrees=True).as_matrix() - rotation_z_up = ( - _Y_UP_TO_Z_UP_ROTATION @ rotation_y_up @ _Y_UP_TO_Z_UP_ROTATION.T + pos_z_up, rotation_z_up = SceneExporter._final_z_up_pose( + scene_object=scene_object, + z_up_rotation=z_up_rotation, + z_up_pivot_xy=z_up_pivot_xy, ) rot_z_up = Rotation.from_matrix(rotation_z_up).as_euler( # RigidObjectCfg.init_rot is interpreted with uppercase XYZ. @@ -166,6 +255,8 @@ def _scene_object_config( return { "uid": scene_object.id, + "category": scene_object.category, + "name": scene_object.name, "description": scene_object.description, "shape": { "shape_type": "Mesh", @@ -179,9 +270,136 @@ def _scene_object_config( # Do not permute this scale: it belongs to the original y-up GLB, # which SimulationManager itself converts to z-up. "body_scale": scale_y_up, + "center_xy": SceneExporter._transformed_optional_xy( + scene_object=scene_object, + field_name="center_xy", + z_up_rotation=z_up_rotation, + z_up_pivot_xy=z_up_pivot_xy, + ), + "support_surface_z": scene_object.support_surface_z, + "support_contour_xy": SceneExporter._transformed_optional_xy_points( + scene_object=scene_object, + field_name="support_contour_xy", + z_up_rotation=z_up_rotation, + z_up_pivot_xy=z_up_pivot_xy, + ), + "support_optimization_rect_xy": ( + SceneExporter._transformed_optional_xy_points( + scene_object=scene_object, + field_name="support_optimization_rect_xy", + z_up_rotation=z_up_rotation, + z_up_pivot_xy=z_up_pivot_xy, + ) + ), "max_convex_hull_num": scene_object.physics.max_convex_hull_num, } + @staticmethod + def _final_z_up_pose( + *, + scene_object: SceneObject, + z_up_rotation: np.ndarray, + z_up_pivot_xy: np.ndarray, + ) -> tuple[np.ndarray, np.ndarray]: + """Convert one y-up pose and apply the export's global z-up rotation.""" + pos_y_up = SceneExporter._scene_vector(scene_object, "pos") + rot_y_up = SceneExporter._scene_vector(scene_object, "rot") + pos_z_up = _Y_UP_TO_Z_UP_ROTATION @ np.asarray(pos_y_up, dtype=float) + pos_z_up[:2] = z_up_pivot_xy + z_up_rotation[:2, :2] @ ( + pos_z_up[:2] - z_up_pivot_xy + ) + rotation_y_up = Rotation.from_euler("xyz", rot_y_up, degrees=True).as_matrix() + rotation_z_up = z_up_rotation @ ( + _Y_UP_TO_Z_UP_ROTATION @ rotation_y_up @ _Y_UP_TO_Z_UP_ROTATION.T + ) + return pos_z_up, rotation_z_up + + @staticmethod + def _scene_object_y_up_dict( + *, + scene_object: SceneObject, + z_up_rotation: np.ndarray, + z_up_pivot_xy: np.ndarray, + ) -> dict[str, object]: + """Serialize the globally rotated export without mutating the input scene.""" + pos_z_up, rotation_z_up = SceneExporter._final_z_up_pose( + scene_object=scene_object, + z_up_rotation=z_up_rotation, + z_up_pivot_xy=z_up_pivot_xy, + ) + result = scene_object.to_dict() + result["pos"] = (_Y_UP_TO_Z_UP_ROTATION.T @ pos_z_up).tolist() + result["rot"] = ( + Rotation.from_matrix( + _Y_UP_TO_Z_UP_ROTATION.T @ rotation_z_up @ _Y_UP_TO_Z_UP_ROTATION + ) + .as_euler("xyz", degrees=True) + .tolist() + ) + result["center_xy"] = SceneExporter._transformed_optional_xy( + scene_object=scene_object, + field_name="center_xy", + z_up_rotation=z_up_rotation, + z_up_pivot_xy=z_up_pivot_xy, + ) + result["support_contour_xy"] = SceneExporter._transformed_optional_xy_points( + scene_object=scene_object, + field_name="support_contour_xy", + z_up_rotation=z_up_rotation, + z_up_pivot_xy=z_up_pivot_xy, + ) + result["support_optimization_rect_xy"] = ( + SceneExporter._transformed_optional_xy_points( + scene_object=scene_object, + field_name="support_optimization_rect_xy", + z_up_rotation=z_up_rotation, + z_up_pivot_xy=z_up_pivot_xy, + ) + ) + return result + + @staticmethod + def _transformed_optional_xy( + *, + scene_object: SceneObject, + field_name: str, + z_up_rotation: np.ndarray, + z_up_pivot_xy: np.ndarray, + ) -> list[float] | None: + """Rotate one optional z-up XY metadata point around the table center.""" + value = getattr(scene_object, field_name) + if value is None: + return None + point = np.asarray(value, dtype=float) + if point.shape != (2,) or not np.all(np.isfinite(point)): + raise ValueError( + f"Scene object {scene_object.id!r} has invalid {field_name!r} metadata." + ) + return ( + z_up_pivot_xy + z_up_rotation[:2, :2] @ (point - z_up_pivot_xy) + ).tolist() + + @staticmethod + def _transformed_optional_xy_points( + *, + scene_object: SceneObject, + field_name: str, + z_up_rotation: np.ndarray, + z_up_pivot_xy: np.ndarray, + ) -> list[list[float]] | None: + """Rotate optional z-up XY support geometry around the table center.""" + value = getattr(scene_object, field_name) + if value is None: + return None + points = np.asarray(value, dtype=float) + if points.ndim != 2 or points.shape[1] != 2 or not np.all(np.isfinite(points)): + raise ValueError( + f"Scene object {scene_object.id!r} has invalid {field_name!r} metadata." + ) + return ( + z_up_pivot_xy + (z_up_rotation[:2, :2] @ (points - z_up_pivot_xy).T).T + ).tolist() + @staticmethod def _scene_vector(scene_object: SceneObject, field_name: str) -> list[float]: """Read one finite final y-up layout vector from a scene object.""" diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py new file mode 100644 index 000000000..e730b88cb --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py @@ -0,0 +1,417 @@ +# ---------------------------------------------------------------------------- +# 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 typing import Any + +import numpy as np +from scipy.spatial.transform import Rotation + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, + SceneGraphRelation, + TABLE_REGIONS, +) +from embodichain.gen_sim.scene_engine.core.scene_object import ( + ObjectPhysics, + SceneObject, +) +from embodichain.utils.logger import log_info + +_Y_UP_TO_Z_UP_ROTATION = np.array( + [ + [1.0, 0.0, 0.0], + [0.0, 0.0, -1.0], + [0.0, 1.0, 0.0], + ], + dtype=float, +) +_Z_UP_TO_Y_UP_ROTATION = _Y_UP_TO_Z_UP_ROTATION.T + + +class SceneExportImporter: + """Import an editable ``Scene`` from an exported Scene Engine directory.""" + + def __init__( + self, + *, + output_root: str | Path, + ) -> None: + self.output_root = Path(output_root).expanduser().resolve() + self.scene_export_root = self.output_root / "scene_export" + self.mesh_assets_root = self.scene_export_root / "mesh_assets" + self.scene_config_path = self.scene_export_root / "scene_config.json" + self.scene_graph_path = self.scene_export_root / "scene_graph.json" + self.scene_json_path = self.scene_export_root / "scene.json" + + def import_scene(self) -> Scene: + """Validate the scene export, write ``scene.json``, and return a ``Scene``.""" + scene = self._load_scene() + self._write_scene_json(scene) + return scene + + def import_scene_and_graph(self) -> tuple[Scene, SceneGraph]: + """Import a scene and graph after validating the complete edit input.""" + scene = self._load_scene() + scene_graph = self._load_scene_graph() + if set(scene_graph.node_by_id()) != { + scene_object.id for scene_object in scene.objects + }: + raise ValueError("Scene graph nodes must match imported scene object ids.") + self._write_scene_json(scene) + return scene, scene_graph + + def _load_scene(self) -> Scene: + """Validate the exported scene files and restore the ``Scene`` data.""" + # Editing only runs on an existing Scene Engine output directory. + if not self.output_root.is_dir() or not any(self.output_root.iterdir()): + raise ValueError( + "Output root must exist and contain files when edit_prompt is provided." + ) + + # The editor consumes the portable scene export and its copied GLB assets. + if not self.scene_export_root.is_dir(): + raise FileNotFoundError( + f"Scene export directory not found: {self.scene_export_root}" + ) + if not self.mesh_assets_root.is_dir(): + raise FileNotFoundError( + f"Scene mesh assets directory not found: {self.mesh_assets_root}" + ) + if not self.scene_config_path.is_file(): + raise FileNotFoundError(f"Scene config not found: {self.scene_config_path}") + + try: + scene_config = json.loads( + self.scene_config_path.read_text(encoding="utf-8") + ) + except json.JSONDecodeError as exc: + raise ValueError( + f"Scene config is not valid JSON: {self.scene_config_path}" + ) from exc + if not isinstance(scene_config, dict): + raise ValueError("Scene config must be a JSON object.") + + return self._scene_from_config(scene_config) + + def _load_scene_graph(self) -> SceneGraph: + """Read and validate the exported scene graph.""" + if not self.scene_graph_path.is_file(): + raise FileNotFoundError(f"Scene graph not found: {self.scene_graph_path}") + try: + scene_graph_data = json.loads( + self.scene_graph_path.read_text(encoding="utf-8") + ) + except json.JSONDecodeError as exc: + raise ValueError( + f"Scene graph is not valid JSON: {self.scene_graph_path}" + ) from exc + return self._scene_graph_from_data(scene_graph_data) + + def _write_scene_json(self, scene: Scene) -> None: + """Write the restored scene debugging artifact after validation succeeds.""" + self.scene_json_path.write_text( + json.dumps(scene.to_dict(), indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + log_info(f"Imported scene JSON: {self.scene_json_path}") + + def _scene_from_config(self, scene_config: dict[str, Any]) -> Scene: + """Build a y-up ``Scene`` from the z-up scene-export config.""" + # The table is the required support object for scene-edit operations. + background = scene_config.get("background", []) + if not isinstance(background, list): + raise ValueError("Scene config background must be a list.") + table_entry = next( + ( + scene_object + for scene_object in background + if isinstance(scene_object, dict) and scene_object.get("uid") == "table" + ), + None, + ) + if table_entry is None: + raise ValueError("Scene config background must contain a table entry.") + + rigid_object_entries = scene_config.get("rigid_object", []) + if not isinstance(rigid_object_entries, list): + raise ValueError("Scene config rigid_object must be a list.") + + return Scene( + objects=[ + self._scene_object_from_export_entry(table_entry, kind="table"), + *[ + self._scene_object_from_export_entry(entry, kind="asset") + for entry in rigid_object_entries + ], + ] + ) + + @staticmethod + def _scene_graph_from_data(value: object) -> SceneGraph: + """Build a validated ``SceneGraph`` from exported graph JSON.""" + if not isinstance(value, dict) or set(value) != {"nodes", "relations"}: + raise ValueError("Scene graph must contain exactly nodes and relations.") + nodes_value = value["nodes"] + relations_value = value["relations"] + if not isinstance(nodes_value, list) or not isinstance(relations_value, list): + raise ValueError("Scene graph nodes and relations must be lists.") + + nodes = [ + SceneExportImporter._scene_graph_node_from_data(node) + for node in nodes_value + ] + relations = [ + SceneExportImporter._scene_graph_relation_from_data(relation) + for relation in relations_value + ] + return SceneGraph(nodes=nodes, relations=relations) + + @staticmethod + def _scene_graph_node_from_data(value: object) -> SceneGraphNode: + if not isinstance(value, dict) or set(value) != { + "object_id", + "parent_id", + "parent_relation", + "table_region", + "orientation_state", + }: + raise ValueError("Scene graph nodes must use the serialized node schema.") + object_id = value["object_id"] + parent_id = value["parent_id"] + parent_relation = value["parent_relation"] + table_region = value["table_region"] + orientation_state = value["orientation_state"] + if not isinstance(object_id, str) or not isinstance( + parent_id, (str, type(None)) + ): + raise ValueError("Scene graph node ids must be strings or null.") + if parent_relation not in {None, "on"}: + raise ValueError("Scene graph parent_relation must be 'on' or null.") + if table_region is not None and table_region not in TABLE_REGIONS: + raise ValueError("Scene graph table_region is invalid.") + if orientation_state not in {None, "standing", "lying"}: + raise ValueError("Scene graph orientation_state is invalid.") + return SceneGraphNode( + object_id=object_id, + parent_id=parent_id, + parent_relation=parent_relation, + table_region=table_region, + orientation_state=orientation_state, + ) + + @staticmethod + def _scene_graph_relation_from_data(value: object) -> SceneGraphRelation: + if not isinstance(value, dict) or set(value) != { + "source_id", + "relation", + "target_id", + }: + raise ValueError( + "Scene graph relations must use the serialized relation schema." + ) + source_id = value["source_id"] + relation = value["relation"] + target_id = value["target_id"] + if not isinstance(source_id, str) or not isinstance(target_id, str): + raise ValueError("Scene graph relation ids must be strings.") + if relation not in {"left_of", "right_of", "in_front_of", "behind"}: + raise ValueError("Scene graph relation is invalid.") + return SceneGraphRelation( + source_id=source_id, + relation=relation, + target_id=target_id, + ) + + def _scene_object_from_export_entry( + self, + entry: object, + *, + kind: str, + ) -> SceneObject: + """Convert one z-up scene-export entry back to a y-up ``SceneObject``.""" + if not isinstance(entry, dict): + raise ValueError("Scene config entries must be objects.") + uid = entry.get("uid") + if not isinstance(uid, str) or not uid: + raise ValueError("Scene config entries must contain a valid uid.") + + glb_path = self._resolve_export_glb_path(entry, uid=uid) + pos_z_up = self._vector3( + entry.get("init_pos", [0.0, 0.0, 0.0]), + field_name=f"{uid}.init_pos", + ) + rot_z_up = self._vector3( + entry.get("init_rot", [0.0, 0.0, 0.0]), + field_name=f"{uid}.init_rot", + ) + scale = self._vector3( + entry.get("body_scale", [1.0, 1.0, 1.0]), + field_name=f"{uid}.body_scale", + ) + center_xy = entry.get("center_xy") + if center_xy is not None: + center_xy = self._vector2(center_xy, field_name=f"{uid}.center_xy") + support_surface_z = entry.get("support_surface_z") + if support_surface_z is not None: + support_surface_z = float(support_surface_z) + support_contour_xy = self._points2( + entry.get("support_contour_xy"), field_name=f"{uid}.support_contour_xy" + ) + support_optimization_rect_xy = self._points2( + entry.get("support_optimization_rect_xy"), + field_name=f"{uid}.support_optimization_rect_xy", + ) + + pos_y_up = _Z_UP_TO_Y_UP_ROTATION @ np.asarray(pos_z_up, dtype=float) + rotation_z_up = Rotation.from_euler("XYZ", rot_z_up, degrees=True).as_matrix() + rotation_y_up = ( + _Z_UP_TO_Y_UP_ROTATION @ rotation_z_up @ _Z_UP_TO_Y_UP_ROTATION.T + ) + rot_y_up = Rotation.from_matrix(rotation_y_up).as_euler("xyz", degrees=True) + + return SceneObject( + id=uid, + kind=kind, # type: ignore[arg-type] + category=self._semantic_text( + entry.get("category"), + field_name=f"{uid}.category", + default=uid, + ), + name=self._semantic_text( + entry.get("name"), + field_name=f"{uid}.name", + default=uid, + ), + description=str(entry.get("description") or uid), + simready_glb_path=str(glb_path), + rot=rot_y_up.tolist(), + pos=pos_y_up.tolist(), + scale=scale, + center_xy=center_xy, + support_surface_z=support_surface_z, + support_contour_xy=support_contour_xy, + support_optimization_rect_xy=support_optimization_rect_xy, + physics=ObjectPhysics( + body_type=str(entry.get("body_type", "dynamic")), # type: ignore[arg-type] + attrs=self._physics_attrs(entry.get("attrs", {"mass": 1.0})), + max_convex_hull_num=max(1, int(entry.get("max_convex_hull_num", 32))), + ), + ) + + def _resolve_export_glb_path( + self, + entry: dict[str, Any], + *, + uid: str, + ) -> Path: + """Validate one exported mesh reference and return its absolute GLB path.""" + shape = entry.get("shape") + if not isinstance(shape, dict) or not isinstance(shape.get("fpath"), str): + raise ValueError(f"Scene object {uid!r} must contain shape.fpath.") + fpath = Path(shape["fpath"]) + if fpath.is_absolute(): + raise ValueError(f"Scene object {uid!r} shape.fpath must be relative.") + if fpath.suffix.lower() != ".glb": + raise ValueError(f"Scene object {uid!r} shape.fpath must point to a GLB.") + expected_fpath = Path("mesh_assets") / uid / f"{uid}.glb" + if fpath != expected_fpath: + raise ValueError( + f"Scene object {uid!r} shape.fpath must be {expected_fpath.as_posix()!r}." + ) + glb_path = (self.scene_export_root / fpath).resolve() + if self.scene_export_root.resolve() not in glb_path.parents: + raise ValueError( + f"Scene object {uid!r} shape.fpath must stay within " + f"{self.scene_export_root.resolve()}." + ) + if not glb_path.is_file(): + raise FileNotFoundError(f"Scene object {uid!r} GLB not found: {glb_path}") + return glb_path + + @staticmethod + def _vector3(value: object, *, field_name: str) -> list[float]: + """Validate one length-3 numeric vector.""" + if not isinstance(value, list) or len(value) != 3: + raise ValueError( + f"Scene config field {field_name!r} must be a length-3 list." + ) + vector = [float(item) for item in value] + if not np.all(np.isfinite(vector)): + raise ValueError(f"Scene config field {field_name!r} must be finite.") + return vector + + @staticmethod + def _semantic_text( + value: object, + *, + field_name: str, + default: str, + ) -> str: + """Read one non-empty semantic label with a legacy-export fallback.""" + if value is None: + return default + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"Scene config field {field_name!r} must be non-empty.") + return value + + @staticmethod + def _vector2(value: object, *, field_name: str) -> list[float]: + """Validate one length-2 numeric vector.""" + if not isinstance(value, list) or len(value) != 2: + raise ValueError( + f"Scene config field {field_name!r} must be a length-2 list." + ) + vector = [float(item) for item in value] + if not np.all(np.isfinite(vector)): + raise ValueError(f"Scene config field {field_name!r} must be finite.") + return vector + + @classmethod + def _points2(cls, value: object, *, field_name: str) -> list[list[float]] | None: + """Validate an optional list of XY points from the scene export.""" + if value is None: + return None + if not isinstance(value, list) or len(value) < 3: + raise ValueError( + f"Scene config field {field_name!r} must contain 3 points." + ) + return [ + cls._vector2(point, field_name=f"{field_name}[{index}]") + for index, point in enumerate(value) + ] + + @staticmethod + def _physics_attrs(value: object) -> dict[str, float | int]: + """Validate exported physics attributes.""" + if not isinstance(value, dict) or not value: + raise ValueError("Scene object attrs must be a non-empty object.") + attrs: dict[str, float | int] = {} + for key, item in value.items(): + if not isinstance(key, str) or not isinstance(item, (float, int)): + raise ValueError("Scene object attrs must map strings to numbers.") + attrs[key] = item + return attrs + + +def import_scene_from_output_root(output_root: str | Path) -> Scene: + """Import an editable ``Scene`` from ``scene_export/scene_config.json``.""" + return SceneExportImporter(output_root=output_root).import_scene() diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py new file mode 100644 index 000000000..95c77d712 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py @@ -0,0 +1,353 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + TABLE_OBJECT_ID, + SceneGraph, +) +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.parent_surface_layout_optimizer import ( + ParentSurfaceLayoutOptimizer, + ParentSurfaceLayoutOptimizerConfig, + ParentSurfaceLayoutProblem, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_utils import ( + translate_scene_object_y_up_by_z_up_delta, + update_scene_object_y_up_pose_from_z_up_support, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.table_surface_layout_optimizer import ( + TableSurfaceLayoutOptimizer, + TableSurfaceLayoutOptimizerConfig, + TableSurfaceLayoutProblem, +) + + +@dataclass(frozen=True) +class SceneLayoutGroup: + """One parent and its direct on-children handled in one layout pass.""" + + parent_id: str + child_ids: list[str] + + +@dataclass(frozen=True) +class SceneLayoutProblem: + """Prepared graph-constrained inputs for one scene-layout construction.""" + + post_edit_scene: Scene + goal_scene_graph: SceneGraph + layout_variable_ids: set[str] + initial_xy_by_id: dict[str, list[float] | None] + groups: list[SceneLayoutGroup] + + +class SceneLayoutConstructor: + """Construct a scene layout from its goal graph. + + ``formal_scene`` may be empty for text-to-scene. In that case every table + and asset object must be supplied through ``generated_scene_objects``. + """ + + def __init__( + self, + *, + formal_scene: Scene, + goal_scene_graph: SceneGraph, + layout_variable_ids: set[str], + generated_scene_objects: list[SceneObject], + output_root: str | Path, + table_surface_config: TableSurfaceLayoutOptimizerConfig | None = None, + parent_surface_config: ParentSurfaceLayoutOptimizerConfig | None = None, + ) -> None: + self.formal_scene = formal_scene + self.goal_scene_graph = goal_scene_graph + self.layout_variable_ids = layout_variable_ids + self.generated_scene_objects = generated_scene_objects + self.output_root = Path(output_root).expanduser().resolve() + # Table surface optimizer. + self.table_surface_layout_optimizer = TableSurfaceLayoutOptimizer( + config=table_surface_config + ) + # Parent surface (on) optimizer. + self.parent_surface_layout_optimizer = ParentSurfaceLayoutOptimizer( + config=parent_surface_config + ) + self._current_xy_by_id: dict[str, list[float] | None] = {} + self._solved_delta_xy_by_id: dict[str, list[float]] = {} + self._updated_object_ids: set[str] = set() + + def construct(self) -> Scene: + """Construct table-root layouts before later stacked-group refinement.""" + # Build layout problem. + layout_problem = self._build_problem() + # Get current XY centers. + self._current_xy_by_id = { + object_id: list(initial_xy) if initial_xy is not None else None + for object_id, initial_xy in layout_problem.initial_xy_by_id.items() + } + self._solved_delta_xy_by_id = {} + self._updated_object_ids = set() + # Check the group. + if ( + layout_problem.groups + and layout_problem.groups[0].parent_id != TABLE_OBJECT_ID + ): + raise ValueError("The first layout group must be rooted at the table.") + + # Optimize each group in BFS order, propagating solved deltas to descendants. + for group in layout_problem.groups: + if group.parent_id == TABLE_OBJECT_ID: + self._optimize_table_group( + layout_problem=layout_problem, + group=group, + ) + continue + self._optimize_parent_group( + layout_problem=layout_problem, + group=group, + ) + + return layout_problem.post_edit_scene + + def _optimize_table_group( + self, + *, + layout_problem: SceneLayoutProblem, + group: SceneLayoutGroup, + ) -> None: + """Optimize all direct on-table children before any stacked child groups.""" + table_surface_problem = TableSurfaceLayoutProblem.from_layout_problem( + layout_problem=layout_problem, + group=group, + current_xy_by_id=self._current_xy_by_id, + ) + solved_root_xy_by_id = self.table_surface_layout_optimizer.optimize( + table_surface_problem + ) + table = layout_problem.post_edit_scene.table + if table is None: + raise ValueError("Table group optimization requires a table.") + # Check the table's z. + if table.support_surface_z is None and any( + root_id in layout_problem.layout_variable_ids for root_id in group.child_ids + ): + raise ValueError("Table group optimization requires support_surface_z.") + assets_by_id = { + asset.id: asset for asset in layout_problem.post_edit_scene.assets + } + for root_id, solved_xy in solved_root_xy_by_id.items(): + seed_xy = table_surface_problem.root_seed_xy_by_id[root_id] + delta_xy = [ + solved_xy[0] - seed_xy[0], + solved_xy[1] - seed_xy[1], + ] + self._current_xy_by_id[root_id] = list(solved_xy) + self._solved_delta_xy_by_id[root_id] = delta_xy + if root_id in layout_problem.layout_variable_ids: + # Direct add/move roots receive a new pose on the table support. + assert table.support_surface_z is not None + update_scene_object_y_up_pose_from_z_up_support( + scene_object=assets_by_id[root_id], + support_region_z=table.support_surface_z, + center_xy=solved_xy, + clearance_m=0.00, # Directly place on the support surface. + ) + self._updated_object_ids.add(root_id) + self._propagate_descendant_delta( + scene=layout_problem.post_edit_scene, + root_id=root_id, + delta_xy=delta_xy, + ) + + def _propagate_descendant_delta( + self, + *, + scene: Scene, + root_id: str, + delta_xy: list[float], + ) -> None: + """Move every positioned descendant by one solved ancestor XY delta.""" + # A zero root delta cannot change any descendant pose, so skip the subtree walk. + if delta_xy == [0.0, 0.0]: + return + assets_by_id = {asset.id: asset for asset in scene.assets} + children_by_parent: dict[str, list[str]] = {} + for node in self.goal_scene_graph.nodes: + if node.parent_id is not None: + children_by_parent.setdefault(node.parent_id, []).append(node.object_id) + + pending = list(children_by_parent.get(root_id, [])) + while pending: + descendant_id = pending.pop(0) + descendant_xy = self._current_xy_by_id[descendant_id] + if descendant_xy is not None: + self._current_xy_by_id[descendant_id] = [ + descendant_xy[0] + delta_xy[0], + descendant_xy[1] + delta_xy[1], + ] + translate_scene_object_y_up_by_z_up_delta( + scene_object=assets_by_id[descendant_id], + delta_xy=delta_xy, + ) + self._updated_object_ids.add(descendant_id) + pending.extend(children_by_parent.get(descendant_id, [])) + + def _optimize_parent_group( + self, + *, + layout_problem: SceneLayoutProblem, + group: SceneLayoutGroup, + ) -> None: + """Optimize one settled parent's direct on-children in local XY coordinates.""" + parent_surface_problem = ParentSurfaceLayoutProblem.from_layout_problem( + layout_problem=layout_problem, + group=group, + current_xy_by_id=self._current_xy_by_id, + ) + # Get results. + solved_child_xy_by_id = self.parent_surface_layout_optimizer.optimize( + parent_surface_problem + ) + for child_id, solved_xy in solved_child_xy_by_id.items(): + seed_xy = parent_surface_problem.child_seed_xy_by_id[child_id] + delta_xy = [ + solved_xy[0] - seed_xy[0], + solved_xy[1] - seed_xy[1], + ] + self._current_xy_by_id[child_id] = list(solved_xy) + self._solved_delta_xy_by_id[child_id] = delta_xy + if child_id in layout_problem.layout_variable_ids: + # Variable children are placed directly above the parent's current top. + update_scene_object_y_up_pose_from_z_up_support( + scene_object=parent_surface_problem.assets_by_id[child_id], + support_region_z=parent_surface_problem.parent_top_z, + center_xy=solved_xy, + clearance_m=0.00, # Directly place on the parent's top surface. + ) + self._updated_object_ids.add(child_id) + self._propagate_descendant_delta( + scene=layout_problem.post_edit_scene, + root_id=child_id, + delta_xy=delta_xy, + ) + + def _build_problem(self) -> SceneLayoutProblem: + """Build post-edit objects and preserve formal-scene centers as seeds.""" + # Validate the graph first. + self.goal_scene_graph.validate() + graph_object_ids = set(self.goal_scene_graph.node_by_id()) + generated_objects_by_id = self._generated_scene_objects_by_id() + + # The goal graph removes deleted formal-scene objects from the layout input. + post_edit_objects = [ + scene_object + for scene_object in self.formal_scene.objects + if scene_object.id in graph_object_ids + ] + imported_object_ids = {scene_object.id for scene_object in post_edit_objects} + if imported_object_ids.intersection(generated_objects_by_id): + raise ValueError( + "Generated scene objects must not reuse formal scene object ids." + ) + post_edit_objects.extend(generated_objects_by_id.values()) + + post_edit_scene = Scene(objects=post_edit_objects) + post_edit_object_ids = { + scene_object.id for scene_object in post_edit_scene.objects + } + if post_edit_object_ids != graph_object_ids: + raise ValueError("Goal scene graph and post-edit scene have different ids.") + # Get the movable asset ids. + if not self.layout_variable_ids.issubset(post_edit_object_ids - {"table"}): + raise ValueError( + "Only post-edit assets may participate in layout optimization." + ) + # Get initial XY centers. + initial_xy_by_id = { + asset.id: self._initial_xy( # The assets' center XY should always be updated whenever changes are made. + asset, + is_generated=asset.id in generated_objects_by_id, + ) + for asset in post_edit_scene.assets + } + for object_id, initial_xy in initial_xy_by_id.items(): + if initial_xy is None and object_id not in self.layout_variable_ids: + raise ValueError( + f"New asset {object_id!r} must participate in layout optimization." + ) + # Build the table-rooted BFS groups. + groups = self._build_groups() + + return SceneLayoutProblem( + post_edit_scene=post_edit_scene, + goal_scene_graph=self.goal_scene_graph, + layout_variable_ids=set(self.layout_variable_ids), + initial_xy_by_id=initial_xy_by_id, + groups=groups, + ) + + def _build_groups(self) -> list[SceneLayoutGroup]: + """Build table-rooted BFS groups of direct on-children.""" + children_by_parent: dict[str, list[str]] = {} + for node in self.goal_scene_graph.nodes: + if node.parent_id is None: + continue + if node.parent_relation != "on": + raise ValueError(f"Node {node.object_id!r} must be on its parent.") + children_by_parent.setdefault(node.parent_id, []).append(node.object_id) + + groups: list[SceneLayoutGroup] = [] + pending = [TABLE_OBJECT_ID] + while pending: + parent_id = pending.pop(0) + child_ids = children_by_parent.get(parent_id, []) + if not child_ids: + continue + groups.append(SceneLayoutGroup(parent_id=parent_id, child_ids=child_ids)) + pending.extend(child_ids) + return groups + + def _generated_scene_objects_by_id(self) -> dict[str, SceneObject]: + """Index generated scene objects before merging them into the formal scene.""" + generated_objects_by_id = { + scene_object.id: scene_object + for scene_object in self.generated_scene_objects + } + if len(generated_objects_by_id) != len(self.generated_scene_objects): + raise ValueError("Generated scene objects must use unique object ids.") + return generated_objects_by_id + + @staticmethod + def _initial_xy( + asset: SceneObject, + *, + is_generated: bool, + ) -> list[float] | None: + """Retain formal-scene centers while generated assets await initialization.""" + if is_generated: + return None + if asset.center_xy is None or len(asset.center_xy) != 2: + raise ValueError( + f"Formal-scene asset {asset.id!r} must have a 2D center_xy." + ) + return [float(value) for value in asset.center_xy] diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_utils.py new file mode 100644 index 000000000..100f5b269 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_utils.py @@ -0,0 +1,155 @@ +# ---------------------------------------------------------------------------- +# 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 numpy as np + +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( + layout_object_to_transform_matrix, + load_glb_mesh, + transform_matrix_to_layout_object, +) + + +def update_scene_object_y_up_pose_from_z_up_support( + *, + scene_object: SceneObject, + support_region_z: float, + center_xy: list[float], + clearance_m: float = 0.02, +) -> None: + """Place a SimReady asset on a horizontal z-up support region.""" + if ( + not np.isfinite(support_region_z) + or clearance_m < 0.0 + or not np.isfinite(clearance_m) + ): + raise ValueError("support_region_z and clearance_m must be finite and valid.") + target_xy = two_floats(center_xy, field_name="center_xy") + rotation_y_up = three_floats_or_default( + scene_object.rot, field_name="rot", default=[0.0, 0.0, 0.0] + ) + mesh = load_scene_object_z_up_mesh( + scene_object=scene_object, rotation_y_up=rotation_y_up + ) + target_position_z_up = np.array( + [ + target_xy[0] - float(mesh.bounds[:, 0].mean()), + target_xy[1] - float(mesh.bounds[:, 1].mean()), + float(support_region_z) + clearance_m - float(mesh.bounds[0, 2]), + ] + ) + scene_object.pos = ( + np.linalg.inv(y_up_to_z_up_matrix())[:3, :3] @ target_position_z_up + ).tolist() + scene_object.rot = rotation_y_up + scene_object.center_xy = target_xy + + +def translate_scene_object_y_up_by_z_up_delta( + *, scene_object: SceneObject, delta_xy: list[float] +) -> None: + """Translate an existing y-up pose by a solved z-up XY delta.""" + dx, dy = two_floats(delta_xy, field_name="delta_xy") + position = three_floats_or_default(scene_object.pos, field_name="pos", default=None) + scene_object.pos = [position[0] + dx, position[1], position[2] - dy] + if scene_object.center_xy is not None: + scene_object.center_xy = [ + scene_object.center_xy[0] + dx, + scene_object.center_xy[1] + dy, + ] + + +def measure_scene_object_z_up_world_aabb( + *, scene_object: SceneObject +) -> list[list[float]]: + """Measure one current SceneObject pose in z-up world coordinates.""" + position_y_up = three_floats_or_default( + scene_object.pos, field_name="pos", default=None + ) + mesh = load_scene_object_z_up_mesh(scene_object=scene_object) + mesh.apply_translation( + y_up_to_z_up_matrix()[:3, :3] @ np.asarray(position_y_up, dtype=float) + ) + return mesh.bounds.tolist() + + +def load_scene_object_z_up_mesh( + *, scene_object: SceneObject, rotation_y_up: list[float] | None = None +): + """Load a SimReady mesh in z-up with orientation and scale but no translation.""" + if scene_object.simready_glb_path is None: + raise ValueError(f"Asset {scene_object.id!r} has no SimReady GLB path.") + y_up_layout = { + "id": scene_object.id, + "rot": ( + rotation_y_up + if rotation_y_up is not None + else three_floats_or_default( + scene_object.rot, field_name="rot", default=[0.0, 0.0, 0.0] + ) + ), + "pos": [0.0, 0.0, 0.0], + "scale": three_floats_or_default( + scene_object.scale, field_name="scale", default=[1.0, 1.0, 1.0] + ), + } + y_up_to_z_up = y_up_to_z_up_matrix() + z_up_layout = transform_matrix_to_layout_object( + scene_object.id, + y_up_to_z_up + @ layout_object_to_transform_matrix(y_up_layout) + @ np.linalg.inv(y_up_to_z_up), + ) + mesh = load_glb_mesh(scene_object.simready_glb_path) + mesh.apply_transform(y_up_to_z_up) + mesh.apply_transform(layout_object_to_transform_matrix(z_up_layout)) + return mesh + + +def y_up_to_z_up_matrix() -> np.ndarray: + """Return the coordinate transform used by layout and export stages.""" + matrix = np.eye(4) + matrix[:3, :3] = np.array([[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]]) + return matrix + + +def two_floats(value: object, *, field_name: str) -> list[float]: + """Validate and return one finite two-value vector.""" + if not isinstance(value, (list, tuple)) or len(value) != 2: + raise ValueError(f"{field_name} must contain two values.") + result = [float(component) for component in value] + if not np.all(np.isfinite(result)): + raise ValueError(f"{field_name} must contain finite values.") + return result + + +def three_floats_or_default( + value: object, *, field_name: str, default: list[float] | None +) -> list[float]: + """Validate three finite values, or return a canonical default.""" + if value is None: + if default is None: + raise ValueError(f"{field_name} must contain three values.") + return list(default) + if not isinstance(value, (list, tuple)) or len(value) != 3: + raise ValueError(f"{field_name} must contain three values.") + result = [float(component) for component in value] + if not np.all(np.isfinite(result)): + raise ValueError(f"{field_name} must contain finite values.") + return result diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_scene_processor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py similarity index 56% rename from embodichain/gen_sim/scene_engine/pipeline/utils/simready_scene_processor.py rename to embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py index 40b4a6f94..864e39a59 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_scene_processor.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py @@ -17,21 +17,34 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path -import re import numpy as np -import open3d as o3d -from scipy.spatial import ConvexHull, QhullError from scipy.spatial.transform import Rotation import trimesh from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import OrientationState from embodichain.gen_sim.scene_engine.core.scene_object import ( ObjectPhysics, SceneObject, ) +from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor_utils import ( + DEFAULT_NEEDED_LAYOUT, + LYING_NEEDED_LAYOUT, + STANDING_NEEDED_LAYOUT, + compute_uniform_xy_scale_for_target, + query_vlm_object_rotation_and_target_size, + render_object_front_top_views, + rotate_glb_about_x_axis, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.table_support_surface import ( + TableSupportSurfaceDetector, +) from embodichain.utils.logger import log_info _TABLE_PHYSICS_ATTRS = { @@ -53,16 +66,17 @@ @dataclass(frozen=True) -class SimReadySceneProcessorConfig: - """Object-category policy for SimReady mesh canonicalization.""" +class SimReadyProcessorConfig: + """SceneGraph-conditioned policy for SimReady mesh canonicalization.""" - upright_container_id_tokens: frozenset[str] = frozenset( - {"bottle", "can", "jar", "flask", "thermos"} - ) # Object-id tokens that enable upright-container standardization. + use_vlm_scale: bool = False # Use the VLM-selected asset scale. + use_vlm_rotation: bool = False # Use the VLM-selected asset rotation. + # Explicit graph orientation overrides the default stable tabletop pose. + orientation_states_by_id: dict[str, OrientationState] = field(default_factory=dict) -class SimReadySceneProcessor: - """Create SimReady GLBs and layouts for one table and its scene assets.""" +class SimReadyProcessor: + """Create SimReady GLBs and layouts for scene objects.""" def __init__( self, @@ -71,7 +85,9 @@ def __init__( coarse_layout_by_id: dict[str, dict[str, object]], coarse_geometry_root: str | Path, simready_geometry_root: str | Path, - config: SimReadySceneProcessorConfig | None = None, + debug_output_root: str | Path | None = None, + config: SimReadyProcessorConfig | None = None, + vlm_client: OpenAICompatibleVLM | None = None, ) -> None: self.scene = scene self.coarse_layout_by_id = coarse_layout_by_id @@ -79,11 +95,22 @@ def __init__( self.simready_geometry_root = ( Path(simready_geometry_root).expanduser().resolve() ) + # Save rendered debug images. + self.debug_output_root = ( + Path(debug_output_root).expanduser().resolve() + if debug_output_root is not None + else None + ) self.simready_table_layout: dict[str, object] | None = None self.simready_assets_layout: list[dict[str, object]] | None = None - self.config = config if config is not None else SimReadySceneProcessorConfig() - if not self.config.upright_container_id_tokens: - raise ValueError("upright_container_id_tokens must not be empty.") + self.config = config if config is not None else SimReadyProcessorConfig() + self.vlm_client = vlm_client + if ( + self.config.use_vlm_scale + or self.config.use_vlm_rotation + or self.config.orientation_states_by_id + ) and vlm_client is None: + raise ValueError("vlm_client is required when VLM transforms are enabled.") def process_table(self) -> dict[str, object]: """Process the required scene table and return its SimReady layout.""" @@ -104,7 +131,13 @@ def process_assets(self) -> list[dict[str, object]]: self.simready_assets_layout = processed_assets return self.simready_assets_layout - def _process_object(self, scene_object: SceneObject) -> dict[str, object]: + def _process_object( + self, + scene_object: SceneObject, + *, + scale: object | None = None, + rot: object | None = None, + ) -> dict[str, object]: """Canonicalize one coarse object and write its SimReady GLB.""" object_id = scene_object.id object_role = scene_object.kind @@ -113,12 +146,18 @@ def _process_object(self, scene_object: SceneObject) -> dict[str, object]: coarse_layout = self.coarse_layout_by_id.get(object_id) if coarse_layout is None: raise ValueError(f"Coarse layout does not contain object {object_id!r}.") + prepared_glb_path, vlm_scale = self._prepare_vlm_rotated_glb(scene_object) + selected_scale = scale + if selected_scale is None: + selected_scale = vlm_scale or coarse_layout.get("scale") simready_mesh, simready_transform = self._canonicalize_object_mesh( - coarse_glb_path=self.coarse_geometry_root / f"{object_id}.glb", + coarse_glb_path=prepared_glb_path, object_id=object_id, - rot=coarse_layout.get("rot"), + # An enabled external rotation replaces the coarse-layout rotation. + rot=coarse_layout.get("rot") if rot is None else rot, pos=coarse_layout.get("pos"), - scale=coarse_layout.get("scale"), + # An enabled VLM scale replaces the coarse-layout scale. + scale=selected_scale, ) output_path = self.simready_geometry_root / f"{object_id}.glb" output_path.parent.mkdir(parents=True, exist_ok=True) @@ -129,9 +168,120 @@ def _process_object(self, scene_object: SceneObject) -> dict[str, object]: ) scene_object.simready_glb_path = str(output_path) scene_object.physics = self._fixed_physics_for_kind(object_role) + # For table. (currently the id is fixed into table) + if object_role == "table": + # Detect and persist all reusable tabletop support geometry at SimReady time. + support_detector = TableSupportSurfaceDetector( + table_world_mesh=self._z_up_table_mesh(simready_mesh), + debug_output_root=self.debug_output_root, + ) + support_region = support_detector.detect() + scene_object.support_surface_z = support_region.top_z + scene_object.support_contour_xy = [ + [float(x), float(y)] + for x, y in support_region.support_polygon.exterior.coords[:-1] + ] + scene_object.support_optimization_rect_xy = [ + [float(x), float(y)] + for x, y in support_region.optimization_rectangle.exterior.coords[:-1] + ] + if self.debug_output_root is not None: + # Keep the 3D selected surface and 2D contour diagnostics beside SimReady output. + support_detector.save_support_surface_debug_images() log_info(f"Created SimReady {object_role}: {object_id!r}.") return {"id": object_id, **simready_transform} + @staticmethod + def _z_up_table_mesh(mesh: trimesh.Trimesh) -> trimesh.Trimesh: + """Convert one canonical y-up GLB mesh into the detector's z-up frame.""" + y_up_to_z_up = np.eye(4) + y_up_to_z_up[:3, :3] = Rotation.from_euler("x", 90.0, degrees=True).as_matrix() + z_up_mesh = mesh.copy() + z_up_mesh.apply_transform(y_up_to_z_up) + return z_up_mesh + + def _prepare_vlm_rotated_glb( + self, scene_object: SceneObject + ) -> tuple[Path, list[float] | None]: + """Render, query, and optionally bake the VLM-selected x-axis rotation.""" + coarse_path = self.coarse_geometry_root / f"{scene_object.id}.glb" + orientation_state = self._orientation_state_for_object(scene_object.id) + orientation_pose_required = orientation_state is not None + if not ( + self.config.use_vlm_scale + or self.config.use_vlm_rotation + or orientation_pose_required + ): + return coarse_path, None + decision = self._vlm_transform_for_object( + scene_object, + needed_layout=self._needed_layout_for_object(scene_object.id), + ) + rotate_about_x = bool(decision["rotate_about_x"]) + vlm_scale = None + if self.config.use_vlm_scale: + # The VLM target describes the final, post-rotation z-up XY footprint. + vlm_scale = compute_uniform_xy_scale_for_target( + glb_path=coarse_path, + target_xy_size_cm=decision["target_xy_size_cm"], + rotate_about_x=rotate_about_x, + ) + rotated_path = rotate_glb_about_x_axis( + input_path=coarse_path, + output_path=self.simready_geometry_root + / "vlm_rotated" + / f"{scene_object.id}.glb", + rotate=rotate_about_x + and (orientation_pose_required or self.config.use_vlm_rotation), + ) + # The scale flag controls whether this VLM-derived isotropic scale is used. + # Apply the same factor on x, y, and z to preserve the asset's proportions. + return ( + rotated_path, + [vlm_scale, vlm_scale, vlm_scale] if self.config.use_vlm_scale else None, + ) + + def _orientation_state_for_object(self, object_id: str) -> OrientationState | None: + """Return the explicit graph orientation requested for one object.""" + return self.config.orientation_states_by_id.get(object_id) + + def _needed_layout_for_object(self, object_id: str) -> str: + """Return the VLM layout instruction for one object's graph semantics.""" + return ( + STANDING_NEEDED_LAYOUT + if self._orientation_state_for_object(object_id) == "standing" + else ( + LYING_NEEDED_LAYOUT + if self._orientation_state_for_object(object_id) == "lying" + else DEFAULT_NEEDED_LAYOUT + ) + ) + + def _vlm_transform_for_object( + self, + scene_object: SceneObject, + *, + needed_layout: str, + ) -> dict[str, object]: + """Render the object and return the validated VLM pose decision.""" + assert self.vlm_client is not None + coarse_path = self.coarse_geometry_root / f"{scene_object.id}.glb" + debug_root = ( + self.debug_output_root or self.simready_geometry_root.parent / "debug" + ) + rendered_path = render_object_front_top_views( + glb_path=coarse_path, + output_path=debug_root / "vlm_views" / f"{scene_object.id}.png", + ) + # Both semantic questions are always answered in one multimodal call. + return query_vlm_object_rotation_and_target_size( + scene_object_description=scene_object.description, + needed_layout=needed_layout, + rendered_views_path=rendered_path, + vlm_client=self.vlm_client, + debug_output_path=debug_root / "vlm_outputs" / f"{scene_object.id}.json", + ) + @staticmethod def _fixed_physics_for_kind(kind: str) -> ObjectPhysics: """Create the fixed initial physics profile for one SimReady object.""" @@ -185,8 +335,6 @@ def _canonicalize_object_mesh( ) if np.any(coarse_scale <= 0): raise ValueError("Coarse object scale values must be positive.") - # We need the object id to determine whether it is a bottle-like object. - # If it does, then we will do a special standardization. (Hard code) if not isinstance(object_id, str) or not object_id: raise ValueError("Scene object id must be a non-empty string.") @@ -197,16 +345,6 @@ def _canonicalize_object_mesh( y_up_to_z_up_transform[:3, :3] = y_up_to_z_up_matrix mesh.apply_transform(y_up_to_z_up_transform) - # Standardize upright containers in temporary z-up coordinates before the - # shared center, scale, and bottom-center preprocessing. - # This is to ensure the action agent can pick up the bottle or can-like objects. - bottle_alignment_matrix = np.eye(3) - if self._is_upright_container_id(object_id): - bottle_alignment_matrix = self._standardize_bottle_z_up(mesh) - bottle_alignment_transform = np.eye(4) - bottle_alignment_transform[:3, :3] = bottle_alignment_matrix - mesh.apply_transform(bottle_alignment_transform) - # First make the object's AABB center at the origin. original_aabb_center = mesh.bounds.mean(axis=0) mesh.apply_translation(-original_aabb_center) @@ -214,13 +352,7 @@ def _canonicalize_object_mesh( # Scale the object with the value in the coarse layout. scale_transform = np.eye(4) scale_transform[:3, :3] = ( - # Actually there's no need to do so, for the scale factor is all equal - # in x, y, z axes. - bottle_alignment_matrix - @ y_up_to_z_up_matrix - @ np.diag(coarse_scale) - @ y_up_to_z_up_matrix.T - @ bottle_alignment_matrix.T + y_up_to_z_up_matrix @ np.diag(coarse_scale) @ y_up_to_z_up_matrix.T ) mesh.apply_transform(scale_transform) @@ -240,17 +372,7 @@ def _canonicalize_object_mesh( z_up_to_y_up_transform[:3, :3] = y_up_to_z_up_matrix.T mesh.apply_transform(z_up_to_y_up_transform) - # Compensate the bottle's local rotation so that its coarse world pose does - # not change. - local_bottle_rotation = Rotation.from_matrix( - y_up_to_z_up_matrix.T @ bottle_alignment_matrix @ y_up_to_z_up_matrix - ) - coarse_rotation_matrix = Rotation.from_euler( - "xyz", coarse_rot, degrees=True - ).as_matrix() - rotation = Rotation.from_matrix( - coarse_rotation_matrix @ local_bottle_rotation.inv().as_matrix() - ) + rotation = Rotation.from_euler("xyz", coarse_rot, degrees=True) # Update the pos. position_offset = y_up_to_z_up_matrix.T @ ( scale_transform[:3, :3] @ original_aabb_center + scaled_aabb_bottom_center @@ -261,87 +383,6 @@ def _canonicalize_object_mesh( "scale": [1.0, 1.0, 1.0], } - def _is_upright_container_id(self, object_id: str) -> bool: - """Return whether object-id tokens indicate a bottle-like container.""" - # Example: soda_can_0 - # tokens: {"soda", "can", "0"} - # upright_container_id_tokens: {"bottle", "can", "jar"} - # So this returns True because "can" is in the configured token set. - tokens = set(re.findall(r"[a-z0-9]+", object_id.lower())) - return bool(tokens & self.config.upright_container_id_tokens) - - @staticmethod - def _standardize_bottle_z_up(mesh: trimesh.Trimesh) -> np.ndarray: - """Return a proper rotation that maps a bottle-like mesh's long axis to z-up. - - Thanks to chenjian for this idea! - """ - if len(mesh.vertices) < 4 or len(mesh.faces) < 4: - raise ValueError( - "Bottle standardization requires a non-degenerate triangle mesh." - ) - open3d_mesh = o3d.geometry.TriangleMesh( - vertices=o3d.utility.Vector3dVector(mesh.vertices), - triangles=o3d.utility.Vector3iVector(mesh.faces), - ) - sampled_points = np.asarray( - open3d_mesh.sample_points_uniformly(number_of_points=10_000).points - ) # (10000, 3) x (x, y, z) - - # Check the number of the points again, and check whether have some - # non-finite values. - if sampled_points.shape[0] < 4 or not np.all(np.isfinite(sampled_points)): - raise ValueError( - "Bottle standardization could not sample valid mesh points." - ) - - centered_points = sampled_points - sampled_points.mean(axis=0) - # SVD find the longest axis. - _, _, principal_axes = np.linalg.svd(centered_points, full_matrices=False) - if np.linalg.det(principal_axes) < 0: - principal_axes[2, :] *= -1 # in case the SVD returns a reflection. - - bottle_rotation = Rotation.from_euler( - "y", 90.0, degrees=True - ).as_matrix() # 3x3 matrix - # The first PCA axis is the longest axis; rotate it onto the temporary z axis. - bottle_rotation = bottle_rotation @ principal_axes - standardized_points = (bottle_rotation @ centered_points.T).T - - axis_min = standardized_points[:, 2].min() - axis_max = standardized_points[:, 2].max() - axis_range = axis_max - axis_min - upper_points = standardized_points[ - standardized_points[:, 2] > axis_min + axis_range * 0.8 - ] - lower_points = standardized_points[ - standardized_points[:, 2] < axis_min + axis_range * 0.2 - ] - upper_volume = SimReadySceneProcessor._convex_hull_volume(upper_points) - lower_volume = SimReadySceneProcessor._convex_hull_volume(lower_points) - - # Bottles usually have a smaller top (neck) than bottom; flip if necessary. - if upper_volume > lower_volume: - bottle_rotation = ( - Rotation.from_euler("x", 180.0, degrees=True).as_matrix() - @ bottle_rotation - ) - return bottle_rotation - - @staticmethod - def _convex_hull_volume(points: np.ndarray) -> float: - """Return the volume of a non-degenerate point set's convex hull.""" - if points.shape[0] < 4: - raise ValueError( - "Bottle standardization needs at least four points per end." - ) - try: - return float(ConvexHull(points).volume) - except QhullError as exc: - raise ValueError( - "Bottle standardization found a degenerate end volume." - ) from exc - @staticmethod def _three_floats(value: object, *, field_name: str) -> list[float]: """Validate and convert a three-value layout field to floats.""" diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py new file mode 100644 index 000000000..4523aa3e6 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.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 + +import json +import os +from pathlib import Path +import sys +from typing import Callable + +import numpy as np +from PIL import Image, ImageDraw, ImageFont +from scipy.spatial.transform import Rotation +import trimesh + +from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) + +_VLM_SYSTEM_PROMPT = """You inspect one isolated 3D object from front and top views. +Use the object description and the rendered views together. + +Use the rendered views and the needed layout to decide whether the object should +be rotated around its own center by +90 degrees around the z-up world's x axis. +The z-up world is right-handed: x is left-right, y is front-back, and z is up. +In the composed image, FRONT VIEW is the left panel: x is horizontal and z is +vertical; the upper-right marker shows the positive z direction. TOP VIEW is +the right panel: x is horizontal and y is vertical; the upper-right markers +show the positive x and y directions. +Do not confuse the top view with looking at the object from above in the image +description: it is a projection along the z axis onto the x-y plane. +After deciding and applying that rotation, estimate the object's desired AABB +footprint on the x-y plane in real-world centimetres. The first value is the x +size and the second value is the y size. + +Return JSON only with exactly this schema: +{ + "rotate_about_x": false, + "target_xy_size_cm": [12.0, 5.0] +} + +Examples: +- Fork lying flat on a table: in FRONT VIEW the fork is mostly a thin + horizontal line; in TOP VIEW its length is visible. Keep it flat with + rotate_about_x=false, and use the tabletop footprint, for example + target_xy_size_cm=[15.0, 3.0]. +- Fork placed in a pen holder: the desired fork is upright, so its long axis is + approximately z. If the input coarse fork is lying in the x-y plane, set + rotate_about_x=true; if the input coarse fork is already upright, set it to + false. The target is the footprint inside the holder, not the fork's full + length, for example target_xy_size_cm=[3.0, 3.0]. +- Fork requested to lie flat on a table even when the input coarse fork is + upright: set rotate_about_x=true and estimate the final flat footprint, for + example target_xy_size_cm=[15.0, 3.0]. +- Bottle already standing on its flat base: keep it upright with + rotate_about_x=false and use target_xy_size_cm=[8.0, 8.0]. +""" + +DEFAULT_NEEDED_LAYOUT = ( + "Place this asset on the table in its natural, physically stable resting " + "orientation. For example, a fork should lie flat on the table rather " + "than stand on an edge." +) +STANDING_NEEDED_LAYOUT = ( + "The scene graph requires this asset to stand vertically on the table, " + "even when its natural stable pose would be lying down. For example, a " + "bottle should stand on its base and a fork should stand upright. If the " + "coarse GLB is lying flat, set rotate_about_x=true so its semantic vertical " + "axis aligns with the z-up world's z axis; if it is already upright, set " + "it to false." +) +LYING_NEEDED_LAYOUT = ( + "The scene graph requires this asset to lie flat on the table, even when " + "its natural stable pose would be standing. For example, a bottle should " + "lie on its side and a fork should lie flat. Choose rotate_about_x so the " + "asset's semantic long axis remains in the tabletop x-y plane rather than " + "along the z-up world's z axis." +) + + +def render_object_front_top_views( + *, + glb_path: str | Path, + output_path: str | Path, + resolution: int = 512, +) -> Path: + """Render fixed z-up front/top views and compose them horizontally.""" + if resolution <= 0: + raise ValueError("resolution must be positive.") + source_path = Path(glb_path).expanduser().resolve() + if not source_path.is_file(): + raise FileNotFoundError(f"GLB for VLM rendering not found: {source_path}") + output_path = Path(output_path).expanduser().resolve() + output_path.parent.mkdir(parents=True, exist_ok=True) + front_path = output_path.with_name(f"{output_path.stem}_front.png") + top_path = output_path.with_name(f"{output_path.stem}_top.png") + try: + import bpy + from mathutils import Vector + except ImportError as exc: + raise RuntimeError( + "Blender's bpy is required for SimReady VLM view rendering." + ) from exc + + _run_blender_operation_silently( + lambda: bpy.ops.wm.read_factory_settings(use_empty=True) + ) + _run_blender_operation_silently( + lambda: bpy.ops.import_scene.gltf(filepath=str(source_path)) + ) + if not any(obj.type == "MESH" for obj in bpy.context.scene.objects): + raise ValueError(f"GLB contains no mesh objects: {source_path}") + scene = bpy.context.scene + # Eevee renders imported GLB materials and textures instead of Workbench previews. + try: + scene.render.engine = "BLENDER_EEVEE_NEXT" + except TypeError: + scene.render.engine = "BLENDER_EEVEE" + scene.render.resolution_x = resolution + scene.render.resolution_y = resolution + scene.render.resolution_percentage = 100 + scene.render.image_settings.file_format = "PNG" + scene.render.film_transparent = False + if scene.world is None: + scene.world = bpy.data.worlds.new("VLM_World") + scene.world.color = (0.08, 0.08, 0.08) + for name, location, energy in ( + ("VLM_Key", (2.0, -2.0, 3.0), 700.0), + ("VLM_Fill", (-2.0, 1.0, 2.0), 400.0), + ): + light_data = bpy.data.lights.new(name, type="AREA") + light_data.energy = energy + light_data.shape = "DISK" + light_data.size = 4.0 + light = bpy.data.objects.new(name, light_data) + light.location = location + light.rotation_euler = ( + (Vector((0.0, 0.0, 0.0)) - light.location) + .to_track_quat("-Z", "Y") + .to_euler() + ) + scene.collection.objects.link(light) + camera_data = bpy.data.cameras.new("VLM_Camera") + camera = bpy.data.objects.new("VLM_Camera", camera_data) + scene.collection.objects.link(camera) + scene.camera = camera + camera.data.type = "ORTHO" + camera.data.ortho_scale = 1.25 + + def render_view(path: Path, location: tuple[float, float, float]) -> None: + camera.location = location + camera.rotation_euler = ( + (Vector((0.0, 0.0, 0.0)) - camera.location) + .to_track_quat("-Z", "Y") + .to_euler() + ) + scene.render.filepath = str(path) + _run_blender_operation_silently(lambda: bpy.ops.render.render(write_still=True)) + + # Blender uses a right-handed z-up world; front is viewed along +y. + render_view(front_path, (0.0, -3.0, 0.0)) + render_view(top_path, (0.0, 0.0, 3.0)) + with Image.open(front_path) as front, Image.open(top_path) as top: + composed = Image.new("RGB", (resolution * 2, resolution), "white") + composed.paste(front.convert("RGB"), (0, 0)) + composed.paste(top.convert("RGB"), (resolution, 0)) + draw = ImageDraw.Draw(composed) + # Use a readable scaled font for the panel labels when available. + try: + font = ImageFont.truetype( + "/usr/share/fonts/truetype/ubuntu/Ubuntu-B.ttf", + max(24, resolution // 16), + ) + except OSError: + font = ImageFont.load_default() + # Label each panel so the VLM and manual debugging can distinguish views. + for label, origin in (("FRONT VIEW", (0, 0)), ("TOP VIEW", (resolution, 0))): + x, y = origin + text_box = draw.textbbox((x + 16, y + 16), label, font=font) + draw.rectangle( + (text_box[0] - 8, text_box[1] - 6, text_box[2] + 8, text_box[3] + 6), + fill="white", + ) + draw.text((x + 16, y + 16), label, fill="black", font=font) + # Mark the positive axes used by each projection for VLM interpretation. + _draw_arrow( + draw, + (resolution - 62, 62), + (resolution - 62, 20), + "+Z", + font, + color="blue", + ) + _draw_arrow( + draw, + (2 * resolution - 92, 62), + (2 * resolution - 42, 62), + "+X", + font, + color="red", + ) + _draw_arrow( + draw, + (2 * resolution - 92, 62), + (2 * resolution - 92, 20), + "+Y", + font, + color="green", + ) + composed.save(output_path) + return output_path + + +def _run_blender_operation_silently(operation: Callable[[], object]) -> object: + """Run one bpy operation without forwarding Blender-native console output.""" + # bpy writes render progress directly to process file descriptors, not Python streams. + sys.stdout.flush() + sys.stderr.flush() + saved_stdout_fd = os.dup(1) + saved_stderr_fd = os.dup(2) + try: + with open(os.devnull, "w", encoding="utf-8") as null_output: + os.dup2(null_output.fileno(), 1) + os.dup2(null_output.fileno(), 2) + return operation() + finally: + os.dup2(saved_stdout_fd, 1) + os.dup2(saved_stderr_fd, 2) + os.close(saved_stdout_fd) + os.close(saved_stderr_fd) + + +def _draw_arrow( + draw: ImageDraw.ImageDraw, + start: tuple[int, int], + end: tuple[int, int], + label: str, + font: ImageFont.FreeTypeFont | ImageFont.ImageFont, + color: str, +) -> None: + """Draw one labeled positive-axis arrow on a rendered view.""" + dx, dy = end[0] - start[0], end[1] - start[1] + length = max(abs(dx), abs(dy)) + if length == 0: + raise ValueError("Axis arrow start and end must differ.") + unit_x, unit_y = dx / length, dy / length + perpendicular_x, perpendicular_y = -unit_y, unit_x + head_length = 14.0 + head_width = 8.0 + tip_x, tip_y = end + base_x = tip_x - unit_x * head_length + base_y = tip_y - unit_y * head_length + arrowhead = ( + (tip_x, tip_y), + ( + base_x + perpendicular_x * head_width, + base_y + perpendicular_y * head_width, + ), + ( + base_x - perpendicular_x * head_width, + base_y - perpendicular_y * head_width, + ), + ) + draw.line((*start, *end), fill=color, width=4) + draw.polygon(arrowhead, fill=color) + # Put each axis label beside its arrowhead so it does not cover the arrow. + draw.text((int(tip_x + 8), int(tip_y - 8)), label, fill=color, font=font) + + +def query_vlm_object_rotation_and_target_size( + *, + scene_object_description: str, + needed_layout: str, + rendered_views_path: str | Path, + vlm_client: OpenAICompatibleVLM, + debug_output_path: str | Path | None = None, + json_max_attempts: int = 3, +) -> dict[str, object]: + """Ask the VLM for a valid rotation and post-rotation tabletop footprint.""" + if json_max_attempts < 1: + raise ValueError("json_max_attempts must be at least 1.") + last_validation_error: ValueError | None = None + for _ in range(json_max_attempts): + response_text = vlm_client.complete( + system_prompt=_VLM_SYSTEM_PROMPT, + user_prompt=( + f"Object description:\n{scene_object_description}\n\n" + f"Needed layout:\n{needed_layout}\n\n" + "The image contains front view on the left and top view on the right." + ), + image_path=rendered_views_path, + ) + try: + value = _parse_vlm_rotation_and_target_size_response(response_text) + break + except ValueError as exc: + last_validation_error = exc + else: + assert last_validation_error is not None + raise ValueError( + "VLM transform response is invalid after " + f"{json_max_attempts} attempts: {last_validation_error}" + ) from last_validation_error + + if debug_output_path is not None: + output_path = Path(debug_output_path).expanduser().resolve() + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text( + json.dumps( + { + "description": scene_object_description, + "needed_layout": needed_layout, + "rendered_views_path": str( + Path(rendered_views_path).expanduser().resolve() + ), + "vlm_output": value, + }, + indent=2, + ensure_ascii=False, + ) + + "\n", + encoding="utf-8", + ) + return value + + +def _parse_vlm_rotation_and_target_size_response( + response_text: str, +) -> dict[str, object]: + """Validate one VLM rotation-and-scale JSON response.""" + try: + value = json.loads(_strip_json_code_fence(response_text)) + except json.JSONDecodeError as exc: + raise ValueError(f"VLM transform response is not valid JSON: {exc}") from exc + if not isinstance(value, dict) or set(value) != { + "rotate_about_x", + "target_xy_size_cm", + }: + raise ValueError( + "VLM transform response must contain exactly rotate_about_x and " + "target_xy_size_cm." + ) + if not isinstance(value["rotate_about_x"], bool): + raise ValueError("VLM rotate_about_x must be boolean.") + target_size = value["target_xy_size_cm"] + if ( + not isinstance(target_size, list) + or len(target_size) != 2 + or not all(isinstance(item, (int, float)) for item in target_size) + or not all(np.isfinite(item) and item > 0 for item in target_size) + ): + raise ValueError("VLM target_xy_size_cm must contain two positive numbers.") + return value + + +def compute_uniform_xy_scale_for_target( + *, + glb_path: str | Path, + target_xy_size_cm: list[float], + rotate_about_x: bool, +) -> float: + """Compute an isotropic scale from the rotated mesh XY AABB and target size.""" + loaded = trimesh.load(Path(glb_path).expanduser().resolve(), process=False) + mesh = ( + loaded.dump(concatenate=True) if isinstance(loaded, trimesh.Scene) else loaded + ) + if not isinstance(mesh, trimesh.Trimesh): + raise ValueError(f"GLB is not a mesh: {glb_path}") + if len(target_xy_size_cm) != 2 or any(value <= 0 for value in target_xy_size_cm): + raise ValueError("target_xy_size_cm must contain two positive values.") + # GLB geometry is y-up, while the target footprint is defined on z-up table XY. + y_up_to_z_up = np.eye(4) + y_up_to_z_up[:3, :3] = Rotation.from_euler("x", 90.0, degrees=True).as_matrix() + mesh.apply_transform(y_up_to_z_up) + if rotate_about_x: + center = mesh.bounds.mean(axis=0) + mesh.apply_translation(-center) + transform = np.eye(4) + transform[:3, :3] = Rotation.from_euler("x", 90.0, degrees=True).as_matrix() + mesh.apply_transform(transform) + mesh.apply_translation(center) + actual_xy_size = mesh.bounds[1, :2] - mesh.bounds[0, :2] + if np.any(actual_xy_size <= 0): + raise ValueError("Rotated mesh must have a positive XY AABB.") + # Convert the VLM's centimetres to metres before comparing with the GLB AABB. + target_xy_size_m = np.asarray(target_xy_size_cm, dtype=float) / 100.0 + axis_scales = target_xy_size_m / actual_xy_size + # Use sqrt(target XY area / actual XY area) as one uniform scale on all axes. + return float(np.sqrt(axis_scales[0] * axis_scales[1])) + + +def rotate_glb_about_x_axis( + *, + input_path: str | Path, + output_path: str | Path, + rotate: bool, +) -> Path: + """Bake an optional +90-degree x-axis rotation around the mesh centre.""" + # Current coarse layouts are either flat on xy with possible random z rotation, + # or upright with almost no random y rotation, so this x-axis toggle is enough. + source_path = Path(input_path).expanduser().resolve() + destination_path = Path(output_path).expanduser().resolve() + destination_path.parent.mkdir(parents=True, exist_ok=True) + loaded = trimesh.load(source_path, process=False) + mesh = ( + loaded.dump(concatenate=True) if isinstance(loaded, trimesh.Scene) else loaded + ) + if not isinstance(mesh, trimesh.Trimesh): + raise ValueError(f"GLB is not a mesh: {source_path}") + if rotate: + center = mesh.bounds.mean(axis=0) + mesh.apply_translation(-center) + transform = np.eye(4) + transform[:3, :3] = Rotation.from_euler("x", 90.0, degrees=True).as_matrix() + mesh.apply_transform(transform) + mesh.apply_translation(center) + mesh.export(destination_path, file_type="glb") + return destination_path + + +def _strip_json_code_fence(response_text: str) -> str: + """Remove one optional Markdown JSON fence from a VLM response.""" + stripped = response_text.strip() + if not stripped.startswith("```"): + return stripped + lines = stripped.splitlines() + if lines and lines[0].startswith("```"): + lines = lines[1:] + if lines and lines[-1].strip() == "```": + lines = lines[:-1] + return "\n".join(lines).strip() diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/table_support_surface.py b/embodichain/gen_sim/scene_engine/pipeline/utils/table_support_surface.py index 6c3ad3e94..983c16098 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/table_support_surface.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/table_support_surface.py @@ -55,6 +55,7 @@ class TableSupportRegion: vertices: np.ndarray # Full z-up table vertex array referenced by ``faces``. faces: np.ndarray # Indices of triangles selected as the main support surface. support_polygon: Polygon # Largest valid outer support contour in z-up XY. + optimization_rectangle: Polygon # Axis-aligned rectangle fully inside the contour. class TableSupportSurfaceDetector: @@ -131,11 +132,13 @@ def detect(self) -> TableSupportRegion: selected_vertices = face_vertices[selected_faces] vertices = mesh.vertices.copy() faces = mesh.faces[selected_faces].copy() + support_polygon = self._extract_largest_support_polygon(vertices[faces, :2]) self.support_region = TableSupportRegion( top_z=float(selected_vertices[:, :, 2].max()), vertices=vertices, faces=faces, - support_polygon=self._extract_largest_support_polygon(vertices[faces, :2]), + support_polygon=support_polygon, + optimization_rectangle=self._largest_inscribed_rectangle(support_polygon), ) return self.support_region @@ -367,6 +370,46 @@ def _extract_largest_support_polygon(cls, triangles_xy: np.ndarray) -> Polygon: raise ValueError("The merged 2D support contour is degenerate.") return Polygon(boundary_xy) + @staticmethod + def _largest_inscribed_rectangle(polygon: Polygon) -> Polygon: + """Find a conservative axis-aligned rectangle contained by the support contour.""" + coordinates = np.asarray(polygon.exterior.coords[:-1], dtype=float) + x_values = np.unique(coordinates[:, 0]) + y_values = np.unique(coordinates[:, 1]) + # Keep the search bounded for highly tessellated support contours. + if len(x_values) > 48: + x_values = x_values[np.linspace(0, len(x_values) - 1, 48, dtype=int)] + if len(y_values) > 48: + y_values = y_values[np.linspace(0, len(y_values) - 1, 48, dtype=int)] + + best_rectangle: Polygon | None = None + best_area = 0.0 + for x_index, minimum_x in enumerate(x_values[:-1]): + for maximum_x in x_values[x_index + 1 :]: + if maximum_x <= minimum_x: + continue + for y_index, minimum_y in enumerate(y_values[:-1]): + for maximum_y in y_values[y_index + 1 :]: + if maximum_y <= minimum_y: + continue + rectangle = Polygon( + [ + (minimum_x, minimum_y), + (maximum_x, minimum_y), + (maximum_x, maximum_y), + (minimum_x, maximum_y), + ] + ) + area = rectangle.area + if area > best_area and polygon.covers(rectangle): + best_rectangle = rectangle + best_area = area + if best_rectangle is None: + raise ValueError( + "Support contour has no non-degenerate inscribed rectangle." + ) + return best_rectangle + @staticmethod def _face_adjacency(mesh: trimesh.Trimesh) -> dict[int, set[int]]: """Build a face adjacency dictionary for the mesh.""" @@ -476,6 +519,14 @@ def _save_support_region_2d_image( linewidth=2.0, label="outer support contour", ) + rectangle_xy = np.asarray(support_region.optimization_rectangle.exterior.coords) + axis.plot( + rectangle_xy[:, 0], + rectangle_xy[:, 1], + color="seagreen", + linewidth=2.0, + label="optimization rectangle", + ) axis.autoscale_view() axis.set_aspect("equal", adjustable="box") axis.set_xlabel("x (z-up world)") diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/table_surface_layout_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/table_surface_layout_optimizer.py new file mode 100644 index 000000000..9d874f239 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/table_surface_layout_optimizer.py @@ -0,0 +1,598 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import numpy as np +from scipy.optimize import minimize + +from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraphRelation +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_utils import ( + load_scene_object_z_up_mesh, +) + +if TYPE_CHECKING: + from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_constructor import ( + SceneLayoutGroup, + SceneLayoutProblem, + ) + + +@dataclass +class TableSurfaceLayoutProblem: + """All scene-graph and geometry inputs for one table-surface solve.""" + + assets_by_id: dict[str, SceneObject] + root_ids: list[str] + root_seed_xy_by_id: dict[str, list[float]] + imported_root_ids: set[str] + fixed_root_xy_by_id: dict[str, list[float] | None] + root_table_regions_by_id: dict[str, str | None] + table_optimization_rect_xy: list[list[float]] + root_relations: list[SceneGraphRelation] + + @classmethod + def from_layout_problem( + cls, + *, + layout_problem: SceneLayoutProblem, + group: SceneLayoutGroup, + current_xy_by_id: dict[str, list[float] | None], + ) -> TableSurfaceLayoutProblem: + """Build one table-surface problem without mutating layout state.""" + table = layout_problem.post_edit_scene.table + if table is None: + raise ValueError("Table group optimization requires a table.") + if table.support_optimization_rect_xy is None: + raise ValueError( + "Table group optimization requires a table support optimization rectangle." + ) + root_ids = set(group.child_ids) + nodes_by_id = layout_problem.goal_scene_graph.node_by_id() + root_seed_xy_by_id = {} + for root_id in group.child_ids: + inherited_xy = current_xy_by_id[root_id] + # New roots begin from the table origin; imported roots retain their seed. + root_seed_xy_by_id[root_id] = ( + [0.0, 0.0] if inherited_xy is None else list(inherited_xy) + ) + return cls( + assets_by_id={ + asset.id: asset for asset in layout_problem.post_edit_scene.assets + }, + root_ids=group.child_ids, + root_seed_xy_by_id=root_seed_xy_by_id, + imported_root_ids={ + root_id + for root_id in group.child_ids + if layout_problem.initial_xy_by_id[root_id] is not None + }, + fixed_root_xy_by_id={ + root_id: ( + None + if root_id in layout_problem.layout_variable_ids + else current_xy_by_id[root_id] + ) + for root_id in group.child_ids + }, + root_table_regions_by_id={ + root_id: nodes_by_id[root_id].table_region + for root_id in group.child_ids + }, + table_optimization_rect_xy=table.support_optimization_rect_xy, + root_relations=[ + relation + for relation in layout_problem.goal_scene_graph.relations + if relation.source_id in root_ids and relation.target_id in root_ids + ], + ) + + +@dataclass(frozen=True) +class TableSurfaceLayoutOptimizerConfig: + """Numerical controls for one direct-table sibling layout solve.""" + + relation_clearance_m: float = 0.03 + collision_margin_m: float = 0.02 + max_slsqp_iterations: int = 500 + slsqp_ftol: float = 1e-6 + max_collision_rounds: int = 8 + max_added_collision_pairs: int = 64 + imported_seed_weight: float = 5.0 + min_center_distance_m: float = 0.01 + min_center_distance_weight: float = 0.05 + + def __post_init__(self) -> None: + """Reject invalid controls before assembling table-surface constraints.""" + if self.relation_clearance_m < 0.0: + raise ValueError("relation_clearance_m must be non-negative.") + if self.collision_margin_m < 0.0: + raise ValueError("collision_margin_m must be non-negative.") + if self.max_slsqp_iterations <= 0: + raise ValueError("max_slsqp_iterations must be positive.") + if self.slsqp_ftol <= 0.0: + raise ValueError("slsqp_ftol must be positive.") + if self.max_collision_rounds <= 0: + raise ValueError("max_collision_rounds must be positive.") + if self.max_added_collision_pairs <= 0: + raise ValueError("max_added_collision_pairs must be positive.") + + +class TableSurfaceLayoutOptimizer: + """Solve direct table children with table, relation, and collision constraints.""" + + def __init__( + self, + *, + config: TableSurfaceLayoutOptimizerConfig | None = None, + ) -> None: + self.config = ( + config if config is not None else TableSurfaceLayoutOptimizerConfig() + ) + + def optimize( + self, + problem: TableSurfaceLayoutProblem, + ) -> dict[str, list[float]]: + """Return the table-frame XY centers satisfying this atomic problem.""" + # Measure only this sibling group from the complete scene-asset index. + root_half_extents_xy = _asset_half_extents_xy( + assets_by_id=problem.assets_by_id, + object_ids=problem.root_ids, + ) + # Equality constraints for fixed roots, and inequality constraints for table-region and planar-relation bounds. + inequality_constraints, equality_constraints = _build_constraints( + problem=problem, + root_half_extents_xy=root_half_extents_xy, + config=self.config, + ) + # Solve with the SLSQP optimizer. + solved_root_xy_by_id = _solve_root_xy( + root_ids=problem.root_ids, + root_seed_xy_by_id=problem.root_seed_xy_by_id, + imported_root_ids=problem.imported_root_ids, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + config=self.config, + ) + return _refine_root_collisions( + root_ids=problem.root_ids, + root_seed_xy_by_id=problem.root_seed_xy_by_id, + imported_root_ids=problem.imported_root_ids, + root_half_extents_xy=root_half_extents_xy, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + fixed_root_xy_by_id=problem.fixed_root_xy_by_id, + solved_root_xy_by_id=solved_root_xy_by_id, + config=self.config, + ) + + +class _LayoutInfeasibleError(ValueError): + """Internal marker for an SLSQP failure while testing one collision direction.""" + + +def _build_constraints( + *, + problem: TableSurfaceLayoutProblem, + root_half_extents_xy: dict[str, np.ndarray], + config: TableSurfaceLayoutOptimizerConfig, +) -> tuple[list[tuple[np.ndarray, float]], list[tuple[np.ndarray, float]]]: + """Build hard table-region, planar-relation, and fixed-root constraints.""" + # Objects which need to be optimized. + root_index = {root_id: index for index, root_id in enumerate(problem.root_ids)} + table_bounds = _bounds_from_points(problem.table_optimization_rect_xy) + # Initi constraints. + inequality_constraints: list[tuple[np.ndarray, float]] = [] + equality_constraints: list[tuple[np.ndarray, float]] = [] + for root_id in problem.root_ids: + # Get the table region bound for this root asset. + region_bounds = _table_region_bounds( + table_bounds=table_bounds, + table_region=problem.root_table_regions_by_id[root_id], + ) + # Add AABB constraints for each root's center inside the table region. + _append_aabb_center_bounds( + constraints=inequality_constraints, + root_index=root_index, + root_id=root_id, + bounds=region_bounds, + half_extents_xy=root_half_extents_xy[root_id], + ) + fixed_xy = problem.fixed_root_xy_by_id[root_id] + if fixed_xy is not None: + # Add fixed-root constraints for each root with a fixed XY center. + _append_fixed_root_constraints( + constraints=equality_constraints, + root_index=root_index, + root_id=root_id, + fixed_xy=fixed_xy, + ) + for relation in problem.root_relations: + # Add planar-relation constraints for each sibling relation in this group. + _append_planar_relation_constraint( + constraints=inequality_constraints, + root_index=root_index, + source_id=relation.source_id, + relation=relation.relation, + target_id=relation.target_id, + source_half_extents_xy=root_half_extents_xy[relation.source_id], + target_half_extents_xy=root_half_extents_xy[relation.target_id], + relation_clearance_m=config.relation_clearance_m, + ) + return inequality_constraints, equality_constraints + + +def _table_region_bounds( + *, + table_bounds: np.ndarray, + table_region: str | None, +) -> np.ndarray: + """Return the requested 3x3 table region, with y increasing toward front.""" + if table_region is None: + return table_bounds.copy() + column_by_region = { + "left_back": 0, + "left_center": 0, + "left_front": 0, + "back_center": 1, + "center": 1, + "front_center": 1, + "right_back": 2, + "right_center": 2, + "right_front": 2, + } + row_by_region = { + "left_back": 0, + "back_center": 0, + "right_back": 0, + "left_center": 1, + "center": 1, + "right_center": 1, + "left_front": 2, + "front_center": 2, + "right_front": 2, + } + if table_region not in column_by_region: + raise ValueError(f"Unsupported table region {table_region!r}.") + minimum, maximum = table_bounds + # 9-grid. + cell_size = (maximum - minimum) / 3.0 + region_minimum = minimum + cell_size * np.array( + [column_by_region[table_region], row_by_region[table_region]] + ) + return np.stack([region_minimum, region_minimum + cell_size]) + + +def _asset_half_extents_xy( + *, assets_by_id: dict[str, SceneObject], object_ids: list[str] +) -> dict[str, np.ndarray]: + """Measure each optimized asset's oriented z-up XY half-extents.""" + result = {} + for object_id in object_ids: + asset = assets_by_id.get(object_id) + if asset is None: + raise ValueError(f"Table root {object_id!r} is not an asset.") + mesh = load_scene_object_z_up_mesh(scene_object=asset) + result[object_id] = (mesh.bounds[1, :2] - mesh.bounds[0, :2]) / 2.0 + return result + + +def _bounds_from_points(points: list[list[float]]) -> np.ndarray: + coordinates = np.asarray(points, dtype=float) + if ( + coordinates.ndim != 2 + or coordinates.shape[1] != 2 + or len(coordinates) < 2 + or not np.all(np.isfinite(coordinates)) + ): + raise ValueError("XY bounds must contain at least two finite points.") + return np.stack([coordinates.min(axis=0), coordinates.max(axis=0)]) + + +def _append_aabb_center_bounds( + *, + constraints: list[tuple[np.ndarray, float]], + root_index: dict[str, int], + root_id: str, + bounds: np.ndarray, + half_extents_xy: np.ndarray, +) -> None: + minimum, maximum = bounds[0] + half_extents_xy, bounds[1] - half_extents_xy + if np.any(minimum > maximum): + raise ValueError( + f"Asset {root_id!r} cannot fit inside its assigned table region." + ) + # root_id is the sibling whose center is constrained in this AABB bound. + offset, count = 2 * root_index[root_id], 2 * len(root_index) + # offset selects this root's XY pair; count is the full flattened XY vector size. + for axis in range(2): + upper, lower = np.zeros(count), np.zeros(count) + upper[offset + axis], lower[offset + axis] = 1.0, -1.0 + constraints.extend( + [(upper, float(maximum[axis])), (lower, -float(minimum[axis]))] + ) + + +def _append_fixed_root_constraints( + *, + constraints: list[tuple[np.ndarray, float]], + root_index: dict[str, int], + root_id: str, + fixed_xy: list[float], +) -> None: + offset, count = 2 * root_index[root_id], 2 * len(root_index) + for axis, coordinate in enumerate(fixed_xy): + row = np.zeros(count) + row[offset + axis] = 1.0 + constraints.append((row, float(coordinate))) + + +def _append_planar_relation_constraint( + *, + constraints: list[tuple[np.ndarray, float]], + root_index: dict[str, int], + source_id: str, + relation: str, + target_id: str, + source_half_extents_xy: np.ndarray, + target_half_extents_xy: np.ndarray, + relation_clearance_m: float, +) -> None: + axis, sign = { + "left_of": (0, 1.0), + "right_of": (0, -1.0), + "behind": (1, 1.0), + "in_front_of": (1, -1.0), + }.get(relation, (None, None)) + if axis is None or source_id not in root_index or target_id not in root_index: + raise ValueError(f"Unsupported table-root planar relation {relation!r}.") + row = np.zeros(2 * len(root_index)) + row[2 * root_index[source_id] + axis] = sign + row[2 * root_index[target_id] + axis] = -sign + constraints.append( + ( + row, + -float( + source_half_extents_xy[axis] + + target_half_extents_xy[axis] + + relation_clearance_m + ), + ) + ) + + +def _solve_root_xy( + *, + root_ids: list[str], + root_seed_xy_by_id: dict[str, list[float]], + imported_root_ids: set[str], + inequality_constraints: list[tuple[np.ndarray, float]], + equality_constraints: list[tuple[np.ndarray, float]], + config: TableSurfaceLayoutOptimizerConfig, +) -> dict[str, list[float]]: + # Init with XY-seeds. + initial = np.asarray( + [root_seed_xy_by_id[root_id] for root_id in root_ids], dtype=float + ) + + def objective(values: np.ndarray) -> float: + xy = values.reshape(-1, 2) + loss = 0.0 + for index, root_id in enumerate(root_ids): + if root_id in imported_root_ids: + delta = xy[index] - initial[index] + loss += config.imported_seed_weight * float(delta @ delta) + return loss + + constraints = [ + { + "type": "ineq", + "fun": lambda values, row=row, bound=bound: bound - float(row @ values), + } + for row, bound in inequality_constraints + ] + [ + { + "type": "eq", + "fun": lambda values, row=row, bound=bound: float(row @ values) - bound, + } + for row, bound in equality_constraints + ] + result = minimize( + objective, + initial.reshape(-1), + method="SLSQP", + constraints=constraints, + options={ + "maxiter": config.max_slsqp_iterations, + "ftol": config.slsqp_ftol, + "disp": False, + }, + ) + if not result.success: + raise _LayoutInfeasibleError( + f"Table layout optimization failed: {result.message}" + ) + return { + root_id: [float(result.x[2 * index]), float(result.x[2 * index + 1])] + for index, root_id in enumerate(root_ids) + } + + +def _refine_root_collisions( + *, + root_ids: list[str], + root_seed_xy_by_id: dict[str, list[float]], + imported_root_ids: set[str], + root_half_extents_xy: dict[str, np.ndarray], + inequality_constraints: list[tuple[np.ndarray, float]], + equality_constraints: list[tuple[np.ndarray, float]], + fixed_root_xy_by_id: dict[str, list[float] | None], + solved_root_xy_by_id: dict[str, list[float]], + config: TableSurfaceLayoutOptimizerConfig, +) -> dict[str, list[float]]: + # Get current SLSQP solution. + current = solved_root_xy_by_id + seen: set[tuple[str, str]] = set() + for _ in range(config.max_collision_rounds): + # Fine overlaps. + overlaps = [ + pair + for pair in _root_aabb_overlaps( + root_ids=root_ids, half_extents=root_half_extents_xy, xy_by_id=current + ) + if fixed_root_xy_by_id[pair[1]] is None + or fixed_root_xy_by_id[pair[2]] is None + ] + if not overlaps: + return current + added = 0 + for _, first_id, second_id in overlaps[: config.max_added_collision_pairs]: + key = tuple(sorted((first_id, second_id))) + if key in seen: + continue + # Earlier pair updates may already have separated this stale overlap. + if key not in { + tuple(sorted((first, second))) + for _, first, second in _root_aabb_overlaps( + root_ids=root_ids, + half_extents=root_half_extents_xy, + xy_by_id=current, + ) + }: + continue + for separation_constraint in _aabb_separation_constraints( + root_ids=root_ids, + first_id=first_id, + second_id=second_id, + half_extents=root_half_extents_xy, + xy_by_id=current, + margin=config.collision_margin_m, + ): + # Keep a candidate only when it is compatible with all hard constraints. + try: + solved_xy_by_id = _solve_root_xy( + root_ids=root_ids, + root_seed_xy_by_id=current, + imported_root_ids=imported_root_ids, + inequality_constraints=[ + *inequality_constraints, + separation_constraint, + ], + equality_constraints=equality_constraints, + config=config, + ) + except _LayoutInfeasibleError: + continue + inequality_constraints.append(separation_constraint) + current = solved_xy_by_id + seen.add(key) + added += 1 + break + else: + raise ValueError( + "Table-root AABB pair has no feasible separation direction: " + f"{first_id!r}, {second_id!r}." + ) + if not added: + break + raise ValueError("Table-root AABB collisions remain after layout refinement.") + + +def _root_aabb_overlaps( + *, + root_ids: list[str], + half_extents: dict[str, np.ndarray], + xy_by_id: dict[str, list[float]], +) -> list[tuple[float, str, str]]: + """Return all overlapping root pairs with their minimum XY overlap distance.""" + result = [] + for index, first_id in enumerate(root_ids): + for second_id in root_ids[index + 1 :]: + overlap = np.minimum( + np.asarray(xy_by_id[first_id]) + half_extents[first_id], + np.asarray(xy_by_id[second_id]) + half_extents[second_id], + ) - np.maximum( + np.asarray(xy_by_id[first_id]) - half_extents[first_id], + np.asarray(xy_by_id[second_id]) - half_extents[second_id], + ) + if np.all(overlap > 1e-9): + result.append((float(np.min(overlap)), first_id, second_id)) + return sorted(result, reverse=True) + + +def _aabb_separation_constraints( + *, + root_ids: list[str], + first_id: str, + second_id: str, + half_extents: dict[str, np.ndarray], + xy_by_id: dict[str, list[float]], + margin: float, +) -> list[tuple[np.ndarray, float]]: + """Return ordered feasible-direction candidates for one overlapping AABB pair.""" + first, second = np.asarray(xy_by_id[first_id]), np.asarray(xy_by_id[second_id]) + # Positive overlap on both axes means these two center-based AABBs intersect. + overlap = np.minimum( + first + half_extents[first_id], second + half_extents[second_id] + ) - np.maximum(first - half_extents[first_id], second - half_extents[second_id]) + # Try the least-penetrating axis first, but permit order reversal if required. + axes = np.argsort(overlap) + constraints = [] + for axis in axes: + current_order = first[axis] < second[axis] or ( + first[axis] == second[axis] and first_id < second_id + ) + for first_is_lower in (current_order, not current_order): + constraints.append( + _aabb_separation_constraint_for_direction( + root_ids=root_ids, + first_id=first_id, + second_id=second_id, + half_extents=half_extents, + axis=int(axis), + first_is_lower=first_is_lower, + margin=margin, + ) + ) + return constraints + + +def _aabb_separation_constraint_for_direction( + *, + root_ids: list[str], + first_id: str, + second_id: str, + half_extents: dict[str, np.ndarray], + axis: int, + first_is_lower: bool, + margin: float, +) -> tuple[np.ndarray, float]: + """Return one directed AABB separation inequality on a selected axis.""" + index = {root_id: i for i, root_id in enumerate(root_ids)} + # One row addresses the x/y variable pair of each root in the flattened solver vector. + row = np.zeros(2 * len(root_ids)) + sign = 1.0 if first_is_lower else -1.0 + row[2 * index[first_id] + axis], row[2 * index[second_id] + axis] = sign, -sign + # row @ values <= bound keeps the selected AABB faces apart by the requested margin. + return row, -float( + half_extents[first_id][axis] + half_extents[second_id][axis] + margin + ) diff --git a/embodichain/gen_sim/task_engine/__init__.py b/embodichain/gen_sim/task_engine/__init__.py new file mode 100644 index 000000000..087983cc4 --- /dev/null +++ b/embodichain/gen_sim/task_engine/__init__.py @@ -0,0 +1,179 @@ +# ---------------------------------------------------------------------------- +# 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 task interpretation and protocol ownership.""" + +from __future__ import annotations + +from typing import Any + +from .agent import ( + TaskAgent, + TaskGenerationError, + derive_scene_request, + derive_success_spec, +) +from .contracts import ( + SCENE_REQUEST_SCHEMA, + SUCCESS_SPEC_SCHEMA, + TASK_CANDIDATE_SET_SCHEMA, + TASK_DRAFT_SCHEMA, + SceneRequest, + SuccessSpec, + TaskCandidate, + TaskCandidateSet, + TaskDraft, + canonical_hash, + validate_scene_request, + validate_success_spec, + validate_task_candidate, + validate_task_candidate_set, + validate_task_draft, +) +from .interpretation import ( + INSTRUCTION_INTENT_SCHEMA, + InstructionCaller, + InstructionDraftResult, + InstructionIntent, + interpret_instruction_draft, + validate_instruction_intent, +) +from .ontology import ( + RELATIONS, + TASK_CONTRACTS, + TERMINAL_BEHAVIORS, + TRANSPORT_DIRECTIONS, + TaskContract, + task_contract, + task_success_type, +) +from .config import ( + TASK_ENGINE_DEFAULTS_SCHEMA, + TaskEngineExecutionCfg, + TaskEnginePlanningCfg, + TaskEngineWorkflowCfg, + load_task_engine_config, +) +from .state_machine import ( + StageStatus, + TaskEngineState, + WorkflowStage, + complete_stage, + fail_stage, + initial_state, + replay_events, + skip_stage, + start_stage, +) +from .workflow_contracts import ( + TASK_RUN_REQUEST_SCHEMA, + SceneInputKind, + TaskRunRequest, + scene_input_kind, + validate_scene_history_root, + validate_scene_output_separation, + validate_task_run_request, +) + +__all__ = [ + "INSTRUCTION_INTENT_SCHEMA", + "InstructionCaller", + "InstructionDraftResult", + "InstructionIntent", + "RELATIONS", + "SCENE_REQUEST_SCHEMA", + "SUCCESS_SPEC_SCHEMA", + "SceneRequest", + "SuccessSpec", + "TASK_CANDIDATE_SET_SCHEMA", + "TASK_CONTRACTS", + "TASK_DRAFT_SCHEMA", + "TERMINAL_BEHAVIORS", + "TRANSPORT_DIRECTIONS", + "TaskAgent", + "TaskCandidate", + "TaskCandidateSet", + "TaskContract", + "TaskDraft", + "TaskGenerationError", + "TASK_RUN_REQUEST_SCHEMA", + "TASK_ENGINE_DEFAULTS_SCHEMA", + "SceneInputKind", + "SceneAnalysis", + "SceneEngineBackend", + "SceneRevision", + "StageStatus", + "TaskEngineState", + "TaskEngineExecutionCfg", + "TaskEnginePlanningCfg", + "TaskEngineWorkflowCfg", + "TaskEngineRunResult", + "TaskEngineWorkflow", + "TaskRunRequest", + "WorkflowStage", + "TASK_ENGINE_RUN_MANIFEST_SCHEMA", + "SubprocessActionExecutor", + "canonical_hash", + "derive_scene_request", + "derive_success_spec", + "complete_stage", + "fail_stage", + "initial_state", + "interpret_instruction_draft", + "load_task_engine_config", + "replay_events", + "task_contract", + "task_success_type", + "scene_input_kind", + "validate_scene_history_root", + "scene_blueprint_objects", + "skip_stage", + "start_stage", + "validate_instruction_intent", + "validate_scene_request", + "validate_scene_output_separation", + "validate_success_spec", + "validate_task_candidate", + "validate_task_candidate_set", + "validate_task_draft", + "validate_task_run_request", +] + +_SCENE_BACKEND_EXPORTS = { + "SceneAnalysis", + "SceneEngineBackend", + "SceneRevision", + "scene_blueprint_objects", +} +_WORKFLOW_EXPORTS = { + "TASK_ENGINE_RUN_MANIFEST_SCHEMA", + "SubprocessActionExecutor", + "TaskEngineRunResult", + "TaskEngineWorkflow", +} + + +def __getattr__(name: str) -> Any: + """Load orchestration entry points lazily to avoid engine import cycles.""" + if name in _SCENE_BACKEND_EXPORTS: + from . import scene_backend + + return getattr(scene_backend, name) + if name in _WORKFLOW_EXPORTS: + from . import workflow + + return getattr(workflow, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/embodichain/gen_sim/task_engine/__main__.py b/embodichain/gen_sim/task_engine/__main__.py new file mode 100644 index 000000000..9e4f06dbe --- /dev/null +++ b/embodichain/gen_sim/task_engine/__main__.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. +# ---------------------------------------------------------------------------- + +"""Module entry point for Task Engine workflows.""" + +from __future__ import annotations + +from .cli import main + +__all__ = ["main"] + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/embodichain/gen_sim/task_engine/_bundle_runner.py b/embodichain/gen_sim/task_engine/_bundle_runner.py new file mode 100644 index 000000000..5bc80bd9b --- /dev/null +++ b/embodichain/gen_sim/task_engine/_bundle_runner.py @@ -0,0 +1,197 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Private subprocess boundary for executing one prepared Task Engine bundle.""" + +from __future__ import annotations + +import argparse +from contextlib import contextmanager +import json +from pathlib import Path +import sys +from typing import Any, Iterator, Sequence + +from embodichain.gen_sim.action_engine.agent import ActionAgent +from embodichain.gen_sim.action_engine.protocol import ( + AGENT_CONFIG_FILENAME, + EXECUTION_PROGRAM_FILENAME, + FAST_GYM_CONFIG_FILENAME, +) +from embodichain.gen_sim.action_engine.runtime import ExecutionReport + +from .orchestration.artifacts import ( + GROUNDED_TASK_PLAN_FILENAME, + STATIC_SCENE_MANIFEST_FILENAME, + write_execution_report, +) +from .orchestration.contracts import validate_grounded_task_plan +from .orchestration.scene_source import verify_scene_source_fingerprint + +__all__ = ["execute_bundle", "main"] + + +def main(argv: Sequence[str] | None = None) -> int: + """Parse the private runner protocol and execute one bundle.""" + parser = argparse.ArgumentParser( + prog="embodichain.gen_sim.task_engine._bundle_runner" + ) + parser.add_argument("--bundle", required=True) + args, forwarded = parser.parse_known_args(argv) + return execute_bundle(args.bundle, forwarded) + + +def execute_bundle( + bundle: str | Path, + forwarded: Sequence[str] = (), +) -> int: + """Execute one prepared bundle through the existing Action Engine launcher.""" + root = Path(bundle).expanduser().resolve() + if not root.is_dir(): + raise FileNotFoundError(f"Bundle directory does not exist: {root}") + agent_config = root / AGENT_CONFIG_FILENAME + gym_config = root / FAST_GYM_CONFIG_FILENAME + for path in (agent_config, gym_config): + if not path.is_file(): + raise FileNotFoundError(f"Bundle is missing required artifact: {path}") + task_id = _bundle_task_id(root, agent_config) + run_args = list(forwarded) + if run_args and run_args[0] == "--": + run_args.pop(0) + rejection = _preflight_bundle( + root, + agent_config=agent_config, + gym_config=gym_config, + forwarded=run_args, + ) + if rejection is not None: + write_execution_report(root, rejection) + _print_json(rejection.as_mapping()) + return 2 + legacy_argv = [ + "--task_name", + task_id, + "--gym_config", + str(gym_config), + "--agent_config", + str(agent_config), + "--task-engine-report", + *run_args, + ] + from embodichain.gen_sim.action_engine.cli import run_agent + + with _temporary_argv(["run_agent", *legacy_argv]): + return int(run_agent.cli() or 0) + + +def _preflight_bundle( + bundle: Path, + *, + agent_config: Path, + gym_config: Path, + forwarded: Sequence[str], +) -> ExecutionReport | None: + static_manifest_path = bundle / STATIC_SCENE_MANIFEST_FILENAME + if static_manifest_path.is_file(): + static_manifest = _read_json(static_manifest_path) + source = static_manifest.get("source", {}) + if isinstance(source, dict) and isinstance( + source.get("source_fingerprint"), dict + ): + verify_scene_source_fingerprint(source["source_fingerprint"]) + grounded_path = bundle / GROUNDED_TASK_PLAN_FILENAME + if not grounded_path.is_file(): + return None + grounded = validate_grounded_task_plan(_read_json(grounded_path)) + agent = _read_json(agent_config) + graph_value = agent.get("seed_task_graph", EXECUTION_PROGRAM_FILENAME) + if not isinstance(graph_value, str) or not graph_value: + raise ValueError("Bundle agent_config.seed_task_graph must be a path string.") + graph_path = Path(graph_value).expanduser() + if not graph_path.is_absolute(): + graph_path = (bundle / graph_path).resolve() + else: + graph_path = graph_path.resolve() + if graph_path != bundle and bundle not in graph_path.parents: + raise ValueError("Bundle SeedGraph path escapes the bundle directory.") + if not graph_path.is_file(): + raise FileNotFoundError(f"Bundle is missing SeedGraph: {graph_path}") + action_agent = ActionAgent() + try: + action_agent.preflight( + graph_path, + scene_manifest=grounded["scene_manifest"], + ) + except (TypeError, ValueError, OSError) as exc: + return action_agent.rejection_report( + graph_path, + exc, + grounded_plan=grounded, + environment_count=_environment_count(gym_config, forwarded), + ) + return None + + +def _environment_count(gym_config: Path, forwarded: Sequence[str]) -> int: + value: Any = _read_json(gym_config).get("num_envs", 1) + for index, argument in enumerate(forwarded): + if argument == "--num_envs" and index + 1 < len(forwarded): + value = forwarded[index + 1] + elif argument.startswith("--num_envs="): + value = argument.partition("=")[2] + try: + return max(1, int(value)) + except (TypeError, ValueError): + return 1 + + +def _bundle_task_id(bundle: Path, agent_config: Path) -> str: + grounded_path = bundle / GROUNDED_TASK_PLAN_FILENAME + if grounded_path.is_file(): + task_id = _read_json(grounded_path).get("task_id") + else: + task_id = _read_json(agent_config).get("task_name") + if not isinstance(task_id, str) or not task_id.strip(): + raise ValueError("Bundle does not declare a non-empty task ID.") + return task_id.strip() + + +def _read_json(path: Path) -> 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 JSON artifact {path}: {exc}") from exc + if not isinstance(value, dict): + raise ValueError(f"JSON artifact must contain an object: {path}") + return value + + +@contextmanager +def _temporary_argv(arguments: list[str]) -> Iterator[None]: + original = sys.argv + sys.argv = arguments + try: + yield + finally: + sys.argv = original + + +def _print_json(value: Any) -> None: + print(json.dumps(value, ensure_ascii=False, indent=2, allow_nan=False)) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/embodichain/gen_sim/task_engine/agent.py b/embodichain/gen_sim/task_engine/agent.py new file mode 100644 index 000000000..f5ab773f0 --- /dev/null +++ b/embodichain/gen_sim/task_engine/agent.py @@ -0,0 +1,289 @@ +# ---------------------------------------------------------------------------- +# 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 semantic candidate generation for Task Engine.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from concurrent.futures import ThreadPoolExecutor, as_completed +from copy import deepcopy +from dataclasses import dataclass +from typing import Any + +from .contracts import ( + SCENE_REQUEST_SCHEMA, + SUCCESS_SPEC_SCHEMA, + TASK_CANDIDATE_SET_SCHEMA, + TASK_DRAFT_SCHEMA, + TaskCandidate, + TaskCandidateSet, + canonical_hash, + validate_task_candidate, + validate_task_candidate_set, +) +from .interpretation import ( + InstructionCaller, + InstructionDraftResult, + interpret_instruction_draft, + validate_instruction_intent, +) +from .ontology import TASK_CONTRACTS, task_success_type + +__all__ = [ + "TaskAgent", + "TaskGenerationError", + "derive_scene_request", + "derive_success_spec", +] + +DraftInterpreter = Callable[..., InstructionDraftResult] + + +class TaskGenerationError(ValueError): + """Raised when every independently generated candidate fails validation.""" + + +@dataclass(frozen=True) +class _CandidateAttempt: + index: int + result: InstructionDraftResult | None = None + error: str = "" + + +class TaskAgent: + """Generate, validate, normalize, and vote on independent task drafts.""" + + def __init__( + self, + *, + caller: InstructionCaller | None = None, + interpreter: DraftInterpreter = interpret_instruction_draft, + ) -> None: + self._caller = caller + self._interpreter = interpreter + + def generate( + self, + task_id: str, + instruction: str, + model: str | None = None, + candidate_count: int = 3, + ) -> TaskCandidateSet: + """Generate candidates concurrently and retain votes after deduplication.""" + normalized_task_id = str(task_id).strip() + normalized_instruction = str(instruction).strip() + if not normalized_task_id or not normalized_instruction: + raise ValueError("task_id and instruction must be non-empty.") + if ( + isinstance(candidate_count, bool) + or not isinstance(candidate_count, int) + or candidate_count < 1 + ): + raise ValueError("candidate_count must be a positive integer.") + + attempts: list[_CandidateAttempt] = [] + with ThreadPoolExecutor( + max_workers=candidate_count, + thread_name_prefix="task-agent", + ) as executor: + futures = { + executor.submit( + self._interpreter, + normalized_instruction, + model=model, + caller=self._caller, + ): index + for index in range(candidate_count) + } + for future in as_completed(futures): + index = futures[future] + try: + attempts.append( + _CandidateAttempt(index=index, result=future.result()) + ) + except Exception as error: # Each candidate is an isolated vote. + attempts.append( + _CandidateAttempt( + index=index, + error=f"candidate_{index + 1:02d}: {type(error).__name__}: {error}", + ) + ) + attempts.sort(key=lambda item: item.index) + errors = [item.error for item in attempts if item.result is None] + unique: dict[str, TaskCandidate] = {} + valid_response_count = 0 + for attempt in attempts: + if attempt.result is None: + continue + assert attempt.result is not None + try: + canonical_intent = _canonicalize_intent(attempt.result.intent) + draft = { + "schema_version": TASK_DRAFT_SCHEMA, + "task_id": normalized_task_id, + "instruction": normalized_instruction, + "steps": canonical_intent["steps"], + } + semantic_hash = canonical_hash(draft["steps"]) + candidate_id = f"candidate_{len(unique) + 1:02d}" + candidate = validate_task_candidate( + { + "candidate_id": candidate_id, + "draft": draft, + "scene_request": derive_scene_request(draft), + "success_spec": derive_success_spec(draft), + "semantic_hash": semantic_hash, + "vote_count": 1, + "attempts": attempt.result.attempts, + "normalizations": deepcopy(list(attempt.result.normalizations)), + } + ) + existing = unique.get(semantic_hash) + if existing is not None: + existing["vote_count"] += 1 + existing["attempts"] = max( + existing["attempts"], attempt.result.attempts + ) + existing["normalizations"].extend(candidate["normalizations"]) + else: + unique[semantic_hash] = candidate + valid_response_count += 1 + except Exception as error: # Post-processing failures stay candidate-local. + errors.append( + f"candidate_{attempt.index + 1:02d}: " + f"{type(error).__name__}: {error}" + ) + + if not unique: + raise TaskGenerationError( + "All Task Agent candidates failed validation: " + "; ".join(errors) + ) + + return validate_task_candidate_set( + { + "schema_version": TASK_CANDIDATE_SET_SCHEMA, + "task_id": normalized_task_id, + "instruction": normalized_instruction, + "candidates": list(unique.values()), + "requested_candidate_count": candidate_count, + "valid_response_count": valid_response_count, + "errors": errors, + } + ) + + +def derive_scene_request(draft: Mapping[str, Any]) -> dict[str, Any]: + """Derive structural scene constraints without classifying reference text.""" + from .contracts import validate_scene_request, validate_task_draft + + normalized = validate_task_draft(draft) + references: list[dict[str, Any]] = [] + for step in normalized["steps"]: + task_type = str(step["task_type"]) + contract = TASK_CONTRACTS[task_type] + for role in ("object", "target"): + selector = step[role] + if selector["kind"] != "scene_ref": + continue + if role == "object": + structure = contract.source_structure + affordances = sorted(contract.scene_affordances) + initial_state = {"orientation": "fallen"} if task_type == "E2" else {} + attributes: dict[str, Any] = {} + else: + structure = _target_structure(task_type, str(step["relation"])) + affordances = _target_affordances(task_type, str(step["relation"])) + initial_state = {} + attributes = {} + references.append( + { + "reference_id": f"{step['id']}.{role}", + "step_id": step["id"], + "role": role, + "reference": selector["reference"], + "quantifier": selector["quantifier"], + "count": selector["count"], + "source_structure": structure, + "affordances": affordances, + "initial_state": initial_state, + "attributes": attributes, + } + ) + if not references: + raise ValueError("A TaskDraft must contain at least one scene_ref selector.") + return validate_scene_request( + { + "schema_version": SCENE_REQUEST_SCHEMA, + "task_id": normalized["task_id"], + "references": references, + } + ) + + +def derive_success_spec(draft: Mapping[str, Any]) -> dict[str, Any]: + """Derive every success term exclusively from the E-task ontology.""" + from .contracts import validate_success_spec, validate_task_draft + + normalized = validate_task_draft(draft) + return validate_success_spec( + { + "schema_version": SUCCESS_SPEC_SCHEMA, + "task_id": normalized["task_id"], + "op": "all", + "terms": [ + { + "step_id": step["id"], + "type": task_success_type(step["task_type"], step), + } + for step in normalized["steps"] + ], + }, + draft=normalized, + ) + + +def _canonicalize_intent(intent: Mapping[str, Any]) -> dict[str, Any]: + """Remove arbitrary model step IDs while preserving the explicit DAG order.""" + normalized = validate_instruction_intent(intent) + id_map = { + step["id"]: f"step_{index + 1:02d}" + for index, step in enumerate(normalized["steps"]) + } + steps = deepcopy(normalized["steps"]) + for step in steps: + old_id = step["id"] + step["id"] = id_map[old_id] + step["depends_on"] = [id_map[item] for item in step["depends_on"]] + for selector_name in ("object", "target"): + selector = step[selector_name] + if selector["kind"] == "step_result": + selector["step_id"] = id_map[selector["step_id"]] + return validate_instruction_intent({"steps": steps}) + + +def _target_affordances(task_type: str, relation: str) -> list[str]: + if task_type == "E3" or (task_type == "E1" and relation == "inside"): + return ["container"] + return [] + + +def _target_structure(task_type: str, relation: str) -> str: + if task_type == "E1" and relation == "on": + return "physical_entity" + if task_type == "E3" or (task_type == "E1" and relation == "inside"): + return "rigid_object" + return "spatial_reference" diff --git a/embodichain/gen_sim/task_engine/cli.py b/embodichain/gen_sim/task_engine/cli.py new file mode 100644 index 000000000..fec6de86a --- /dev/null +++ b/embodichain/gen_sim/task_engine/cli.py @@ -0,0 +1,246 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Unified CLI for complete Task Engine workflows.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +import sys +from typing import Any, Final, Sequence + +from .config import load_task_engine_config +from .orchestration.scene_adapter import SceneAdapter +from .run_directory import reserve_run_directory +from .workflow import SubprocessActionExecutor, TaskEngineWorkflow +from .workflow_contracts import ( + TASK_RUN_REQUEST_SCHEMA, + validate_scene_history_root, + validate_scene_output_separation, +) + +__all__ = ["build_parser", "main"] + + +_ROBOT_PROFILES = ( + "ur5", + "ur10", + "dual_ur5", + "dual_ur10", + "franka", + "dual_franka", +) +_MODES: Final = ("image", "image-edit", "scene", "scene-edit") + + +def build_parser() -> argparse.ArgumentParser: + """Build the Task Engine parser.""" + parser = argparse.ArgumentParser( + prog="embodichain task-engine", + description="Prepare, run, or complete one Scene and Action workflow.", + ) + subparsers = parser.add_subparsers(dest="command", required=True) + prepare_parser = subparsers.add_parser( + "prepare", help="Prepare a bundle without simulator execution." + ) + _add_workflow_arguments(prepare_parser) + run_all_parser = subparsers.add_parser( + "run-all", help="Prepare and execute one complete workflow." + ) + _add_workflow_arguments(run_all_parser) + run_parser = subparsers.add_parser( + "run", help="Execute an already prepared Task Engine bundle." + ) + run_parser.add_argument("--bundle", required=True) + run_parser.add_argument("--output-root", required=True) + run_parser.add_argument("--config", default=None) + run_parser.add_argument("--seed", type=int, default=0) + run_parser.add_argument("--num-envs", type=int, default=None) + run_parser.add_argument("--dataset-saving", action="store_true") + return parser + + +def _add_workflow_arguments(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--mode", choices=_MODES, required=True) + parser.add_argument("--task-id", "--task_id", required=True) + instruction = parser.add_mutually_exclusive_group(required=True) + instruction.add_argument("--instruction") + instruction.add_argument("--task-file", "--task_file") + parser.add_argument("--image") + parser.add_argument("--scene") + parser.add_argument("--scene-edit", "--scene_edit", default=None) + parser.add_argument("--output-root", required=True) + parser.add_argument("--config", default=None) + parser.add_argument("--model", default=None) + parser.add_argument("--vlm-model", default=None) + parser.add_argument("--base-seed", type=int, default=0) + parser.add_argument( + "--dataset_saving", + action="store_true", + help="Opt in to the Gym project's dataset recorder during execution.", + ) + parser.add_argument( + "--robot-profile", + choices=_ROBOT_PROFILES, + default="franka", + ) + + +def main(argv: Sequence[str] | None = None) -> int: + """Dispatch one Task Engine workflow command.""" + parser = build_parser() + arguments = list(sys.argv[1:] if argv is None else argv) + if arguments and arguments[0] not in { + "prepare", + "run", + "run-all", + "-h", + "--help", + }: + arguments.insert(0, "run-all") + args = parser.parse_args(arguments) + if args.command == "run": + return _run_prepared_bundle(args) + return _run_workflow( + args, + execute=args.command == "run-all", + parser=parser, + ) + + +def _run_workflow( + args: argparse.Namespace, + *, + execute: bool, + parser: argparse.ArgumentParser, +) -> int: + try: + image, scene, edit = _mode_inputs(args) + except ValueError as exc: + parser.error(str(exc)) + if scene is not None: + validate_scene_history_root(scene, args.output_root) + instruction = _instruction(args) + adapter = SceneAdapter(model=args.model, robot_profile=args.robot_profile) + workflow = TaskEngineWorkflow(scene_adapter=adapter) + with reserve_run_directory(args.output_root) as allocation: + if scene is not None: + validate_scene_output_separation(scene, allocation.path) + result = workflow.run( + { + "schema_version": TASK_RUN_REQUEST_SCHEMA, + "task_id": args.task_id, + "task_instruction": instruction, + "image_path": image, + "gym_project": scene, + "scene_edit_prompt": edit, + "output_dir": allocation.path.as_posix(), + }, + config_path=args.config, + model=args.model, + vlm_model=args.vlm_model, + base_seed=args.base_seed, + dataset_saving=args.dataset_saving, + run_id=allocation.run_id, + created_at=allocation.created_at, + execute=execute, + ) + _print_json( + { + "run_id": allocation.run_id, + "status": result.status, + "failure_class": result.failure_class, + "output_dir": result.output_dir.as_posix(), + "manifest": result.manifest_path.as_posix(), + "final_bundle": ( + None if result.final_bundle is None else result.final_bundle.as_posix() + ), + } + ) + accepted = result.succeeded if execute else result.status == "prepared" + return 0 if accepted else 2 + + +def _run_prepared_bundle(args: argparse.Namespace) -> int: + _, _, execution_cfg = load_task_engine_config(args.config) + num_envs = execution_cfg.num_envs if args.num_envs is None else int(args.num_envs) + if num_envs < 1: + raise ValueError("num_envs must be positive.") + with reserve_run_directory(args.output_root) as allocation: + report = SubprocessActionExecutor()( + args.bundle, + allocation.path, + seed=int(args.seed), + num_envs=num_envs, + dataset_saving=bool(args.dataset_saving), + ) + environments = report.get("environments", ()) + successes = [ + bool(item.get("success")) for item in environments if isinstance(item, dict) + ] + accepted = ( + str(report.get("status")) not in {"rejected", "aborted"} + and len(successes) == num_envs + and sum(successes) >= execution_cfg.required_successes + ) + _print_json( + { + "run_id": allocation.run_id, + "status": "succeeded" if accepted else "failed", + "output_dir": allocation.path.as_posix(), + "execution_report": report, + } + ) + return 0 if accepted else 2 + + +def _instruction(args: argparse.Namespace) -> str: + instruction = ( + str(args.instruction).strip() + if args.instruction is not None + else Path(args.task_file).expanduser().read_text(encoding="utf-8").strip() + ) + if not instruction: + raise ValueError("Task instruction must not be empty.") + return instruction + + +def _mode_inputs(args: argparse.Namespace) -> tuple[str | None, str | None, str | None]: + image = None if args.image is None else str(args.image).strip() + scene = None if args.scene is None else str(args.scene).strip() + edit = None if args.scene_edit is None else str(args.scene_edit).strip() + expected = { + "image": (True, False, False), + "image-edit": (True, False, True), + "scene": (False, True, False), + "scene-edit": (False, True, True), + }[args.mode] + actual = (bool(image), bool(scene), bool(edit)) + if actual != expected: + raise ValueError( + f"mode={args.mode!r} requires image/scene/edit={expected}, got {actual}." + ) + return image, scene, edit + + +def _print_json(value: Any) -> None: + print(json.dumps(value, ensure_ascii=False, indent=2, allow_nan=False)) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/embodichain/gen_sim/task_engine/config.py b/embodichain/gen_sim/task_engine/config.py new file mode 100644 index 000000000..60cda86ef --- /dev/null +++ b/embodichain/gen_sim/task_engine/config.py @@ -0,0 +1,190 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Configuration owned by Task Engine orchestration.""" + +from __future__ import annotations + +from collections.abc import Mapping +from importlib.resources import files +from pathlib import Path +from typing import Any, Final + +import yaml + +from embodichain.utils import configclass + +__all__ = [ + "TASK_ENGINE_DEFAULTS_SCHEMA", + "TaskEngineExecutionCfg", + "TaskEnginePlanningCfg", + "TaskEngineWorkflowCfg", + "load_task_engine_config", +] + +TASK_ENGINE_DEFAULTS_SCHEMA: Final = "embodichain.task-engine-defaults/v1" + + +@configclass +class TaskEngineExecutionCfg: + """Success policy for vectorized simulator execution.""" + + num_envs: int = 1 + success_policy: str = "any" + min_successful_envs: int = 1 + + def __post_init__(self) -> None: + if ( + isinstance(self.num_envs, bool) + or not isinstance(self.num_envs, int) + or self.num_envs < 1 + ): + raise ValueError("num_envs must be a positive integer.") + if self.success_policy not in {"any", "all", "at_least"}: + raise ValueError("success_policy must be any, all, or at_least.") + if ( + isinstance(self.min_successful_envs, bool) + or not isinstance(self.min_successful_envs, int) + or not 1 <= self.min_successful_envs <= self.num_envs + ): + raise ValueError("min_successful_envs must be in [1, num_envs].") + if self.success_policy == "any" and self.min_successful_envs != 1: + raise ValueError("success_policy=any requires min_successful_envs=1.") + if self.success_policy == "all" and self.min_successful_envs != self.num_envs: + raise ValueError( + "success_policy=all requires min_successful_envs=num_envs." + ) + + @property + def required_successes(self) -> int: + """Return the number of successful replicas required for acceptance.""" + if self.success_policy == "all": + return self.num_envs + if self.success_policy == "any": + return 1 + return self.min_successful_envs + + +@configclass +class TaskEngineWorkflowCfg: + """Conservative first-version orchestration limits. + + The packaged YAML owns retry limits so deployment testing can tune them + without changing the orchestration implementation. + """ + + max_parallel_workers: int = 2 + max_scene_attempts: int = 2 + max_action_attempts: int = 3 + + def __post_init__(self) -> None: + for field_name in ( + "max_parallel_workers", + "max_scene_attempts", + "max_action_attempts", + ): + value = getattr(self, field_name) + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise ValueError(f"{field_name} must be a positive integer.") + + +@configclass +class TaskEnginePlanningCfg: + """Task interpretation and Action bundle generation defaults.""" + + candidate_count: int = 3 + planning_mode: str = "offline" + max_episodes: int = 1 + max_episode_steps: int = 4000 + + def __post_init__(self) -> None: + for field_name in ( + "candidate_count", + "max_episodes", + "max_episode_steps", + ): + value = getattr(self, field_name) + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise ValueError(f"{field_name} must be a positive integer.") + if self.planning_mode not in {"offline", "ab"}: + raise ValueError("planning_mode must be offline or ab.") + + +def load_task_engine_config( + path: str | Path | None = None, +) -> tuple[ + TaskEngineWorkflowCfg, + TaskEnginePlanningCfg, + TaskEngineExecutionCfg, +]: + """Load strict Task Engine defaults from YAML. + + Args: + path: Optional override YAML. The packaged defaults are used when omitted. + + Returns: + Validated workflow, planning, and execution configurations. + + Raises: + TypeError: If a configuration section is not a mapping. + ValueError: If the YAML schema or fields are invalid. + """ + content = ( + Path(path).expanduser().resolve().read_text(encoding="utf-8") + if path is not None + else files(__package__).joinpath("defaults.yaml").read_text(encoding="utf-8") + ) + raw = yaml.safe_load(content) + if not isinstance(raw, Mapping): + raise TypeError("Task Engine configuration must be a mapping.") + expected = {"schema_version", "workflow", "planning", "execution"} + if set(raw) != expected: + raise ValueError("Task Engine configuration fields are invalid.") + if raw.get("schema_version") != TASK_ENGINE_DEFAULTS_SCHEMA: + raise ValueError("Task Engine configuration schema_version is invalid.") + workflow = _mapping(raw.get("workflow"), "workflow") + planning = _mapping(raw.get("planning"), "planning") + execution = _mapping(raw.get("execution"), "execution") + if set(workflow) != { + "max_parallel_workers", + "max_scene_attempts", + "max_action_attempts", + }: + raise ValueError("Task Engine workflow configuration fields are invalid.") + if set(planning) != { + "candidate_count", + "planning_mode", + "max_episodes", + "max_episode_steps", + }: + raise ValueError("Task Engine planning configuration fields are invalid.") + if set(execution) != { + "num_envs", + "success_policy", + "min_successful_envs", + }: + raise ValueError("Task Engine execution configuration fields are invalid.") + return ( + TaskEngineWorkflowCfg(**workflow), + TaskEnginePlanningCfg(**planning), + TaskEngineExecutionCfg(**execution), + ) + + +def _mapping(value: Any, field_name: str) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise TypeError(f"Task Engine {field_name} configuration must be a mapping.") + return dict(value) diff --git a/embodichain/gen_sim/task_engine/contracts.py b/embodichain/gen_sim/task_engine/contracts.py new file mode 100644 index 000000000..237423e55 --- /dev/null +++ b/embodichain/gen_sim/task_engine/contracts.py @@ -0,0 +1,410 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Strict, JSON-safe public contracts owned by Task Engine.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +import hashlib +import json +from typing import Any, TypeAlias + +from .interpretation import validate_instruction_intent +from .ontology import TASK_CONTRACTS, task_success_type + +__all__ = [ + "SCENE_REQUEST_SCHEMA", + "SUCCESS_SPEC_SCHEMA", + "TASK_CANDIDATE_SET_SCHEMA", + "TASK_DRAFT_SCHEMA", + "SceneRequest", + "SuccessSpec", + "TaskCandidate", + "TaskCandidateSet", + "TaskDraft", + "canonical_hash", + "validate_scene_request", + "validate_success_spec", + "validate_task_candidate", + "validate_task_candidate_set", + "validate_task_draft", +] + +TASK_DRAFT_SCHEMA = "action_engine_task_draft_v1" +SCENE_REQUEST_SCHEMA = "action_engine_scene_request_v1" +SUCCESS_SPEC_SCHEMA = "action_engine_success_spec_v1" +TASK_CANDIDATE_SET_SCHEMA = "action_engine_task_candidate_set_v1" + +TaskDraft: TypeAlias = dict[str, Any] +SceneRequest: TypeAlias = dict[str, Any] +SuccessSpec: TypeAlias = dict[str, Any] +TaskCandidate: TypeAlias = dict[str, Any] +TaskCandidateSet: TypeAlias = dict[str, Any] + +_SUCCESS_TYPES = frozenset( + {contract.success_type for contract in TASK_CONTRACTS.values()} | {"semantic_goal"} +) +_DRAFT_KEYS = frozenset({"schema_version", "task_id", "instruction", "steps"}) +_SCENE_REQUEST_KEYS = frozenset({"schema_version", "task_id", "references"}) +_REFERENCE_KEYS = frozenset( + { + "reference_id", + "step_id", + "role", + "reference", + "quantifier", + "count", + "source_structure", + "affordances", + "initial_state", + "attributes", + } +) +_SUCCESS_KEYS = frozenset({"schema_version", "task_id", "op", "terms"}) +_SUCCESS_TERM_KEYS = frozenset({"step_id", "type"}) +_CANDIDATE_KEYS = frozenset( + { + "candidate_id", + "draft", + "scene_request", + "success_spec", + "semantic_hash", + "vote_count", + "attempts", + "normalizations", + } +) +_CANDIDATE_SET_KEYS = frozenset( + { + "schema_version", + "task_id", + "instruction", + "candidates", + "requested_candidate_count", + "valid_response_count", + "errors", + } +) + + +def canonical_hash(value: Any) -> str: + """Return the stable SHA-256 of one JSON-safe protocol value.""" + payload = json.dumps( + value, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def validate_task_draft(value: Mapping[str, Any]) -> TaskDraft: + result = _mapping(value, "TaskDraft") + _keys(result, _DRAFT_KEYS, "TaskDraft") + _schema(result, TASK_DRAFT_SCHEMA, "TaskDraft") + result["task_id"] = _nonempty(result.get("task_id"), "TaskDraft.task_id") + result["instruction"] = _nonempty( + result.get("instruction"), "TaskDraft.instruction" + ) + intent = validate_instruction_intent({"steps": result.get("steps")}) + result["steps"] = intent["steps"] + return result + + +def validate_scene_request(value: Mapping[str, Any]) -> SceneRequest: + result = _mapping(value, "SceneRequest") + _keys(result, _SCENE_REQUEST_KEYS, "SceneRequest") + _schema(result, SCENE_REQUEST_SCHEMA, "SceneRequest") + task_id = _nonempty(result.get("task_id"), "SceneRequest.task_id") + references: list[dict[str, Any]] = [] + for index, raw in enumerate( + _sequence(result.get("references"), "SceneRequest.references") + ): + context = f"SceneRequest.references[{index}]" + reference = _mapping(raw, context) + _keys(reference, _REFERENCE_KEYS, context) + for key in ("reference_id", "step_id", "role", "reference", "source_structure"): + reference[key] = _nonempty(reference.get(key), f"{context}.{key}") + reference["role"] = _enum( + reference["role"], {"object", "target"}, f"{context}.role" + ) + reference["quantifier"] = _enum( + reference.get("quantifier"), + {"one", "all", "count"}, + f"{context}.quantifier", + ) + reference["count"] = _integer( + reference.get("count"), f"{context}.count", minimum=0 + ) + if reference["quantifier"] in {"one", "all"} and reference["count"] != 0: + raise ValueError( + f"{context} quantifier={reference['quantifier']} requires count=0." + ) + if reference["quantifier"] == "count" and reference["count"] < 1: + raise ValueError(f"{context} quantifier=count requires count>=1.") + reference["affordances"] = _strings( + reference.get("affordances"), f"{context}.affordances" + ) + reference["initial_state"] = _mapping( + reference.get("initial_state"), f"{context}.initial_state" + ) + reference["attributes"] = _mapping( + reference.get("attributes"), f"{context}.attributes" + ) + references.append(reference) + _unique([item["reference_id"] for item in references], "SceneRequest reference IDs") + result["task_id"] = task_id + result["references"] = references + _json_safe(result, "SceneRequest") + return result + + +def validate_success_spec( + value: Mapping[str, Any], + *, + draft: Mapping[str, Any] | None = None, +) -> SuccessSpec: + result = _mapping(value, "SuccessSpec") + _keys(result, _SUCCESS_KEYS, "SuccessSpec") + _schema(result, SUCCESS_SPEC_SCHEMA, "SuccessSpec") + task_id = _nonempty(result.get("task_id"), "SuccessSpec.task_id") + if result.get("op") != "all": + raise ValueError("SuccessSpec.op must be 'all'.") + terms: list[dict[str, str]] = [] + for index, raw in enumerate(_sequence(result.get("terms"), "SuccessSpec.terms")): + context = f"SuccessSpec.terms[{index}]" + term = _mapping(raw, context) + _keys(term, _SUCCESS_TERM_KEYS, context) + terms.append( + { + "step_id": _nonempty(term.get("step_id"), f"{context}.step_id"), + "type": _enum(term.get("type"), set(_SUCCESS_TYPES), f"{context}.type"), + } + ) + if not terms: + raise ValueError("SuccessSpec.terms must not be empty.") + _unique([term["step_id"] for term in terms], "SuccessSpec step IDs") + if draft is not None: + normalized_draft = validate_task_draft(draft) + if normalized_draft["task_id"] != task_id: + raise ValueError("SuccessSpec.task_id must match TaskDraft.task_id.") + expected = [ + { + "step_id": step["id"], + "type": task_success_type(step["task_type"], step), + } + for step in normalized_draft["steps"] + ] + if terms != expected: + raise ValueError( + "SuccessSpec terms must be ordered, complete, and derived from " + "task_success_type." + ) + result["task_id"] = task_id + result["terms"] = terms + return result + + +def validate_task_candidate(value: Mapping[str, Any]) -> TaskCandidate: + result = _mapping(value, "TaskCandidate") + _keys(result, _CANDIDATE_KEYS, "TaskCandidate") + result["candidate_id"] = _nonempty( + result.get("candidate_id"), "TaskCandidate.candidate_id" + ) + result["draft"] = validate_task_draft(result.get("draft")) + result["scene_request"] = validate_scene_request(result.get("scene_request")) + result["success_spec"] = validate_success_spec( + result.get("success_spec"), draft=result["draft"] + ) + for name in ("scene_request", "success_spec"): + if result[name]["task_id"] != result["draft"]["task_id"]: + raise ValueError(f"TaskCandidate {name}.task_id must match its draft.") + from .agent import derive_scene_request + + if result["scene_request"] != derive_scene_request(result["draft"]): + raise ValueError( + "TaskCandidate.scene_request must be derived exactly from its draft." + ) + result["semantic_hash"] = _digest( + result.get("semantic_hash"), "TaskCandidate.semantic_hash" + ) + if result["semantic_hash"] != canonical_hash(result["draft"]["steps"]): + raise ValueError( + "TaskCandidate.semantic_hash does not match its canonical steps." + ) + result["vote_count"] = _integer( + result.get("vote_count"), "TaskCandidate.vote_count", minimum=1 + ) + result["attempts"] = _integer( + result.get("attempts"), "TaskCandidate.attempts", minimum=1, maximum=2 + ) + result["normalizations"] = _mapping_sequence( + result.get("normalizations"), "TaskCandidate.normalizations" + ) + return result + + +def validate_task_candidate_set(value: Mapping[str, Any]) -> TaskCandidateSet: + result = _mapping(value, "TaskCandidateSet") + _keys(result, _CANDIDATE_SET_KEYS, "TaskCandidateSet") + _schema(result, TASK_CANDIDATE_SET_SCHEMA, "TaskCandidateSet") + task_id = _nonempty(result.get("task_id"), "TaskCandidateSet.task_id") + instruction = _nonempty(result.get("instruction"), "TaskCandidateSet.instruction") + requested = _integer( + result.get("requested_candidate_count"), + "TaskCandidateSet.requested_candidate_count", + minimum=1, + ) + valid = _integer( + result.get("valid_response_count"), + "TaskCandidateSet.valid_response_count", + minimum=1, + maximum=requested, + ) + candidates = [ + validate_task_candidate(item) + for item in _sequence(result.get("candidates"), "TaskCandidateSet.candidates") + ] + if not candidates: + raise ValueError("TaskCandidateSet.candidates must not be empty.") + _unique([item["candidate_id"] for item in candidates], "TaskCandidate IDs") + _unique( + [item["semantic_hash"] for item in candidates], "TaskCandidate semantic hashes" + ) + if sum(item["vote_count"] for item in candidates) != valid: + raise ValueError( + "TaskCandidate vote_count values must sum to valid_response_count." + ) + for candidate in candidates: + if ( + candidate["draft"]["task_id"] != task_id + or candidate["draft"]["instruction"] != instruction + ): + raise ValueError("Every TaskCandidate draft must match its candidate set.") + errors = _strings(result.get("errors"), "TaskCandidateSet.errors", allow_empty=True) + if valid + len(errors) != requested: + raise ValueError( + "Valid responses plus errors must equal requested_candidate_count." + ) + result.update( + { + "task_id": task_id, + "instruction": instruction, + "requested_candidate_count": requested, + "valid_response_count": valid, + "candidates": candidates, + "errors": errors, + } + ) + return result + + +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] | frozenset[str], context: str +) -> None: + if set(value) != set(expected): + raise ValueError( + f"{context} requires exactly fields {sorted(expected)}; received {sorted(value)}." + ) + + +def _schema(value: Mapping[str, Any], expected: str, context: str) -> None: + if value.get("schema_version") != expected: + raise ValueError(f"{context}.schema_version must be {expected!r}.") + + +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, maximum: int | None = None +) -> int: + if ( + not isinstance(value, int) + or isinstance(value, bool) + or value < minimum + or (maximum is not None and value > maximum) + ): + raise ValueError(f"{context} must be an integer in the allowed range.") + return value + + +def _strings(value: Any, context: str, *, allow_empty: bool = False) -> list[str]: + result = [_string(item, context) for item in _sequence(value, context)] + if not allow_empty and any(not item for item in result): + raise ValueError(f"{context} values must not be empty.") + if len(result) != len(set(result)): + raise ValueError(f"{context} values must be unique.") + 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 _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 _unique(values: Sequence[str], context: str) -> None: + if len(values) != len(set(values)): + raise ValueError(f"{context} must be unique.") + + +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/task_engine/defaults.yaml b/embodichain/gen_sim/task_engine/defaults.yaml new file mode 100644 index 000000000..33169e461 --- /dev/null +++ b/embodichain/gen_sim/task_engine/defaults.yaml @@ -0,0 +1,33 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +schema_version: embodichain.task-engine-defaults/v1 + +workflow: + max_parallel_workers: 2 + max_scene_attempts: 2 + max_action_attempts: 3 + +planning: + candidate_count: 3 + planning_mode: offline + max_episodes: 1 + max_episode_steps: 4000 + +execution: + num_envs: 1 + success_policy: any + min_successful_envs: 1 diff --git a/embodichain/gen_sim/task_engine/interpretation.py b/embodichain/gen_sim/task_engine/interpretation.py new file mode 100644 index 000000000..68ee25c7c --- /dev/null +++ b/embodichain/gen_sim/task_engine/interpretation.py @@ -0,0 +1,1240 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Scene-independent structured interpretation for Task Engine.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass +import json +import os +from pathlib import Path +import re +from time import perf_counter +from typing import Any, TypeAlias + +from .ontology import ( + RELATIONS, + TASK_CONTRACTS, + TERMINAL_BEHAVIORS, + TRANSPORT_DIRECTIONS, +) + +__all__ = [ + "INSTRUCTION_INTENT_SCHEMA", + "InstructionDraftResult", + "InstructionIntent", + "InstructionCaller", + "interpret_instruction_draft", + "validate_instruction_intent", +] + +InstructionCaller = Callable[..., Mapping[str, Any]] +InstructionIntent: TypeAlias = dict[str, Any] +TASK_TYPES = frozenset(TASK_CONTRACTS) + +_RELATIONS = RELATIONS +_ARMS = frozenset({"none", "auto", "left_arm", "right_arm"}) +_ORIENTATIONS = frozenset({"none", "preserve", "upright"}) +_TARGET_STATES = frozenset({"none", "open", "closed", "activated"}) +_LAYOUTS = frozenset({"none", "line"}) +_AXES = frozenset({"none", "world_x", "world_y"}) +_DIRECTIONS = TRANSPORT_DIRECTIONS +_TERMINAL_BEHAVIORS = TERMINAL_BEHAVIORS +_SELECTOR_KINDS = frozenset({"none", "scene_ref", "step_result"}) +_QUANTIFIERS = frozenset({"one", "all", "count"}) +_STEP_KEYS = frozenset( + { + "id", + "task_type", + "object", + "target", + "relation", + "required_arm", + "transfer_arm", + "receive_arm", + "orientation_goal", + "target_state", + "target_setting", + "layout", + "axis", + "direction", + "terminal_behavior", + "depends_on", + } +) +_INTENT_TASK_FIELD_REGISTRY = { + task_type: contract.applicable_intent_fields + for task_type, contract in TASK_CONTRACTS.items() +} +_INTENT_FIELD_DEFAULTS: dict[str, Any] = { + "target": None, + "relation": "none", + "required_arm": "none", + "transfer_arm": "none", + "receive_arm": "none", + "orientation_goal": "none", + "target_state": "none", + "target_setting": 0, + "layout": "none", + "axis": "none", + "direction": "none", + "terminal_behavior": "none", +} +_SELECTOR_KEYS = frozenset( + { + "kind", + "step_id", + "reference", + "quantifier", + "count", + } +) +_FORBIDDEN_FIELDS = frozenset( + { + "atomic_action", + "atomic_actions", + "atomicaction", + "coordinates", + "bbox", + "bboxes", + "grasp_pose", + "keypoint", + "keypoints", + "joint_positions", + "joints", + "pose", + "position", + "qpos", + "rotation", + "target_pose", + "translation", + "trajectory", + "waypoints", + } +) +# MiMo's OpenAI-compatible endpoint can spend the whole completion budget in +# hidden reasoning when the request leaves thinking enabled. A sparse final +# JSON object then looks like a schema failure to the deterministic verifier. +# Keep the budget bounded and turn reasoning off for the text interpretation +# call; the parser must return an auditable object rather than a thought trace. +_MIMO_MAX_COMPLETION_TOKENS = 4096 +_GEN_SIM_DIR = Path(__file__).resolve().parents[1] +_GEN_SIM_ENV_PATH = _GEN_SIM_DIR / ".env" +_GEN_CONFIG_PATH = _GEN_SIM_DIR / "simready_pipeline" / "configs" / "gen_config.json" + + +class _MissingRequiredTargetError(ValueError): + """Identify a validation failure that receives targeted repair guidance.""" + + +@dataclass(frozen=True) +class InstructionDraftResult: + """One validated, scene-independent interpretation and its audit metadata.""" + + intent: InstructionIntent + model: str + attempts: int + latency_seconds: float + normalizations: tuple[dict[str, Any], ...] + + +# Object semantics remain open natural-language references until the dedicated +# scene-grounding phase resolves them. All other values are strict protocol +# enums; non-canonical model output is repaired by the model, never guessed by +# a local language alias table. + +_SELECTOR_SCHEMA = { + "type": "object", + "additionalProperties": False, + "required": sorted(_SELECTOR_KEYS), + "properties": { + "kind": {"type": "string", "enum": sorted(_SELECTOR_KINDS)}, + "step_id": {"type": "string"}, + "reference": {"type": "string"}, + "quantifier": {"type": "string", "enum": sorted(_QUANTIFIERS)}, + "count": {"type": "integer", "minimum": 0}, + }, +} + +_INTENT_OUTPUT_SCHEMA = { + "title": "ActionEngineInstructionIntent", + "type": "object", + "additionalProperties": False, + "required": ["steps"], + "properties": { + "steps": { + "type": "array", + "minItems": 1, + "items": { + "type": "object", + "additionalProperties": False, + "required": sorted(_STEP_KEYS), + "properties": { + "id": {"type": "string"}, + "task_type": {"type": "string", "enum": sorted(TASK_TYPES)}, + "object": _SELECTOR_SCHEMA, + "target": _SELECTOR_SCHEMA, + "relation": {"type": "string", "enum": sorted(_RELATIONS)}, + "required_arm": {"type": "string", "enum": sorted(_ARMS)}, + "transfer_arm": {"type": "string", "enum": sorted(_ARMS)}, + "receive_arm": {"type": "string", "enum": sorted(_ARMS)}, + "orientation_goal": { + "type": "string", + "enum": sorted(_ORIENTATIONS), + }, + "target_state": { + "type": "string", + "enum": sorted(_TARGET_STATES), + }, + "target_setting": {"type": "integer"}, + "layout": {"type": "string", "enum": sorted(_LAYOUTS)}, + "axis": {"type": "string", "enum": sorted(_AXES)}, + "direction": { + "type": "string", + "enum": sorted(_DIRECTIONS), + }, + "terminal_behavior": { + "type": "string", + "enum": sorted(_TERMINAL_BEHAVIORS), + }, + "depends_on": { + "type": "array", + "items": {"type": "string"}, + }, + }, + }, + } + }, +} + +# Keep a read-only-by-convention public copy for callers that need to configure +# a structured client. The schema is an input contract, not a persisted task +# graph; ``validate_instruction_intent`` remains the authoritative verifier. +INSTRUCTION_INTENT_SCHEMA = deepcopy(_INTENT_OUTPUT_SCHEMA) + + +def interpret_instruction_draft( + instruction: str, + *, + model: str | None = None, + caller: InstructionCaller | None = None, +) -> InstructionDraftResult: + """Interpret one instruction without reading or grounding a scene.""" + instruction_text = str(instruction).strip() + if not instruction_text: + raise ValueError("instruction must be non-empty.") + prompt = _instruction_prompt(instruction_text) + invoke = caller or _default_instruction_caller + # An injected caller owns its transport and does not need provider config. + selected_model = model if caller is not None else _instruction_model(model) + if caller is None and selected_model is None: + raise ValueError( + "A text LLM model is required through --llm-model, " + "ACTION_ENGINE_LLM_MODEL, or OPENAI_MODEL." + ) + started = perf_counter() + first_error: Exception | None = None + for attempt in range(2): + current_prompt = prompt + if first_error is not None: + current_prompt += ( + "\n\nREPAIR OVERRIDE: the previous JSON was invalid. Return a corrected " + "JSON object only; do not repeat the sparse response. Every step " + "must contain all 16 step keys and every selector all 5 selector " + "keys. Keep semantic fields explicit: E4 requires transfer_arm " + "and receive_arm, and E1/E3 require target plus relation (unless " + "E1 layout=line). Use canonical defaults only for fields that do " + "not apply. Validation error: " + f"{first_error}\n" + "Copy this complete shape before filling values (shape only; do " + "not copy its values or step count):\n" + f"{json.dumps(_instruction_shape_example(), ensure_ascii=False, sort_keys=True)}\n" + "Selector kind rules:\n" + f"{_instruction_selector_rules()}" + f"{_instruction_repair_guidance(first_error)}" + ) + try: + response = invoke( + prompt=current_prompt, + schema=deepcopy(INSTRUCTION_INTENT_SCHEMA), + model=selected_model, + ) + normalized, normalizations = _normalize_instruction_intent_fields( + _coerce_instruction_response(response) + ) + intent = validate_instruction_intent(normalized) + return InstructionDraftResult( + intent=intent, + model=selected_model or "injected_caller", + attempts=attempt + 1, + latency_seconds=perf_counter() - started, + normalizations=tuple(normalizations), + ) + except (TypeError, ValueError) as error: + if attempt: + raise ValueError( + "Instruction intent failed validation after one repair: " f"{error}" + ) from error + first_error = error + raise AssertionError("unreachable") + + +def _normalize_instruction_intent_fields( + value: Mapping[str, Any], +) -> tuple[dict[str, Any], list[dict[str, Any]]]: + """Canonicalize defaults and uniquely constrained cross-step continuity. + + The strict public validator deliberately remains unchanged. This pass is + confined to the LLM boundary, where weak JSON-mode providers sometimes + copy a meaningful value into an inapplicable slot such as E4.required_arm. + Required scene facts and ambiguous arm assignments are never inferred here + and still fail closed. + """ + result = deepcopy(dict(value)) + raw_steps = result.get("steps") + if not isinstance(raw_steps, list): + return result, [] + changes: list[dict[str, Any]] = [] + for index, raw_step in enumerate(raw_steps): + if not isinstance(raw_step, dict) or set(raw_step) != _STEP_KEYS: + continue + task_type = raw_step.get("task_type") + applicable = _INTENT_TASK_FIELD_REGISTRY.get(task_type) + if applicable is None: + continue + for field, configured_default in _INTENT_FIELD_DEFAULTS.items(): + field_applies = field in applicable + if task_type == "E1" and field in {"target", "relation"}: + field_applies = raw_step.get("layout") != "line" + if task_type == "E1" and field == "axis": + field_applies = raw_step.get("layout") == "line" + if field_applies: + continue + default = ( + _empty_selector() + if field == "target" and configured_default is None + else deepcopy(configured_default) + ) + if raw_step[field] == default: + continue + previous = deepcopy(raw_step[field]) + raw_step[field] = default + changes.append( + { + "path": f"steps[{index}].{field}", + "from": previous, + "to": deepcopy(default), + "reason": f"inapplicable_for_{task_type}", + } + ) + target = raw_step.get("target") + if ( + task_type == "E5" + and isinstance(target, Mapping) + and target.get("kind") == "none" + and raw_step.get("relation") == "none" + and raw_step.get("direction") == "none" + and raw_step.get("terminal_behavior") == "hold" + ): + raw_step["direction"] = "up" + changes.append( + { + "path": f"steps[{index}].direction", + "from": "none", + "to": "up", + "reason": "e5_hold_defaults_to_lift", + } + ) + _normalize_handover_arm_continuity(raw_steps, changes) + return result, changes + + +def _normalize_handover_arm_continuity( + steps: Sequence[Any], + changes: list[dict[str, Any]], +) -> None: + """Repair a same-arm E4 only when adjacent ownership fixes both roles.""" + by_id: dict[str, Mapping[str, Any]] = {} + for step in steps: + if not isinstance(step, dict) or set(step) != _STEP_KEYS: + return + step_id = step.get("id") + if not isinstance(step_id, str) or not step_id or step_id in by_id: + return + by_id[step_id] = step + + explicit_arms = {"left_arm", "right_arm"} + for index, step in enumerate(steps): + assert isinstance(step, dict) + transfer = step.get("transfer_arm") + receive = step.get("receive_arm") + if ( + step.get("task_type") != "E4" + or transfer not in explicit_arms + or transfer != receive + ): + continue + + object_key = _object_lineage_key(step, by_id) + upstream_arm: str | None = None + for producer in reversed(steps[:index]): + assert isinstance(producer, Mapping) + if object_key is None or _object_lineage_key(producer, by_id) != object_key: + continue + candidate = ( + producer.get("receive_arm") + if producer.get("task_type") == "E4" + else producer.get("required_arm") + ) + if candidate in explicit_arms: + upstream_arm = str(candidate) + break + + downstream_arm: str | None = None + for consumer in steps[index + 1 :]: + assert isinstance(consumer, Mapping) + if object_key is None or _object_lineage_key(consumer, by_id) != object_key: + continue + candidate = ( + consumer.get("transfer_arm") + if consumer.get("task_type") == "E4" + else consumer.get("required_arm") + ) + if candidate in explicit_arms: + downstream_arm = str(candidate) + break + + desired_transfer = upstream_arm or str(transfer) + desired_receive = downstream_arm or str(receive) + if desired_transfer == desired_receive: + continue + for field, desired in ( + ("transfer_arm", desired_transfer), + ("receive_arm", desired_receive), + ): + if step[field] == desired: + continue + previous = step[field] + step[field] = desired + changes.append( + { + "path": f"steps[{index}].{field}", + "from": previous, + "to": desired, + "reason": "handover_arm_continuity", + } + ) + + +def _object_lineage_key( + step: Mapping[str, Any], + by_id: Mapping[str, Mapping[str, Any]], + seen: frozenset[str] = frozenset(), +) -> tuple[str, str] | None: + """Resolve object identity only through explicit step-result lineage.""" + selector = step.get("object") + if not isinstance(selector, Mapping): + return None + kind = selector.get("kind") + if kind == "scene_ref": + step_id = step.get("id") + if not isinstance(step_id, str) or not step_id: + return None + return ("step_result", step_id) + if kind != "step_result": + return None + producer_id = selector.get("step_id") + if not isinstance(producer_id, str) or producer_id in seen: + return None + producer = by_id.get(producer_id) + if producer is None: + return None + return _object_lineage_key(producer, by_id, seen | {producer_id}) + + +def _empty_selector() -> dict[str, Any]: + """Return the canonical selector value for an inapplicable target.""" + return { + "kind": "none", + "step_id": "", + "reference": "", + "quantifier": "one", + "count": 0, + } + + +def validate_instruction_intent(value: Mapping[str, Any]) -> dict[str, Any]: + """Validate the private, non-graph instruction interpretation contract.""" + if not isinstance(value, Mapping): + raise TypeError("Instruction intent must be a mapping.") + _reject_forbidden_fields(value) + if set(value) != {"steps"}: + raise ValueError("Instruction intent may contain only 'steps'.") + raw_steps = value.get("steps") + if not isinstance(raw_steps, Sequence) or isinstance(raw_steps, (str, bytes)): + raise ValueError("Instruction intent steps must be a list.") + if not raw_steps: + raise ValueError("Instruction intent steps must not be empty.") + steps = [] + ids: set[str] = set() + dependencies: dict[str, list[str]] = {} + for index, raw in enumerate(raw_steps): + context = f"InstructionIntent.steps[{index}]" + if not isinstance(raw, Mapping): + raise ValueError(f"{context} must be a mapping.") + if set(raw) != _STEP_KEYS: + raise ValueError( + f"{context} requires exactly fields {sorted(_STEP_KEYS)}; " + f"received {sorted(raw)}." + ) + step = deepcopy(dict(raw)) + step_id = _nonempty(step["id"], f"{context}.id") + if step_id in ids: + raise ValueError(f"Duplicate instruction step ID {step_id!r}.") + ids.add(step_id) + step["id"] = step_id + step["task_type"] = _choice( + step["task_type"], TASK_TYPES, f"{context}.task_type" + ) + step["object"] = _validate_selector(step["object"], f"{context}.object") + step["target"] = _validate_selector(step["target"], f"{context}.target") + step["relation"] = _canonical_relation(step["relation"], f"{context}.relation") + for key in ("required_arm", "transfer_arm", "receive_arm"): + step[key] = _canonical_arm(step[key], f"{context}.{key}") + step["orientation_goal"] = _canonical_orientation( + step["orientation_goal"], f"{context}.orientation_goal" + ) + step["target_state"] = _choice( + step["target_state"], _TARGET_STATES, f"{context}.target_state" + ) + if isinstance(step["target_setting"], bool) or not isinstance( + step["target_setting"], int + ): + raise ValueError(f"{context}.target_setting must be an integer.") + step["layout"] = _choice(step["layout"], _LAYOUTS, f"{context}.layout") + step["axis"] = _choice(step["axis"], _AXES, f"{context}.axis") + step["direction"] = _choice( + step["direction"], _DIRECTIONS, f"{context}.direction" + ) + step["terminal_behavior"] = _choice( + step["terminal_behavior"], + _TERMINAL_BEHAVIORS, + f"{context}.terminal_behavior", + ) + raw_depends = step["depends_on"] + if not isinstance(raw_depends, Sequence) or isinstance( + raw_depends, (str, bytes) + ): + raise ValueError(f"{context}.depends_on must be a list.") + step["depends_on"] = [ + _nonempty(item, f"{context}.depends_on") for item in raw_depends + ] + if step_id in step["depends_on"]: + raise ValueError(f"{context}.depends_on cannot contain its own ID.") + dependencies[step_id] = step["depends_on"] + _validate_task_fields(step, context) + steps.append(step) + positions = {str(step["id"]): index for index, step in enumerate(steps)} + for index, step in enumerate(steps): + for selector_name in ("object", "target"): + selector = step[selector_name] + if selector["kind"] != "step_result": + continue + reference = str(selector["step_id"]) + if reference not in positions: + raise ValueError( + f"Instruction step {step['id']!r} {selector_name} references " + f"unknown step {reference!r}." + ) + if positions[reference] >= index: + raise ValueError( + f"Instruction step {step['id']!r} {selector_name} must reference " + f"a preceding step, not {reference!r}." + ) + for step_id, depends_on in dependencies.items(): + unknown = set(depends_on) - ids + if unknown: + raise ValueError( + f"Instruction step {step_id!r} has unknown dependencies " + f"{sorted(unknown)}." + ) + _validate_dag(dependencies) + return {"steps": steps} + + +def _validate_selector(value: Any, context: str) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise ValueError(f"{context} must be a mapping.") + if set(value) != _SELECTOR_KEYS: + raise ValueError( + f"{context} requires exactly fields {sorted(_SELECTOR_KEYS)}; " + f"received {sorted(value)}." + ) + selector = deepcopy(dict(value)) + selector["kind"] = _choice(selector["kind"], _SELECTOR_KINDS, f"{context}.kind") + selector["step_id"] = _selector_string(selector["step_id"], f"{context}.step_id") + selector["reference"] = _selector_string( + selector["reference"], f"{context}.reference" + ) + selector["quantifier"] = _canonical_quantifier( + selector["quantifier"], f"{context}.quantifier" + ) + if isinstance(selector["count"], bool) or not isinstance(selector["count"], int): + raise ValueError(f"{context}.count must be an integer.") + if selector["count"] < 0: + raise ValueError(f"{context}.count must be non-negative.") + kind = selector["kind"] + if kind == "scene_ref" and not selector["reference"]: + raise ValueError(f"{context} scene_ref requires a reference.") + if kind == "step_result": + if not selector["step_id"]: + raise ValueError(f"{context} step_result requires step_id.") + if selector["reference"]: + raise ValueError( + f"{context} step_result may identify only a prior step_id." + ) + if selector["quantifier"] != "one" or selector["count"] != 0: + raise ValueError( + f"{context} step_result requires quantifier=one and count=0." + ) + if kind == "scene_ref" and selector["step_id"]: + raise ValueError(f"{context} scene_ref cannot carry step_id.") + if kind == "none" and (selector["step_id"] or selector["reference"]): + raise ValueError(f"{context} kind=none cannot carry constraints.") + if kind == "none" and (selector["quantifier"] != "one" or selector["count"] != 0): + raise ValueError(f"{context} kind=none requires quantifier=one and count=0.") + if selector["quantifier"] == "one" and selector["count"] != 0: + raise ValueError(f"{context} quantifier=one requires count=0.") + if selector["quantifier"] == "all" and selector["count"] != 0: + raise ValueError(f"{context} quantifier=all requires count=0.") + if selector["quantifier"] == "count" and selector["count"] < 1: + raise ValueError(f"{context} quantifier=count requires count>=1.") + return selector + + +def _validate_task_fields(step: Mapping[str, Any], context: str) -> None: + task_type = str(step["task_type"]) + target_kind = str(step["target"]["kind"]) + if task_type not in {"E1", "E3", "E5"} and step["relation"] != "none": + raise ValueError(f"{context} {task_type} does not accept relation.") + if task_type == "E3" and step["relation"] != "above": + raise ValueError(f"{context} E3 relation must be above.") + target_setting = int(step["target_setting"]) + if task_type != "E8" and target_setting != 0: + raise ValueError(f"{context} target_setting is only valid for E8.") + if task_type != "E1" and step["axis"] != "none": + raise ValueError(f"{context} axis is only valid for E1 line arrangement.") + if task_type == "E1" and step["layout"] != "line" and step["axis"] != "none": + raise ValueError(f"{context} axis is only valid for E1 line arrangement.") + if task_type not in {"E6", "E7", "E9"} and step["target_state"] != "none": + raise ValueError(f"{context} target_state is not valid for {task_type}.") + if task_type != "E4" and step["transfer_arm"] != "none": + raise ValueError(f"{context} transfer_arm is only valid for E4.") + if task_type != "E4" and step["receive_arm"] != "none": + raise ValueError(f"{context} receive_arm is only valid for E4.") + orientation_goal = str(step["orientation_goal"]) + if task_type == "E2" and orientation_goal != "upright": + raise ValueError(f"{context} E2 orientation_goal must be upright.") + if task_type not in {"E1", "E2", "E4"} and orientation_goal != "none": + raise ValueError( + f"{context} orientation_goal is only valid for E1, E2, and E4." + ) + if task_type == "E1" and step["layout"] == "line": + if target_kind != "none": + raise ValueError(f"{context} E1 line arrangement cannot carry a target.") + if step["relation"] != "none": + raise ValueError(f"{context} E1 line arrangement cannot carry a relation.") + elif task_type in {"E1", "E3"}: + if target_kind == "none": + raise _MissingRequiredTargetError( + f"{context} {task_type} requires a target selector." + ) + if step["relation"] == "none" and task_type == "E3": + raise ValueError(f"{context} {task_type} requires a symbolic relation.") + elif task_type == "E5": + direction = str(step["direction"]) + terminal = str(step["terminal_behavior"]) + if terminal not in _TERMINAL_BEHAVIORS - {"none"}: + raise ValueError(f"{context} E5 requires terminal_behavior hold/place.") + if target_kind == "none": + if step["relation"] != "none": + raise ValueError(f"{context} E5 relation requires a target selector.") + if direction == "none" and terminal != "place": + raise ValueError( + f"{context} E5 requires a direction or target relation." + ) + else: + if step["relation"] == "none": + raise ValueError(f"{context} E5 target requires a relation.") + if direction != "none": + raise ValueError( + f"{context} E5 target relation cannot also carry direction." + ) + elif target_kind != "none": + raise ValueError(f"{context} {task_type} does not accept a target selector.") + if task_type != "E5": + if step["direction"] != "none": + raise ValueError(f"{context} direction is only valid for E5.") + if step["terminal_behavior"] != "none": + raise ValueError(f"{context} terminal_behavior is only valid for E5.") + if task_type == "E4": + transfer = str(step["transfer_arm"]) + receive = str(step["receive_arm"]) + if transfer not in {"left_arm", "right_arm"} or receive not in { + "left_arm", + "right_arm", + }: + raise ValueError(f"{context} E4 requires two explicit arms.") + if transfer == receive: + raise ValueError(f"{context} E4 transfer and receive arms must differ.") + if step["required_arm"] not in {"none", "auto"}: + raise ValueError( + f"{context} E4 uses transfer_arm/receive_arm, not required_arm." + ) + if task_type == "E5" and step["required_arm"] not in {"none", "auto"}: + raise ValueError(f"{context} E5 always uses both arms, not required_arm.") + if task_type == "E6" and step["target_state"] != "open": + raise ValueError(f"{context} E6 target_state must be open.") + if task_type == "E7" and step["target_state"] != "closed": + raise ValueError(f"{context} E7 target_state must be closed.") + if task_type == "E9" and step["target_state"] != "activated": + raise ValueError(f"{context} E9 target_state must be activated.") + if step["layout"] == "line" and task_type != "E1": + raise ValueError(f"{context} only E1 supports layout=line.") + + +def _instruction_prompt(instruction: str) -> str: + return ( + "Convert the user's explicit L1-L3 instruction into typed E1-E9 task " + "intent. Understand synonyms, ellipsis, and pronouns, but " + "do not invent missing objects. Use step_result for cross-step pronouns " + "and explicit references to the result of an earlier manipulation. Keep " + "an independently selected repeated noun as scene_ref; identical text " + "alone does not prove object identity. " + "Object directions are robot-relative; arm names are robot body sides. " + "Preserve each concrete object or target phrase from the instruction as " + "an open scene_ref.reference. Do not classify it or emit a scene UID. " + "Emit no AtomicAction, category label, affordance, coordinates, poses, " + "paths, or reasoning. Encode explicit ordering with depends_on; same-action set " + "members may remain independent. Use empty strings and 'none' for " + "inapplicable required fields. A request to retract the transfer arm " + "immediately after an E4 handover is a mandatory runtime retreat/home " + "barrier for that E4; do " + "not emit a separate task step for it. The exact output keys are steps -> id, " + "task_type, object, target, relation, required_arm, transfer_arm, " + "receive_arm, orientation_goal, target_state, target_setting, layout, " + "axis, direction, terminal_behavior, depends_on; each selector has kind, " + "step_id, reference, quantifier, count.\n\n" + "Use orientation_goal=none unless the instruction explicitly requests " + "upright orientation or preserving the original orientation. Spatial " + "placement and handover alone do not imply preserve. " + f"Instruction:\n{instruction}\n\n" + f"E1-E9 catalog:\n{json.dumps(_intent_capability_catalog(), ensure_ascii=False, sort_keys=True)}\n\n" + "Shape-only complete JSON example (do not copy its step count or values; " + "copy every key, including keys whose value is none/empty/0):\n" + f"{json.dumps(_instruction_shape_example(), ensure_ascii=False, sort_keys=True)}\n\n" + "Selector kind rules (these are not extra output fields):\n" + f"{_instruction_selector_rules()}\n\n" + "For E5, use target+relation for moving an object relative to another " + "object, or direction for a small robot-relative move. A dual-arm pick, " + "lift, raise, or hold request without another target uses direction=up " + "and terminal_behavior=hold. Use hold unless the instruction explicitly " + "says to put/release the object. For pick " + "and release at the original location, use direction=none and place. A dual-arm " + "pick/move/transport request is E5, not E1. Final checklist: every step " + "has all 16 step keys; every object and target " + "has all 5 selector keys. For an inapplicable field use the canonical " + "default shown in the example, never omit the field. E4 must explicitly " + "state transfer_arm and receive_arm. E1/E3 must explicitly state target " + "and relation (except E1 layout=line)." + ) + + +def _instruction_shape_example() -> dict[str, Any]: + """Return a compact field-complete example for providers with weak schemas.""" + selector = { + "kind": "scene_ref", + "step_id": "", + "reference": "example object A", + "quantifier": "one", + "count": 0, + } + empty_selector = { + "kind": "none", + "step_id": "", + "reference": "", + "quantifier": "one", + "count": 0, + } + return { + "steps": [ + { + "id": "step_1", + "task_type": "E2", + "object": selector, + "target": empty_selector, + "relation": "none", + "required_arm": "auto", + "transfer_arm": "none", + "receive_arm": "none", + "orientation_goal": "upright", + "target_state": "none", + "target_setting": 0, + "layout": "none", + "axis": "none", + "direction": "none", + "terminal_behavior": "none", + "depends_on": [], + } + ] + } + + +def _instruction_selector_rules() -> str: + """Return the mutually exclusive selector encodings for model prompts.""" + step_result = { + "kind": "step_result", + "step_id": "step_1", + "reference": "", + "quantifier": "one", + "count": 0, + } + return ( + "- kind=none: step_id and reference are empty strings; " + "quantifier='one'; count=0.\n" + "- kind=scene_ref: step_id is empty and reference preserves the concrete " + "object phrase from the user's instruction. Repeated scene_ref text does " + "not establish cross-step identity.\n" + "- kind=step_result: use it only for a pronoun that means exactly one " + "object, or an explicit continuation of the result of an earlier " + "instruction step. Set step_id to that prior " + "step ID and set reference='', quantifier='one', count=0. Do not copy " + "the prior object's phrase into this selector. Replace step_1 in this " + f"complete shape with the actual prior step ID: {json.dumps(step_result, sort_keys=True)}\n" + "A step_result may identify only a prior step_id; it cannot carry any " + "other object constraint." + ) + + +def _instruction_repair_guidance(error: Exception) -> str: + """Add narrow semantic guidance for errors weak JSON-mode models repeat.""" + if "E4 transfer and receive arms must differ" in str(error): + return ( + "\nSame-arm handover repair rule: transfer_arm and receive_arm must " + "name different arms. Preserve the explicitly stated transfer arm. " + "When a later clause clearly continues with the handed object using " + "the other arm, use that arm as receive_arm. Resolve coreference from " + "the instruction semantics; identical scene_ref text alone does not " + "prove that two independently selected objects are the same.\n" + ) + if not isinstance(error, _MissingRequiredTargetError): + return "" + return ( + "\nMissing-target repair rule: for a non-line E1 placement, object is " + "the item being moved and target is the explicit reference object " + "after the spatial relation in the original instruction. For example, " + "in 'place it to the left of the striped pedestal', object is the earlier " + "step_result for 'it', while target selects the striped pedestal; target " + "must not use kind=none. Use target kind=step_result only when the " + "reference object itself is exactly the result of a prior step.\n" + ) + + +def _intent_capability_catalog() -> dict[str, dict[str, Any]]: + """Return the LLM's thin, import-safe E1-E9 capability view. + + Action Engine's online planning catalog also reports runtime availability + and therefore imports simulator action classes. Text interpretation only + needs symbolic E semantics and must remain testable before a simulator + backend is installed. + """ + return { + task_type: { + "semantics": contract.semantics, + "applicable_fields": sorted(_INTENT_TASK_FIELD_REGISTRY[task_type]), + } + for task_type, contract in TASK_CONTRACTS.items() + } + + +def _default_instruction_caller( + *, + prompt: str, + schema: Mapping[str, Any], + model: str | None, +) -> Mapping[str, Any]: + from langchain_core.messages import HumanMessage, SystemMessage + from langchain_openai import ChatOpenAI + + settings = _load_llm_settings(model=model) + kwargs: dict[str, Any] = { + "api_key": settings["api_key"], + "model": settings["model"], + "temperature": 0, + "http_socket_options": (), + } + for key in ("base_url", "default_query"): + if settings[key]: + kwargs[key] = settings[key] + if _is_mimo_compatible(settings): + # MiMo documents ``thinking`` as a provider extension carried in the + # OpenAI client's extra body. Disabling it is important here: hidden + # reasoning can consume the completion and leave only id/object/type. + kwargs.update( + { + "max_completion_tokens": _MIMO_MAX_COMPLETION_TOKENS, + "extra_body": {"thinking": {"type": "disabled"}}, + } + ) + client = ChatOpenAI(**kwargs) + # The full schema remains in the prompt and the local validator is still + # authoritative even when the provider only offers JSON mode. + structured = _structured_output_runnable( + client, + schema, + settings=settings, + ) + schema_prompt = ( + f"{prompt}\n\nReturn one JSON object conforming exactly to this JSON " + f"Schema:\n{json.dumps(schema, ensure_ascii=False, sort_keys=True)}" + ) + response = structured.invoke( + [ + SystemMessage( + content=( + "Return only the requested structured JSON response. Never " + "return reasoning, coordinates, or AtomicAction nodes." + ) + ), + HumanMessage(content=schema_prompt), + ] + ) + return _coerce_instruction_response(response) + + +def _instruction_model(explicit: str | None) -> str | None: + if isinstance(explicit, str) and explicit.strip(): + return explicit.strip() + # Keep model selection separate from credential loading. Reading the local + # dotenv file is side-effect free and gives generation the documented + # priority without leaking credentials into TaskSpec metadata. + for name in ("TASK_ENGINE_LLM_MODEL", "ACTION_ENGINE_LLM_MODEL", "OPENAI_MODEL"): + for source in ( + os.environ, + _load_local_env(), + ): + value = source.get(name) + if isinstance(value, str) and value.strip(): + return value.strip() + return None + + +def _load_local_env() -> dict[str, str]: + """Read Task Engine model configuration without mutating the environment.""" + return _load_env_file(_GEN_SIM_ENV_PATH) + + +def _is_mimo_compatible(settings: Mapping[str, Any]) -> bool: + model = str(settings.get("model", "")).casefold() + base_url = str(settings.get("base_url", "")).casefold() + return "mimo" in model or "xiaomimimo.com" in base_url + + +def _structured_output_runnable( + client: Any, + schema: Mapping[str, Any], + *, + settings: Mapping[str, Any], +) -> Any: + if not hasattr(client, "with_structured_output"): + return client + method = "json_mode" if _is_mimo_compatible(settings) else "json_schema" + try: + return client.with_structured_output(schema, method=method) + except (TypeError, ValueError): + if method == "json_mode" and hasattr(client, "bind"): + from langchain_core.output_parsers import JsonOutputParser + + return ( + client.bind(response_format={"type": "json_object"}) + | JsonOutputParser() + ) + return client.with_structured_output(schema) + + +def _load_llm_settings(*, model: str | None) -> dict[str, Any]: + local_env = _load_local_env() + config: dict[str, Any] = {} + if _GEN_CONFIG_PATH.is_file(): + raw = json.loads(_GEN_CONFIG_PATH.read_text(encoding="utf-8")) + if isinstance(raw, Mapping): + llm = raw.get("llm", {}) + if isinstance(llm, Mapping): + configured = llm.get("openai_compatible", {}) + if isinstance(configured, Mapping): + config = dict(configured) + api_key, base_url = _resolve_transport_settings(local_env, config) + selected_model = ( + (model.strip() if isinstance(model, str) else "") + or _first_env_value( + local_env, + "TASK_ENGINE_LLM_MODEL", + "ACTION_ENGINE_LLM_MODEL", + "OPENAI_MODEL", + "LLM_MODEL", + ) + or str(config.get("model", "")).strip() + ) + default_query = config.get("default_query", {}) or {} + if not api_key: + raise ValueError( + "OPENAI_API_KEY is required for Task Engine interpretation. Set it " + f"in the process environment or {_GEN_SIM_ENV_PATH}." + ) + if not selected_model: + raise ValueError( + "A text LLM model is required through model=, TASK_ENGINE_LLM_MODEL, " + f"OPENAI_MODEL, or {_GEN_CONFIG_PATH}." + ) + if not isinstance(default_query, Mapping): + raise ValueError("LLM default_query must be a mapping.") + return { + "api_key": api_key, + "model": selected_model, + "base_url": base_url, + "default_query": dict(default_query), + } + + +def _resolve_transport_settings( + local_env: Mapping[str, str], + config: Mapping[str, Any], +) -> tuple[str, str]: + """Resolve an API key and endpoint from one configuration source.""" + transports = ( + ( + _mapping_value(os.environ, "OPENAI_API_KEY"), + _mapping_value( + os.environ, + "OPENAI_BASE_URL", + "OPENAI_API_BASE", + "LLM_URL", + ), + ), + ( + _mapping_value(local_env, "OPENAI_API_KEY"), + _mapping_value( + local_env, + "OPENAI_BASE_URL", + "OPENAI_API_BASE", + "LLM_URL", + ), + ), + ( + _mapping_value(config, "api_key"), + _mapping_value(config, "base_url"), + ), + ) + for api_key, base_url in transports: + if api_key and base_url: + return api_key, base_url.rstrip("/") + for api_key, base_url in transports: + if api_key: + return api_key, base_url.rstrip("/") + return "", "" + + +def _mapping_value(source: Mapping[str, Any], *names: str) -> str: + for name in names: + value = source.get(name) + if isinstance(value, str) and value.strip(): + return value.strip() + return "" + + +def _load_env_file(path: Path) -> dict[str, str]: + if not path.is_file(): + return {} + values: dict[str, str] = {} + for line_number, raw_line in enumerate( + path.read_text(encoding="utf-8").splitlines(), start=1 + ): + line = raw_line.strip() + if not line or line.startswith("#"): + continue + if line.startswith("export "): + line = line[len("export ") :].lstrip() + if "=" not in line: + continue + key, raw_value = line.split("=", 1) + key = key.strip() + if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", key): + raise ValueError(f"Invalid dotenv key at {path}:{line_number}.") + value = raw_value.strip() + if len(value) >= 2 and value[0] == value[-1] and value[0] in {"'", '"'}: + value = value[1:-1] + elif " #" in value: + value = value.split(" #", 1)[0].rstrip() + values[key] = value + return values + + +def _first_env_value(local_env: Mapping[str, str], *names: str) -> str | None: + for source in (os.environ, local_env): + for name in names: + value = source.get(name) + if isinstance(value, str) and value.strip(): + return value.strip() + return None + + +def _choice(value: Any, allowed: set[str] | frozenset[str], context: str) -> str: + if not isinstance(value, str) or value not in allowed: + raise ValueError(f"{context} must be one of {sorted(allowed)}.") + return value + + +def _selector_string(value: Any, context: str) -> str: + if not isinstance(value, str): + raise ValueError(f"{context} must be a string.") + return value.strip() + + +def _canonical_quantifier(value: Any, context: str) -> str: + return _choice(value, _QUANTIFIERS, context) + + +def _canonical_arm(value: Any, context: str) -> str: + return _choice(value, _ARMS, context) + + +def _canonical_relation(value: Any, context: str) -> str: + return _choice(value, _RELATIONS, context) + + +def _canonical_orientation(value: Any, context: str) -> str: + return _choice(value, _ORIENTATIONS, context) + + +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() + + +def _topological_steps(steps: Sequence[Mapping[str, Any]]) -> list[dict[str, Any]]: + """Return a stable topological ordering for validated intent steps.""" + by_id = {str(step["id"]): dict(step) for step in steps} + effective_dependencies: dict[str, tuple[str, ...]] = {} + for step_id, step in by_id.items(): + deps = list(str(dep) for dep in step["depends_on"]) + for selector in (step["object"], step["target"]): + if selector["kind"] == "step_result": + reference = str(selector["step_id"]) + if reference not in deps: + deps.append(reference) + effective_dependencies[step_id] = tuple(deps) + pending = set(by_id) + ordered: list[dict[str, Any]] = [] + original = [str(step["id"]) for step in steps] + while pending: + ready = [ + step_id + for step_id in original + if step_id in pending + and all(str(dep) not in pending for dep in effective_dependencies[step_id]) + ] + if not ready: + raise ValueError("Instruction intent dependencies contain a cycle.") + for step_id in ready: + ordered.append(by_id[step_id]) + pending.remove(step_id) + return ordered + + +def _coerce_instruction_response(response: Any) -> Mapping[str, Any]: + """Coerce common structured-client response wrappers without accepting prose.""" + if isinstance(response, Mapping): + return dict(response) + model_dump = getattr(response, "model_dump", None) + if callable(model_dump): + dumped = model_dump() + if isinstance(dumped, Mapping): + return dict(dumped) + content = getattr(response, "content", response) + if isinstance(content, Mapping): + return dict(content) + if isinstance(content, list): + content = "\n".join( + str(item.get("text", "")) + for item in content + if isinstance(item, Mapping) and item.get("type") == "text" + ) + if not isinstance(content, str): + raise ValueError( + f"Instruction model output has unsupported type {type(content).__name__}." + ) + text = content.strip() + if text.startswith("```"): + lines = text.splitlines() + if lines: + lines = lines[1:] + if lines and lines[-1].strip().startswith("```"): + lines = lines[:-1] + text = "\n".join(lines).strip() + try: + parsed = json.loads(text) + except json.JSONDecodeError as exc: + raise ValueError(f"Instruction model output is not valid JSON: {exc}") from exc + if not isinstance(parsed, Mapping): + raise ValueError("Instruction model output must decode to a JSON object.") + return dict(parsed) + + +def _validate_dag(dependencies: Mapping[str, Sequence[str]]) -> None: + visiting: set[str] = set() + visited: set[str] = set() + + def visit(node: str) -> None: + if node in visiting: + raise ValueError("Instruction intent dependencies contain a cycle.") + if node in visited: + return + visiting.add(node) + for dependency in dependencies[node]: + visit(str(dependency)) + visiting.remove(node) + visited.add(node) + + for node in dependencies: + visit(node) + + +def _reject_forbidden_fields(value: Any) -> None: + if isinstance(value, Mapping): + forbidden = _FORBIDDEN_FIELDS & {str(key).strip().lower() for key in value} + if forbidden: + raise ValueError( + f"Instruction intent contains forbidden fields {sorted(forbidden)}." + ) + for item in value.values(): + _reject_forbidden_fields(item) + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + for item in value: + _reject_forbidden_fields(item) diff --git a/embodichain/gen_sim/task_engine/ontology.py b/embodichain/gen_sim/task_engine/ontology.py new file mode 100644 index 000000000..add919683 --- /dev/null +++ b/embodichain/gen_sim/task_engine/ontology.py @@ -0,0 +1,220 @@ +# ---------------------------------------------------------------------------- +# 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 semantic ontology for the canonical E1-E9 tasks.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any + +__all__ = [ + "RELATIONS", + "TASK_CONTRACTS", + "TERMINAL_BEHAVIORS", + "TRANSPORT_DIRECTIONS", + "TaskContract", + "task_contract", + "task_success_type", +] + + +# These are protocol values consumed by executable planners. They are not a +# vocabulary for matching words in user instructions. +RELATIONS = frozenset( + { + "none", + "on", + "inside", + "above", + "left_of", + "right_of", + "front_of", + "behind", + "front_left_of", + "front_right_of", + "back_left_of", + "back_right_of", + } +) +TRANSPORT_DIRECTIONS = frozenset( + { + "none", + "world_x", + "world_y", + "front", + "back", + "left", + "right", + "front_left", + "front_right", + "back_left", + "back_right", + "up", + "down", + } +) +TERMINAL_BEHAVIORS = frozenset({"none", "hold", "place"}) + + +@dataclass(frozen=True, slots=True) +class TaskContract: + """One scene-independent semantic E-task contract.""" + + task_type: str + semantics: str + applicable_intent_fields: frozenset[str] + source_structure: str + required_affordances: frozenset[str] + success_type: str + scene_affordances: frozenset[str] + + +def _contract( + task_type: str, + semantics: str, + applicable_intent_fields: frozenset[str], + source_structure: str, + required_affordances: frozenset[str], + success_type: str, + *, + scene_affordances: frozenset[str] | None = None, +) -> TaskContract: + return TaskContract( + task_type=task_type, + semantics=semantics, + applicable_intent_fields=applicable_intent_fields, + source_structure=source_structure, + required_affordances=required_affordances, + success_type=success_type, + scene_affordances=scene_affordances or required_affordances, + ) + + +TASK_CONTRACTS: Mapping[str, TaskContract] = MappingProxyType( + { + "E1": _contract( + "E1", + "Pick, move, and place one object at a symbolic relation.", + frozenset( + { + "target", + "relation", + "required_arm", + "orientation_goal", + "layout", + "axis", + } + ), + "rigid_object", + frozenset({"graspable", "placeable"}), + "semantic_goal", + ), + "E2": _contract( + "E2", + "Make one fallen object upright and place it stably.", + frozenset({"required_arm", "orientation_goal"}), + "rigid_object", + frozenset({"graspable", "orientable"}), + "object_upright", + ), + "E3": _contract( + "E3", + "Pour from a held source container into a target container.", + frozenset({"target", "relation", "required_arm"}), + "rigid_object", + frozenset({"graspable", "pourable"}), + "poured", + ), + "E4": _contract( + "E4", + "Transfer one held object from one arm to the other.", + frozenset({"transfer_arm", "receive_arm", "orientation_goal"}), + "rigid_object", + frozenset({"graspable", "handover"}), + "handover_complete", + ), + "E5": _contract( + "E5", + "Use both arms to pick, move, and optionally release one shared rigid object.", + frozenset({"target", "relation", "direction", "terminal_behavior"}), + "rigid_object", + frozenset({"dual_graspable"}), + "held_by_both_grippers", + scene_affordances=frozenset({"dual_graspable", "rigid"}), + ), + "E6": _contract( + "E6", + "Pull an articulated part to its requested state.", + frozenset({"required_arm", "target_state"}), + "articulation", + frozenset({"pullable"}), + "articulation_joint_near", + scene_affordances=frozenset({"articulated", "pullable"}), + ), + "E7": _contract( + "E7", + "Push an articulated part to its requested state.", + frozenset({"required_arm", "target_state"}), + "articulation", + frozenset({"pushable"}), + "articulation_joint_near", + scene_affordances=frozenset({"articulated", "pushable"}), + ), + "E8": _contract( + "E8", + "Turn one knob to a requested setting.", + frozenset({"required_arm", "target_setting"}), + "articulation", + frozenset({"turnable"}), + "articulation_joint_near", + ), + "E9": _contract( + "E9", + "Press one button until its requested terminal state.", + frozenset({"required_arm", "target_state"}), + "articulation", + frozenset({"pressable"}), + "pressed", + ), + } +) + + +def task_contract(task_type: str) -> TaskContract: + """Return the canonical contract or reject an unknown E-task type.""" + try: + return TASK_CONTRACTS[str(task_type)] + except KeyError as exc: + raise ValueError(f"Unsupported task type {task_type!r}.") from exc + + +def task_success_type( + task_type: str, + params: Mapping[str, Any] | None = None, +) -> str: + """Resolve a TaskSpec success type, including E5's terminal behavior.""" + contract = task_contract(task_type) + if contract.task_type != "E5": + return contract.success_type + terminal_behavior = str((params or {}).get("terminal_behavior", "hold")) + if terminal_behavior == "hold": + return "held_by_both_grippers" + if terminal_behavior == "place": + return "semantic_goal" + raise ValueError("E5 terminal_behavior must be 'hold' or 'place'.") diff --git a/embodichain/gen_sim/task_engine/orchestration/__init__.py b/embodichain/gen_sim/task_engine/orchestration/__init__.py new file mode 100644 index 000000000..ed3a9539c --- /dev/null +++ b/embodichain/gen_sim/task_engine/orchestration/__init__.py @@ -0,0 +1,109 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Task-owned orchestration across task, scene, and action engines.""" + +from __future__ import annotations + +from embodichain.gen_sim.action_engine.agent import ActionAgent, ActionGraph +from embodichain.gen_sim.action_engine.runtime import ExecutionReport + +from .artifacts import ( + ArtifactTransaction, + CONSERVATIVE_SCENE_GRAPH_FILENAME, + TaskEngineArtifactPaths, + FEASIBILITY_REPORT_FILENAME, + PREPARATION_FAILURE_FILENAME, + STATIC_SCENE_MANIFEST_FILENAME, + task_engine_artifact_paths, + write_execution_report, + write_preparation_failure, +) +from .contracts import ( + BINDING_REPORT_SCHEMA, + EXECUTION_REPORT_SCHEMA, + GROUNDED_TASK_PLAN_SCHEMA, + ROLE_BINDINGS_SCHEMA, + SCENE_MANIFEST_SCHEMA, + BindingReport, + GroundedTaskPlan, + RoleBindings, + SceneManifest, +) +from .coordinator import ( + TaskEngineCoordinator, + PreparationResult, + build_grounded_task_plan, + lower_task_candidate, +) +from .scene_adapter import ( + CandidateSelection, + SceneAdaptation, + SceneAdapter, + SceneAdapterProtocolError, +) +from .scene_source import ( + SceneSourceFingerprint, + SceneSourceRef, + fingerprint_scene_source, + verify_scene_source_fingerprint, +) +from .legacy_scene import ( + LEGACY_SCENE_CONVERSION_SCHEMA, + LegacySceneRevision, + convert_legacy_gym_project, + restore_locked_scene_entities, +) + +__all__ = [ + "ActionAgent", + "ActionGraph", + "ArtifactTransaction", + "CONSERVATIVE_SCENE_GRAPH_FILENAME", + "BINDING_REPORT_SCHEMA", + "BindingReport", + "TaskEngineArtifactPaths", + "TaskEngineCoordinator", + "EXECUTION_REPORT_SCHEMA", + "ExecutionReport", + "FEASIBILITY_REPORT_FILENAME", + "GROUNDED_TASK_PLAN_SCHEMA", + "GroundedTaskPlan", + "PREPARATION_FAILURE_FILENAME", + "PreparationResult", + "ROLE_BINDINGS_SCHEMA", + "RoleBindings", + "SCENE_MANIFEST_SCHEMA", + "STATIC_SCENE_MANIFEST_FILENAME", + "SceneAdaptation", + "CandidateSelection", + "SceneAdapter", + "SceneAdapterProtocolError", + "SceneManifest", + "SceneSourceFingerprint", + "SceneSourceRef", + "build_grounded_task_plan", + "task_engine_artifact_paths", + "fingerprint_scene_source", + "LEGACY_SCENE_CONVERSION_SCHEMA", + "LegacySceneRevision", + "convert_legacy_gym_project", + "restore_locked_scene_entities", + "verify_scene_source_fingerprint", + "lower_task_candidate", + "write_execution_report", + "write_preparation_failure", +] diff --git a/embodichain/gen_sim/task_engine/orchestration/artifacts.py b/embodichain/gen_sim/task_engine/orchestration/artifacts.py new file mode 100644 index 000000000..b941c86dc --- /dev/null +++ b/embodichain/gen_sim/task_engine/orchestration/artifacts.py @@ -0,0 +1,331 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Transactional publication for Task Engine artifacts.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +import json +import os +from pathlib import Path +import shutil +import tempfile +from typing import Any + +from embodichain.gen_sim.action_engine.runtime import ( + EXECUTION_REPORT_FILENAME, + write_execution_report as _write_execution_report, +) + +__all__ = [ + "BINDING_REPORT_FILENAME", + "CONSERVATIVE_SCENE_GRAPH_FILENAME", + "EXECUTION_REPORT_FILENAME", + "GROUNDED_TASK_PLAN_FILENAME", + "FEASIBILITY_REPORT_FILENAME", + "FINAL_SCENE_INSPECTION_FILENAME", + "PREPARATION_FAILURE_FILENAME", + "ROLE_BINDINGS_FILENAME", + "SCENE_MANIFEST_FILENAME", + "STATIC_SCENE_MANIFEST_FILENAME", + "SUCCESS_SPEC_FILENAME", + "TASK_CANDIDATE_SET_FILENAME", + "TASK_DRAFT_FILENAME", + "SCENE_REQUEST_FILENAME", + "ArtifactTransaction", + "TaskEngineArtifactPaths", + "task_engine_artifact_paths", + "write_task_engine_artifacts", + "write_execution_report", + "write_preparation_failure", +] + + +TASK_CANDIDATE_SET_FILENAME = "task_candidate_set.json" +TASK_DRAFT_FILENAME = "task_draft.json" +SCENE_REQUEST_FILENAME = "scene_request.json" +SUCCESS_SPEC_FILENAME = "success_spec.json" +SCENE_MANIFEST_FILENAME = "scene_manifest.json" +STATIC_SCENE_MANIFEST_FILENAME = "static_scene_manifest.json" +CONSERVATIVE_SCENE_GRAPH_FILENAME = "conservative_scene_graph.json" +ROLE_BINDINGS_FILENAME = "role_bindings.json" +BINDING_REPORT_FILENAME = "binding_report.json" +FEASIBILITY_REPORT_FILENAME = "feasibility_report.json" +FINAL_SCENE_INSPECTION_FILENAME = "final_scene_inspection.json" +GROUNDED_TASK_PLAN_FILENAME = "grounded_task_plan.json" +PREPARATION_FAILURE_FILENAME = "preparation_failure.json" + + +@dataclass(frozen=True) +class TaskEngineArtifactPaths: + """Canonical Task Engine paths rooted at one published bundle.""" + + root: Path + task_candidate_set: Path + task_draft: Path + scene_request: Path + success_spec: Path + scene_manifest: Path + static_scene_manifest: Path + conservative_scene_graph: Path + role_bindings: Path + binding_report: Path + feasibility_report: Path + final_scene_inspection: Path + grounded_task_plan: Path + preparation_failure: Path + execution_report: Path + + +def task_engine_artifact_paths( + output_dir: str | Path, +) -> TaskEngineArtifactPaths: + """Return all Task Engine paths without creating the directory.""" + root = Path(output_dir).expanduser().resolve() + return TaskEngineArtifactPaths( + root=root, + task_candidate_set=root / TASK_CANDIDATE_SET_FILENAME, + task_draft=root / TASK_DRAFT_FILENAME, + scene_request=root / SCENE_REQUEST_FILENAME, + success_spec=root / SUCCESS_SPEC_FILENAME, + scene_manifest=root / SCENE_MANIFEST_FILENAME, + static_scene_manifest=root / STATIC_SCENE_MANIFEST_FILENAME, + conservative_scene_graph=root / CONSERVATIVE_SCENE_GRAPH_FILENAME, + role_bindings=root / ROLE_BINDINGS_FILENAME, + binding_report=root / BINDING_REPORT_FILENAME, + feasibility_report=root / FEASIBILITY_REPORT_FILENAME, + final_scene_inspection=root / FINAL_SCENE_INSPECTION_FILENAME, + grounded_task_plan=root / GROUNDED_TASK_PLAN_FILENAME, + preparation_failure=root / PREPARATION_FAILURE_FILENAME, + execution_report=root / EXECUTION_REPORT_FILENAME, + ) + + +class ArtifactTransaction: + """Build a complete bundle beside its destination and publish it by rename.""" + + def __init__(self, output_dir: str | Path, *, overwrite: bool = False) -> None: + raw = Path(output_dir).expanduser() + self.output_dir = ( + (Path.cwd() / raw).resolve() if not raw.is_absolute() else raw.resolve() + ) + self.overwrite = bool(overwrite) + self.staging_dir: Path | None = None + self._committed = False + + def __enter__(self) -> "ArtifactTransaction": + destination = self.output_dir + destination.parent.mkdir(parents=True, exist_ok=True) + if destination.exists() and not self.overwrite: + raise FileExistsError( + f"Output directory already exists: {destination}. " + "Pass overwrite=True to replace it." + ) + self.staging_dir = Path( + tempfile.mkdtemp( + prefix=f".{destination.name}.staging-", + dir=destination.parent, + ) + ) + return self + + def commit(self) -> Path: + """Rewrite staging-local absolute paths, then atomically publish.""" + if self.staging_dir is None: + raise RuntimeError("ArtifactTransaction has not been entered.") + if self._committed: + raise RuntimeError("ArtifactTransaction has already been committed.") + staging = self.staging_dir + destination = self.output_dir + _relocate_json_paths(staging, destination) + + backup: Path | None = None + if destination.exists(): + if not self.overwrite: + raise FileExistsError( + f"Output directory already exists: {destination}." + ) + backup = Path( + tempfile.mkdtemp( + prefix=f".{destination.name}.backup-", + dir=destination.parent, + ) + ) + backup.rmdir() + os.replace(destination, backup) + try: + os.replace(staging, destination) + except BaseException: + if backup is not None and backup.exists() and not destination.exists(): + os.replace(backup, destination) + raise + else: + self._committed = True + self.staging_dir = None + if backup is not None: + _remove_path(backup) + _fsync_directory(destination.parent) + return destination + + def __exit__(self, exc_type: Any, exc: Any, traceback: Any) -> bool: + if self.staging_dir is not None and self.staging_dir.exists(): + shutil.rmtree(self.staging_dir) + return False + + +def write_task_engine_artifacts( + output_dir: str | Path, + *, + candidate_set: Mapping[str, Any], + scene_manifest: Mapping[str, Any] | None, + role_bindings: Mapping[str, Any] | None, + binding_report: Mapping[str, Any], + grounded_task_plan: Mapping[str, Any] | None = None, + static_scene_manifest: Mapping[str, Any] | None = None, + conservative_scene_graph: Mapping[str, Any] | None = None, + feasibility_report: Mapping[str, Any] | None = None, + final_scene_inspection: Mapping[str, Any] | None = None, +) -> TaskEngineArtifactPaths: + """Write Task Engine protocols into an unpublished staging directory. + + An unsuccessful adaptation can omit SceneManifest and RoleBindings rather + than publishing protocol filenames whose payloads do not satisfy their + schemas. + """ + paths = task_engine_artifact_paths(output_dir) + paths.root.mkdir(parents=True, exist_ok=True) + _write_json(paths.task_candidate_set, candidate_set) + if scene_manifest is not None: + _write_json(paths.scene_manifest, scene_manifest) + if static_scene_manifest is not None: + _write_json(paths.static_scene_manifest, static_scene_manifest) + if conservative_scene_graph is not None: + _write_json(paths.conservative_scene_graph, conservative_scene_graph) + if role_bindings is not None: + _write_json(paths.role_bindings, role_bindings) + _write_json(paths.binding_report, binding_report) + if feasibility_report is not None: + _write_json(paths.feasibility_report, feasibility_report) + if final_scene_inspection is not None: + _write_json(paths.final_scene_inspection, final_scene_inspection) + + if grounded_task_plan is not None: + _write_json(paths.grounded_task_plan, grounded_task_plan) + _write_json(paths.task_draft, grounded_task_plan["task_draft"]) + candidate_id = grounded_task_plan["selected_candidate_id"] + selected = next( + candidate + for candidate in candidate_set["candidates"] + if candidate["candidate_id"] == candidate_id + ) + _write_json(paths.scene_request, selected["scene_request"]) + _write_json(paths.success_spec, grounded_task_plan["success_spec"]) + return paths + + +def write_execution_report(output_dir: str | Path, value: Any) -> Path: + """Publish through the Action Engine-owned report boundary.""" + return _write_execution_report(output_dir, value) + + +def write_preparation_failure(output_dir: str | Path, value: Any) -> Path: + """Write a strict-JSON audit for a failed candidate planning transaction.""" + path = task_engine_artifact_paths(output_dir).preparation_failure + path.parent.mkdir(parents=True, exist_ok=True) + _write_json(path, value) + return path + + +def _write_json(path: Path, value: Any) -> None: + try: + payload = ( + json.dumps( + value, + ensure_ascii=False, + indent=2, + sort_keys=False, + allow_nan=False, + ) + + "\n" + ).encode("utf-8") + except (TypeError, ValueError) as exc: + raise ValueError(f"Artifact {path.name} is not strict JSON data.") from exc + temporary: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="wb", + dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + delete=False, + ) as stream: + temporary = Path(stream.name) + stream.write(payload) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path) + temporary = None + finally: + if temporary is not None: + temporary.unlink(missing_ok=True) + + +def _relocate_json_paths(staging: Path, destination: Path) -> None: + """Replace staging-root absolute paths embedded by the legacy generator.""" + source_prefix = staging.resolve().as_posix() + destination_prefix = destination.resolve().as_posix() + + def relocate(value: Any) -> Any: + if isinstance(value, str): + if value == source_prefix: + return destination_prefix + if value.startswith(source_prefix + "/"): + return destination_prefix + value[len(source_prefix) :] + return value + if isinstance(value, list): + return [relocate(item) for item in value] + if isinstance(value, dict): + return {key: relocate(item) for key, item in value.items()} + return value + + for path in staging.rglob("*.json"): + try: + value = json.loads(path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise ValueError(f"Generated artifact is invalid JSON: {path}") from exc + relocated = relocate(value) + if relocated != value: + _write_json(path, relocated) + + +def _remove_path(path: Path) -> None: + if path.is_dir() and not path.is_symlink(): + shutil.rmtree(path) + else: + path.unlink(missing_ok=True) + + +def _fsync_directory(path: Path) -> None: + try: + descriptor = os.open(path, os.O_RDONLY) + except OSError: + return + try: + os.fsync(descriptor) + finally: + os.close(descriptor) diff --git a/embodichain/gen_sim/task_engine/orchestration/contracts.py b/embodichain/gen_sim/task_engine/orchestration/contracts.py new file mode 100644 index 000000000..5c38f284c --- /dev/null +++ b/embodichain/gen_sim/task_engine/orchestration/contracts.py @@ -0,0 +1,647 @@ +# ---------------------------------------------------------------------------- +# 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 cross-engine contracts for scene binding and orchestration.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +import json +import math +from typing import Any, TypeAlias + +from embodichain.gen_sim.action_engine.domain import ( + validate_scene_requirements, + validate_task_spec, +) +from embodichain.gen_sim.action_engine.runtime import ( + EXECUTION_REPORT_SCHEMA, + validate_execution_report, +) +from embodichain.gen_sim.task_engine import ( + SCENE_REQUEST_SCHEMA, + SUCCESS_SPEC_SCHEMA, + TASK_CANDIDATE_SET_SCHEMA, + TASK_DRAFT_SCHEMA, + SceneRequest, + SuccessSpec, + TaskCandidate, + TaskCandidateSet, + TaskDraft, + canonical_hash, + task_success_type, + validate_scene_request, + validate_success_spec, + validate_task_candidate, + validate_task_candidate_set, + validate_task_draft, +) + +__all__ = [ + "BINDING_REPORT_SCHEMA", + "EXECUTION_REPORT_SCHEMA", + "GROUNDED_TASK_PLAN_SCHEMA", + "ROLE_BINDINGS_SCHEMA", + "SCENE_MANIFEST_SCHEMA", + "SCENE_REQUEST_SCHEMA", + "SUCCESS_SPEC_SCHEMA", + "TASK_CANDIDATE_SET_SCHEMA", + "TASK_DRAFT_SCHEMA", + "BindingReport", + "ExecutionReport", + "GroundedTaskPlan", + "RoleBindings", + "SceneManifest", + "SceneRequest", + "SuccessSpec", + "TaskCandidate", + "TaskCandidateSet", + "TaskDraft", + "canonical_hash", + "validate_binding_report", + "validate_execution_report", + "validate_grounded_task_plan", + "validate_role_bindings", + "validate_scene_manifest", + "validate_scene_request", + "validate_success_spec", + "validate_task_candidate", + "validate_task_candidate_set", + "validate_task_draft", +] + +SCENE_MANIFEST_SCHEMA = "action_engine_scene_manifest_v1" +ROLE_BINDINGS_SCHEMA = "action_engine_role_bindings_v1" +BINDING_REPORT_SCHEMA = "action_engine_binding_report_v1" +GROUNDED_TASK_PLAN_SCHEMA = "action_engine_grounded_task_plan_v1" +SceneManifest: TypeAlias = dict[str, Any] +RoleBindings: TypeAlias = dict[str, Any] +BindingReport: TypeAlias = dict[str, Any] +GroundedTaskPlan: TypeAlias = dict[str, Any] +ExecutionReport: TypeAlias = dict[str, Any] + + +def validate_scene_manifest(value: Mapping[str, Any]) -> SceneManifest: + result = _mapping(value, "SceneManifest") + _keys( + result, + {"schema_version", "scene_id", "source_format", "robot_profile", "objects"}, + "SceneManifest", + ) + _schema(result, SCENE_MANIFEST_SCHEMA, "SceneManifest") + for key in ("scene_id", "source_format", "robot_profile"): + result[key] = _nonempty(result.get(key), f"SceneManifest.{key}") + object_keys = { + "uid", + "role", + "name", + "description", + "category", + "color", + "affordances", + "initial_state", + "attributes", + } + objects = [] + for index, raw in enumerate( + _sequence(result.get("objects"), "SceneManifest.objects") + ): + context = f"SceneManifest.objects[{index}]" + item = _mapping(raw, context) + _keys(item, object_keys, context) + item["uid"] = _nonempty(item.get("uid"), f"{context}.uid") + for key in ("role", "name", "description", "category"): + item[key] = _string(item.get(key), f"{context}.{key}") + if item.get("color") is not None: + item["color"] = _string(item.get("color"), f"{context}.color") + item["affordances"] = _strings( + item.get("affordances"), f"{context}.affordances" + ) + item["initial_state"] = _mapping( + item.get("initial_state"), f"{context}.initial_state" + ) + item["attributes"] = _mapping(item.get("attributes"), f"{context}.attributes") + objects.append(item) + _unique([item["uid"] for item in objects], "SceneManifest object UIDs") + result["objects"] = objects + _json_safe(result, "SceneManifest") + return result + + +def validate_role_bindings(value: Mapping[str, Any]) -> RoleBindings: + result = _mapping(value, "RoleBindings") + _keys( + result, + { + "schema_version", + "task_id", + "candidate_id", + "reference_bindings", + "role_bindings", + }, + "RoleBindings", + ) + _schema(result, ROLE_BINDINGS_SCHEMA, "RoleBindings") + for key in ("task_id", "candidate_id"): + result[key] = _nonempty(result.get(key), f"RoleBindings.{key}") + result["reference_bindings"] = _string_lists( + result.get("reference_bindings"), "RoleBindings.reference_bindings" + ) + if any(not uids for uids in result["reference_bindings"].values()): + raise ValueError("RoleBindings.reference_bindings values must not be empty.") + result["role_bindings"] = _string_map( + result.get("role_bindings"), "RoleBindings.role_bindings" + ) + return result + + +def validate_binding_report(value: Mapping[str, Any]) -> BindingReport: + result = _mapping(value, "BindingReport") + _keys( + result, + { + "schema_version", + "task_id", + "status", + "selected_candidate_id", + "selection_reason", + "candidates", + }, + "BindingReport", + ) + _schema(result, BINDING_REPORT_SCHEMA, "BindingReport") + result["task_id"] = _nonempty(result.get("task_id"), "BindingReport.task_id") + result["status"] = _enum( + result.get("status"), + {"bound", "ambiguous", "unsatisfied"}, + "BindingReport.status", + ) + result["selected_candidate_id"] = _string( + result.get("selected_candidate_id"), "BindingReport.selected_candidate_id" + ) + result["selection_reason"] = _string( + result.get("selection_reason"), "BindingReport.selection_reason" + ) + if result["status"] == "bound" and not result["selected_candidate_id"]: + raise ValueError("A bound BindingReport requires selected_candidate_id.") + candidate_keys = { + "candidate_id", + "semantic_hash", + "status", + "references", + "reasons", + } + reference_keys = { + "reference_id", + "status", + "confidence", + "candidate_uids", + "selected_uids", + "reasons", + } + candidates = [] + for index, raw in enumerate( + _sequence(result.get("candidates"), "BindingReport.candidates") + ): + context = f"BindingReport.candidates[{index}]" + candidate = _mapping(raw, context) + _keys(candidate, candidate_keys, context) + candidate["candidate_id"] = _nonempty( + candidate.get("candidate_id"), f"{context}.candidate_id" + ) + candidate["semantic_hash"] = _digest( + candidate.get("semantic_hash"), f"{context}.semantic_hash" + ) + candidate["status"] = _enum( + candidate.get("status"), + {"resolved", "ambiguous", "not_found", "incompatible"}, + f"{context}.status", + ) + references = [] + for ref_index, ref_raw in enumerate( + _sequence(candidate.get("references"), f"{context}.references") + ): + ref_context = f"{context}.references[{ref_index}]" + reference = _mapping(ref_raw, ref_context) + _keys(reference, reference_keys, ref_context) + reference["reference_id"] = _nonempty( + reference.get("reference_id"), f"{ref_context}.reference_id" + ) + reference["status"] = _enum( + reference.get("status"), + {"resolved", "ambiguous", "not_found", "incompatible"}, + f"{ref_context}.status", + ) + reference["confidence"] = _number( + reference.get("confidence"), + f"{ref_context}.confidence", + minimum=0.0, + maximum=1.0, + ) + reference["candidate_uids"] = _strings( + reference.get("candidate_uids"), + f"{ref_context}.candidate_uids", + allow_empty=True, + ) + reference["selected_uids"] = _strings( + reference.get("selected_uids"), + f"{ref_context}.selected_uids", + allow_empty=True, + ) + reference["reasons"] = _strings( + reference.get("reasons"), f"{ref_context}.reasons", allow_empty=True + ) + selected = set(reference["selected_uids"]) + candidates_for_reference = set(reference["candidate_uids"]) + if not selected <= candidates_for_reference: + raise ValueError( + f"{ref_context}.selected_uids must be a subset of candidate_uids." + ) + if reference["status"] == "resolved" and not selected: + raise ValueError( + f"{ref_context} status=resolved requires selected_uids." + ) + if reference["status"] != "resolved" and selected: + raise ValueError( + f"{ref_context} non-resolved status cannot select UIDs." + ) + if reference["status"] == "not_found" and candidates_for_reference: + raise ValueError( + f"{ref_context} status=not_found cannot carry candidate_uids." + ) + references.append(reference) + if not references: + raise ValueError(f"{context}.references must not be empty.") + _unique( + [item["reference_id"] for item in references], + f"{context} reference IDs", + ) + expected_status = _candidate_binding_status(references) + if candidate["status"] != expected_status: + raise ValueError( + f"{context}.status must be {expected_status!r} for its references." + ) + candidate["references"] = references + candidate["reasons"] = _strings( + candidate.get("reasons"), f"{context}.reasons", allow_empty=True + ) + candidates.append(candidate) + if not candidates: + raise ValueError("BindingReport.candidates must not be empty.") + _unique( + [item["candidate_id"] for item in candidates], "BindingReport candidate IDs" + ) + if result["selected_candidate_id"] and result["selected_candidate_id"] not in { + item["candidate_id"] for item in candidates + }: + raise ValueError("BindingReport.selected_candidate_id is unknown.") + if result["status"] != "bound" and result["selected_candidate_id"]: + raise ValueError( + "A non-bound BindingReport cannot carry selected_candidate_id." + ) + selected = next( + ( + candidate + for candidate in candidates + if candidate["candidate_id"] == result["selected_candidate_id"] + ), + None, + ) + if result["status"] == "bound" and ( + selected is None or selected["status"] != "resolved" + ): + raise ValueError( + "A bound BindingReport must select a resolved candidate audit." + ) + if result["status"] == "unsatisfied" and any( + candidate["status"] in {"resolved", "ambiguous"} for candidate in candidates + ): + raise ValueError( + "An unsatisfied BindingReport cannot contain resolved or ambiguous candidates." + ) + result["candidates"] = candidates + return result + + +def validate_grounded_task_plan(value: Mapping[str, Any]) -> GroundedTaskPlan: + result = _mapping(value, "GroundedTaskPlan") + keys = { + "schema_version", + "task_id", + "instruction", + "selected_candidate_id", + "task_draft", + "task_spec", + "scene_requirements", + "success_spec", + "scene_manifest", + "role_bindings", + "binding_report", + "hashes", + } + _keys(result, keys, "GroundedTaskPlan") + _schema(result, GROUNDED_TASK_PLAN_SCHEMA, "GroundedTaskPlan") + task_id = _nonempty(result.get("task_id"), "GroundedTaskPlan.task_id") + result["instruction"] = _nonempty( + result.get("instruction"), "GroundedTaskPlan.instruction" + ) + result["selected_candidate_id"] = _nonempty( + result.get("selected_candidate_id"), "GroundedTaskPlan.selected_candidate_id" + ) + result["task_draft"] = validate_task_draft(result.get("task_draft")) + # Candidate SuccessSpec terms use draft step IDs. Once count/all selectors + # are lowered, the grounded plan carries one term per concrete TaskSpec + # instance instead, and is checked against the v2 recipe below. + result["success_spec"] = validate_success_spec(result.get("success_spec")) + result["scene_manifest"] = validate_scene_manifest(result.get("scene_manifest")) + result["role_bindings"] = validate_role_bindings(result.get("role_bindings")) + result["binding_report"] = validate_binding_report(result.get("binding_report")) + result["task_spec"] = validate_task_spec( + _mapping(result.get("task_spec"), "GroundedTaskPlan.task_spec") + ) + result["scene_requirements"] = validate_scene_requirements( + _mapping( + result.get("scene_requirements"), + "GroundedTaskPlan.scene_requirements", + ) + ) + hashes = _mapping(result.get("hashes"), "GroundedTaskPlan.hashes") + _keys( + hashes, + {"task_draft", "task_spec", "scene_manifest", "role_bindings", "plan"}, + "GroundedTaskPlan.hashes", + ) + for key in hashes: + hashes[key] = _digest(hashes[key], f"GroundedTaskPlan.hashes.{key}") + if any( + part["task_id"] != task_id + for part in ( + result["task_draft"], + result["task_spec"], + result["scene_requirements"], + result["success_spec"], + result["role_bindings"], + result["binding_report"], + ) + ): + raise ValueError("GroundedTaskPlan task IDs must agree.") + if result["task_draft"]["instruction"] != result["instruction"]: + raise ValueError("GroundedTaskPlan instruction must match TaskDraft.") + if result["task_spec"]["instruction"] != result["instruction"]: + raise ValueError("GroundedTaskPlan instruction must match TaskSpec.") + if ( + result["role_bindings"]["candidate_id"] != result["selected_candidate_id"] + or result["binding_report"]["selected_candidate_id"] + != result["selected_candidate_id"] + ): + raise ValueError("GroundedTaskPlan selected candidate IDs must agree.") + if result["binding_report"]["status"] != "bound": + raise ValueError("GroundedTaskPlan requires a bound BindingReport.") + selected_audit = next( + candidate + for candidate in result["binding_report"]["candidates"] + if candidate["candidate_id"] == result["selected_candidate_id"] + ) + if selected_audit["semantic_hash"] != canonical_hash(result["task_draft"]["steps"]): + raise ValueError( + "GroundedTaskPlan selected candidate hash must match TaskDraft." + ) + task_metadata = result["task_spec"].get("metadata", {}) + task_oracle = result["task_spec"].get("oracle", {}) + serialized_bindings = ( + task_metadata.get("role_bindings") + if isinstance(task_metadata, Mapping) + and task_metadata.get("role_bindings") is not None + else ( + task_oracle.get("role_bindings") + if isinstance(task_oracle, Mapping) + else None + ) + ) + if serialized_bindings != result["role_bindings"]["role_bindings"]: + raise ValueError( + "GroundedTaskPlan RoleBindings must match the TaskSpec binding hand-off." + ) + requirement_roles = { + str(item["role_id"]) for item in result["scene_requirements"]["objects"] + } + if requirement_roles != set(result["role_bindings"]["role_bindings"]): + raise ValueError( + "GroundedTaskPlan SceneRequirements roles must match RoleBindings." + ) + manifest_uids = {item["uid"] for item in result["scene_manifest"]["objects"]} + missing_uids = sorted( + set(result["role_bindings"]["role_bindings"].values()) - manifest_uids + ) + if missing_uids: + raise ValueError( + f"GroundedTaskPlan RoleBindings reference unknown scene UIDs {missing_uids}." + ) + reference_ids = { + f"{step['id']}.{role}" + for step in result["task_draft"]["steps"] + for role in ("object", "target") + if step[role]["kind"] == "scene_ref" + } + reference_bindings = result["role_bindings"]["reference_bindings"] + if set(reference_bindings) != reference_ids: + raise ValueError( + "GroundedTaskPlan reference bindings must cover every draft scene_ref exactly." + ) + bound_uids = {uid for uids in reference_bindings.values() for uid in uids} | set( + result["role_bindings"]["role_bindings"].values() + ) + unknown_uids = sorted(bound_uids - manifest_uids) + if unknown_uids: + raise ValueError( + "GroundedTaskPlan bindings reference unknown SceneManifest UIDs: " + f"{unknown_uids}." + ) + audited_bindings = { + reference["reference_id"]: reference["selected_uids"] + for reference in selected_audit["references"] + } + if audited_bindings != reference_bindings: + raise ValueError( + "GroundedTaskPlan RoleBindings must match the selected candidate audit." + ) + ontology_success = [ + { + "step_id": instance["id"], + "type": task_success_type(instance["task_type"], instance["params"]), + } + for instance in result["task_spec"]["task_instances"] + ] + recipe_success = [ + { + "step_id": term["task_instance_id"], + "type": term["type"], + } + for term in result["task_spec"]["success"]["terms"] + ] + if recipe_success != ontology_success: + raise ValueError( + "GroundedTaskPlan TaskSpec success recipe must be derived from " + "task_success_type." + ) + if result["success_spec"]["terms"] != ontology_success: + raise ValueError( + "GroundedTaskPlan SuccessSpec must exactly match the lowered " + "TaskSpec success recipe." + ) + expected_hashes = { + "task_draft": canonical_hash(result["task_draft"]), + "task_spec": canonical_hash(result["task_spec"]), + "scene_manifest": canonical_hash(result["scene_manifest"]), + "role_bindings": canonical_hash(result["role_bindings"]), + } + base = {key: value for key, value in result.items() if key != "hashes"} + expected_hashes["plan"] = canonical_hash(base) + if hashes != expected_hashes: + raise ValueError("GroundedTaskPlan hashes do not match their contents.") + result["task_id"] = task_id + result["hashes"] = hashes + _json_safe(result, "GroundedTaskPlan") + return result + + +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 _candidate_binding_status(references: Sequence[Mapping[str, Any]]) -> str: + statuses = {str(reference["status"]) for reference in references} + if statuses == {"resolved"}: + return "resolved" + if "ambiguous" in statuses: + return "ambiguous" + if "incompatible" in statuses: + return "incompatible" + return "not_found" + + +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] | frozenset[str], context: str +) -> None: + if set(value) != set(expected): + raise ValueError( + f"{context} requires exactly fields {sorted(expected)}; received {sorted(value)}." + ) + + +def _schema(value: Mapping[str, Any], expected: str, context: str) -> None: + if value.get("schema_version") != expected: + raise ValueError(f"{context}.schema_version must be {expected!r}.") + + +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, maximum: int | None = None +) -> int: + if ( + not isinstance(value, int) + or isinstance(value, bool) + or value < minimum + or (maximum is not None and value > maximum) + ): + raise ValueError(f"{context} must be an integer in the allowed range.") + return value + + +def _number(value: Any, context: str, *, minimum: float, maximum: float) -> float: + if ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + or not minimum <= float(value) <= maximum + ): + raise ValueError( + f"{context} must be a finite number between {minimum} and {maximum}." + ) + return float(value) + + +def _strings(value: Any, context: str, *, allow_empty: bool = False) -> list[str]: + result = [_string(item, context) for item in _sequence(value, context)] + if not allow_empty and any(not item for item in result): + raise ValueError(f"{context} values must not be empty.") + if len(result) != len(set(result)): + raise ValueError(f"{context} values must be unique.") + return result + + +def _string_map(value: Any, context: str) -> dict[str, str]: + result = _mapping(value, context) + return { + _nonempty(key, context): _nonempty(item, context) + for key, item in result.items() + } + + +def _string_lists(value: Any, context: str) -> dict[str, list[str]]: + result = _mapping(value, context) + return { + _nonempty(key, context): _strings(item, context) for key, item in result.items() + } + + +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 _unique(values: Sequence[str], context: str) -> None: + if len(values) != len(set(values)): + raise ValueError(f"{context} must be unique.") + + +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/task_engine/orchestration/coordinator.py b/embodichain/gen_sim/task_engine/orchestration/coordinator.py new file mode 100644 index 000000000..c6c9f0282 --- /dev/null +++ b/embodichain/gen_sim/task_engine/orchestration/coordinator.py @@ -0,0 +1,862 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""End-to-end Task Engine preparation for an existing scene source.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass, replace +import json +from pathlib import Path +import shutil +from typing import Any + +from embodichain.gen_sim.action_engine.generation import ( + GeneratedConfigPaths, + generate_action_engine_config, +) +from embodichain.gen_sim.action_engine.generation.artifacts import artifact_paths +from embodichain.gen_sim.action_engine.agent import ActionAgent +from embodichain.gen_sim.action_engine.unbound import ActionCapabilityError +from embodichain.gen_sim.action_engine.domain.task_contracts import ( + TASK_CONTRACTS as ACTION_TASK_CONTRACTS, +) +from embodichain.gen_sim.action_engine.tasks import ( + GroundedTaskSpec, + ground_instruction_draft, +) +from embodichain.gen_sim.task_engine import ( + TaskAgent, + TaskCandidate, + TaskCandidateSet, + validate_scene_output_separation, + validate_task_candidate, + validate_task_candidate_set, +) +from embodichain.gen_sim.task_engine.scene import FeasibilityBroker, FeasibilityReport + +from .artifacts import ( + ArtifactTransaction, + TaskEngineArtifactPaths, + task_engine_artifact_paths, + write_task_engine_artifacts, + write_preparation_failure, +) +from .contracts import ( + GROUNDED_TASK_PLAN_SCHEMA, + GroundedTaskPlan, + RoleBindings, + canonical_hash, + validate_grounded_task_plan, + validate_binding_report, + validate_role_bindings, +) +from .scene_adapter import SceneAdaptation, SceneAdapter +from .scene_source import SceneSourceRef + +__all__ = [ + "TaskEngineCoordinator", + "PreparationResult", + "build_grounded_task_plan", + "lower_task_candidate", +] + + +BundleGenerator = Callable[..., GeneratedConfigPaths] +_PREPARATION_FAILURE_SCHEMA = "action_engine_preparation_failure_v1" + + +def lower_task_candidate( + candidate: Mapping[str, Any], + reference_bindings: Mapping[str, Any], + scene_objects: Sequence[Mapping[str, Any]], + robot_profile: str, +) -> GroundedTaskSpec: + """Lower a selected TaskCandidate across the Task/Action boundary.""" + normalized = validate_task_candidate(candidate) + if reference_bindings.get("schema_version") is not None: + role_bindings = validate_role_bindings(reference_bindings) + if role_bindings["task_id"] != normalized["draft"]["task_id"]: + raise ValueError("RoleBindings.task_id must match the TaskCandidate.") + if role_bindings["candidate_id"] != normalized["candidate_id"]: + raise ValueError("RoleBindings.candidate_id must match the TaskCandidate.") + raw_bindings = role_bindings["reference_bindings"] + else: + raw_bindings = reference_bindings + bindings = { + str(reference_id): [str(uid) for uid in uids] + for reference_id, uids in raw_bindings.items() + } + grounded = ground_instruction_draft( + normalized["draft"]["task_id"], + normalized["draft"]["instruction"], + {"steps": normalized["draft"]["steps"]}, + scene_objects, + robot_profile=robot_profile, + reference_bindings=bindings, + ) + _validate_lowered_success(normalized, bindings, grounded) + return grounded + + +@dataclass(frozen=True) +class PreparationResult: + """Published result of one Task -> Scene -> Action preparation attempt.""" + + status: str + output_dir: Path + candidate_set: TaskCandidateSet + adaptation: SceneAdaptation + artifacts: TaskEngineArtifactPaths + grounded_task_plan: GroundedTaskPlan | None = None + action_graph: dict[str, Any] | None = None + generated_paths: GeneratedConfigPaths | None = None + feasibility_report: FeasibilityReport | None = None + planning_attempts: tuple[dict[str, Any], ...] = () + unbound_action_plan: dict[str, Any] | None = None + + @property + def bound(self) -> bool: + return self.status == "bound" + + @property + def selected_candidate_id(self) -> str | None: + return self.adaptation.selected_candidate_id + + +@dataclass(frozen=True) +class _PlannedCandidate: + adaptation: SceneAdaptation + selected: TaskCandidate + role_bindings: RoleBindings + feasibility_report: FeasibilityReport | None + grounded: GroundedTaskSpec + grounded_plan: GroundedTaskPlan + action_graph: dict[str, Any] + unbound_action_plan: dict[str, Any] | None + + +class TaskEngineCoordinator: + """Run Task Agent, Scene Adapter, and Action Agent as one transaction.""" + + def __init__( + self, + *, + task_agent: TaskAgent | None = None, + scene_adapter: SceneAdapter | None = None, + action_agent: ActionAgent | None = None, + bundle_generator: BundleGenerator = generate_action_engine_config, + feasibility_broker: FeasibilityBroker | None = None, + ) -> None: + self.task_agent = task_agent or TaskAgent() + self.scene_adapter = scene_adapter or SceneAdapter() + self.action_agent = action_agent or ActionAgent() + self.bundle_generator = bundle_generator + self.feasibility_broker = feasibility_broker or FeasibilityBroker() + + def prepare( + self, + task_id: str, + instruction: str, + source: SceneSourceRef | str | Path, + output_dir: str | Path, + *, + model: str | None = None, + candidate_count: int = 3, + overwrite: bool = False, + planning_mode: str = "offline", + vlm_model: str | None = None, + max_episodes: int | None = None, + max_episode_steps: int | None = None, + randomize_scene: bool = False, + randomize_table_material: bool = False, + candidate_set: TaskCandidateSet | Mapping[str, Any] | None = None, + force_most_likely: bool = False, + final_inspection: Mapping[str, Any] | None = None, + unbound_action_plan: Mapping[str, Any] | None = None, + ) -> PreparationResult: + """Prepare and atomically publish a Task Engine bundle. + + Ambiguous and unsatisfied scene adaptations are valid terminal results. + They publish the complete audit hand-off but never publish a TaskSpec, + SeedGraph, Gym configuration, or GroundedTaskPlan. + """ + normalized_source = self._coerce_source(source) + validate_scene_output_separation(normalized_source.path, output_dir) + with ArtifactTransaction(output_dir, overwrite=overwrite) as transaction: + staging_dir = transaction.staging_dir + assert staging_dir is not None + if candidate_set is None: + normalized_candidates = self.task_agent.generate( + task_id, + instruction, + model=model, + candidate_count=candidate_count, + ) + else: + normalized_candidates = validate_task_candidate_set(candidate_set) + if normalized_candidates["task_id"] != str(task_id).strip(): + raise ValueError("TaskCandidateSet.task_id must match task_id.") + if normalized_candidates["instruction"] != str(instruction).strip(): + raise ValueError( + "TaskCandidateSet.instruction must match instruction." + ) + candidate_set = normalized_candidates + adaptation_kwargs: dict[str, Any] = {"force_most_likely": force_most_likely} + if final_inspection is not None: + adaptation_kwargs["final_inspection"] = final_inspection + adaptation = self.scene_adapter.adapt( + candidate_set, + normalized_source, + **adaptation_kwargs, + ) + status = str(adaptation.binding_report["status"]) + + if status != "bound": + write_task_engine_artifacts( + staging_dir, + candidate_set=candidate_set, + scene_manifest=None, + role_bindings=None, + binding_report=adaptation.binding_report, + static_scene_manifest=adaptation.static_scene_manifest, + conservative_scene_graph=adaptation.conservative_scene_graph, + final_scene_inspection=final_inspection, + ) + published = transaction.commit() + return PreparationResult( + status=status, + output_dir=published, + candidate_set=deepcopy(candidate_set), + adaptation=adaptation, + artifacts=task_engine_artifact_paths(published), + planning_attempts=(), + ) + + selected = adaptation.selected_candidate + raw_role_bindings = adaptation.role_bindings + if selected is None or raw_role_bindings is None: + raise ValueError( + "A bound SceneAdaptation must include a selected candidate " + "and RoleBindings." + ) + feasibility_report = self._assess_feasibility( + selected, + raw_role_bindings, + adaptation, + ) + if ( + feasibility_report is not None + and feasibility_report["status"] == "contradicted" + and feasibility_report["remediation_class"] != "action_capability" + ): + adaptation, selected, raw_role_bindings, feasibility_report = ( + self._fallback_feasible_candidate( + candidate_set, + adaptation, + selected, + raw_role_bindings, + feasibility_report, + ) + ) + if ( + feasibility_report is not None + and feasibility_report["status"] == "contradicted" + ): + write_task_engine_artifacts( + staging_dir, + candidate_set=candidate_set, + scene_manifest=adaptation.scene_manifest, + role_bindings=raw_role_bindings, + binding_report=adaptation.binding_report, + static_scene_manifest=adaptation.static_scene_manifest, + conservative_scene_graph=adaptation.conservative_scene_graph, + feasibility_report=feasibility_report, + final_scene_inspection=final_inspection, + ) + published = transaction.commit() + return PreparationResult( + status="infeasible", + output_dir=published, + candidate_set=deepcopy(candidate_set), + adaptation=adaptation, + artifacts=task_engine_artifact_paths(published), + feasibility_report=deepcopy(feasibility_report), + planning_attempts=(), + ) + robot_profile = str(adaptation.scene_manifest["robot_profile"]) + planned, planning_failures = self._plan_with_candidate_fallback( + candidate_set, + adaptation, + selected, + raw_role_bindings, + feasibility_report, + robot_profile=robot_profile, + unbound_action_plan=unbound_action_plan, + ) + if planned is None: + write_task_engine_artifacts( + staging_dir, + candidate_set=candidate_set, + scene_manifest=adaptation.scene_manifest, + role_bindings=raw_role_bindings, + binding_report=adaptation.binding_report, + static_scene_manifest=adaptation.static_scene_manifest, + conservative_scene_graph=adaptation.conservative_scene_graph, + feasibility_report=feasibility_report, + final_scene_inspection=final_inspection, + ) + write_preparation_failure( + staging_dir, + { + "schema_version": _PREPARATION_FAILURE_SCHEMA, + "task_id": str(candidate_set["task_id"]), + "status": "planning_failed", + "selected_candidate_id": str(selected["candidate_id"]), + "attempts": planning_failures, + }, + ) + published = transaction.commit() + return PreparationResult( + status="planning_failed", + output_dir=published, + candidate_set=deepcopy(candidate_set), + adaptation=adaptation, + artifacts=task_engine_artifact_paths(published), + feasibility_report=deepcopy(feasibility_report), + planning_attempts=tuple(deepcopy(planning_failures)), + ) + + adaptation = planned.adaptation + selected = planned.selected + role_bindings = planned.role_bindings + feasibility_report = planned.feasibility_report + grounded = planned.grounded + grounded_plan = planned.grounded_plan + action_graph = planned.action_graph + + generator_kwargs: dict[str, Any] = { + "task_name": grounded_plan["task_id"], + "task_spec": grounded_plan["task_spec"], + "robot_profile": robot_profile, + "source_scene_z_rotation_degrees": ( + adaptation.prepared_scene.z_rotation_degrees + ), + "body_scale_policy": adaptation.prepared_scene.body_scale_policy, + "body_scale": adaptation.prepared_scene.body_scale, + "overwrite": False, + "randomize_scene": randomize_scene, + "randomize_table_material": randomize_table_material, + "planning_mode": planning_mode, + "vlm_model": vlm_model, + } + if max_episodes is not None: + generator_kwargs["max_episodes"] = max_episodes + if max_episode_steps is not None: + generator_kwargs["max_episode_steps"] = max_episode_steps + compatibility_input = staging_dir / ".task_engine_input" + compatibility_input.mkdir() + task_spec_path = compatibility_input / "task_spec.json" + requirements_path = compatibility_input / "scene_requirements.json" + _write_compatibility_input(task_spec_path, grounded.task_spec) + _write_compatibility_input( + requirements_path, + grounded.scene_requirements, + ) + generator_kwargs["task_spec"] = task_spec_path + try: + generated = self.bundle_generator( + adaptation.source_config_path, + staging_dir, + **generator_kwargs, + ) + finally: + shutil.rmtree(compatibility_input, ignore_errors=True) + _require_matching_generated_graph(generated, action_graph) + write_task_engine_artifacts( + staging_dir, + candidate_set=candidate_set, + scene_manifest=adaptation.scene_manifest, + role_bindings=role_bindings, + binding_report=adaptation.binding_report, + grounded_task_plan=grounded_plan, + static_scene_manifest=adaptation.static_scene_manifest, + conservative_scene_graph=adaptation.conservative_scene_graph, + feasibility_report=feasibility_report, + final_scene_inspection=final_inspection, + ) + published = transaction.commit() + return PreparationResult( + status="bound", + output_dir=published, + candidate_set=deepcopy(candidate_set), + adaptation=adaptation, + artifacts=task_engine_artifact_paths(published), + grounded_task_plan=grounded_plan, + action_graph=deepcopy(action_graph), + generated_paths=artifact_paths( + published, + planning_mode=planning_mode, + ), + feasibility_report=deepcopy(feasibility_report), + planning_attempts=tuple(deepcopy(planning_failures)), + unbound_action_plan=deepcopy(planned.unbound_action_plan), + ) + + def _plan_with_candidate_fallback( + self, + candidate_set: Mapping[str, Any], + adaptation: SceneAdaptation, + selected: TaskCandidate, + role_bindings: RoleBindings, + feasibility_report: FeasibilityReport | None, + *, + robot_profile: str, + unbound_action_plan: Mapping[str, Any] | None, + ) -> tuple[_PlannedCandidate | None, list[dict[str, Any]]]: + """Treat lowering and Action planning failures as candidate-local.""" + candidates = { + str(candidate["candidate_id"]): candidate + for candidate in candidate_set.get("candidates", ()) + if isinstance(candidate, Mapping) and candidate.get("candidate_id") + } + resolved = { + str(audit["candidate_id"]) + for audit in adaptation.binding_report["candidates"] + if audit["status"] == "resolved" + } + selected_id = str(selected["candidate_id"]) + ordered_ids = [selected_id] + [ + candidate_id + for candidate_id in candidates + if candidate_id != selected_id and candidate_id in resolved + ] + failures: list[dict[str, Any]] = [] + + for candidate_id in ordered_ids: + candidate = candidates.get(candidate_id) + raw_bindings = ( + role_bindings + if candidate_id == selected_id + else adaptation.candidate_bindings.get(candidate_id) + ) + if candidate is None or raw_bindings is None: + continue + report = ( + feasibility_report + if candidate_id == selected_id + else self._assess_feasibility(candidate, raw_bindings, adaptation) + ) + if report is not None and report["status"] == "contradicted": + failures.append( + _candidate_failure( + candidate, + raw_bindings, + stage="static_feasibility", + error_type="FeasibilityContradiction", + error_message="Static feasibility contradicted this candidate.", + feasibility_report=report, + ) + ) + continue + + candidate_adaptation = _select_candidate_adaptation( + adaptation, + candidate, + raw_bindings, + failures, + ) + grounded: GroundedTaskSpec | None = None + grounded_plan: GroundedTaskPlan | None = None + candidate_unbound: Mapping[str, Any] | None = None + action_graph: Mapping[str, Any] | None = None + stage = "lowering" + try: + grounded = lower_task_candidate( + candidate, + raw_bindings, + adaptation.prepared_scene.planner_objects, + robot_profile, + ) + canonical_bindings = validate_role_bindings( + { + **deepcopy(raw_bindings), + "role_bindings": deepcopy(grounded.role_bindings), + } + ) + candidate_adaptation = replace( + candidate_adaptation, + role_bindings=deepcopy(canonical_bindings), + ) + stage = "grounded_plan" + grounded_plan = build_grounded_task_plan( + candidate=candidate, + task_spec=grounded.task_spec, + scene_requirements=grounded.scene_requirements, + scene_manifest=adaptation.scene_manifest, + role_bindings=canonical_bindings, + binding_report=candidate_adaptation.binding_report, + ) + stage = "action_planning" + bind_and_plan = getattr(self.action_agent, "bind_and_plan", None) + if callable(bind_and_plan): + candidate_unbound = ( + unbound_action_plan + if unbound_action_plan is not None + and str(unbound_action_plan.get("candidate_id")) == candidate_id + else self.action_agent.draft(candidate) + ) + action_graph = bind_and_plan(candidate_unbound, grounded_plan) + else: + action_graph = self.action_agent.plan(grounded_plan) + stage = "preflight" + preflight = getattr(self.action_agent, "preflight", None) + if callable(preflight): + preflight( + action_graph, + scene_manifest=adaptation.scene_manifest, + ) + except ActionCapabilityError: + raise + except (TypeError, ValueError, OSError) as error: + failures.append( + _candidate_failure( + candidate, + raw_bindings, + stage=stage, + error_type=type(error).__name__, + error_message=str(error), + feasibility_report=report, + grounded_task_plan=grounded_plan, + unbound_action_plan=candidate_unbound, + action_graph=action_graph, + ) + ) + continue + + assert grounded is not None and grounded_plan is not None + return ( + _PlannedCandidate( + adaptation=candidate_adaptation, + selected=deepcopy(candidate), + role_bindings=canonical_bindings, + feasibility_report=deepcopy(report), + grounded=grounded, + grounded_plan=grounded_plan, + action_graph=deepcopy(action_graph), + unbound_action_plan=( + None + if candidate_unbound is None + else deepcopy(dict(candidate_unbound)) + ), + ), + failures, + ) + return None, failures + + def _fallback_feasible_candidate( + self, + candidate_set: Mapping[str, Any], + adaptation: SceneAdaptation, + selected: TaskCandidate, + role_bindings: RoleBindings, + report: FeasibilityReport, + ) -> tuple[ + SceneAdaptation, + TaskCandidate, + RoleBindings, + FeasibilityReport | None, + ]: + """Try other resolved semantic candidates after a static contradiction.""" + candidates = { + str(candidate["candidate_id"]): candidate + for candidate in candidate_set.get("candidates", ()) + if isinstance(candidate, Mapping) and candidate.get("candidate_id") + } + selected_id = str(selected["candidate_id"]) + for audit in adaptation.binding_report["candidates"]: + candidate_id = str(audit["candidate_id"]) + if candidate_id == selected_id or audit["status"] != "resolved": + continue + candidate = candidates.get(candidate_id) + alternative_bindings = adaptation.candidate_bindings.get(candidate_id) + if candidate is None or alternative_bindings is None: + continue + alternative_report = self._assess_feasibility( + candidate, + alternative_bindings, + adaptation, + ) + if ( + alternative_report is not None + and alternative_report["status"] == "contradicted" + ): + continue + binding_report = validate_binding_report( + { + **deepcopy(adaptation.binding_report), + "selected_candidate_id": candidate_id, + "selection_reason": ( + "Selected the next resolved candidate after static " + f"feasibility contradicted {selected_id}." + ), + } + ) + chosen = deepcopy(candidate) + updated = replace( + adaptation, + selected_candidate=chosen, + role_bindings=deepcopy(alternative_bindings), + binding_report=binding_report, + ) + return ( + updated, + chosen, + deepcopy(alternative_bindings), + alternative_report, + ) + return adaptation, selected, role_bindings, report + + def _assess_feasibility( + self, + candidate: Mapping[str, Any], + role_bindings: Mapping[str, Any], + adaptation: SceneAdaptation, + ) -> FeasibilityReport | None: + """Intersect task requirements with scene and Action Engine capabilities.""" + manifest = adaptation.static_scene_manifest + registry = getattr(self.action_agent, "registry", None) + if manifest is None or registry is None: + return None + catalog = getattr(registry, "catalog", None) + if not callable(catalog): + return None + return self.feasibility_broker.assess( + candidate, + role_bindings, + manifest, + capability_catalog=catalog(), + task_actions={ + task_type: contract.core_actions + for task_type, contract in ACTION_TASK_CONTRACTS.items() + }, + ) + + def _coerce_source( + self, + source: SceneSourceRef | str | Path, + ) -> SceneSourceRef: + if isinstance(source, SceneSourceRef): + return source + path = Path(source).expanduser() + return SceneSourceRef( + path, + robot_profile=self.scene_adapter.robot_profile, + ) + + +def _select_candidate_adaptation( + adaptation: SceneAdaptation, + candidate: Mapping[str, Any], + role_bindings: Mapping[str, Any], + prior_failures: Sequence[Mapping[str, Any]], +) -> SceneAdaptation: + candidate_id = str(candidate["candidate_id"]) + current_id = adaptation.selected_candidate_id + if candidate_id == current_id and not prior_failures: + return adaptation + failed = ", ".join( + f"{failure['candidate_id']} failed {failure['stage']}" + for failure in prior_failures + ) + reason = f"Selected {candidate_id} after {failed}." + binding_report = validate_binding_report( + { + **deepcopy(adaptation.binding_report), + "selected_candidate_id": candidate_id, + "selection_reason": reason, + } + ) + return replace( + adaptation, + selected_candidate=deepcopy(candidate), + role_bindings=deepcopy(role_bindings), + binding_report=binding_report, + ) + + +def _candidate_failure( + candidate: Mapping[str, Any], + role_bindings: Mapping[str, Any], + *, + stage: str, + error_type: str, + error_message: str, + feasibility_report: Mapping[str, Any] | None = None, + grounded_task_plan: Mapping[str, Any] | None = None, + unbound_action_plan: Mapping[str, Any] | None = None, + action_graph: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + return { + "candidate_id": str(candidate["candidate_id"]), + "stage": stage, + "draft": deepcopy(candidate["draft"]), + "bindings": deepcopy(dict(role_bindings)), + "grounded_task_plan": ( + None if grounded_task_plan is None else deepcopy(dict(grounded_task_plan)) + ), + "unbound_action_plan": ( + None if unbound_action_plan is None else deepcopy(dict(unbound_action_plan)) + ), + "action_graph": None if action_graph is None else deepcopy(dict(action_graph)), + "feasibility_report": ( + None if feasibility_report is None else deepcopy(dict(feasibility_report)) + ), + "error": {"type": error_type, "message": error_message}, + } + + +# Short public name used in the phase-one design document. +def build_grounded_task_plan( + *, + candidate: Mapping[str, Any], + task_spec: Mapping[str, Any], + scene_requirements: Mapping[str, Any], + scene_manifest: Mapping[str, Any], + role_bindings: RoleBindings, + binding_report: Mapping[str, Any], +) -> GroundedTaskPlan: + """Assemble a validated plan with hashes over every authoritative hand-off.""" + draft = deepcopy(candidate["draft"]) + # A scene-ref quantifier may expand one draft step into several concrete + # task instances. The grounded plan records the executable success terms, + # while the selected TaskCandidate retains the pre-grounding SuccessSpec. + success_spec = { + **deepcopy(candidate["success_spec"]), + "terms": [ + { + "step_id": str(term["task_instance_id"]), + "type": str(term["type"]), + } + for term in task_spec["success"]["terms"] + ], + } + task = deepcopy(dict(task_spec)) + requirements = deepcopy(dict(scene_requirements)) + manifest = deepcopy(dict(scene_manifest)) + bindings = deepcopy(dict(role_bindings)) + report = deepcopy(dict(binding_report)) + base = { + "schema_version": GROUNDED_TASK_PLAN_SCHEMA, + "task_id": draft["task_id"], + "instruction": draft["instruction"], + "selected_candidate_id": candidate["candidate_id"], + "task_draft": draft, + "task_spec": task, + "scene_requirements": requirements, + "success_spec": success_spec, + "scene_manifest": manifest, + "role_bindings": bindings, + "binding_report": report, + } + plan = { + **base, + "hashes": { + "task_draft": canonical_hash(draft), + "task_spec": canonical_hash(task), + "scene_manifest": canonical_hash(manifest), + "role_bindings": canonical_hash(bindings), + "plan": canonical_hash(base), + }, + } + return validate_grounded_task_plan(plan) + + +def _validate_lowered_success( + candidate: TaskCandidate, + bindings: Mapping[str, list[str]], + grounded: GroundedTaskSpec, +) -> None: + success_by_step = { + term["step_id"]: term["type"] for term in candidate["success_spec"]["terms"] + } + expected: list[str] = [] + multiplicity_by_step: dict[str, int] = {} + for step in _topological_steps(candidate["draft"]["steps"]): + selector = step["object"] + multiplicity = 1 + if selector["kind"] == "scene_ref": + multiplicity = len(bindings.get(f"{step['id']}.object", ())) + elif selector["kind"] == "step_result": + multiplicity = multiplicity_by_step[str(selector["step_id"])] + multiplicity_by_step[str(step["id"])] = multiplicity + expected.extend([success_by_step[step["id"]]] * multiplicity) + actual = [term.get("type") for term in grounded.task_spec["success"]["terms"]] + if actual != expected: + raise ValueError( + "Lowered TaskSpec success terms do not match the expanded SuccessSpec." + ) + + +def _topological_steps( + steps: list[Mapping[str, Any]], +) -> list[Mapping[str, Any]]: + positions = {str(step["id"]): index for index, step in enumerate(steps)} + pending = {str(step["id"]): set(step["depends_on"]) for step in steps} + result: list[Mapping[str, Any]] = [] + emitted: set[str] = set() + while len(result) < len(steps): + ready = [ + step + for step in steps + if step["id"] not in emitted and pending[str(step["id"])] <= emitted + ] + if not ready: + raise ValueError("TaskDraft step dependencies contain a cycle.") + step = min(ready, key=lambda item: positions[str(item["id"])]) + result.append(step) + emitted.add(str(step["id"])) + return result + + +def _require_matching_generated_graph( + generated: GeneratedConfigPaths, + expected: Mapping[str, Any], +) -> None: + """Catch a compatibility-generator drift before publishing the bundle.""" + graph_path = getattr(generated, "seed_task_graph", None) + if graph_path is None or not Path(graph_path).is_file(): + # Injected generators used by API consumers may publish by other means. + # Task Engine's independently planned graph remains authoritative. + return + try: + actual = json.loads(Path(graph_path).read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError(f"Generated SeedGraph is unreadable: {graph_path}") from exc + if canonical_hash(actual) != canonical_hash(expected): + raise ValueError( + "Legacy bundle generation produced a SeedGraph different from " + "ActionAgent.plan." + ) + + +def _write_compatibility_input(path: Path, value: Mapping[str, Any]) -> None: + path.write_text( + json.dumps(value, ensure_ascii=False, indent=2, allow_nan=False) + "\n", + encoding="utf-8", + ) diff --git a/embodichain/gen_sim/task_engine/orchestration/legacy_scene.py b/embodichain/gen_sim/task_engine/orchestration/legacy_scene.py new file mode 100644 index 000000000..4cb17ff7d --- /dev/null +++ b/embodichain/gen_sim/task_engine/orchestration/legacy_scene.py @@ -0,0 +1,381 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Read-only conversion of legacy Gym projects into editable scene revisions.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass +import json +from pathlib import Path +import shutil +from typing import Any, Final + +import numpy as np +from scipy.spatial.transform import Rotation +import trimesh + +from embodichain.gen_sim.action_engine.generation.source_scene import ( + prepare_scene, + resolve_source_scene, +) + +from .scene_source import ( + SceneSourceFingerprint, + fingerprint_scene_source, + verify_scene_source_fingerprint, +) + +__all__ = [ + "LEGACY_SCENE_CONVERSION_SCHEMA", + "LegacySceneRevision", + "convert_legacy_gym_project", + "restore_locked_scene_entities", +] + +LEGACY_SCENE_CONVERSION_SCHEMA: Final = "embodichain.legacy-scene-conversion/v1" +_CONVERSION_MANIFEST = "legacy_conversion.json" + + +@dataclass(frozen=True) +class LegacySceneRevision: + """A new editable revision derived without modifying its legacy source.""" + + output_root: Path + scene_config_path: Path + scene_graph_path: Path + manifest_path: Path + source_fingerprint: SceneSourceFingerprint + locked_entity_uids: tuple[str, ...] + + +def convert_legacy_gym_project( + source: str | Path, + output_root: str | Path, +) -> LegacySceneRevision: + """Convert a supported legacy Gym project into a Scene Engine revision. + + Args: + source: Legacy Gym project directory or explicit configuration path. + output_root: Empty destination owned by the new scene revision. + + Returns: + Paths and provenance for the converted revision. + + Raises: + ValueError: If the source is not legacy or the destination already exists. + FileNotFoundError: If a referenced source asset is missing. + """ + resolved = resolve_source_scene(source) + if resolved.source_format != "legacy_gym_config": + raise ValueError("Legacy conversion requires a legacy Gym configuration.") + destination = Path(output_root).expanduser().resolve() + if destination.exists(): + if not destination.is_dir() or any(destination.iterdir()): + raise ValueError("Legacy scene revision output_root must be empty.") + source_fingerprint = fingerprint_scene_source(source) + prepared = prepare_scene(source) + export_root = destination / "scene_export" + assets_root = export_root / "mesh_assets" + assets_root.mkdir(parents=True, exist_ok=True) + semantics = {str(item.get("uid")): item for item in prepared.planner_objects} + + background = [ + _editable_entry(item, semantics=semantics, assets_root=assets_root) + for item in prepared.background + ] + rigid_objects = [ + _editable_entry(item, semantics=semantics, assets_root=assets_root) + for item in prepared.rigid_objects + ] + articulations = [ + _locked_articulation( + item, + source_root=resolved.path.parent, + destination_root=export_root / "locked_assets", + ) + for item in prepared.articulations + ] + table = next((item for item in background if item.get("uid") == "table"), None) + if table is None: + raise ValueError("Legacy conversion requires one table support object.") + _measure_support_metadata(table, export_root=export_root) + for item in rigid_objects: + _measure_center(item, export_root=export_root) + + scene_config = { + "format": "embodichain.scene-export/v1", + "scene_id": f"legacy-revision-{source_fingerprint.config_sha256[:16]}", + "background": background, + "rigid_object": rigid_objects, + "articulation": articulations, + } + scene_config_path = export_root / "scene_config.json" + _write_json(scene_config_path, scene_config) + scene_graph = { + "nodes": [ + { + "object_id": "table", + "parent_id": None, + "parent_relation": None, + "table_region": None, + "orientation_state": None, + }, + *[ + { + "object_id": str(item["uid"]), + "parent_id": "table", + "parent_relation": "on", + "table_region": None, + "orientation_state": None, + } + for item in rigid_objects + ], + ], + "relations": [], + } + scene_graph_path = export_root / "scene_graph.json" + _write_json(scene_graph_path, scene_graph) + locked_uids = tuple( + sorted(str(item["uid"]) for item in [*background, *articulations]) + ) + manifest = { + "schema_version": LEGACY_SCENE_CONVERSION_SCHEMA, + "source": source_fingerprint.to_dict(), + "scene_config": scene_config_path.as_posix(), + "audit_hierarchy": "unknown", + "operational_hierarchy": "assumed_on_table", + "assumptions": [ + { + "uid": str(item["uid"]), + "relation": "on", + "parent_uid": "table", + "confidence": None, + "source": "operational_assumption", + } + for item in rigid_objects + ], + "locked_entity_uids": list(locked_uids), + "locked_articulations": deepcopy(articulations), + "locked_background": deepcopy( + [item for item in background if item.get("uid") != "table"] + ), + } + manifest_path = destination / _CONVERSION_MANIFEST + _write_json(manifest_path, manifest) + verify_scene_source_fingerprint(source_fingerprint.to_dict()) + return LegacySceneRevision( + output_root=destination, + scene_config_path=scene_config_path, + scene_graph_path=scene_graph_path, + manifest_path=manifest_path, + source_fingerprint=source_fingerprint, + locked_entity_uids=locked_uids, + ) + + +def restore_locked_scene_entities(revision_root: str | Path) -> Path: + """Restore collision-only legacy entities after Scene Engine export. + + Args: + revision_root: Converted revision root containing ``legacy_conversion.json``. + + Returns: + Updated scene configuration path. + + Raises: + FileNotFoundError: If the conversion manifest or scene config is absent. + ValueError: If a generated scene attempts to reuse a locked UID. + """ + root = Path(revision_root).expanduser().resolve() + manifest_path = root / _CONVERSION_MANIFEST + if not manifest_path.is_file(): + raise FileNotFoundError( + f"Legacy conversion manifest not found: {manifest_path}" + ) + manifest = _read_mapping(manifest_path) + if manifest.get("schema_version") != LEGACY_SCENE_CONVERSION_SCHEMA: + raise ValueError("Legacy conversion manifest schema is invalid.") + config_path = root / "scene_export" / "scene_config.json" + config = _read_mapping(config_path) + existing = { + str(item.get("uid")) + for section in ("background", "rigid_object", "articulation") + for item in config.get(section, ()) + if isinstance(item, Mapping) and item.get("uid") + } + for section, key in ( + ("background", "locked_background"), + ("articulation", "locked_articulations"), + ): + values = manifest.get(key, ()) + if not isinstance(values, Sequence) or isinstance(values, (str, bytes)): + raise TypeError(f"Legacy conversion manifest {key} must be a sequence.") + target = list(config.get(section, ())) + for raw in values: + item = deepcopy(dict(raw)) + uid = str(item.get("uid", "")) + if uid in existing: + raise ValueError(f"Generated scene reused locked entity UID {uid!r}.") + existing.add(uid) + target.append(item) + config[section] = target + _write_json(config_path, config) + verify_scene_source_fingerprint(manifest["source"]) + return config_path + + +def _editable_entry( + value: Mapping[str, Any], + *, + semantics: Mapping[str, Mapping[str, Any]], + assets_root: Path, +) -> dict[str, Any]: + item = deepcopy(dict(value)) + uid = str(item.get("uid", "")).strip() + if not uid: + raise ValueError("Converted scene entities require a UID.") + semantic = semantics.get(uid, {}) + for key in ("category", "name", "description"): + item[key] = str(semantic.get(key) or item.get(key) or uid) + shape = item.get("shape") + if not isinstance(shape, Mapping): + raise ValueError(f"Legacy scene entity {uid!r} has no supported shape.") + destination = assets_root / uid / f"{uid}.glb" + destination.parent.mkdir(parents=True, exist_ok=True) + _shape_to_glb(shape, destination) + item["shape"] = { + "shape_type": "Mesh", + "fpath": destination.relative_to(assets_root.parent).as_posix(), + "compute_uv": False, + } + item.setdefault("body_scale", [1.0, 1.0, 1.0]) + item.setdefault("init_pos", [0.0, 0.0, 0.0]) + item.setdefault("init_rot", [0.0, 0.0, 0.0]) + item.setdefault("attrs", {"mass": 1.0}) + item.setdefault("body_type", "kinematic" if uid == "table" else "dynamic") + item.setdefault("max_convex_hull_num", 1 if uid == "table" else 16) + return item + + +def _shape_to_glb(shape: Mapping[str, Any], destination: Path) -> None: + shape_type = str(shape.get("shape_type", "")) + if shape_type == "Mesh": + source = Path(str(shape.get("fpath", ""))).expanduser().resolve() + if not source.is_file(): + raise FileNotFoundError(f"Legacy mesh asset not found: {source}") + mesh = trimesh.load(source, force="scene") + elif shape_type == "Cube": + size = _vector(shape.get("size", [1.0, 1.0, 1.0]), length=3) + mesh = trimesh.Scene(trimesh.creation.box(extents=size)) + elif shape_type == "Sphere": + radius = float(shape.get("radius", 1.0)) + if not np.isfinite(radius) or radius <= 0.0: + raise ValueError("Legacy sphere radius must be positive and finite.") + mesh = trimesh.Scene(trimesh.creation.icosphere(radius=radius)) + else: + raise ValueError(f"Unsupported legacy shape_type {shape_type!r}.") + mesh.export(destination, file_type="glb") + + +def _locked_articulation( + value: Mapping[str, Any], + *, + source_root: Path, + destination_root: Path, +) -> dict[str, Any]: + item = deepcopy(dict(value)) + uid = str(item.get("uid", "")).strip() + raw = Path(str(item.get("fpath", ""))).expanduser() + source = raw.resolve() if raw.is_absolute() else (source_root / raw).resolve() + if not source.is_file(): + raise FileNotFoundError(f"Legacy articulation asset not found: {source}") + target_root = destination_root / uid + shutil.copytree(source.parent, target_root, dirs_exist_ok=True) + copied = target_root / source.name + item["fpath"] = copied.resolve().as_posix() + return item + + +def _measure_support_metadata(entry: dict[str, Any], *, export_root: Path) -> None: + bounds = _world_bounds(entry, export_root=export_root) + entry["support_surface_z"] = float(bounds[1, 2]) + rectangle = [ + [float(bounds[0, 0]), float(bounds[0, 1])], + [float(bounds[1, 0]), float(bounds[0, 1])], + [float(bounds[1, 0]), float(bounds[1, 1])], + [float(bounds[0, 0]), float(bounds[1, 1])], + ] + entry["support_contour_xy"] = rectangle + entry["support_optimization_rect_xy"] = deepcopy(rectangle) + entry["center_xy"] = [ + float((bounds[0, 0] + bounds[1, 0]) / 2.0), + float((bounds[0, 1] + bounds[1, 1]) / 2.0), + ] + + +def _measure_center(entry: dict[str, Any], *, export_root: Path) -> None: + bounds = _world_bounds(entry, export_root=export_root) + entry["center_xy"] = [ + float((bounds[0, 0] + bounds[1, 0]) / 2.0), + float((bounds[0, 1] + bounds[1, 1]) / 2.0), + ] + + +def _world_bounds(entry: Mapping[str, Any], *, export_root: Path) -> np.ndarray: + shape = dict(entry["shape"]) + mesh_path = (export_root / str(shape["fpath"])).resolve() + loaded = trimesh.load(mesh_path, force="scene") + mesh = loaded.to_geometry() + scale = np.asarray(_vector(entry.get("body_scale", [1.0] * 3), length=3)) + mesh.apply_scale(scale) + transform = np.eye(4) + transform[:3, :3] = Rotation.from_euler( + "XYZ", + _vector(entry.get("init_rot", [0.0] * 3), length=3), + degrees=True, + ).as_matrix() + transform[:3, 3] = _vector(entry.get("init_pos", [0.0] * 3), length=3) + mesh.apply_transform(transform) + return np.asarray(mesh.bounds, dtype=float) + + +def _vector(value: Any, *, length: int) -> list[float]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise TypeError("Legacy scene vector must be a sequence.") + result = [float(item) for item in value] + if len(result) != length or not np.all(np.isfinite(result)): + raise ValueError(f"Legacy scene vector must contain {length} finite values.") + return result + + +def _read_mapping(path: Path) -> dict[str, Any]: + if not path.is_file(): + raise FileNotFoundError(path) + value = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, Mapping): + raise TypeError(f"JSON document must contain an object: {path}") + return dict(value) + + +def _write_json(path: Path, value: 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/task_engine/orchestration/scene_adapter.py b/embodichain/gen_sim/task_engine/orchestration/scene_adapter.py new file mode 100644 index 000000000..58ee084f8 --- /dev/null +++ b/embodichain/gen_sim/task_engine/orchestration/scene_adapter.py @@ -0,0 +1,1049 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Bind Task Agent candidates to a redacted, authoritative scene inventory.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass, field +import hashlib +import json +from pathlib import Path +from typing import Any + +from embodichain.gen_sim.action_engine.generation.models import PreparedScene +from embodichain.gen_sim.action_engine.generation.source_scene import ( + prepare_scene, + resolve_source_scene, +) +from embodichain.gen_sim.action_engine.tasks.assembly import ( + SceneInventory, + validate_source_compatibility, + validate_target_compatibility, +) +from embodichain.gen_sim.action_engine.tasks.grounding import ( + GroundingCaller, + ground_scene_references, +) +from embodichain.gen_sim.task_engine import TaskCandidate, TaskCandidateSet +from embodichain.gen_sim.task_engine.interpretation import ( + _default_instruction_caller, +) +from embodichain.gen_sim.task_engine.scene import ( + ConservativeSceneGraph, + SceneEngineV1Adapter, + StaticSceneManifest, + build_conservative_scene_graph, + validate_static_scene_manifest, +) +from embodichain.gen_sim.task_engine.scene.final_inspection import ( + apply_final_inspection, + validate_final_scene_inspection, +) + +from .contracts import ( + BINDING_REPORT_SCHEMA, + ROLE_BINDINGS_SCHEMA, + SCENE_MANIFEST_SCHEMA, + BindingReport, + RoleBindings, + SceneManifest, + validate_binding_report, + validate_role_bindings, + validate_scene_manifest, + validate_task_candidate, + validate_task_candidate_set, +) +from .scene_source import ( + SceneSourceRef, + fingerprint_scene_source, + scene_revision_id, +) + +__all__ = [ + "Adjudicator", + "CandidateSelection", + "SceneAdaptation", + "SceneAdapter", + "SceneAdapterProtocolError", +] + + +Adjudicator = Callable[..., Mapping[str, Any]] + +_REDACTED_KEYS = frozenset( + { + "absolute_position", + "bbox", + "bboxes", + "bounding_box", + "center", + "centroid", + "coordinates", + "dimensions", + "extrinsics", + "grasp_pose", + "init_local_pose", + "init_pos", + "init_rot", + "intrinsics", + "joint_positions", + "joints", + "keypoint", + "keypoints", + "location", + "matrix", + "pose", + "position", + "position_xyz", + "qpos", + "quaternion", + "rotation", + "scale", + "target_pose", + "trajectory", + "transform", + "translation", + "waypoints", + "world_x", + "world_y", + "world_z", + "x", + "y", + "z", + } +) + + +class SceneAdapterProtocolError(ValueError): + """The grounding or adjudication transport violated its JSON protocol.""" + + +@dataclass(frozen=True) +class CandidateSelection: + """Candidate binding against semantic scene data before materialization.""" + + scene_manifest: SceneManifest + role_bindings: RoleBindings | None + binding_report: BindingReport + selected_candidate: TaskCandidate | None + candidate_bindings: dict[str, RoleBindings] = field(default_factory=dict) + + @property + def selected_candidate_id(self) -> str | None: + """Return the chosen candidate identifier, when one was bindable.""" + return ( + str(self.selected_candidate["candidate_id"]) + if self.selected_candidate is not None + else None + ) + + +@dataclass(frozen=True) +class SceneAdaptation: + """Complete Scene Adapter result, including the reusable prepared scene.""" + + scene_manifest: SceneManifest + role_bindings: RoleBindings | None + binding_report: BindingReport + selected_candidate: TaskCandidate | None + prepared_scene: PreparedScene + source_config_path: Path + conservative_scene_graph: ConservativeSceneGraph + static_scene_manifest: StaticSceneManifest | None = None + candidate_bindings: dict[str, RoleBindings] = field(default_factory=dict) + + @property + def selected_candidate_id(self) -> str | None: + return ( + str(self.selected_candidate["candidate_id"]) + if self.selected_candidate is not None + else None + ) + + @property + def reference_bindings(self) -> dict[str, list[str]]: + if self.role_bindings is None: + return {} + return deepcopy(self.role_bindings["reference_bindings"]) + + +class SceneAdapter: + """Adapt one existing or packaged scene to a set of task candidates.""" + + def __init__( + self, + *, + model: str | None = None, + grounding_caller: GroundingCaller | None = None, + adjudicator: Adjudicator | None = None, + robot_profile: str = "franka", + scene_engine_adapter: SceneEngineV1Adapter | None = None, + ) -> None: + self.model = model + self.grounding_caller = grounding_caller + self.adjudicator = adjudicator + self.robot_profile = robot_profile + self.scene_engine_adapter = scene_engine_adapter or SceneEngineV1Adapter() + + def adapt( + self, + candidate_set: TaskCandidateSet | Sequence[Mapping[str, Any]], + source: SceneSourceRef | str | Path, + *, + grounding_caller: GroundingCaller | None = None, + adjudicator: Adjudicator | None = None, + force_most_likely: bool = False, + final_inspection: Mapping[str, Any] | None = None, + ) -> SceneAdaptation: + """Ground all candidates, then deterministically choose a bindable one.""" + task_id, instruction, candidates = _coerce_candidates(candidate_set) + source_ref = self._resolve_source(source) + source_fingerprint = fingerprint_scene_source(source_ref) + prepared = prepare_scene( + source_ref.path, + z_rotation_degrees=source_ref.z_rotation_degrees, + body_scale_policy=source_ref.body_scale_policy, + body_scale=source_ref.body_scale, + ) + if final_inspection is not None: + normalized_inspection = validate_final_scene_inspection(final_inspection) + if normalized_inspection["scene_revision_id"] != scene_revision_id( + source_ref + ): + raise ValueError( + "FinalSceneInspection does not describe the adapted scene revision." + ) + prepared = apply_final_inspection(prepared, normalized_inspection) + inventory = SceneInventory( + prepared.planner_objects, + robot_profile=source_ref.robot_profile, + ) + resolved_source = resolve_source_scene(source_ref.path) + manifest = _build_manifest( + prepared, + inventory, + source_format=resolved_source.source_format, + ) + static_manifest = self.scene_engine_adapter.adapt_prepared_scene( + prepared, + source_format=resolved_source.source_format, + robot_profile=inventory.profile, + ) + static_manifest["source"]["source_fingerprint"] = source_fingerprint.to_dict() + static_manifest = validate_static_scene_manifest(static_manifest) + conservative_scene_graph = build_conservative_scene_graph( + prepared, + scene_id=static_manifest["scene_id"], + ) + if fingerprint_scene_source(source_ref) != source_fingerprint: + raise RuntimeError("Source Gym project changed while it was being adapted.") + + selection = self._select_candidates( + task_id, + instruction, + candidates, + manifest=manifest, + inventory=inventory, + scene_objects=prepared.planner_objects, + grounding_caller=grounding_caller, + adjudicator=adjudicator, + force_most_likely=force_most_likely, + ) + return SceneAdaptation( + scene_manifest=manifest, + role_bindings=selection.role_bindings, + binding_report=selection.binding_report, + selected_candidate=selection.selected_candidate, + prepared_scene=prepared, + source_config_path=prepared.source_config_path, + conservative_scene_graph=conservative_scene_graph, + static_scene_manifest=static_manifest, + candidate_bindings=selection.candidate_bindings, + ) + + def select_objects( + self, + candidate_set: TaskCandidateSet | Sequence[Mapping[str, Any]], + scene_objects: Sequence[Mapping[str, Any]], + *, + source_format: str = "embodichain.scene-blueprint/v1", + robot_profile: str | None = None, + grounding_caller: GroundingCaller | None = None, + adjudicator: Adjudicator | None = None, + force_most_likely: bool = False, + ) -> CandidateSelection: + """Bind candidates to semantic objects before assets are generated. + + Args: + candidate_set: Validated Task Engine candidate set. + scene_objects: Blueprint-level semantic object records. + source_format: Provenance label included in the semantic manifest. + robot_profile: Optional robot profile override. + grounding_caller: Optional structured grounding transport. + adjudicator: Optional candidate tie-breaker. + force_most_likely: Resolve ranked UID hypotheses instead of rejecting + low-confidence or ambiguous responses. + + Returns: + Audited candidate selection without requiring generated assets. + """ + task_id, instruction, candidates = _coerce_candidates(candidate_set) + inventory = SceneInventory( + scene_objects, + robot_profile=robot_profile or self.robot_profile, + ) + manifest = _build_semantic_manifest( + inventory, + source_format=source_format, + ) + return self._select_candidates( + task_id, + instruction, + candidates, + manifest=manifest, + inventory=inventory, + scene_objects=scene_objects, + grounding_caller=grounding_caller, + adjudicator=adjudicator, + force_most_likely=force_most_likely, + ) + + def _select_candidates( + self, + task_id: str, + instruction: str, + candidates: Sequence[TaskCandidate], + *, + manifest: SceneManifest, + inventory: SceneInventory, + scene_objects: Sequence[Mapping[str, Any]], + grounding_caller: GroundingCaller | None, + adjudicator: Adjudicator | None, + force_most_likely: bool, + ) -> CandidateSelection: + invoke = grounding_caller or self.grounding_caller + use_default_adjudicator = invoke is None + if invoke is None: + invoke = _default_grounding_caller() + choose = adjudicator or self.adjudicator + if choose is None and use_default_adjudicator: + choose = _default_adjudicator(model=self.model) + audits: list[dict[str, Any]] = [] + bindings_by_candidate: dict[str, dict[str, tuple[str, ...]]] = {} + for candidate in candidates: + audit, bindings = _ground_candidate( + candidate, + instruction=instruction, + inventory=inventory, + scene_objects=scene_objects, + model=self.model, + caller=invoke, + force_most_likely=force_most_likely, + ) + audits.append(audit) + if bindings is not None: + bindings_by_candidate[str(candidate["candidate_id"])] = bindings + + selected_id, status, reason = _select_candidate( + candidates, + audits, + manifest=manifest, + instruction=instruction, + adjudicator=choose, + ) + report = validate_binding_report( + { + "schema_version": BINDING_REPORT_SCHEMA, + "task_id": task_id, + "status": status, + "selected_candidate_id": selected_id or "", + "selection_reason": reason, + "candidates": audits, + } + ) + selected = next( + ( + deepcopy(candidate) + for candidate in candidates + if candidate["candidate_id"] == selected_id + ), + None, + ) + candidate_bindings = { + candidate_id: validate_role_bindings( + { + "schema_version": ROLE_BINDINGS_SCHEMA, + "task_id": task_id, + "candidate_id": candidate_id, + "reference_bindings": { + key: list(value) for key, value in sorted(raw_bindings.items()) + }, + "role_bindings": {}, + } + ) + for candidate_id, raw_bindings in bindings_by_candidate.items() + } + role_bindings = None if selected_id is None else candidate_bindings[selected_id] + return CandidateSelection( + scene_manifest=manifest, + role_bindings=role_bindings, + binding_report=report, + selected_candidate=selected, + candidate_bindings=candidate_bindings, + ) + + def _resolve_source( + self, + source: SceneSourceRef | str | Path, + ) -> SceneSourceRef: + if isinstance(source, SceneSourceRef): + return source + return SceneSourceRef(source, robot_profile=self.robot_profile) + + +def _coerce_candidates( + value: TaskCandidateSet | Sequence[Mapping[str, Any]], +) -> tuple[str, str, list[TaskCandidate]]: + if isinstance(value, Mapping): + normalized = validate_task_candidate_set(value) + return ( + normalized["task_id"], + normalized["instruction"], + normalized["candidates"], + ) + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + candidates = [validate_task_candidate(candidate) for candidate in value] + if not candidates: + raise ValueError("SceneAdapter requires at least one TaskCandidate.") + task_ids = {candidate["draft"]["task_id"] for candidate in candidates} + instructions = {candidate["draft"]["instruction"] for candidate in candidates} + if len(task_ids) != 1 or len(instructions) != 1: + raise ValueError("All TaskCandidates must describe the same task.") + return task_ids.pop(), instructions.pop(), candidates + raise TypeError("candidate_set must be a TaskCandidateSet or candidate sequence.") + + +def _default_grounding_caller() -> GroundingCaller: + # Keep provider setup lazy so package import and offline tests never load an + # LLM client. This is the same structured transport used by interpretation. + return _default_instruction_caller + + +def _default_adjudicator(*, model: str | None) -> Adjudicator: + caller = _default_grounding_caller() + + def adjudicate(**kwargs: Any) -> Mapping[str, Any]: + candidates = [ + { + key: deepcopy(candidate[key]) + for key in ( + "candidate_id", + "draft", + "scene_request", + "success_spec", + "vote_count", + ) + } + for candidate in kwargs["candidates"] + ] + allowed = [str(candidate["candidate_id"]) for candidate in candidates] + schema = { + "title": "ActionEngineTaskAdjudication", + "type": "object", + "additionalProperties": False, + "required": ["candidate_id"], + "properties": { + "candidate_id": {"type": "string", "enum": allowed}, + }, + } + prompt = ( + "Select exactly one already verified, fully bindable task candidate " + "that best matches the instruction and redacted scene manifest. Do " + "not alter a candidate or invent a new interpretation. Return only " + "candidate_id.\n\n" + f"Instruction:\n{kwargs['instruction']}\n\n" + "Candidates:\n" + f"{json.dumps(candidates, ensure_ascii=False, sort_keys=True)}\n\n" + "Redacted scene manifest:\n" + f"{json.dumps(kwargs['scene_manifest'], ensure_ascii=False, sort_keys=True)}" + ) + try: + return caller(prompt=prompt, schema=schema, model=model) + except (TypeError, ValueError) as exc: + raise SceneAdapterProtocolError( + f"Task adjudication returned invalid structured output: {exc}" + ) from exc + + return adjudicate + + +def _ground_candidate( + candidate: TaskCandidate, + *, + instruction: str, + inventory: SceneInventory, + scene_objects: Sequence[Mapping[str, Any]], + model: str | None, + caller: GroundingCaller, + force_most_likely: bool, +) -> tuple[dict[str, Any], dict[str, tuple[str, ...]] | None]: + responses: list[Any] = [] + + def audited_caller(**kwargs: Any) -> Mapping[str, Any]: + call_kwargs = dict(kwargs) + if force_most_likely: + call_kwargs["prompt"] = ( + f"{kwargs['prompt']}\n\nFINAL BINDING OVERRIDE: do not return " + "ambiguous merely because confidence is low. Choose the most " + "likely existing UID that satisfies the supplied structured " + "role, affordance, state, and attribute metadata. Return " + "candidate UIDs in descending likelihood order. Do not invent, " + "add, delete, move, or modify any scene object. Use not_found " + "when no structurally compatible existing object is plausible." + ) + response = caller(**call_kwargs) + responses.append(deepcopy(response)) + if force_most_likely: + return _force_most_likely_response(response, candidate=candidate) + return response + + candidate_id = str(candidate["candidate_id"]) + try: + result = ground_scene_references( + instruction=instruction, + intent=candidate["draft"], + inventory=inventory, + scene_objects=scene_objects, + model=model, + caller=audited_caller, + ) + except (TypeError, ValueError) as exc: + if responses: + audits = _audit_unresolved_response( + responses[-1], + candidate=candidate, + inventory=inventory, + error=str(exc), + ) + status = _candidate_status(audits) + return ( + _candidate_audit(candidate, status, audits, [str(exc)]), + None, + ) + raise SceneAdapterProtocolError( + f"Grounding candidate {candidate_id!r} failed before returning JSON: {exc}" + ) from exc + + raw_bindings = result.bindings + response_by_id = _response_bindings(responses[-1], candidate=candidate) + self_reference_reasons = _self_reference_reasons(candidate["draft"], raw_bindings) + reference_audits = [] + incompatible: set[str] = set() + request_by_id = { + str(request["reference_id"]): request + for request in candidate["scene_request"]["references"] + } + for reference_id, uids in raw_bindings.items(): + compatibility_reasons = _compatibility_reasons( + request_by_id[reference_id], + uids, + inventory=inventory, + draft=candidate["draft"], + ) + compatibility_reasons.extend(self_reference_reasons.get(reference_id, ())) + compatibility_reasons = sorted(set(compatibility_reasons)) + if compatibility_reasons: + incompatible.add(reference_id) + response = response_by_id[reference_id] + audit_reasons = list(compatibility_reasons) + if ( + force_most_likely + and response.get("status") == "ambiguous" + and response.get("uids") + ): + audit_reasons.append( + "Forced the highest-ranked structurally compatible UID from an " + "ambiguous low-confidence response." + ) + reference_audits.append( + { + "reference_id": reference_id, + "status": ("incompatible" if compatibility_reasons else "resolved"), + "confidence": float(response["confidence"]), + "candidate_uids": list(response["uids"]), + "selected_uids": ([] if compatibility_reasons else list(uids)), + "reasons": audit_reasons, + } + ) + if incompatible: + reasons = [ + f"Reference {reference_id!r} conflicts with authoritative scene semantics." + for reference_id in sorted(incompatible) + ] + return ( + _candidate_audit(candidate, "incompatible", reference_audits, reasons), + None, + ) + return _candidate_audit(candidate, "resolved", reference_audits, []), dict( + raw_bindings + ) + + +def _force_most_likely_response( + response: Mapping[str, Any], + *, + candidate: TaskCandidate, +) -> Mapping[str, Any]: + """Turn ranked low-confidence UID hypotheses into explicit selections.""" + if not isinstance(response, Mapping) or set(response) != {"bindings"}: + return response + requests = { + str(item["reference_id"]): item + for item in candidate["scene_request"]["references"] + } + raw_bindings = response.get("bindings") + if not isinstance(raw_bindings, Sequence) or isinstance(raw_bindings, (str, bytes)): + return response + result = deepcopy(dict(response)) + values = [] + for raw in raw_bindings: + if not isinstance(raw, Mapping): + return response + item = deepcopy(dict(raw)) + request = requests.get(str(item.get("reference_id", ""))) + uids = item.get("uids") + if ( + request is not None + and item.get("status") in {"resolved", "ambiguous"} + and isinstance(uids, Sequence) + and not isinstance(uids, (str, bytes)) + and uids + ): + quantifier = str(request["quantifier"]) + count = int(request["count"]) + if quantifier == "one": + item["uids"] = list(uids[:1]) + elif quantifier == "count": + item["uids"] = list(uids[:count]) + item["status"] = "resolved" + confidence = item.get("confidence") + if isinstance(confidence, (int, float)) and not isinstance( + confidence, bool + ): + item["confidence"] = max(0.5, float(confidence)) + values.append(item) + result["bindings"] = values + return result + + +def _response_bindings( + response: Any, + *, + candidate: TaskCandidate, +) -> dict[str, Mapping[str, Any]]: + expected = { + str(request["reference_id"]) + for request in candidate["scene_request"]["references"] + } + if not isinstance(response, Mapping) or set(response) != {"bindings"}: + raise SceneAdapterProtocolError( + "Grounding response must contain only bindings." + ) + values = response["bindings"] + if not isinstance(values, Sequence) or isinstance(values, (str, bytes)): + raise SceneAdapterProtocolError("Grounding response bindings must be a list.") + result: dict[str, Mapping[str, Any]] = {} + for raw in values: + if not isinstance(raw, Mapping): + raise SceneAdapterProtocolError( + "Every grounding binding must be a mapping." + ) + reference_id = raw.get("reference_id") + if not isinstance(reference_id, str) or reference_id not in expected: + raise SceneAdapterProtocolError( + "Grounding response contains an unknown reference ID." + ) + if reference_id in result: + raise SceneAdapterProtocolError( + "Grounding response contains duplicate reference IDs." + ) + result[reference_id] = raw + if set(result) != expected: + raise SceneAdapterProtocolError( + "Grounding response omitted requested reference IDs." + ) + return result + + +def _audit_unresolved_response( + response: Any, + *, + candidate: TaskCandidate, + inventory: SceneInventory, + error: str, +) -> list[dict[str, Any]]: + by_id = _response_bindings(response, candidate=candidate) + audits: list[dict[str, Any]] = [] + for request in candidate["scene_request"]["references"]: + reference_id = str(request["reference_id"]) + raw = by_id[reference_id] + if set(raw) != {"reference_id", "status", "uids", "confidence"}: + raise SceneAdapterProtocolError( + f"Grounding binding {reference_id!r} has unsupported fields." + ) + status = raw["status"] + if status not in {"resolved", "ambiguous", "not_found"}: + raise SceneAdapterProtocolError( + f"Grounding binding {reference_id!r} has invalid status." + ) + uids = raw["uids"] + confidence = raw["confidence"] + if ( + not isinstance(uids, Sequence) + or isinstance(uids, (str, bytes)) + or any( + not isinstance(uid, str) or uid not in inventory.by_uid for uid in uids + ) + or len(set(uids)) != len(uids) + ): + raise SceneAdapterProtocolError( + f"Grounding binding {reference_id!r} has invalid candidate UIDs." + ) + if ( + isinstance(confidence, bool) + or not isinstance(confidence, (int, float)) + or not 0.0 <= float(confidence) <= 1.0 + ): + raise SceneAdapterProtocolError( + f"Grounding binding {reference_id!r} has invalid confidence." + ) + audit_status = status + reasons: list[str] = [] + if status == "not_found" and uids: + raise SceneAdapterProtocolError( + f"Grounding binding {reference_id!r} status=not_found requires no UIDs." + ) + if status == "resolved": + audit_status = "incompatible" + reasons.append(error) + else: + reasons.append(f"Grounding returned status={status}.") + audits.append( + { + "reference_id": reference_id, + "status": audit_status, + "confidence": float(confidence), + "candidate_uids": list(uids), + "selected_uids": [], + "reasons": reasons, + } + ) + return audits + + +def _compatibility_reasons( + request: Mapping[str, Any], + uids: Sequence[str], + *, + inventory: SceneInventory, + draft: Mapping[str, Any], +) -> list[str]: + entities = [inventory.by_uid[uid] for uid in uids] + reasons: list[str] = [] + role = str(request["role"]) + step = next(item for item in draft["steps"] if item["id"] == request["step_id"]) + try: + if role == "object": + validate_source_compatibility(str(step["task_type"]), entities) + else: + for entity in entities: + validate_target_compatibility( + str(step["task_type"]), + entity, + relation=str(step["relation"]), + ) + except ValueError as exc: + reasons.append(str(exc)) + + expected_structure = str(request["source_structure"]) + for entity in entities: + # Source structure is strict for manipulated objects. Target structure + # is relation-dependent and is already checked by + # validate_target_compatibility; a table support surface must not be + # rejected merely because it is passive rather than a rigid object. + if role == "object": + if expected_structure == "articulation" and entity.role != "articulation": + reasons.append( + f"UID {entity.uid!r} is not an articulation as requested." + ) + if expected_structure in { + "rigid_object", + "movable", + } and entity.role not in { + "object", + "rigid_object", + }: + reasons.append( + f"UID {entity.uid!r} is not a movable rigid object as requested." + ) + required_affordances = set(request["affordances"]) + if entity.affordances: + missing = required_affordances - set(entity.affordances) + if missing: + reasons.append( + f"UID {entity.uid!r} explicitly lacks affordances {sorted(missing)}." + ) + for key, expected in request["initial_state"].items(): + if key in entity.initial_state and entity.initial_state[key] != expected: + reasons.append( + f"UID {entity.uid!r} state {key!r} conflicts with the request." + ) + for key, expected in request["attributes"].items(): + if key in entity.attributes and entity.attributes[key] != expected: + reasons.append( + f"UID {entity.uid!r} attribute {key!r} conflicts with the request." + ) + return sorted(set(reasons)) + + +def _self_reference_reasons( + draft: Mapping[str, Any], + bindings: Mapping[str, Sequence[str]], +) -> dict[str, list[str]]: + """Reject object/target identity overlap, including step_result selectors.""" + objects_by_step: dict[str, tuple[str, ...]] = {} + reasons: dict[str, list[str]] = {} + for step in draft["steps"]: + step_id = str(step["id"]) + object_uids = _selector_uids( + step["object"], + reference_id=f"{step_id}.object", + bindings=bindings, + objects_by_step=objects_by_step, + ) + target_uids = _selector_uids( + step["target"], + reference_id=f"{step_id}.target", + bindings=bindings, + objects_by_step=objects_by_step, + ) + overlap = sorted(set(object_uids) & set(target_uids)) + if overlap: + reason = ( + f"Grounding step {step_id!r} uses the same UID as object and " + f"target: {overlap}." + ) + for role in ("object", "target"): + selector = step[role] + if selector["kind"] == "scene_ref": + reasons.setdefault(f"{step_id}.{role}", []).append(reason) + objects_by_step[step_id] = object_uids + return reasons + + +def _selector_uids( + selector: Mapping[str, Any], + *, + reference_id: str, + bindings: Mapping[str, Sequence[str]], + objects_by_step: Mapping[str, tuple[str, ...]], +) -> tuple[str, ...]: + kind = str(selector["kind"]) + if kind == "scene_ref": + return tuple(str(uid) for uid in bindings[reference_id]) + if kind == "step_result": + return objects_by_step[str(selector["step_id"])] + return () + + +def _candidate_audit( + candidate: TaskCandidate, + status: str, + references: Sequence[Mapping[str, Any]], + reasons: Sequence[str], +) -> dict[str, Any]: + return { + "candidate_id": candidate["candidate_id"], + "semantic_hash": candidate["semantic_hash"], + "status": status, + "references": [deepcopy(dict(reference)) for reference in references], + "reasons": list(reasons), + } + + +def _candidate_status(references: Sequence[Mapping[str, Any]]) -> str: + statuses = {str(reference["status"]) for reference in references} + if statuses == {"resolved"}: + return "resolved" + if "ambiguous" in statuses: + return "ambiguous" + if "incompatible" in statuses: + return "incompatible" + return "not_found" + + +def _select_candidate( + candidates: Sequence[TaskCandidate], + audits: Sequence[Mapping[str, Any]], + *, + manifest: SceneManifest, + instruction: str, + adjudicator: Adjudicator | None, +) -> tuple[str | None, str, str]: + audit_by_id = {str(audit["candidate_id"]): audit for audit in audits} + bound = [ + candidate + for candidate in candidates + if audit_by_id[str(candidate["candidate_id"])]["status"] == "resolved" + ] + majority = [candidate for candidate in bound if int(candidate["vote_count"]) >= 2] + if len(majority) == 1: + return str(majority[0]["candidate_id"]), "bound", "majority_bindable" + if not majority and len(bound) == 1: + return str(bound[0]["candidate_id"]), "bound", "unique_bindable" + + choices = majority if majority else bound + if len(choices) > 1: + if adjudicator is None: + return None, "ambiguous", "multiple_conflicting_bindable_candidates" + raw = adjudicator( + instruction=instruction, + candidates=deepcopy(list(choices)), + scene_manifest=deepcopy(manifest), + ) + if not isinstance(raw, Mapping) or set(raw) != {"candidate_id"}: + raise SceneAdapterProtocolError( + "Adjudicator response must contain only candidate_id." + ) + selected_id = raw["candidate_id"] + allowed = {str(candidate["candidate_id"]) for candidate in choices} + if not isinstance(selected_id, str) or selected_id not in allowed: + raise SceneAdapterProtocolError( + "Adjudicator must select the candidate_id of a verified bindable candidate." + ) + return selected_id, "bound", "adjudicated_bindable" + if any(audit["status"] == "ambiguous" for audit in audits): + return None, "ambiguous", "no_fully_bound_candidate" + return None, "unsatisfied", "no_fully_bound_candidate" + + +def _build_manifest( + prepared: PreparedScene, + inventory: SceneInventory, + *, + source_format: str, +) -> SceneManifest: + objects = [ + { + "uid": entity.uid, + "role": entity.role, + "name": entity.name, + "description": entity.description, + "category": entity.category, + "color": entity.color, + "affordances": sorted(entity.affordances), + "initial_state": _redact_semantics(entity.initial_state), + "attributes": _redact_semantics(entity.attributes), + } + for entity in sorted(inventory.entities, key=lambda item: item.uid) + ] + scene_id = _canonical_hash( + { + "source_format": source_format, + "objects": objects, + "asset_hashes": prepared.asset_hashes, + "rotation": prepared.z_rotation_degrees, + "body_scale_policy": prepared.body_scale_policy, + "body_scale": prepared.body_scale, + } + ) + return validate_scene_manifest( + { + "schema_version": SCENE_MANIFEST_SCHEMA, + "scene_id": scene_id, + "source_format": source_format, + "robot_profile": inventory.profile, + "objects": objects, + } + ) + + +def _build_semantic_manifest( + inventory: SceneInventory, + *, + source_format: str, +) -> SceneManifest: + objects = [ + { + "uid": entity.uid, + "role": entity.role, + "name": entity.name, + "description": entity.description, + "category": entity.category, + "color": entity.color, + "affordances": sorted(entity.affordances), + "initial_state": _redact_semantics(entity.initial_state), + "attributes": _redact_semantics(entity.attributes), + } + for entity in sorted(inventory.entities, key=lambda item: item.uid) + ] + return validate_scene_manifest( + { + "schema_version": SCENE_MANIFEST_SCHEMA, + "scene_id": _canonical_hash( + {"source_format": source_format, "objects": objects} + ), + "source_format": source_format, + "robot_profile": inventory.profile, + "objects": objects, + } + ) + + +def _redact_semantics(value: Mapping[str, Any]) -> dict[str, Any]: + result: dict[str, Any] = {} + for key, child in value.items(): + name = str(key) + normalized = name.strip().lower().replace("-", "_") + if normalized in _REDACTED_KEYS: + continue + if isinstance(child, Mapping): + nested = _redact_semantics(child) + if nested: + result[name] = nested + elif isinstance(child, (str, int, float, bool)) and not isinstance( + child, complex + ): + result[name] = child + elif isinstance(child, Sequence) and not isinstance(child, (str, bytes)): + simple = [item for item in child if isinstance(item, (str, bool))] + if len(simple) == len(child): + result[name] = simple + return result + + +def _canonical_hash(value: Any) -> str: + payload = json.dumps( + value, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() diff --git a/embodichain/gen_sim/task_engine/orchestration/scene_source.py b/embodichain/gen_sim/task_engine/orchestration/scene_source.py new file mode 100644 index 000000000..54941bf7c --- /dev/null +++ b/embodichain/gen_sim/task_engine/orchestration/scene_source.py @@ -0,0 +1,276 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Read-only references and integrity checks for existing Gym projects.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +import hashlib +import json +import os +from pathlib import Path +from typing import Any +from urllib.parse import unquote, urlparse +import xml.etree.ElementTree as ET + +from embodichain.gen_sim.action_engine.generation.source_scene import ( + resolve_source_scene, +) + +__all__ = [ + "SceneSourceFingerprint", + "SceneSourceRef", + "fingerprint_scene_source", + "scene_revision_id", + "verify_scene_source_fingerprint", +] + +_SCENE_SECTIONS = ("background", "rigid_object", "articulation") + + +@dataclass(frozen=True) +class SceneSourceRef: + """Reference an existing scene without copying or owning its files.""" + + path: Path | str + robot_profile: str = "franka" + z_rotation_degrees: float | None = None + body_scale_policy: str = "preserve" + body_scale: tuple[float, float, float] = (1.0, 1.0, 1.0) + + def __post_init__(self) -> None: + object.__setattr__(self, "path", Path(self.path).expanduser()) + + +@dataclass(frozen=True) +class SceneSourceFingerprint: + """Content evidence for one externally owned scene source.""" + + source_format: str + config_path: Path + config_sha256: str + asset_sha256: dict[str, str] + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-safe audit view.""" + return { + "source_format": self.source_format, + "config_path": self.config_path.as_posix(), + "config_sha256": self.config_sha256, + "asset_sha256": dict(sorted(self.asset_sha256.items())), + } + + +def fingerprint_scene_source( + source: SceneSourceRef | str | Path, +) -> SceneSourceFingerprint: + """Hash a source config and referenced assets without copying either.""" + source_path = source.path if isinstance(source, SceneSourceRef) else source + resolved = resolve_source_scene(source_path) + config_bytes = resolved.path.read_bytes() + try: + config = json.loads(config_bytes) + except json.JSONDecodeError as exc: + raise ValueError(f"Scene config is not valid JSON: {resolved.path}") from exc + if not isinstance(config, Mapping): + raise ValueError(f"Scene config must contain an object: {resolved.path}") + + asset_hashes: dict[str, str] = {} + for section in _SCENE_SECTIONS: + entries = config.get(section, ()) + if not isinstance(entries, Sequence) or isinstance(entries, (str, bytes)): + continue + for index, entry in enumerate(entries): + if not isinstance(entry, Mapping): + continue + references: list[tuple[str, Any]] = [] + shape = entry.get("shape") + if isinstance(shape, Mapping) and shape.get("fpath"): + references.append(("shape.fpath", shape["fpath"])) + if section == "articulation" and entry.get("fpath"): + references.append(("fpath", entry["fpath"])) + for field_name, reference in references: + asset_path = Path(str(reference)).expanduser() + if not asset_path.is_absolute(): + asset_path = resolved.path.parent / asset_path + asset_path = asset_path.resolve() + if not asset_path.is_file(): + raise FileNotFoundError( + f"Scene asset does not exist: {asset_path} " + f"({section}[{index}].{field_name})." + ) + for dependency in _asset_dependency_files(asset_path): + asset_hashes[dependency.as_posix()] = _sha256( + dependency.read_bytes() + ) + return SceneSourceFingerprint( + source_format=resolved.source_format, + config_path=resolved.path, + config_sha256=_sha256(config_bytes), + asset_sha256=asset_hashes, + ) + + +def verify_scene_source_fingerprint(expected: Mapping[str, Any]) -> None: + """Raise when an externally owned source changed after preparation.""" + required = {"source_format", "config_path", "config_sha256", "asset_sha256"} + if set(expected) != required: + raise ValueError("Scene source fingerprint fields are invalid.") + actual = fingerprint_scene_source(str(expected["config_path"])).to_dict() + normalized = { + "source_format": str(expected["source_format"]), + "config_path": Path(str(expected["config_path"])).resolve().as_posix(), + "config_sha256": str(expected["config_sha256"]), + "asset_sha256": dict(expected["asset_sha256"]), + } + if actual != normalized: + raise RuntimeError( + "Source Gym project changed after Task Engine preparation; " + "prepare a new bundle before running it." + ) + + +def scene_revision_id(source: SceneSourceRef | str | Path) -> str: + """Return a location-independent content identity for one scene revision. + + Volatile exporter IDs and absolute asset paths are excluded. Referenced + asset content remains part of the identity through SHA-256 placeholders. + + Args: + source: Scene project, configuration path, or Task Engine source reference. + + Returns: + Stable SHA-256 identity of scene semantics and referenced asset content. + """ + source_path = source.path if isinstance(source, SceneSourceRef) else source + resolved = resolve_source_scene(source_path) + try: + config = json.loads(resolved.path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise ValueError(f"Scene config is not valid JSON: {resolved.path}") from exc + if not isinstance(config, Mapping): + raise ValueError(f"Scene config must contain an object: {resolved.path}") + normalized = _normalize_revision_value( + dict(config), + config_root=resolved.path.parent, + ) + normalized.pop("scene_id", None) + payload = json.dumps( + normalized, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ).encode("utf-8") + return _sha256(payload) + + +def _normalize_revision_value(value: Any, *, config_root: Path) -> Any: + if isinstance(value, Mapping): + result = { + str(key): _normalize_revision_value(item, config_root=config_root) + for key, item in value.items() + } + for key in ("fpath",): + raw = result.get(key) + if not isinstance(raw, str) or not raw: + continue + path = Path(raw).expanduser() + if not path.is_absolute(): + path = config_root / path + path = path.resolve() + if path.is_file(): + files = _asset_dependency_files(path) + result[key] = { + "sha256": _sha256(path.read_bytes()), + "dependency_sha256": { + Path( + os.path.relpath(dependency, start=path.parent) + ).as_posix(): (_sha256(dependency.read_bytes())) + for dependency in files + if dependency != path + }, + } + return result + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + return [ + _normalize_revision_value(item, config_root=config_root) for item in value + ] + return value + + +def _asset_dependency_files(asset_path: Path) -> tuple[Path, ...]: + """Return one asset and every local XML-declared dependency transitively.""" + pending = [asset_path.resolve()] + visited: set[Path] = set() + while pending: + path = pending.pop() + if path in visited: + continue + if not path.is_file(): + raise FileNotFoundError(f"Scene asset dependency does not exist: {path}") + visited.add(path) + if path.suffix.lower() not in {".urdf", ".xml", ".mjcf", ".xacro"}: + continue + try: + root = ET.parse(path).getroot() + except ET.ParseError: + # Opaque articulation assets remain valid direct dependencies even + # when their extension suggests XML. + continue + for element in root.iter(): + tag = element.tag.rsplit("}", maxsplit=1)[-1] + if tag not in {"mesh", "texture", "include"}: + continue + for attribute in ("filename", "file", "url"): + reference = element.attrib.get(attribute) + if reference: + pending.append(_resolve_asset_reference(path, reference)) + return tuple(sorted(visited)) + + +def _resolve_asset_reference(owner: Path, reference: str) -> Path: + """Resolve a local filesystem or ROS package URI without global state.""" + parsed = urlparse(reference) + if parsed.scheme in {"http", "https", "data"}: + raise ValueError( + f"Remote scene asset dependencies cannot be integrity-hashed: {reference}" + ) + if parsed.scheme == "file": + return Path(unquote(parsed.path)).expanduser().resolve() + if parsed.scheme == "package": + package_name = parsed.netloc + relative = Path(unquote(parsed.path.lstrip("/"))) + candidates = [ + ancestor / package_name / relative + for ancestor in (owner.parent, *owner.parents) + ] + candidates.append(owner.parent / relative) + for candidate in candidates: + if candidate.is_file(): + return candidate.resolve() + raise FileNotFoundError( + f"Unable to resolve package asset {reference!r} from {owner}." + ) + if parsed.scheme: + raise ValueError(f"Unsupported scene asset URI scheme: {reference}") + return (owner.parent / unquote(reference)).expanduser().resolve() + + +def _sha256(value: bytes) -> str: + return hashlib.sha256(value).hexdigest() diff --git a/embodichain/gen_sim/task_engine/run_directory.py b/embodichain/gen_sim/task_engine/run_directory.py new file mode 100644 index 000000000..0092cffc0 --- /dev/null +++ b/embodichain/gen_sim/task_engine/run_directory.py @@ -0,0 +1,87 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Collision-safe allocation of human-readable Task Engine run directories.""" + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from typing import Iterator + +__all__ = ["RunDirectory", "reserve_run_directory"] + + +@dataclass(frozen=True) +class RunDirectory: + """One reserved run identifier and its not-yet-published destination.""" + + run_id: str + output_root: Path + path: Path + created_at: datetime + + +@contextmanager +def reserve_run_directory( + output_root: str | Path, + *, + now: datetime | None = None, +) -> Iterator[RunDirectory]: + """Reserve a timestamped child name without creating its destination. + + Args: + output_root: Persistent task-history directory. + now: Optional timezone-aware timestamp used by deterministic tests. + + Yields: + A run directory allocation safe to publish through ArtifactTransaction. + """ + created_at = now or datetime.now().astimezone() + if created_at.tzinfo is None or created_at.utcoffset() is None: + raise ValueError("Task Engine run timestamps must include a timezone.") + root = Path(output_root).expanduser().resolve() + root.mkdir(parents=True, exist_ok=True) + if not root.is_dir(): + raise NotADirectoryError(root) + + base = created_at.strftime("%Y%m%d_%H%M%S") + for collision_index in range(10_000): + run_id = base if collision_index == 0 else f"{base}_{collision_index:02d}" + destination = root / run_id + reservation = root / f".{run_id}.reserve" + if destination.exists(): + continue + try: + reservation.mkdir() + except FileExistsError: + continue + if destination.exists(): + reservation.rmdir() + continue + try: + yield RunDirectory( + run_id=run_id, + output_root=root, + path=destination, + created_at=created_at, + ) + finally: + reservation.rmdir() + return + raise RuntimeError("Unable to reserve a Task Engine run directory.") diff --git a/embodichain/gen_sim/task_engine/scene/__init__.py b/embodichain/gen_sim/task_engine/scene/__init__.py new file mode 100644 index 000000000..40a218b57 --- /dev/null +++ b/embodichain/gen_sim/task_engine/scene/__init__.py @@ -0,0 +1,53 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Task Engine ownership of scene adaptation and static feasibility.""" + +from __future__ import annotations + +from .contracts import ( + ASSESSMENT_STATUSES, + FEASIBILITY_REPORT_SCHEMA, + STATIC_SCENE_MANIFEST_SCHEMA, + FeasibilityReport, + StaticSceneManifest, + validate_feasibility_report, + validate_static_scene_manifest, +) +from .feasibility import FeasibilityBroker +from .scene_engine_v1 import SceneEngineV1Adapter +from .conservative_graph import ( + CONSERVATIVE_SCENE_GRAPH_SCHEMA, + ConservativeSceneGraph, + build_conservative_scene_graph, + validate_conservative_scene_graph, +) + +__all__ = [ + "ASSESSMENT_STATUSES", + "CONSERVATIVE_SCENE_GRAPH_SCHEMA", + "FEASIBILITY_REPORT_SCHEMA", + "STATIC_SCENE_MANIFEST_SCHEMA", + "FeasibilityBroker", + "FeasibilityReport", + "SceneEngineV1Adapter", + "StaticSceneManifest", + "ConservativeSceneGraph", + "build_conservative_scene_graph", + "validate_feasibility_report", + "validate_static_scene_manifest", + "validate_conservative_scene_graph", +] diff --git a/embodichain/gen_sim/task_engine/scene/conservative_graph.py b/embodichain/gen_sim/task_engine/scene/conservative_graph.py new file mode 100644 index 000000000..ac3f2a634 --- /dev/null +++ b/embodichain/gen_sim/task_engine/scene/conservative_graph.py @@ -0,0 +1,247 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Conservative hierarchy evidence for imported scenes.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +import json +from pathlib import Path +from typing import Any, Final, TypeAlias + +__all__ = [ + "CONSERVATIVE_SCENE_GRAPH_SCHEMA", + "ConservativeSceneGraph", + "build_conservative_scene_graph", + "validate_conservative_scene_graph", +] + +CONSERVATIVE_SCENE_GRAPH_SCHEMA: Final = "embodichain.conservative-scene-graph/v1" +ConservativeSceneGraph: TypeAlias = dict[str, Any] + + +def build_conservative_scene_graph( + prepared_scene: Any, + *, + scene_id: str, +) -> ConservativeSceneGraph: + """Use exported hierarchy when available and mark every gap as unknown.""" + source_path = Path(getattr(prepared_scene, "source_config_path")).resolve() + uid_map = dict(getattr(prepared_scene, "uid_map", {}) or {}) + exported = _read_exported_graph(source_path.with_name("scene_graph.json")) + operational_assumptions = _legacy_operational_assumption_uids(source_path) + exported_nodes = { + str(node.get("object_id")): node + for node in exported.get("nodes", ()) + if isinstance(node, Mapping) and node.get("object_id") + } + + nodes: list[dict[str, Any]] = [] + for raw in getattr(prepared_scene, "planner_objects"): + uid = str(raw.get("uid", "")) + source_uid = str(raw.get("source_uid", uid)) + known = exported_nodes.get(source_uid) or exported_nodes.get(uid) + attributes = raw.get("attributes", {}) + final_support = ( + attributes.get("final_support") if isinstance(attributes, Mapping) else None + ) + initial_state = raw.get("initial_state", {}) + final_orientation = ( + initial_state.get("orientation") + if isinstance(initial_state, Mapping) + else None + ) + if uid == "table": + node = { + "uid": uid, + "parent_uid": None, + "parent_relation": "root", + "orientation": "unknown", + "source": "structural_root", + } + elif isinstance(final_support, Mapping): + relation = str(final_support.get("relation", "unknown")) + parent_uid = final_support.get("parent_uid", "unknown") + node = { + "uid": uid, + "parent_uid": ( + str(parent_uid) + if isinstance(parent_uid, str) and parent_uid + else "unknown" + ), + "parent_relation": relation if relation == "on" else "unknown", + "orientation": ( + "standing" + if final_orientation == "upright" + else ("lying" if final_orientation == "fallen" else "unknown") + ), + "source": "final_inspection", + } + elif ( + known is None + or uid in operational_assumptions + or source_uid in operational_assumptions + ): + node = { + "uid": uid, + "parent_uid": "unknown", + "parent_relation": "unknown", + "orientation": "unknown", + "source": "conservative_import", + } + else: + raw_parent = known.get("parent_id") + parent_uid = ( + uid_map.get(str(raw_parent), str(raw_parent)) + if raw_parent is not None + else "unknown" + ) + relation = known.get("parent_relation") + orientation = known.get("orientation_state") + node = { + "uid": uid, + "parent_uid": parent_uid, + "parent_relation": relation if relation == "on" else "unknown", + "orientation": ( + orientation if orientation in {"standing", "lying"} else "unknown" + ), + "source": "scene_graph", + } + nodes.append(node) + + relations = [] + for raw in exported.get("relations", ()): + if not isinstance(raw, Mapping): + continue + source_uid = uid_map.get(str(raw.get("source_id")), str(raw.get("source_id"))) + target_uid = uid_map.get(str(raw.get("target_id")), str(raw.get("target_id"))) + relation = str(raw.get("relation", "")) + if source_uid and target_uid and relation: + relations.append( + { + "source_uid": source_uid, + "relation": relation, + "target_uid": target_uid, + "source": "scene_graph", + } + ) + return validate_conservative_scene_graph( + { + "schema_version": CONSERVATIVE_SCENE_GRAPH_SCHEMA, + "scene_id": str(scene_id), + "nodes": nodes, + "relations": relations, + } + ) + + +def _legacy_operational_assumption_uids(source_path: Path) -> set[str]: + manifest_path = source_path.parent.parent / "legacy_conversion.json" + if not manifest_path.is_file(): + return set() + try: + value = json.loads(manifest_path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise ValueError( + f"Legacy conversion manifest is invalid JSON: {manifest_path}" + ) from exc + if not isinstance(value, Mapping): + raise ValueError("Legacy conversion manifest must contain an object.") + assumptions = value.get("assumptions", ()) + if not isinstance(assumptions, Sequence) or isinstance(assumptions, (str, bytes)): + raise ValueError("Legacy conversion assumptions must be a sequence.") + return { + str(item["uid"]) + for item in assumptions + if isinstance(item, Mapping) and isinstance(item.get("uid"), str) + } + + +def validate_conservative_scene_graph( + value: Mapping[str, Any], +) -> ConservativeSceneGraph: + """Validate and detach one conservative graph.""" + if not isinstance(value, Mapping): + raise TypeError("ConservativeSceneGraph must be a mapping.") + result = deepcopy(dict(value)) + expected = {"schema_version", "scene_id", "nodes", "relations"} + if set(result) != expected: + raise ValueError("ConservativeSceneGraph fields are invalid.") + if result.get("schema_version") != CONSERVATIVE_SCENE_GRAPH_SCHEMA: + raise ValueError("ConservativeSceneGraph schema version is invalid.") + if not isinstance(result.get("scene_id"), str) or not result["scene_id"]: + raise ValueError("ConservativeSceneGraph.scene_id must not be empty.") + nodes = _sequence(result.get("nodes"), "nodes") + normalized_nodes = [] + for index, raw in enumerate(nodes): + if not isinstance(raw, Mapping): + raise TypeError(f"ConservativeSceneGraph.nodes[{index}] must be a mapping.") + node = dict(raw) + if set(node) != { + "uid", + "parent_uid", + "parent_relation", + "orientation", + "source", + }: + raise ValueError( + f"ConservativeSceneGraph.nodes[{index}] fields are invalid." + ) + if not isinstance(node["uid"], str) or not node["uid"]: + raise ValueError(f"ConservativeSceneGraph.nodes[{index}].uid is invalid.") + if node["parent_uid"] is not None and not isinstance(node["parent_uid"], str): + raise TypeError( + f"ConservativeSceneGraph.nodes[{index}].parent_uid is invalid." + ) + if node["parent_relation"] not in {"root", "on", "unknown"}: + raise ValueError( + f"ConservativeSceneGraph.nodes[{index}].parent_relation is invalid." + ) + if node["orientation"] not in {"standing", "lying", "unknown"}: + raise ValueError( + f"ConservativeSceneGraph.nodes[{index}].orientation is invalid." + ) + if not isinstance(node["source"], str) or not node["source"]: + raise ValueError( + f"ConservativeSceneGraph.nodes[{index}].source is invalid." + ) + normalized_nodes.append(node) + if len({node["uid"] for node in normalized_nodes}) != len(normalized_nodes): + raise ValueError("ConservativeSceneGraph node UIDs must be unique.") + result["nodes"] = normalized_nodes + result["relations"] = [ + dict(item) for item in _sequence(result.get("relations"), "relations") + ] + json.dumps(result, allow_nan=False) + return result + + +def _read_exported_graph(path: Path) -> dict[str, Any]: + if not path.is_file(): + return {} + try: + value = json.loads(path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise ValueError(f"Scene graph is not valid JSON: {path}") from exc + return dict(value) if isinstance(value, Mapping) else {} + + +def _sequence(value: Any, field_name: str) -> list[Any]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise TypeError(f"ConservativeSceneGraph.{field_name} must be a sequence.") + return list(value) diff --git a/embodichain/gen_sim/task_engine/scene/contracts.py b/embodichain/gen_sim/task_engine/scene/contracts.py new file mode 100644 index 000000000..974864785 --- /dev/null +++ b/embodichain/gen_sim/task_engine/scene/contracts.py @@ -0,0 +1,304 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""JSON contracts owned by the Scene Engine anti-corruption boundary.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +import json +import math +from typing import Any, TypeAlias + +__all__ = [ + "ASSESSMENT_STATUSES", + "FEASIBILITY_REPORT_SCHEMA", + "REMEDIATION_CLASSES", + "STATIC_SCENE_MANIFEST_SCHEMA", + "FeasibilityReport", + "StaticSceneManifest", + "validate_feasibility_report", + "validate_static_scene_manifest", +] + + +STATIC_SCENE_MANIFEST_SCHEMA = "embodichain.static-scene-manifest/v1" +FEASIBILITY_REPORT_SCHEMA = "embodichain.scene-action-feasibility/v2" +ASSESSMENT_STATUSES = frozenset({"proven", "runtime_probe", "unknown", "contradicted"}) +REMEDIATION_CLASSES = frozenset( + {"none", "scene_remediable", "action_capability", "input_conflict", "terminal"} +) +_EVIDENCE_STATUSES = frozenset({"declared", "inferred", "verified", "contradicted"}) + +StaticSceneManifest: TypeAlias = dict[str, Any] +FeasibilityReport: TypeAlias = dict[str, Any] + + +def validate_static_scene_manifest(value: Mapping[str, Any]) -> StaticSceneManifest: + """Validate and detach one static scene manifest.""" + result = _mapping(value, "StaticSceneManifest") + _exact_keys( + result, + { + "schema_version", + "scene_id", + "source_format", + "robot_profile", + "source", + "adapter_capabilities", + "objects", + }, + "StaticSceneManifest", + ) + _schema(result, STATIC_SCENE_MANIFEST_SCHEMA, "StaticSceneManifest") + for key in ("scene_id", "source_format", "robot_profile"): + result[key] = _nonempty(result.get(key), f"StaticSceneManifest.{key}") + result["source"] = _mapping(result.get("source"), "StaticSceneManifest.source") + result["adapter_capabilities"] = _bool_mapping( + result.get("adapter_capabilities"), + "StaticSceneManifest.adapter_capabilities", + ) + + objects: list[dict[str, Any]] = [] + for index, raw in enumerate(_sequence(result.get("objects"), "objects")): + context = f"StaticSceneManifest.objects[{index}]" + item = _mapping(raw, context) + _exact_keys( + item, + { + "uid", + "source_uid", + "role", + "name", + "description", + "category", + "color", + "geometry", + "initial_pose", + "physics", + "articulation", + "affordances", + "initial_state", + "attributes", + "provenance", + }, + context, + ) + item["uid"] = _nonempty(item.get("uid"), f"{context}.uid") + item["source_uid"] = _string(item.get("source_uid"), f"{context}.source_uid") + item["role"] = _nonempty(item.get("role"), f"{context}.role") + for key in ("name", "description", "category"): + item[key] = _string(item.get(key), f"{context}.{key}") + color = item.get("color") + if color is not None: + color = _string(color, f"{context}.color") + item["color"] = color + for key in ( + "geometry", + "initial_pose", + "physics", + "articulation", + "initial_state", + "attributes", + "provenance", + ): + item[key] = _mapping(item.get(key), f"{context}.{key}") + item["affordances"] = [ + _validate_affordance(evidence, f"{context}.affordances[{evidence_index}]") + for evidence_index, evidence in enumerate( + _sequence(item.get("affordances"), f"{context}.affordances") + ) + ] + objects.append(item) + uids = [item["uid"] for item in objects] + if len(set(uids)) != len(uids): + raise ValueError("StaticSceneManifest object UIDs must be unique.") + result["objects"] = objects + _json_safe(result, "StaticSceneManifest") + return result + + +def validate_feasibility_report(value: Mapping[str, Any]) -> FeasibilityReport: + """Validate and detach one scene/action feasibility report.""" + result = _mapping(value, "FeasibilityReport") + _exact_keys( + result, + { + "schema_version", + "task_id", + "candidate_id", + "scene_id", + "status", + "remediation_class", + "checks", + "blockers", + "summary", + }, + "FeasibilityReport", + ) + _schema(result, FEASIBILITY_REPORT_SCHEMA, "FeasibilityReport") + for key in ("task_id", "candidate_id", "scene_id"): + result[key] = _nonempty(result.get(key), f"FeasibilityReport.{key}") + result["status"] = _status(result.get("status"), "FeasibilityReport.status") + remediation_class = result.get("remediation_class") + if remediation_class not in REMEDIATION_CLASSES: + raise ValueError( + "FeasibilityReport.remediation_class must be one of " + f"{sorted(REMEDIATION_CLASSES)}." + ) + result["remediation_class"] = str(remediation_class) + if result["status"] != "contradicted" and remediation_class != "none": + raise ValueError( + "A non-contradicted FeasibilityReport requires remediation_class=none." + ) + checks: list[dict[str, Any]] = [] + for index, raw in enumerate(_sequence(result.get("checks"), "checks")): + context = f"FeasibilityReport.checks[{index}]" + item = _mapping(raw, context) + _exact_keys( + item, + {"kind", "subject", "status", "reason", "evidence"}, + context, + ) + item["kind"] = _nonempty(item.get("kind"), f"{context}.kind") + item["subject"] = _nonempty(item.get("subject"), f"{context}.subject") + item["status"] = _status(item.get("status"), f"{context}.status") + item["reason"] = _nonempty(item.get("reason"), f"{context}.reason") + item["evidence"] = _mapping(item.get("evidence"), f"{context}.evidence") + checks.append(item) + result["checks"] = checks + blockers = _sequence(result.get("blockers"), "FeasibilityReport.blockers") + if any(not isinstance(item, str) or not item for item in blockers): + raise ValueError("FeasibilityReport.blockers must contain non-empty strings.") + result["blockers"] = list(blockers) + summary = _mapping(result.get("summary"), "FeasibilityReport.summary") + expected = set(ASSESSMENT_STATUSES) + if set(summary) != expected: + raise ValueError( + "FeasibilityReport.summary must count every assessment status." + ) + if any( + isinstance(value, bool) or not isinstance(value, int) or value < 0 + for value in summary.values() + ): + raise ValueError( + "FeasibilityReport.summary counts must be non-negative integers." + ) + if sum(summary.values()) != len(checks): + raise ValueError("FeasibilityReport.summary must match the check count.") + result["summary"] = dict(summary) + _json_safe(result, "FeasibilityReport") + return result + + +def _validate_affordance(value: Any, context: str) -> dict[str, Any]: + item = _mapping(value, context) + _exact_keys( + item, + { + "type", + "status", + "confidence", + "source", + "link_uid", + "frame", + "parameters", + }, + context, + ) + item["type"] = _nonempty(item.get("type"), f"{context}.type") + status = item.get("status") + if status not in _EVIDENCE_STATUSES: + raise ValueError( + f"{context}.status must be one of {sorted(_EVIDENCE_STATUSES)}." + ) + confidence = item.get("confidence") + if confidence is not None: + if ( + isinstance(confidence, bool) + or not isinstance(confidence, (int, float)) + or not math.isfinite(float(confidence)) + or not 0.0 <= float(confidence) <= 1.0 + ): + raise ValueError(f"{context}.confidence must be null or in [0, 1].") + confidence = float(confidence) + item["confidence"] = confidence + item["source"] = _nonempty(item.get("source"), f"{context}.source") + item["link_uid"] = _string(item.get("link_uid"), f"{context}.link_uid") + item["frame"] = _mapping(item.get("frame"), f"{context}.frame") + item["parameters"] = _mapping(item.get("parameters"), f"{context}.parameters") + return item + + +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 _exact_keys(value: Mapping[str, Any], expected: set[str], context: str) -> None: + if set(value) != expected: + missing = sorted(expected - set(value)) + extra = sorted(set(value) - expected) + raise ValueError(f"{context} fields differ; missing={missing}, extra={extra}.") + + +def _schema(value: Mapping[str, Any], expected: str, context: str) -> None: + if value.get("schema_version") != expected: + raise ValueError(f"{context}.schema_version must be {expected!r}.") + + +def _nonempty(value: Any, context: str) -> str: + result = _string(value, context).strip() + if not result: + raise ValueError(f"{context} must not be empty.") + return result + + +def _string(value: Any, context: str) -> str: + if not isinstance(value, str): + raise TypeError(f"{context} must be a string.") + return value + + +def _status(value: Any, context: str) -> str: + if value not in ASSESSMENT_STATUSES: + raise ValueError(f"{context} must be one of {sorted(ASSESSMENT_STATUSES)}.") + return str(value) + + +def _bool_mapping(value: Any, context: str) -> dict[str, bool]: + result = _mapping(value, context) + if any( + not isinstance(key, str) or not isinstance(item, bool) + for key, item in result.items() + ): + raise TypeError(f"{context} must map strings to booleans.") + return result + + +def _json_safe(value: Any, context: str) -> None: + try: + json.dumps(value, allow_nan=False) + except (TypeError, ValueError) as exc: + raise ValueError(f"{context} must contain strict JSON data.") from exc diff --git a/embodichain/gen_sim/task_engine/scene/feasibility.py b/embodichain/gen_sim/task_engine/scene/feasibility.py new file mode 100644 index 000000000..7eaefc162 --- /dev/null +++ b/embodichain/gen_sim/task_engine/scene/feasibility.py @@ -0,0 +1,667 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Deterministic task, scene, robot, and action-capability intersection.""" + +from __future__ import annotations + +from collections import Counter +from collections.abc import Mapping, Sequence +import math +from typing import Any + +from .contracts import ( + ASSESSMENT_STATUSES, + FEASIBILITY_REPORT_SCHEMA, + FeasibilityReport, + validate_feasibility_report, + validate_static_scene_manifest, +) + +__all__ = ["FeasibilityBroker"] + + +_STATUS_PRIORITY = { + "proven": 0, + "runtime_probe": 1, + "unknown": 2, + "contradicted": 3, +} + + +class FeasibilityBroker: + """Produce an auditable compatibility report without repairing inputs.""" + + def assess( + self, + candidate: Mapping[str, Any], + role_bindings: Mapping[str, Any], + scene_manifest: Mapping[str, Any], + *, + capability_catalog: Mapping[str, Mapping[str, Any]], + task_actions: Mapping[str, Sequence[str]], + ) -> FeasibilityReport: + """Assess one grounded candidate against static and runtime capabilities.""" + manifest = validate_static_scene_manifest(scene_manifest) + draft = _mapping(candidate.get("draft"), "candidate.draft") + scene_request = _mapping( + candidate.get("scene_request"), "candidate.scene_request" + ) + bindings = role_bindings.get("reference_bindings", role_bindings) + bindings = _mapping(bindings, "role_bindings.reference_bindings") + objects = {item["uid"]: item for item in manifest["objects"]} + steps = { + str(item["id"]): item + for item in _sequence(draft.get("steps"), "candidate.draft.steps") + } + checks: list[dict[str, Any]] = [] + + for step_id, step in steps.items(): + task_type = str(step.get("task_type", "")) + actions = task_actions.get(task_type) + if not actions: + checks.append( + _check( + "task_capability", + step_id, + "contradicted", + f"Task type {task_type!r} has no registered action recipe.", + ) + ) + continue + for action_name in actions: + capability = capability_catalog.get(str(action_name)) + if capability is None: + checks.append( + _check( + "atomic_capability", + f"{step_id}:{action_name}", + "contradicted", + f"AtomicAction {action_name!r} is not registered.", + ) + ) + elif not bool(capability.get("runtime_available", False)): + checks.append( + _check( + "atomic_capability", + f"{step_id}:{action_name}", + "contradicted", + str( + capability.get("unavailable_reason") + or "Action is planning-only." + ), + evidence={"action": str(action_name)}, + ) + ) + else: + checks.append( + _check( + "atomic_capability", + f"{step_id}:{action_name}", + "proven", + "AtomicAction is registered and executable.", + evidence={"action": str(action_name)}, + ) + ) + + for request in _sequence( + scene_request.get("references"), "candidate.scene_request.references" + ): + reference_id = str(request.get("reference_id", "")) + raw_uids = bindings.get(reference_id, ()) + if not isinstance(raw_uids, Sequence) or isinstance(raw_uids, (str, bytes)): + raw_uids = () + uids = [str(uid) for uid in raw_uids] + if not uids: + checks.append( + _check( + "binding", + reference_id, + "contradicted", + "Reference has no grounded scene entity.", + ) + ) + continue + for uid in uids: + entity = objects.get(uid) + if entity is None: + checks.append( + _check( + "binding", + f"{reference_id}:{uid}", + "contradicted", + "Binding references an entity absent from the static manifest.", + ) + ) + continue + checks.extend(self._entity_checks(request, entity, reference_id)) + + checks.extend(self._workspace_checks(steps, bindings, objects)) + + statuses = Counter(check["status"] for check in checks) + status = max( + (check["status"] for check in checks), + key=_STATUS_PRIORITY.__getitem__, + default="unknown", + ) + blockers = sorted( + { + f"{check['subject']}: {check['reason']}" + for check in checks + if check["status"] == "contradicted" + } + ) + return validate_feasibility_report( + { + "schema_version": FEASIBILITY_REPORT_SCHEMA, + "task_id": str(draft.get("task_id", "")), + "candidate_id": str(candidate.get("candidate_id", "")), + "scene_id": manifest["scene_id"], + "status": status, + "remediation_class": _remediation_class(checks), + "checks": checks, + "blockers": blockers, + "summary": { + name: int(statuses.get(name, 0)) + for name in sorted(ASSESSMENT_STATUSES) + }, + } + ) + + def _entity_checks( + self, + request: Mapping[str, Any], + entity: Mapping[str, Any], + reference_id: str, + ) -> list[dict[str, Any]]: + uid = str(entity["uid"]) + subject = f"{reference_id}:{uid}" + checks = [self._structure_check(request, entity, subject)] + evidence_by_type: dict[str, list[Mapping[str, Any]]] = {} + for evidence in entity["affordances"]: + evidence_by_type.setdefault(str(evidence["type"]), []).append(evidence) + for affordance in request.get("affordances", ()): + name = str(affordance) + checks.append( + self._affordance_check(name, evidence_by_type.get(name, ()), subject) + ) + for field_name in ("initial_state", "attributes"): + required = request.get(field_name, {}) + actual = entity.get(field_name, {}) + if isinstance(required, Mapping) and isinstance(actual, Mapping): + for key, expected in required.items(): + if key not in actual: + status = "unknown" + reason = f"Required {field_name} field {key!r} is not declared." + elif actual[key] != expected: + status = "contradicted" + reason = f"Required {field_name} field {key!r} conflicts with the scene." + else: + status = "proven" + reason = f"Required {field_name} field {key!r} matches." + checks.append( + _check( + field_name, + subject, + status, + reason, + evidence={"field": str(key)}, + ) + ) + if str(request.get("role")) == "object": + checks.append( + _check( + "runtime_reachability", + subject, + "runtime_probe", + "Reachability, collision, and grasp geometry require live planning.", + ) + ) + if ( + str(request.get("role")) == "target" + and str(request.get("source_structure")) == "physical_entity" + ): + checks.append( + _check( + "placement_support", + subject, + "runtime_probe", + "Support depends on the payload, candidate pose, live geometry, " + "and post-release stability.", + evidence={ + "runtime_obligations": [ + "placement_candidates", + "object_supported_by", + "stable_for", + "final_support_revalidation", + ] + }, + ) + ) + return checks + + def _workspace_checks( + self, + steps: Mapping[str, Mapping[str, Any]], + bindings: Mapping[str, Any], + objects: Mapping[str, Mapping[str, Any]], + ) -> list[dict[str, Any]]: + """Defer arm-side compatibility to the live robot frame.""" + checks: list[dict[str, Any]] = [] + object_uids_by_step: dict[str, tuple[str, ...]] = {} + phases: list[dict[str, Any]] = [] + for step_id, step in steps.items(): + object_uids = _step_selector_uids( + step_id, + "object", + step.get("object"), + bindings, + object_uids_by_step, + ) + object_uids_by_step[step_id] = object_uids + target_uids = _step_selector_uids( + step_id, + "target", + step.get("target"), + bindings, + object_uids_by_step, + ) + task_type = str(step.get("task_type", "")) + required_arm = str(step.get("required_arm", "auto")) + if task_type == "E4": + required_arm = str(step.get("transfer_arm", "none")) + if required_arm in {"left_arm", "right_arm"}: + for uid in object_uids: + entity = objects.get(uid) + position = ( + entity.get("initial_pose", {}).get("position", ()) + if isinstance(entity, Mapping) + and isinstance(entity.get("initial_pose"), Mapping) + else () + ) + if ( + not isinstance(position, Sequence) + or isinstance(position, (str, bytes, bytearray)) + or len(position) < 2 + ): + continue + checks.append( + _check( + "arm_layout_risk", + f"{step_id}:{uid}", + "runtime_probe", + "Arm-side compatibility requires live left/right arm-base " + "poses and workspace geometry.", + evidence={ + "required_arm": required_arm, + "object_world_position": [ + float(position[0]), + float(position[1]), + ], + "arm_side_frame": "live_robot", + "mismatch_risk": None, + "geometry_certificate": False, + }, + ) + ) + + phases.extend( + _workflow_phases( + step_id, + task_type, + object_uids, + target_uids, + transfer_arm=str(step.get("transfer_arm", "none")), + receive_arm=str(step.get("receive_arm", "none")), + ) + ) + if phases: + checks.append( + _check( + "task_workspace", + "task_workflow", + "runtime_probe", + "Scene layout must satisfy pickup, transfer, placement, and " + "safety-clearance phases across the complete task workflow.", + evidence={ + "arm_side_frame": "live_robot", + "phases": phases, + "geometry_certificate": False, + }, + ) + ) + return checks + + @staticmethod + def _structure_check( + request: Mapping[str, Any], + entity: Mapping[str, Any], + subject: str, + ) -> dict[str, Any]: + expected = str(request.get("source_structure", "")) + role = str(entity.get("role", "")) + if expected in {"scene_entity", "spatial_reference"}: + if role in {"camera", "light", "robot", "sensor"}: + return _check( + "structure", + subject, + "contradicted", + f"Scene entity role {role!r} cannot be a spatial action target.", + evidence={"static_pose": _has_static_pose(entity)}, + ) + if not _has_static_pose(entity): + return _check( + "structure", + subject, + "unknown", + "Static scene evidence does not provide a finite spatial pose.", + evidence={"static_pose": False}, + ) + if role == "articulation": + return _check( + "structure", + subject, + "runtime_probe", + "Articulation has a static pose, but live spatial target lookup " + "must be validated at runtime.", + evidence={ + "static_pose": True, + "runtime_entity_kind": "articulation", + }, + ) + has_runtime_body = bool(entity.get("physics")) + if ( + role + in { + "background", + "object", + "rigid_object", + "support_surface", + "table", + } + and has_runtime_body + ): + return _check( + "structure", + subject, + "proven", + "Scene entity has a static pose and a rigid runtime body.", + evidence={ + "static_pose": True, + "runtime_entity_kind": "rigid_object", + }, + ) + return _check( + "structure", + subject, + "runtime_probe", + "Scene entity has a static pose, but its live target interface is " + "not proven by the static manifest.", + evidence={"static_pose": True, "runtime_entity_kind": "unknown"}, + ) + if expected == "physical_entity": + geometry = entity.get("geometry", {}) + shape = geometry.get("shape", {}) if isinstance(geometry, Mapping) else {} + asset_sha256 = ( + geometry.get("asset_sha256", "") + if isinstance(geometry, Mapping) + else "" + ) + physics = entity.get("physics", {}) + articulation = entity.get("articulation", {}) + has_physical_geometry = bool(shape) or bool(asset_sha256) + has_runtime_body = bool(physics) or bool(articulation) + if role in {"camera", "light", "sensor"}: + return _check( + "structure", + subject, + "contradicted", + f"Scene entity role {role!r} is not a physical collision body.", + evidence={"physical_geometry": False, "runtime_body": False}, + ) + if role == "articulation" or bool(articulation): + return _check( + "structure", + subject, + "contradicted", + "Placement on an articulation requires a link-level target " + "interface that the current runtime does not provide.", + evidence={ + "physical_geometry": has_physical_geometry, + "runtime_body": bool(articulation), + "runtime_entity_kind": "articulation", + "runtime_target_interface": False, + }, + ) + if has_physical_geometry and has_runtime_body: + return _check( + "structure", + subject, + "proven", + "Scene entity has physical geometry and a runtime body.", + evidence={"physical_geometry": True, "runtime_body": True}, + ) + return _check( + "structure", + subject, + "unknown", + "Static scene evidence does not prove physical geometry and a " + "runtime body required for placement.", + evidence={ + "physical_geometry": has_physical_geometry, + "runtime_body": has_runtime_body, + }, + ) + accepted_by_structure = { + "articulation": {"articulation"}, + "rigid_object": {"object", "rigid_object"}, + "movable": {"object", "rigid_object"}, + "support_surface": {"background", "support_surface", "table"}, + } + accepted = accepted_by_structure.get(expected) + if accepted is None: + return _check( + "structure", + subject, + "unknown", + f"Structure contract {expected!r} is not recognized by the broker.", + evidence={"scene_role": role}, + ) + if role in accepted: + return _check( + "structure", + subject, + "proven", + f"Scene role {role!r} satisfies structure {expected!r}.", + ) + return _check( + "structure", + subject, + "contradicted", + f"Scene role {role!r} does not satisfy structure {expected!r}.", + ) + + @staticmethod + def _affordance_check( + name: str, + evidence: Sequence[Mapping[str, Any]], + subject: str, + ) -> dict[str, Any]: + if not evidence: + return _check( + "affordance", + subject, + "unknown", + f"Affordance {name!r} has no evidence.", + evidence={"affordance": name}, + ) + statuses = {str(item.get("status")) for item in evidence} + if statuses == {"contradicted"}: + status = "contradicted" + reason = f"Affordance {name!r} is explicitly contradicted." + elif "verified" in statuses: + status = "proven" + reason = f"Affordance {name!r} has verified evidence." + else: + status = "runtime_probe" + reason = ( + f"Affordance {name!r} is declared but requires physical validation." + ) + return _check( + "affordance", + subject, + status, + reason, + evidence={ + "affordance": name, + "sources": sorted({str(item.get("source")) for item in evidence}), + }, + ) + + +def _has_static_pose(entity: Mapping[str, Any]) -> bool: + pose = entity.get("initial_pose") + if not isinstance(pose, Mapping): + return False + position = pose.get("position") + if ( + not isinstance(position, Sequence) + or isinstance(position, (str, bytes, bytearray)) + or len(position) != 3 + ): + return False + return all( + not isinstance(value, bool) + and isinstance(value, (int, float)) + and math.isfinite(float(value)) + for value in position + ) + + +def _step_selector_uids( + step_id: str, + role: str, + selector: Any, + bindings: Mapping[str, Any], + object_uids_by_step: Mapping[str, tuple[str, ...]], +) -> tuple[str, ...]: + """Resolve direct and prior-step selectors for static workspace advice.""" + raw = bindings.get(f"{step_id}.{role}", ()) + if isinstance(raw, Sequence) and not isinstance(raw, (str, bytes, bytearray)): + direct = tuple(str(uid) for uid in raw if str(uid)) + if direct: + return direct + if not isinstance(selector, Mapping) or selector.get("kind") != "step_result": + return () + source_step = str(selector.get("step_id", "")) + return tuple(object_uids_by_step.get(source_step, ())) + + +def _workflow_phases( + step_id: str, + task_type: str, + object_uids: Sequence[str], + target_uids: Sequence[str], + *, + transfer_arm: str, + receive_arm: str, +) -> list[dict[str, Any]]: + """Describe whole-task layout anchors without inventing geometry bounds.""" + phases: list[dict[str, Any]] = [] + if object_uids: + phases.append( + { + "step_id": step_id, + "phase": "pickup", + "object_uids": list(object_uids), + } + ) + if task_type == "E4": + phases.append( + { + "step_id": step_id, + "phase": "handover_shared_workspace", + "object_uids": list(object_uids), + "transfer_arm": transfer_arm, + "receive_arm": receive_arm, + } + ) + if target_uids: + phases.append( + { + "step_id": step_id, + "phase": "target_interaction", + "object_uids": list(object_uids), + "target_uids": list(target_uids), + } + ) + if task_type in {"E1", "E2", "E3", "E4", "E5"}: + phases.append( + { + "step_id": step_id, + "phase": "safety_clearance", + "object_uids": list(object_uids), + } + ) + return phases + + +def _check( + kind: str, + subject: str, + status: str, + reason: str, + *, + evidence: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + return { + "kind": kind, + "subject": subject, + "status": status, + "reason": reason, + "evidence": dict(evidence or {}), + } + + +def _remediation_class(checks: Sequence[Mapping[str, Any]]) -> str: + """Classify contradictions by the subsystem capable of changing them.""" + contradicted = [check for check in checks if check.get("status") == "contradicted"] + if not contradicted: + return "none" + kinds = {str(check.get("kind", "")) for check in contradicted} + if kinds.intersection({"task_capability", "atomic_capability"}): + return "action_capability" + # A new materialization seed can change observed pose/orientation, but it + # cannot change task semantics, entity roles, bindings, or declared affordances. + if kinds <= {"initial_state"}: + return "scene_remediable" + if kinds.intersection({"binding", "structure", "affordance", "attributes"}): + return "input_conflict" + return "terminal" + + +def _mapping(value: Any, context: str) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise TypeError(f"{context} must be a mapping.") + return dict(value) + + +def _sequence(value: Any, context: str) -> list[Mapping[str, Any]]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise TypeError(f"{context} must be a sequence.") + if any(not isinstance(item, Mapping) for item in value): + raise TypeError(f"{context} must contain mappings.") + return list(value) diff --git a/embodichain/gen_sim/task_engine/scene/final_inspection.py b/embodichain/gen_sim/task_engine/scene/final_inspection.py new file mode 100644 index 000000000..f7bf3b68f --- /dev/null +++ b/embodichain/gen_sim/task_engine/scene/final_inspection.py @@ -0,0 +1,428 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Geometry-derived evidence from one completed scene revision.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +from dataclasses import replace +import json +from pathlib import Path +from typing import Any, Final, TypeAlias + +import numpy as np +from scipy.spatial.transform import Rotation +import trimesh + +from embodichain.gen_sim.action_engine.generation.models import PreparedScene +from embodichain.gen_sim.action_engine.generation.source_scene import ( + prepare_scene, + resolve_source_scene, +) + +__all__ = [ + "FINAL_SCENE_INSPECTION_SCHEMA", + "FinalSceneInspection", + "apply_final_inspection", + "inspect_final_scene", + "validate_final_scene_inspection", +] + +FINAL_SCENE_INSPECTION_SCHEMA: Final = "embodichain.final-scene-inspection/v1" +FinalSceneInspection: TypeAlias = dict[str, Any] + +_Y_UP_TO_Z_UP = np.array( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]], + dtype=float, +) + + +def inspect_final_scene( + source: str | Path, + *, + revision_id: str, + contact_tolerance_m: float = 0.03, +) -> FinalSceneInspection: + """Measure final AABBs, orientation, and support from exported geometry. + + Args: + source: Completed scene project or configuration path. + revision_id: Content identity already assigned to the completed revision. + contact_tolerance_m: Maximum support-surface contact gap in meters. + + Returns: + Strict geometry-derived final inspection document. + """ + tolerance = float(contact_tolerance_m) + if not np.isfinite(tolerance) or tolerance <= 0.0: + raise ValueError("contact_tolerance_m must be positive and finite.") + normalized_revision_id = str(revision_id) + if len(normalized_revision_id) != 64: + raise ValueError("revision_id must be a SHA-256 hexadecimal digest.") + try: + int(normalized_revision_id, 16) + except ValueError as exc: + raise ValueError("revision_id must be a SHA-256 hexadecimal digest.") from exc + resolved = resolve_source_scene(source) + prepared = prepare_scene(source) + runtime = { + str(item.get("uid")): item + for item in ( + *prepared.background, + *prepared.rigid_objects, + *prepared.articulations, + ) + if isinstance(item, Mapping) and item.get("uid") + } + measured: dict[str, dict[str, Any]] = {} + for raw in prepared.planner_objects: + uid = str(raw.get("uid", "")) + role = str(raw.get("role", "")) + geometry = _measure_geometry( + runtime.get(uid, raw), + convert_y_up=resolved.is_prompt2scene, + ) + measured[uid] = { + "uid": uid, + "role": role, + "orientation": _orientation(geometry), + "support": { + "parent_uid": None if uid == "table" else "unknown", + "relation": "root" if uid == "table" else "unknown", + "confidence": 1.0 if uid == "table" else None, + "gap_m": None, + "xy_overlap_ratio": None, + }, + "world_aabb": ( + None + if geometry is None + else { + "min": geometry["bounds"][0].tolist(), + "max": geometry["bounds"][1].tolist(), + } + ), + "evidence": { + "source": "final_geometry" if geometry is not None else "unmeasured", + "method": "world_aabb_and_dominant_axis", + }, + } + + for uid, item in measured.items(): + child_geometry = _geometry_from_record(item) + if uid == "table" or child_geometry is None: + continue + support = _support_for( + uid, + child_geometry, + measured, + tolerance=tolerance, + ) + if support is not None: + item["support"] = support + + return validate_final_scene_inspection( + { + "schema_version": FINAL_SCENE_INSPECTION_SCHEMA, + "scene_revision_id": normalized_revision_id, + "source_config_path": prepared.source_config_path.as_posix(), + "contact_tolerance_m": tolerance, + "objects": [measured[uid] for uid in sorted(measured)], + } + ) + + +def apply_final_inspection( + prepared_scene: PreparedScene, + inspection: Mapping[str, Any], +) -> PreparedScene: + """Return a detached PreparedScene enriched with measured final evidence. + + Args: + prepared_scene: Normalized scene to enrich without mutation. + inspection: Validated or raw final inspection mapping. + + Returns: + Prepared scene whose semantic state reflects measured final geometry. + """ + normalized = validate_final_scene_inspection(inspection) + by_uid = {str(item["uid"]): item for item in normalized["objects"]} + planner_objects = [] + for raw in prepared_scene.planner_objects: + item = deepcopy(raw) + evidence = by_uid.get(str(item.get("uid"))) + if evidence is not None: + initial_state = deepcopy(dict(item.get("initial_state", {}))) + initial_state.pop("orientation", None) + if evidence["orientation"] == "standing": + initial_state["orientation"] = "upright" + elif evidence["orientation"] == "lying": + initial_state["orientation"] = "fallen" + attributes = deepcopy(dict(item.get("attributes", {}))) + attributes["final_support"] = deepcopy(evidence["support"]) + attributes["final_world_aabb"] = deepcopy(evidence["world_aabb"]) + item["initial_state"] = initial_state + item["attributes"] = attributes + planner_objects.append(item) + return replace(prepared_scene, planner_objects=tuple(planner_objects)) + + +def validate_final_scene_inspection( + value: Mapping[str, Any], +) -> FinalSceneInspection: + """Validate and detach one final scene inspection document. + + Args: + value: Inspection mapping to validate. + + Returns: + Detached, normalized inspection document. + """ + if not isinstance(value, Mapping): + raise TypeError("FinalSceneInspection must be a mapping.") + result = deepcopy(dict(value)) + expected = { + "schema_version", + "scene_revision_id", + "source_config_path", + "contact_tolerance_m", + "objects", + } + if set(result) != expected: + raise ValueError("FinalSceneInspection fields are invalid.") + if result.get("schema_version") != FINAL_SCENE_INSPECTION_SCHEMA: + raise ValueError("FinalSceneInspection schema version is invalid.") + revision_id = result.get("scene_revision_id") + if not isinstance(revision_id, str) or len(revision_id) != 64: + raise ValueError("FinalSceneInspection.scene_revision_id is invalid.") + try: + int(revision_id, 16) + except ValueError as exc: + raise ValueError("FinalSceneInspection.scene_revision_id is invalid.") from exc + source_path = result.get("source_config_path") + if not isinstance(source_path, str) or not source_path: + raise ValueError("FinalSceneInspection.source_config_path is invalid.") + tolerance = result.get("contact_tolerance_m") + if ( + isinstance(tolerance, bool) + or not isinstance(tolerance, (int, float)) + or not np.isfinite(float(tolerance)) + or float(tolerance) <= 0.0 + ): + raise ValueError("FinalSceneInspection.contact_tolerance_m is invalid.") + objects = result.get("objects") + if not isinstance(objects, Sequence) or isinstance(objects, (str, bytes)): + raise TypeError("FinalSceneInspection.objects must be a sequence.") + normalized = [_validate_object(item, index) for index, item in enumerate(objects)] + if len({item["uid"] for item in normalized}) != len(normalized): + raise ValueError("FinalSceneInspection object UIDs must be unique.") + result["objects"] = normalized + result["contact_tolerance_m"] = float(tolerance) + json.dumps(result, ensure_ascii=False, allow_nan=False) + return result + + +def _validate_object(value: Any, index: int) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise TypeError(f"FinalSceneInspection.objects[{index}] must be a mapping.") + item = deepcopy(dict(value)) + expected = {"uid", "role", "orientation", "support", "world_aabb", "evidence"} + if set(item) != expected: + raise ValueError(f"FinalSceneInspection.objects[{index}] fields are invalid.") + if not isinstance(item["uid"], str) or not item["uid"]: + raise ValueError(f"FinalSceneInspection.objects[{index}].uid is invalid.") + if not isinstance(item["role"], str) or not item["role"]: + raise ValueError(f"FinalSceneInspection.objects[{index}].role is invalid.") + if item["orientation"] not in {"standing", "lying", "unknown"}: + raise ValueError( + f"FinalSceneInspection.objects[{index}].orientation is invalid." + ) + if not isinstance(item["support"], Mapping) or not isinstance( + item["evidence"], Mapping + ): + raise TypeError("FinalSceneInspection support and evidence must be mappings.") + support = deepcopy(dict(item["support"])) + if set(support) != { + "parent_uid", + "relation", + "confidence", + "gap_m", + "xy_overlap_ratio", + }: + raise ValueError("FinalSceneInspection support fields are invalid.") + if support["parent_uid"] is not None and not isinstance(support["parent_uid"], str): + raise TypeError("FinalSceneInspection support parent_uid is invalid.") + if support["relation"] not in {"root", "on", "unknown"}: + raise ValueError("FinalSceneInspection support relation is invalid.") + for field_name in ("confidence", "gap_m", "xy_overlap_ratio"): + field_value = support[field_name] + if field_value is not None and ( + isinstance(field_value, bool) + or not isinstance(field_value, (int, float)) + or not np.isfinite(float(field_value)) + ): + raise ValueError(f"FinalSceneInspection support {field_name} is invalid.") + if ( + support["confidence"] is not None + and not 0.0 <= float(support["confidence"]) <= 1.0 + ): + raise ValueError("FinalSceneInspection support confidence is invalid.") + if ( + support["xy_overlap_ratio"] is not None + and not 0.0 <= float(support["xy_overlap_ratio"]) <= 1.0 + 1.0e-6 + ): + raise ValueError("FinalSceneInspection support overlap is invalid.") + item["support"] = support + aabb = item["world_aabb"] + if aabb is not None: + if not isinstance(aabb, Mapping) or set(aabb) != {"min", "max"}: + raise ValueError("FinalSceneInspection world_aabb is invalid.") + if aabb["min"] is None or aabb["max"] is None: + raise ValueError("FinalSceneInspection world_aabb vectors are invalid.") + minimum = _vector(aabb["min"], default=(0.0, 0.0, 0.0)) + maximum = _vector(aabb["max"], default=(0.0, 0.0, 0.0)) + if np.any(np.asarray(maximum) < np.asarray(minimum)): + raise ValueError("FinalSceneInspection world_aabb bounds are inverted.") + item["world_aabb"] = {"min": minimum, "max": maximum} + evidence = deepcopy(dict(item["evidence"])) + if set(evidence) != {"source", "method"} or any( + not isinstance(evidence[key], str) or not evidence[key] for key in evidence + ): + raise ValueError("FinalSceneInspection evidence is invalid.") + item["evidence"] = evidence + return item + + +def _measure_geometry( + entry: Mapping[str, Any], + *, + convert_y_up: bool, +) -> dict[str, Any] | None: + shape = entry.get("shape") + if not isinstance(shape, Mapping): + return None + shape_type = str(shape.get("shape_type", "")) + if shape_type == "Mesh": + path = Path(str(shape.get("fpath", ""))).expanduser().resolve() + if not path.is_file(): + return None + loaded = trimesh.load(path, force="scene") + mesh = loaded.to_geometry() + elif shape_type == "Cube": + mesh = trimesh.creation.box( + extents=_vector(shape.get("size"), default=(1, 1, 1)) + ) + elif shape_type == "Sphere": + radius = float(shape.get("radius", 1.0)) + mesh = trimesh.creation.icosphere(radius=radius) + else: + return None + scale = np.asarray(_vector(entry.get("body_scale"), default=(1, 1, 1))) + mesh.apply_scale(scale) + local_extents = np.asarray(mesh.extents, dtype=float) + conversion = _Y_UP_TO_Z_UP if convert_y_up else np.eye(3) + rotation = Rotation.from_euler( + "XYZ", + _vector(entry.get("init_rot"), default=(0, 0, 0)), + degrees=True, + ).as_matrix() + transform = np.eye(4) + transform[:3, :3] = rotation @ conversion + transform[:3, 3] = _vector(entry.get("init_pos"), default=(0, 0, 0)) + mesh.apply_transform(transform) + return { + "bounds": np.asarray(mesh.bounds, dtype=float), + "local_extents": local_extents, + "axis_transform": transform[:3, :3], + "shape_type": shape_type, + } + + +def _orientation(geometry: Mapping[str, Any] | None) -> str: + if geometry is None or geometry["shape_type"] == "Sphere": + return "unknown" + extents = np.asarray(geometry["local_extents"], dtype=float) + ordered = np.sort(extents) + if ordered[-1] <= 0.0 or ordered[-1] / max(ordered[-2], 1.0e-9) < 1.2: + return "unknown" + dominant = int(np.argmax(extents)) + axis = np.asarray(geometry["axis_transform"], dtype=float)[:, dominant] + vertical = abs(float(axis[2])) / max(float(np.linalg.norm(axis)), 1.0e-9) + if vertical >= 0.75: + return "standing" + if vertical <= 0.35: + return "lying" + return "unknown" + + +def _geometry_from_record(item: Mapping[str, Any]) -> np.ndarray | None: + aabb = item.get("world_aabb") + if not isinstance(aabb, Mapping): + return None + return np.asarray([aabb["min"], aabb["max"]], dtype=float) + + +def _support_for( + uid: str, + child: np.ndarray, + objects: Mapping[str, Mapping[str, Any]], + *, + tolerance: float, +) -> dict[str, Any] | None: + child_bottom = float(child[0, 2]) + child_area = max( + float((child[1, 0] - child[0, 0]) * (child[1, 1] - child[0, 1])), + 1.0e-9, + ) + candidates = [] + for parent_uid, parent_item in objects.items(): + if parent_uid == uid: + continue + parent = _geometry_from_record(parent_item) + if parent is None: + continue + overlap_x = max( + 0.0, min(child[1, 0], parent[1, 0]) - max(child[0, 0], parent[0, 0]) + ) + overlap_y = max( + 0.0, min(child[1, 1], parent[1, 1]) - max(child[0, 1], parent[0, 1]) + ) + overlap_ratio = float(overlap_x * overlap_y / child_area) + gap = child_bottom - float(parent[1, 2]) + if overlap_ratio >= 0.1 and -tolerance <= gap <= tolerance: + candidates.append((overlap_ratio, -abs(gap), parent_uid, gap)) + if not candidates: + return None + overlap_ratio, _, parent_uid, gap = max(candidates) + confidence = min(1.0, overlap_ratio * max(0.0, 1.0 - abs(gap) / tolerance)) + return { + "parent_uid": parent_uid, + "relation": "on", + "confidence": float(confidence), + "gap_m": float(gap), + "xy_overlap_ratio": float(overlap_ratio), + } + + +def _vector(value: Any, *, default: tuple[float, float, float]) -> list[float]: + raw = default if value is None else value + if not isinstance(raw, Sequence) or isinstance(raw, (str, bytes)): + raise TypeError("Scene geometry vectors must be sequences.") + result = [float(item) for item in raw] + if len(result) != 3 or not np.all(np.isfinite(result)): + raise ValueError("Scene geometry vectors must contain three finite values.") + return result diff --git a/embodichain/gen_sim/task_engine/scene/scene_engine_v1.py b/embodichain/gen_sim/task_engine/scene/scene_engine_v1.py new file mode 100644 index 000000000..af0c4a47c --- /dev/null +++ b/embodichain/gen_sim/task_engine/scene/scene_engine_v1.py @@ -0,0 +1,242 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Adapt existing Scene Engine exports without changing their source schema.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +import hashlib +import json +from pathlib import Path +from typing import Any + +from .contracts import ( + STATIC_SCENE_MANIFEST_SCHEMA, + StaticSceneManifest, + validate_static_scene_manifest, +) + +__all__ = ["SceneEngineV1Adapter"] + + +class SceneEngineV1Adapter: + """Convert a normalized Scene Engine v1 export to the neutral manifest.""" + + def adapt_prepared_scene( + self, + prepared_scene: Any, + *, + source_format: str, + robot_profile: str, + ) -> StaticSceneManifest: + """Adapt the existing prepared-scene view through a duck-typed boundary.""" + planner_objects = tuple(getattr(prepared_scene, "planner_objects")) + runtime_objects = ( + tuple(getattr(prepared_scene, "background", ())) + + tuple(getattr(prepared_scene, "rigid_objects", ())) + + tuple(getattr(prepared_scene, "articulations", ())) + ) + runtime_by_uid = { + str(item.get("uid")): item + for item in runtime_objects + if isinstance(item, Mapping) and item.get("uid") + } + asset_hashes = dict(getattr(prepared_scene, "asset_hashes", {}) or {}) + objects = [ + self._object_manifest( + raw, + runtime=runtime_by_uid.get(str(raw.get("uid")), {}), + asset_sha256=str(asset_hashes.get(str(raw.get("uid")), "")), + ) + for raw in planner_objects + ] + identity = { + "source_format": str(source_format), + "robot_profile": str(robot_profile), + "objects": [_identity_object(item) for item in objects], + } + source_path = Path(getattr(prepared_scene, "source_config_path")) + return validate_static_scene_manifest( + { + "schema_version": STATIC_SCENE_MANIFEST_SCHEMA, + "scene_id": _canonical_hash(identity), + "source_format": str(source_format), + "robot_profile": str(robot_profile), + "source": { + "adapter": f"{type(self).__module__}.{type(self).__qualname__}", + "config_path": source_path.expanduser().resolve().as_posix(), + "config_sha256": _file_hash(source_path), + "asset_hashes": asset_hashes, + }, + "adapter_capabilities": { + "task_conditioned_generation": False, + "structured_affordances": any( + bool(item["affordances"]) for item in objects + ), + "articulation_instances": any( + item["role"] == "articulation" for item in objects + ), + "runtime_scene_observation": False, + }, + "objects": objects, + } + ) + + def _object_manifest( + self, + raw: Mapping[str, Any], + *, + runtime: Mapping[str, Any], + asset_sha256: str, + ) -> dict[str, Any]: + uid = str(raw.get("uid", "")).strip() + role = str(raw.get("role", "")).strip() + shape = raw.get("shape", runtime.get("shape", {})) + shape = deepcopy(dict(shape)) if isinstance(shape, Mapping) else {} + physics_keys = ("attrs", "body_type", "max_convex_hull_num") + physics = { + key: deepcopy(runtime[key]) for key in physics_keys if key in runtime + } + articulation = deepcopy(dict(runtime)) if role == "articulation" else {} + affordances = _affordance_evidence(raw.get("affordances", ())) + if role in {"background", "table", "support_surface"}: + affordances = _with_structural_evidence( + affordances, + "support_surface", + ) + if role in {"object", "rigid_object"}: + affordances = _with_structural_evidence(affordances, "rigid") + return { + "uid": uid, + "source_uid": str(raw.get("source_uid", "")), + "role": role, + "name": str(raw.get("name", "")), + "description": str(raw.get("description", "")), + "category": str(raw.get("category", "")), + "color": raw.get("color") if isinstance(raw.get("color"), str) else None, + "geometry": { + "shape": shape, + "asset_sha256": asset_sha256, + }, + "initial_pose": { + "position": deepcopy(list(raw.get("init_pos", ()))), + "rotation": deepcopy(list(raw.get("init_rot", ()))), + "scale": deepcopy(list(raw.get("body_scale", ()))), + }, + "physics": physics, + "articulation": articulation, + "affordances": affordances, + "initial_state": _mapping_or_empty(raw.get("initial_state")), + "attributes": _mapping_or_empty(raw.get("attributes")), + "provenance": { + "semantic_source": "scene_export", + "geometry_source": "prepared_scene", + "physics_source": "prepared_scene_runtime", + }, + } + + +def _affordance_evidence(value: Any) -> list[dict[str, Any]]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + return [] + result: list[dict[str, Any]] = [] + for raw in value: + if isinstance(raw, str) and raw.strip(): + result.append(_evidence(raw.strip(), status="declared")) + continue + if not isinstance(raw, Mapping): + continue + affordance_type = str(raw.get("type", raw.get("name", ""))).strip() + if not affordance_type: + continue + status = str(raw.get("status", "declared")) + result.append( + { + "type": affordance_type, + "status": status, + "confidence": raw.get("confidence"), + "source": str(raw.get("source", "scene_export")), + "link_uid": str(raw.get("link_uid", "")), + "frame": _mapping_or_empty(raw.get("frame")), + "parameters": _mapping_or_empty(raw.get("parameters")), + } + ) + return sorted(result, key=lambda item: (item["type"], item["source"])) + + +def _with_structural_evidence( + evidence: list[dict[str, Any]], affordance_type: str +) -> list[dict[str, Any]]: + if any(item["type"] == affordance_type for item in evidence): + return evidence + return sorted( + [ + *evidence, + _evidence(affordance_type, status="verified", source="adapter_structure"), + ], + key=lambda item: (item["type"], item["source"]), + ) + + +def _evidence( + affordance_type: str, + *, + status: str, + source: str = "scene_export", +) -> dict[str, Any]: + return { + "type": affordance_type, + "status": status, + "confidence": None, + "source": source, + "link_uid": "", + "frame": {}, + "parameters": {}, + } + + +def _mapping_or_empty(value: Any) -> dict[str, Any]: + return deepcopy(dict(value)) if isinstance(value, Mapping) else {} + + +def _identity_object(value: Mapping[str, Any]) -> dict[str, Any]: + result = deepcopy(dict(value)) + geometry = result.get("geometry") + if isinstance(geometry, dict): + shape = geometry.get("shape") + if isinstance(shape, dict) and geometry.get("asset_sha256"): + shape.pop("fpath", None) + return result + + +def _canonical_hash(value: Any) -> str: + payload = json.dumps( + value, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def _file_hash(path: Path) -> str: + resolved = path.expanduser().resolve() + if not resolved.is_file(): + return "" + return hashlib.sha256(resolved.read_bytes()).hexdigest() diff --git a/embodichain/gen_sim/task_engine/scene_backend.py b/embodichain/gen_sim/task_engine/scene_backend.py new file mode 100644 index 000000000..dd4a665c5 --- /dev/null +++ b/embodichain/gen_sim/task_engine/scene_backend.py @@ -0,0 +1,474 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Task-owned adapter for Scene Engine analysis, revisions, and edits.""" + +from __future__ import annotations + +from collections.abc import Mapping +from copy import deepcopy +from dataclasses import dataclass, replace +from pathlib import Path +import json +import shutil +from typing import Any + +from embodichain.gen_sim.action_engine.generation.source_scene import ( + resolve_source_scene, +) +from embodichain.gen_sim.scene_engine.pipeline import ( + SceneBlueprintPackage, + SceneMaterialization, + analyze_edit, + analyze_image, + materialize_blueprint, + materialize_edit, +) + +from .orchestration.legacy_scene import ( + convert_legacy_gym_project, + restore_locked_scene_entities, +) +from .orchestration.scene_adapter import CandidateSelection, SceneAdapter +from .orchestration.scene_source import ( + SceneSourceFingerprint, + fingerprint_scene_source, + scene_revision_id, + verify_scene_source_fingerprint, +) +from .scene.final_inspection import FinalSceneInspection, inspect_final_scene +from .workflow_contracts import TaskRunRequest, scene_input_kind + +__all__ = [ + "SceneRemediableError", + "SceneAnalysis", + "SceneEngineBackend", + "SceneRevision", + "scene_blueprint_objects", +] + +_LOCKED_SCENE_MANIFEST = "locked_scene_entities.json" + + +class SceneRemediableError(RuntimeError): + """A Scene output failure that permits a fresh materialization attempt.""" + + +@dataclass(frozen=True) +class SceneAnalysis: + """Scene semantics available before asset materialization.""" + + input_kind: str + source: Path + blueprint: SceneBlueprintPackage | None + source_fingerprint: SceneSourceFingerprint | None + + +@dataclass(frozen=True) +class SceneRevision: + """One immutable scene source selected for final Action preparation.""" + + source: Path + output_root: Path | None + revision_id: str + seed: int + edit_plan: dict[str, Any] | None + source_fingerprint: SceneSourceFingerprint | None + + +class SceneEngineBackend: + """Expose Scene Engine stages without giving it workflow ownership.""" + + def analyze( + self, + request: TaskRunRequest, + output_root: str | Path, + ) -> SceneAnalysis: + """Analyze an image or fingerprint an existing read-only project. + + Args: + request: Validated Task Engine run request. + output_root: Directory for image-understanding artifacts. + + Returns: + Scene semantics and immutable source provenance. + """ + root = Path(output_root).expanduser().resolve() + if scene_input_kind(request) == "image": + image_path = Path(str(request["image_path"])).resolve() + blueprint = analyze_image(image_path, root) + return SceneAnalysis( + input_kind="image", + source=image_path, + blueprint=blueprint, + source_fingerprint=None, + ) + source = Path(str(request["gym_project"])).resolve() + return SceneAnalysis( + input_kind="gym_project", + source=source, + blueprint=None, + source_fingerprint=fingerprint_scene_source(source), + ) + + def select( + self, + analysis: SceneAnalysis, + candidate_set: Mapping[str, Any], + scene_adapter: SceneAdapter, + *, + force_most_likely: bool, + ) -> CandidateSelection: + """Select a task candidate from blueprint or existing-scene semantics. + + Args: + analysis: Pre-materialization scene analysis. + candidate_set: Task candidates to ground and vote. + scene_adapter: Task-owned semantic binding adapter. + force_most_likely: Whether ranked UID hypotheses must be resolved. + + Returns: + Audited initial candidate selection. + """ + if analysis.blueprint is not None: + return scene_adapter.select_objects( + candidate_set, + scene_blueprint_objects(analysis.blueprint), + force_most_likely=force_most_likely, + ) + adaptation = scene_adapter.adapt( + candidate_set, + analysis.source, + force_most_likely=force_most_likely, + ) + return CandidateSelection( + scene_manifest=adaptation.scene_manifest, + role_bindings=adaptation.role_bindings, + binding_report=adaptation.binding_report, + selected_candidate=adaptation.selected_candidate, + candidate_bindings=adaptation.candidate_bindings, + ) + + def materialize( + self, + analysis: SceneAnalysis, + request: TaskRunRequest, + output_root: str | Path, + *, + seed: int, + ) -> SceneRevision: + """Produce a new revision, or return the untouched existing source. + + Args: + analysis: Pre-materialization scene analysis. + request: Validated Task Engine run request. + output_root: Fresh directory for this scene attempt. + seed: Attempt seed recorded for recovery audit. + + Returns: + Final scene source for binding and Action Engine generation. + """ + root = Path(output_root).expanduser().resolve() + edit_prompt = request["scene_edit_prompt"] + if analysis.input_kind == "image": + assert analysis.blueprint is not None + root.mkdir(parents=True, exist_ok=False) + blueprint = replace(analysis.blueprint, output_root=root) + materialization = materialize_blueprint(blueprint, seed=seed) + edit_plan = None + if edit_prompt is not None: + edit_blueprint = analyze_edit( + output_root=root, + edit_prompt=str(edit_prompt), + ) + edit_plan = edit_blueprint.scene_edit_plan.to_dict() + materialization = materialize_edit(edit_blueprint, seed=seed) + revision = _revision(materialization, seed=seed, edit_plan=edit_plan) + _write_revision_audit( + root, + revision_id=revision.revision_id, + seed=seed, + edit_plan=edit_plan, + ) + return revision + + fingerprint = analysis.source_fingerprint + assert fingerprint is not None + if edit_prompt is None: + verify_scene_source_fingerprint(fingerprint.to_dict()) + return SceneRevision( + source=analysis.source, + output_root=None, + revision_id=scene_revision_id(analysis.source), + seed=seed, + edit_plan=None, + source_fingerprint=fingerprint, + ) + + resolved = resolve_source_scene(analysis.source) + if resolved.source_format == "legacy_gym_config": + converted = convert_legacy_gym_project(analysis.source, root) + editable_root = converted.output_root + else: + editable_root = _copy_scene_export_revision(resolved.path, root) + edit_blueprint = analyze_edit( + output_root=editable_root, + edit_prompt=str(edit_prompt), + ) + edit_plan = edit_blueprint.scene_edit_plan.to_dict() + materialization = materialize_edit(edit_blueprint, seed=seed) + if resolved.source_format == "legacy_gym_config": + restore_locked_scene_entities(editable_root) + else: + _restore_scene_export_locked_entities(editable_root) + verify_scene_source_fingerprint(fingerprint.to_dict()) + _write_revision_audit( + editable_root, + revision_id=scene_revision_id(materialization.scene_config_path), + seed=seed, + edit_plan=edit_plan, + source_fingerprint=fingerprint, + ) + return SceneRevision( + source=materialization.scene_config_path, + output_root=editable_root, + revision_id=scene_revision_id(materialization.scene_config_path), + seed=seed, + edit_plan=edit_plan, + source_fingerprint=fingerprint, + ) + + def inspect( + self, + revision: SceneRevision, + output_path: str | Path, + ) -> FinalSceneInspection: + """Inspect final geometry and publish support/orientation evidence. + + Args: + revision: Completed immutable scene revision. + output_path: JSON path receiving the inspection document. + + Returns: + Validated final scene inspection. + """ + try: + actual_revision_id = scene_revision_id(revision.source) + except (OSError, TypeError, ValueError) as exc: + raise SceneRemediableError( + f"Final scene content could not be hashed: {exc}" + ) from exc + if actual_revision_id != revision.revision_id: + raise RuntimeError("Final scene changed before geometry inspection.") + try: + inspection = inspect_final_scene( + revision.source, + revision_id=actual_revision_id, + ) + except (OSError, TypeError, ValueError) as exc: + raise SceneRemediableError( + f"Final scene assets could not be inspected: {exc}" + ) from exc + path = Path(output_path).expanduser().resolve() + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(inspection, ensure_ascii=False, indent=2, allow_nan=False) + + "\n", + encoding="utf-8", + ) + return inspection + + +def scene_blueprint_objects(blueprint: SceneBlueprintPackage) -> list[dict[str, Any]]: + """Convert image semantics into the redacted grounding inventory shape. + + Args: + blueprint: Scene Engine image-understanding package. + + Returns: + Semantic objects with unknown physical fields represented conservatively. + """ + nodes = blueprint.scene_graph.node_by_id() + result = [] + for item in blueprint.scene.objects: + node = nodes.get(item.id) + orientation = None if node is None else node.orientation_state + initial_state = {} + if orientation == "lying": + initial_state["orientation"] = "fallen" + elif orientation == "standing": + initial_state["orientation"] = "upright" + result.append( + { + "uid": item.id, + "source_uid": item.id, + "role": "table" if item.kind == "table" else "rigid_object", + "name": item.name, + "description": item.description, + "category": item.category, + "color": None, + "init_pos": [0.0, 0.0, 0.0], + "affordances": [], + "initial_state": initial_state, + "attributes": {}, + } + ) + return result + + +def _copy_scene_export_revision(source_config: Path, output_root: Path) -> Path: + if output_root.exists(): + if not output_root.is_dir() or any(output_root.iterdir()): + raise ValueError("Scene revision output_root must be empty.") + source_root = source_config.parent + destination = output_root / "scene_export" + shutil.copytree(source_root, destination) + config_path = destination / "scene_config.json" + config = _read_json_mapping(config_path) + background = list(config.get("background", ())) + rigid_objects = list(config.get("rigid_object", ())) + articulations = list(config.get("articulation", ())) + editable_rigid = [ + item + for item in rigid_objects + if isinstance(item, Mapping) and _scene_editable_rigid(item) + ] + locked_rigid = [item for item in rigid_objects if item not in editable_rigid] + table = [ + item + for item in background + if isinstance(item, Mapping) and item.get("uid") == "table" + ] + if len(table) != 1: + raise ValueError("Scene export revision requires exactly one table.") + locked = { + "schema_version": "embodichain.locked-scene-entities/v1", + "background": [item for item in background if item not in table], + "rigid_object": locked_rigid, + "articulation": articulations, + } + config["background"] = table + config["rigid_object"] = editable_rigid + config["articulation"] = [] + _write_json_mapping(config_path, config) + _write_json_mapping(output_root / _LOCKED_SCENE_MANIFEST, locked) + graph_path = destination / "scene_graph.json" + if graph_path.is_file(): + graph = _read_json_mapping(graph_path) + editable_uids = {str(item.get("uid")) for item in [*table, *editable_rigid]} + graph["nodes"] = [ + item + for item in graph.get("nodes", ()) + if isinstance(item, Mapping) and item.get("object_id") in editable_uids + ] + graph["relations"] = [ + item + for item in graph.get("relations", ()) + if isinstance(item, Mapping) + and item.get("source_id") in editable_uids + and item.get("target_id") in editable_uids + ] + _write_json_mapping(graph_path, graph) + return output_root + + +def _restore_scene_export_locked_entities(output_root: Path) -> None: + manifest = _read_json_mapping(output_root / _LOCKED_SCENE_MANIFEST) + if manifest.get("schema_version") != "embodichain.locked-scene-entities/v1": + raise ValueError("Locked scene entity manifest schema is invalid.") + config_path = output_root / "scene_export" / "scene_config.json" + config = _read_json_mapping(config_path) + existing = { + str(item.get("uid")) + for section in ("background", "rigid_object", "articulation") + for item in config.get(section, ()) + if isinstance(item, Mapping) and item.get("uid") + } + for section in ("background", "rigid_object", "articulation"): + target = list(config.get(section, ())) + for raw in manifest.get(section, ()): + item = deepcopy(dict(raw)) + uid = str(item.get("uid", "")) + if not uid or uid in existing: + raise ValueError(f"Scene edit reused locked entity UID {uid!r}.") + target.append(item) + existing.add(uid) + config[section] = target + _write_json_mapping(config_path, config) + + +def _scene_editable_rigid(value: Mapping[str, Any]) -> bool: + shape = value.get("shape") + return ( + isinstance(shape, Mapping) + and shape.get("shape_type") == "Mesh" + and isinstance(shape.get("fpath"), str) + and bool(shape["fpath"]) + ) + + +def _read_json_mapping(path: Path) -> dict[str, Any]: + value = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, Mapping): + raise TypeError(f"JSON artifact must contain an object: {path}") + return dict(value) + + +def _write_json_mapping(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", + ) + + +def _revision( + value: SceneMaterialization, + *, + seed: int, + edit_plan: dict[str, Any] | None, +) -> SceneRevision: + return SceneRevision( + source=value.scene_config_path, + output_root=value.output_root, + revision_id=scene_revision_id(value.scene_config_path), + seed=seed, + edit_plan=edit_plan, + source_fingerprint=None, + ) + + +def _write_revision_audit( + output_root: Path, + *, + revision_id: str, + seed: int, + edit_plan: Mapping[str, Any] | None, + source_fingerprint: SceneSourceFingerprint | None = None, +) -> None: + payload = { + "schema_version": "embodichain.scene-revision-attempt/v1", + "revision_id": revision_id, + "seed": int(seed), + "edit_plan": None if edit_plan is None else dict(edit_plan), + "source_fingerprint": ( + None if source_fingerprint is None else source_fingerprint.to_dict() + ), + } + (output_root / "scene_revision_attempt.json").write_text( + json.dumps(payload, ensure_ascii=False, indent=2, allow_nan=False) + "\n", + encoding="utf-8", + ) diff --git a/embodichain/gen_sim/task_engine/state_machine.py b/embodichain/gen_sim/task_engine/state_machine.py new file mode 100644 index 000000000..ffdaaba15 --- /dev/null +++ b/embodichain/gen_sim/task_engine/state_machine.py @@ -0,0 +1,294 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Typed, replayable state transitions for cross-engine orchestration.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass, field +from enum import Enum +from types import MappingProxyType +from typing import Any + +from .workflow_contracts import TaskRunRequest, validate_task_run_request + +__all__ = [ + "StageStatus", + "TaskEngineState", + "WorkflowStage", + "complete_stage", + "fail_stage", + "initial_state", + "replay_events", + "skip_stage", + "start_stage", +] + + +class WorkflowStage(str, Enum): + """Stable stages shared by all four supported input combinations.""" + + INPUT = "input" + TASK_CANDIDATES = "task_candidates" + SCENE_PREPARATION = "scene_preparation" + SCENE_EDIT = "scene_edit" + CANDIDATE_SELECTION = "candidate_selection" + SCENE_FINALIZATION = "scene_finalization" + UNBOUND_ACTION = "unbound_action" + FINAL_INSPECTION = "final_inspection" + FINAL_BINDING = "final_binding" + STATIC_FEASIBILITY = "static_feasibility" + GROUNDED_ACTION = "grounded_action" + EXECUTION = "execution" + + +class StageStatus(str, Enum): + """Lifecycle of one independently schedulable workflow stage.""" + + PENDING = "pending" + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + SKIPPED = "skipped" + + +_DEPENDENCIES: dict[WorkflowStage, frozenset[WorkflowStage]] = { + WorkflowStage.INPUT: frozenset(), + WorkflowStage.TASK_CANDIDATES: frozenset({WorkflowStage.INPUT}), + WorkflowStage.SCENE_PREPARATION: frozenset({WorkflowStage.INPUT}), + WorkflowStage.SCENE_EDIT: frozenset({WorkflowStage.SCENE_PREPARATION}), + WorkflowStage.CANDIDATE_SELECTION: frozenset( + { + WorkflowStage.TASK_CANDIDATES, + WorkflowStage.SCENE_PREPARATION, + } + ), + WorkflowStage.SCENE_FINALIZATION: frozenset( + {WorkflowStage.CANDIDATE_SELECTION, WorkflowStage.SCENE_EDIT} + ), + WorkflowStage.UNBOUND_ACTION: frozenset({WorkflowStage.CANDIDATE_SELECTION}), + WorkflowStage.FINAL_INSPECTION: frozenset({WorkflowStage.SCENE_FINALIZATION}), + WorkflowStage.FINAL_BINDING: frozenset( + {WorkflowStage.FINAL_INSPECTION, WorkflowStage.UNBOUND_ACTION} + ), + WorkflowStage.STATIC_FEASIBILITY: frozenset({WorkflowStage.FINAL_BINDING}), + WorkflowStage.GROUNDED_ACTION: frozenset({WorkflowStage.STATIC_FEASIBILITY}), + WorkflowStage.EXECUTION: frozenset({WorkflowStage.GROUNDED_ACTION}), +} + +_SKIPPABLE_STAGES = frozenset({WorkflowStage.SCENE_EDIT}) + + +@dataclass(frozen=True) +class TaskEngineState: + """Immutable state snapshot plus an append-only transition audit.""" + + request: Mapping[str, Any] + stages: Mapping[WorkflowStage, StageStatus] + events: tuple[Mapping[str, Any], ...] = field(default_factory=tuple) + + def __post_init__(self) -> None: + object.__setattr__( + self, + "request", + MappingProxyType(deepcopy(dict(self.request))), + ) + object.__setattr__( + self, + "stages", + MappingProxyType(dict(self.stages)), + ) + object.__setattr__( + self, + "events", + tuple(MappingProxyType(deepcopy(dict(event))) for event in self.events), + ) + + @property + def terminal(self) -> bool: + """Return whether execution succeeded or any stage failed.""" + return ( + self.stages[WorkflowStage.EXECUTION] == StageStatus.SUCCEEDED + or StageStatus.FAILED in self.stages.values() + ) + + def to_dict(self) -> dict[str, Any]: + """Return one JSON-safe audit snapshot.""" + return { + "request": deepcopy(dict(self.request)), + "stages": { + stage.value: self.stages[stage].value for stage in WorkflowStage + }, + "events": deepcopy([dict(event) for event in self.events]), + } + + +def initial_state(request: TaskRunRequest) -> TaskEngineState: + """Create a validated state with the optional edit stage resolved.""" + normalized = validate_task_run_request(request) + stages = {stage: StageStatus.PENDING for stage in WorkflowStage} + stages[WorkflowStage.INPUT] = StageStatus.SUCCEEDED + events = ( + { + "sequence": 1, + "stage": WorkflowStage.INPUT.value, + "from": StageStatus.PENDING.value, + "to": StageStatus.SUCCEEDED.value, + }, + ) + state = TaskEngineState(request=normalized, stages=stages, events=events) + if normalized["scene_edit_prompt"] is None: + state = skip_stage(state, WorkflowStage.SCENE_EDIT) + return state + + +def start_stage(state: TaskEngineState, stage: WorkflowStage) -> TaskEngineState: + """Start a pending stage only after every dependency has completed.""" + if state.terminal: + raise ValueError("A terminal TaskEngineState cannot start another stage.") + if state.stages[stage] != StageStatus.PENDING: + raise ValueError(f"Stage {stage.value!r} is not pending.") + incomplete = [ + dependency.value + for dependency in _DEPENDENCIES[stage] + if state.stages[dependency] not in {StageStatus.SUCCEEDED, StageStatus.SKIPPED} + ] + if incomplete: + raise ValueError( + f"Stage {stage.value!r} has incomplete dependencies: {incomplete}." + ) + return _transition(state, stage, StageStatus.RUNNING) + + +def complete_stage(state: TaskEngineState, stage: WorkflowStage) -> TaskEngineState: + """Complete one running stage.""" + if state.stages[stage] != StageStatus.RUNNING: + raise ValueError(f"Stage {stage.value!r} is not running.") + return _transition(state, stage, StageStatus.SUCCEEDED) + + +def fail_stage( + state: TaskEngineState, + stage: WorkflowStage, + *, + reason: str, +) -> TaskEngineState: + """Fail a stage, including a later retry of a previously successful stage.""" + if state.terminal: + raise ValueError("A terminal TaskEngineState cannot fail another stage.") + if state.stages[stage] not in { + StageStatus.PENDING, + StageStatus.RUNNING, + StageStatus.SUCCEEDED, + }: + raise ValueError(f"Stage {stage.value!r} cannot be failed now.") + normalized_reason = str(reason).strip() + if not normalized_reason: + raise ValueError("A failed stage requires a non-empty reason.") + return _transition( + state, + stage, + StageStatus.FAILED, + details={"reason": normalized_reason}, + ) + + +def skip_stage(state: TaskEngineState, stage: WorkflowStage) -> TaskEngineState: + """Skip one optional pending stage.""" + if stage not in _SKIPPABLE_STAGES: + raise ValueError("Only the optional scene_edit stage can be skipped.") + if state.stages[stage] != StageStatus.PENDING: + raise ValueError(f"Stage {stage.value!r} is not pending.") + return _transition(state, stage, StageStatus.SKIPPED) + + +def replay_events( + request: TaskRunRequest, + events: Sequence[Mapping[str, Any]], +) -> TaskEngineState: + """Rebuild a state by validating and applying its transition audit. + + Args: + request: Original workflow request used to create the state. + events: Complete ordered event audit to validate and replay. + + Returns: + The immutable state reconstructed from the supplied audit. + + Raises: + TypeError: If the audit is not a sequence of event mappings. + ValueError: If any event is missing, altered, or not a valid transition. + """ + if not isinstance(events, Sequence) or isinstance(events, (str, bytes)): + raise TypeError("Task Engine events must be a sequence of mappings.") + recorded = [] + for event in events: + if not isinstance(event, Mapping): + raise TypeError("Each Task Engine event must be a mapping.") + recorded.append(deepcopy(dict(event))) + + state = initial_state(request) + initial_events = [dict(event) for event in state.events] + if recorded[: len(initial_events)] != initial_events: + raise ValueError("Replay event does not match the canonical initial state.") + + for expected in recorded[len(initial_events) :]: + try: + stage = WorkflowStage(expected["stage"]) + target = StageStatus(expected["to"]) + if target == StageStatus.RUNNING: + replayed = start_stage(state, stage) + elif target == StageStatus.SUCCEEDED: + replayed = complete_stage(state, stage) + elif target == StageStatus.FAILED: + replayed = fail_stage(state, stage, reason=expected["reason"]) + elif target == StageStatus.SKIPPED: + replayed = skip_stage(state, stage) + else: + raise ValueError(f"Unsupported replay target: {target.value!r}.") + except (KeyError, TypeError, ValueError) as exc: + raise ValueError("Replay event does not match a valid transition.") from exc + if dict(replayed.events[-1]) != expected: + raise ValueError("Replay event does not match the generated transition.") + state = replayed + return state + + +def _transition( + state: TaskEngineState, + stage: WorkflowStage, + status: StageStatus, + *, + details: dict[str, Any] | None = None, +) -> TaskEngineState: + previous = state.stages[stage] + stages = dict(state.stages) + stages[stage] = status + event = { + "sequence": len(state.events) + 1, + "stage": stage.value, + "from": previous.value, + "to": status.value, + } + if details: + event.update(deepcopy(details)) + return TaskEngineState( + request=dict(state.request), + stages=stages, + events=(*state.events, event), + ) diff --git a/embodichain/gen_sim/task_engine/workflow.py b/embodichain/gen_sim/task_engine/workflow.py new file mode 100644 index 000000000..fca9778c4 --- /dev/null +++ b/embodichain/gen_sim/task_engine/workflow.py @@ -0,0 +1,1275 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Parallel Task Engine workflow with bounded, fully audited recovery.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from copy import deepcopy +from dataclasses import dataclass, replace +from datetime import datetime +import json +from pathlib import Path +import shutil +import subprocess +import sys +from typing import Any, Final + +from embodichain.gen_sim.action_engine.agent import ActionAgent +from embodichain.gen_sim.action_engine.unbound import ActionCapabilityError +from embodichain.gen_sim.action_engine.runtime import ( + EXECUTION_REPORT_FILENAME, + validate_execution_report, +) +from embodichain.gen_sim.scene_engine.errors import SceneServiceError + +from .agent import TaskAgent +from .config import ( + TaskEngineExecutionCfg, + TaskEnginePlanningCfg, + TaskEngineWorkflowCfg, + load_task_engine_config, +) +from .contracts import canonical_hash +from .orchestration.artifacts import ArtifactTransaction +from .orchestration.coordinator import PreparationResult, TaskEngineCoordinator +from .orchestration.scene_adapter import CandidateSelection, SceneAdapter +from .orchestration.scene_source import SceneSourceRef +from .scene_backend import ( + SceneAnalysis, + SceneEngineBackend, + SceneRemediableError, + SceneRevision, +) +from .state_machine import ( + TaskEngineState, + WorkflowStage, + complete_stage, + fail_stage, + initial_state, + start_stage, +) +from .workflow_contracts import TaskRunRequest, validate_task_run_request + +__all__ = [ + "TASK_ENGINE_RUN_MANIFEST_SCHEMA", + "ActionExecutor", + "SubprocessActionExecutor", + "TaskEngineRunResult", + "TaskEngineWorkflow", +] + +TASK_ENGINE_RUN_MANIFEST_SCHEMA: Final = "embodichain.task-engine-run/v1" +ActionExecutor = Callable[..., Mapping[str, Any]] + + +@dataclass(frozen=True) +class TaskEngineRunResult: + """Published outcome of one isolated cross-engine workflow run.""" + + status: str + output_dir: Path + manifest_path: Path + state_path: Path + final_bundle: Path | None + failure_class: str | None = None + + @property + def succeeded(self) -> bool: + """Return whether real simulator execution met the configured policy.""" + return self.status == "succeeded" + + +class SubprocessActionExecutor: + """Execute a prepared bundle through Task Engine's private runner.""" + + def __call__( + self, + bundle: str | Path, + output_root: str | Path, + *, + seed: int, + num_envs: int, + dataset_saving: bool = False, + ) -> Mapping[str, Any]: + """Run one simulator attempt and preserve its report and trajectory. + + Args: + bundle: Prepared Action Engine bundle. + output_root: Fresh directory for this execution attempt. + seed: Action Engine random seed. + num_envs: Number of vectorized scene replicas. + dataset_saving: Whether to enable the Gym project's dataset recorder. + + Returns: + Validated Action Engine execution report. + """ + bundle_root = Path(bundle).expanduser().resolve() + attempt_root = Path(output_root).expanduser().resolve() + attempt_root.mkdir(parents=True, exist_ok=False) + command = [ + sys.executable, + "-u", + "-m", + "embodichain.gen_sim.task_engine._bundle_runner", + "--bundle", + bundle_root.as_posix(), + "--num_envs", + str(num_envs), + "--seed", + str(seed), + "--headless", + ] + if not dataset_saving: + command.append("--filter_dataset_saving") + log_path = attempt_root / "action.log" + print( + "[Task Engine] Starting " + f"{attempt_root.name}: seed={seed}, num_envs={num_envs}, " + f"dataset_saving={dataset_saving}", + flush=True, + ) + completed = _run_streaming_process(command, log_path) + print( + f"[Task Engine] Completed {attempt_root.name}: " + f"returncode={completed.returncode}", + flush=True, + ) + report_path = bundle_root / EXECUTION_REPORT_FILENAME + process_record = { + "command": command, + "returncode": completed.returncode, + "combined_log": log_path.name, + "stdout": completed.stdout, + "stderr": completed.stderr, + } + _write_json(attempt_root / "process.json", process_record) + if not report_path.is_file(): + raise RuntimeError( + "Action execution did not publish execution_report.json; " + f"returncode={completed.returncode}." + ) + report = validate_execution_report(_read_json(report_path)) + shutil.copy2(report_path, attempt_root / EXECUTION_REPORT_FILENAME) + trajectory_copy = _copy_trajectory_record(report, attempt_root) + if report["action_count"] > 0 and trajectory_copy is None: + raise RuntimeError( + "Action execution report did not expose a readable trajectory record." + ) + _write_json( + attempt_root / "execution_attempt.json", + { + "seed": seed, + "num_envs": num_envs, + "dataset_saving": dataset_saving, + "returncode": completed.returncode, + "trajectory_copy": trajectory_copy, + "report": report, + }, + ) + return report + + +def _run_streaming_process( + command: list[str], + log_path: str | Path, +) -> subprocess.CompletedProcess[str]: + """Run a child while teeing its combined output to the terminal and disk. + + Args: + command: Argument vector passed directly to the child process. + log_path: File receiving the exact combined stdout and stderr bytes. + + Returns: + Completed process metadata with a decoded copy of the combined output. + """ + resolved_log = Path(log_path).expanduser().resolve() + resolved_log.parent.mkdir(parents=True, exist_ok=True) + captured = bytearray() + process: subprocess.Popen[bytes] | None = None + try: + with resolved_log.open("wb") as log_stream: + process = subprocess.Popen( + command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + bufsize=0, + ) + assert process.stdout is not None + while True: + chunk = process.stdout.read(64 * 1024) + if not chunk: + break + captured.extend(chunk) + log_stream.write(chunk) + log_stream.flush() + _write_terminal_chunk(chunk) + returncode = process.wait() + except BaseException: + if process is not None and process.poll() is None: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + raise + output = captured.decode("utf-8", errors="replace") + return subprocess.CompletedProcess( + args=command, + returncode=returncode, + stdout=output, + stderr="", + ) + + +def _write_terminal_chunk(chunk: bytes) -> None: + """Best-effort write of raw child output to the parent terminal.""" + try: + stream = getattr(sys.stdout, "buffer", None) + if stream is not None: + stream.write(chunk) + stream.flush() + return + sys.stdout.write(chunk.decode("utf-8", errors="replace")) + sys.stdout.flush() + except (BrokenPipeError, OSError, ValueError): + return + + +class TaskEngineWorkflow: + """Run Scene and Action work concurrently under Task Engine ownership.""" + + def __init__( + self, + *, + task_agent: TaskAgent | None = None, + scene_adapter: SceneAdapter | None = None, + action_agent: ActionAgent | None = None, + coordinator: TaskEngineCoordinator | None = None, + scene_backend: SceneEngineBackend | None = None, + action_executor: ActionExecutor | None = None, + ) -> None: + self.task_agent = task_agent or TaskAgent() + self.scene_adapter = scene_adapter or SceneAdapter() + self.action_agent = action_agent or ActionAgent() + self.coordinator = coordinator or TaskEngineCoordinator( + task_agent=self.task_agent, + scene_adapter=self.scene_adapter, + action_agent=self.action_agent, + ) + self.scene_backend = scene_backend or SceneEngineBackend() + self.action_executor = action_executor or SubprocessActionExecutor() + + def run( + self, + request: TaskRunRequest | Mapping[str, Any], + *, + workflow_cfg: TaskEngineWorkflowCfg | None = None, + planning_cfg: TaskEnginePlanningCfg | None = None, + execution_cfg: TaskEngineExecutionCfg | None = None, + config_path: str | Path | None = None, + model: str | None = None, + vlm_model: str | None = None, + base_seed: int = 0, + dataset_saving: bool = False, + run_id: str | None = None, + created_at: datetime | None = None, + overwrite: bool = False, + execute: bool = True, + ) -> TaskEngineRunResult: + """Run all stages and publish success only after simulator acceptance. + + Args: + request: One of the four image/project plus optional-edit inputs. + workflow_cfg: Optional retry and concurrency configuration. + planning_cfg: Optional interpretation and bundle generation defaults. + execution_cfg: Optional vectorized success policy. + config_path: YAML used for omitted workflow or execution config. + model: Optional Task and grounding model override. + vlm_model: Optional Action Engine VLM override. + base_seed: First audited scene and action attempt seed. + dataset_saving: Whether Action attempts may initialize dataset recording. + run_id: Optional externally allocated run identifier. + created_at: Optional timezone-aware run creation timestamp. + overwrite: Whether to atomically replace an existing run directory. + execute: Whether to execute the prepared bundle in the simulator. + + Returns: + Published run status, manifest, state audit, and final bundle path. + """ + normalized = validate_task_run_request(request) + if not isinstance(dataset_saving, bool): + raise TypeError("dataset_saving must be a boolean.") + if workflow_cfg is None or planning_cfg is None or execution_cfg is None: + loaded_workflow, loaded_planning, loaded_execution = ( + load_task_engine_config(config_path) + ) + workflow_cfg = workflow_cfg or loaded_workflow + planning_cfg = planning_cfg or loaded_planning + execution_cfg = execution_cfg or loaded_execution + effective_candidate_count = planning_cfg.candidate_count + effective_run_id = str(run_id or Path(normalized["output_dir"]).name).strip() + if not effective_run_id or Path(effective_run_id).name != effective_run_id: + raise ValueError("run_id must be one non-empty path component.") + effective_created_at = created_at or datetime.now().astimezone() + if ( + effective_created_at.tzinfo is None + or effective_created_at.utcoffset() is None + ): + raise ValueError("created_at must include a timezone.") + run_metadata = { + "run_id": effective_run_id, + "created_at": effective_created_at.isoformat(), + "dataset_saving": bool(dataset_saving), + } + state = initial_state(normalized) + attempts: list[dict[str, Any]] = [] + output_dir = Path(normalized["output_dir"]) + + with ArtifactTransaction(output_dir, overwrite=overwrite) as transaction: + staging = transaction.staging_dir + assert staging is not None + analysis_root = staging / "scene_analysis" + state = start_stage(state, WorkflowStage.TASK_CANDIDATES) + state = start_stage(state, WorkflowStage.SCENE_PREPARATION) + with ThreadPoolExecutor( + max_workers=workflow_cfg.max_parallel_workers, + thread_name_prefix="task-engine-input", + ) as executor: + candidate_future = executor.submit( + self.task_agent.generate, + normalized["task_id"], + normalized["task_instruction"], + model, + effective_candidate_count, + ) + analysis_future = executor.submit( + self.scene_backend.analyze, + normalized, + analysis_root, + ) + try: + candidate_set = candidate_future.result() + except Exception as exc: + analysis_future.cancel() + state = fail_stage( + state, + WorkflowStage.TASK_CANDIDATES, + reason=str(exc), + ) + return self._publish( + transaction, + staging, + normalized, + workflow_cfg, + planning_cfg, + execution_cfg, + run_metadata, + state, + attempts, + status="failed", + failure_class="task_generation", + ) + state = complete_stage(state, WorkflowStage.TASK_CANDIDATES) + try: + analysis = analysis_future.result() + except Exception as exc: + state = fail_stage( + state, + WorkflowStage.SCENE_PREPARATION, + reason=str(exc), + ) + return self._publish( + transaction, + staging, + normalized, + workflow_cfg, + planning_cfg, + execution_cfg, + run_metadata, + state, + attempts, + status="failed", + failure_class="scene_analysis", + ) + state = complete_stage(state, WorkflowStage.SCENE_PREPARATION) + + state = start_stage(state, WorkflowStage.CANDIDATE_SELECTION) + try: + selection = self.scene_backend.select( + analysis, + candidate_set, + self.scene_adapter, + force_most_likely=True, + ) + except Exception as exc: + state = fail_stage( + state, + WorkflowStage.CANDIDATE_SELECTION, + reason=str(exc), + ) + return self._publish( + transaction, + staging, + normalized, + workflow_cfg, + planning_cfg, + execution_cfg, + run_metadata, + state, + attempts, + status="input_conflict", + failure_class="candidate_selection", + ) + _write_json( + staging / "initial_binding_report.json", selection.binding_report + ) + if selection.selected_candidate is None: + if normalized["scene_edit_prompt"] is None: + state = fail_stage( + state, + WorkflowStage.CANDIDATE_SELECTION, + reason=str(selection.binding_report["selection_reason"]), + ) + return self._publish( + transaction, + staging, + normalized, + workflow_cfg, + planning_cfg, + execution_cfg, + run_metadata, + state, + attempts, + status="input_conflict", + failure_class="unbound_scene_reference", + ) + provisional = _highest_vote_candidate(candidate_set) + selection = replace( + selection, + selected_candidate=deepcopy(provisional), + ) + _write_json( + staging / "provisional_candidate.json", + { + "candidate_id": provisional["candidate_id"], + "reason": "explicit_scene_edit_may_materialize_missing_reference", + "binding_status": selection.binding_report["status"], + }, + ) + state = complete_stage(state, WorkflowStage.CANDIDATE_SELECTION) + + state = start_stage(state, WorkflowStage.UNBOUND_ACTION) + if normalized["scene_edit_prompt"] is not None: + state = start_stage(state, WorkflowStage.SCENE_EDIT) + else: + state = start_stage(state, WorkflowStage.SCENE_FINALIZATION) + + unbound_plan: Mapping[str, Any] | None = None + unbound_failures: list[dict[str, Any]] = [] + unbound_error: Exception | None = None + scene_error: Exception | None = None + inspection_error = False + preparation_error: Exception | None = None + preparation: PreparationResult | None = None + scene_attempt_limit = ( + 1 + if analysis.input_kind == "gym_project" + and normalized["scene_edit_prompt"] is None + else workflow_cfg.max_scene_attempts + ) + for scene_index in range(1, scene_attempt_limit + 1): + inspection_error = False + scene_seed = int(base_seed) + scene_index - 1 + attempt_root = staging / "attempts" / f"scene_{scene_index:04d}" + attempt_root.mkdir(parents=True) + attempt = { + "scene_attempt": scene_index, + "scene_seed": scene_seed, + "status": "running", + "scene_revision": None, + "final_inspection": None, + "unbound_action_plan": None, + "final_unbound_action_plan": None, + "unbound_transition": None, + "unbound_failures": [], + "preparation": None, + "planning_attempts": [], + "action_attempts": [], + "parallel_errors": [], + "error": None, + } + attempts.append(attempt) + revision: SceneRevision | None = None + try: + if unbound_plan is None: + with ThreadPoolExecutor( + max_workers=workflow_cfg.max_parallel_workers, + thread_name_prefix="task-engine-parallel", + ) as executor: + scene_future = executor.submit( + self.scene_backend.materialize, + analysis, + normalized, + attempt_root / "scene_revision", + seed=scene_seed, + ) + draft_future = executor.submit( + self._draft_with_fallback, + candidate_set, + selection, + ) + try: + revision = scene_future.result() + except Exception as exc: + scene_error = exc + revision = None + try: + unbound_plan, unbound_failures = draft_future.result() + except Exception as exc: + unbound_error = exc + if unbound_error is not None: + raise unbound_error + state = complete_stage(state, WorkflowStage.UNBOUND_ACTION) + if scene_error is not None: + raise scene_error + assert revision is not None + else: + revision = self.scene_backend.materialize( + analysis, + normalized, + attempt_root / "scene_revision", + seed=scene_seed, + ) + scene_error = None + except Exception as exc: + if revision is not None: + attempt["scene_revision"] = _revision_record(revision) + state = _complete_materialized_scene( + state, + has_edit=normalized["scene_edit_prompt"] is not None, + ) + if unbound_error is not None: + attempt["status"] = "unbound_action_failed" + attempt["error"] = _error_record(unbound_error) + if scene_error is not None: + attempt["parallel_errors"].append( + { + "branch": "scene", + **_error_record(scene_error), + } + ) + _write_json(attempt_root / "attempt.json", attempt) + break + scene_error = exc + if unbound_plan is not None: + attempt["unbound_action_plan"] = deepcopy(dict(unbound_plan)) + attempt["unbound_failures"] = deepcopy(unbound_failures) + _write_json( + attempt_root / "unbound_action_plan.json", unbound_plan + ) + attempt["status"] = "scene_failed" + attempt["error"] = _error_record(exc) + _write_json(attempt_root / "attempt.json", attempt) + if ( + scene_index < scene_attempt_limit + and _is_scene_remediable_error(exc) + ): + continue + break + + attempt["scene_revision"] = _revision_record(revision) + attempt["unbound_action_plan"] = deepcopy(dict(unbound_plan)) + attempt["unbound_failures"] = deepcopy(unbound_failures) + _write_json(attempt_root / "unbound_action_plan.json", unbound_plan) + state = _complete_materialized_scene( + state, + has_edit=normalized["scene_edit_prompt"] is not None, + ) + if state.stages[WorkflowStage.FINAL_INSPECTION].value == "pending": + state = start_stage(state, WorkflowStage.FINAL_INSPECTION) + try: + final_inspection = self.scene_backend.inspect( + revision, + attempt_root / "final_scene_inspection.json", + ) + except Exception as exc: + scene_error = exc + inspection_error = True + attempt["status"] = "scene_inspection_failed" + attempt["error"] = _error_record(exc) + _write_json(attempt_root / "attempt.json", attempt) + if ( + scene_index < scene_attempt_limit + and _is_scene_remediable_error(exc) + ): + continue + break + attempt["final_inspection"] = deepcopy(dict(final_inspection)) + if state.stages[WorkflowStage.FINAL_INSPECTION].value == "running": + state = complete_stage(state, WorkflowStage.FINAL_INSPECTION) + + bundle_root = attempt_root / "bundle" + try: + preparation = self.coordinator.prepare( + normalized["task_id"], + normalized["task_instruction"], + SceneSourceRef( + revision.source, + robot_profile=self.scene_adapter.robot_profile, + ), + bundle_root, + model=model, + candidate_count=effective_candidate_count, + planning_mode=planning_cfg.planning_mode, + vlm_model=vlm_model, + max_episodes=planning_cfg.max_episodes, + max_episode_steps=planning_cfg.max_episode_steps, + candidate_set=candidate_set, + force_most_likely=True, + final_inspection=final_inspection, + unbound_action_plan=unbound_plan, + ) + except Exception as exc: + preparation_error = exc + attempt["status"] = "preparation_error" + attempt["error"] = _error_record(exc) + _write_json(attempt_root / "attempt.json", attempt) + break + attempt["preparation"] = preparation.status + attempt["planning_attempts"] = deepcopy( + list(preparation.planning_attempts) + ) + if preparation.status == "bound": + attempt["status"] = "prepared" + _write_json(attempt_root / "attempt.json", attempt) + break + attempt["status"] = "preparation_failed" + attempt["error"] = { + "type": "PreparationFailure", + "message": preparation.status, + } + _write_json(attempt_root / "attempt.json", attempt) + if not _scene_remediable( + preparation, + analysis=analysis, + request=normalized, + ): + break + + if preparation is None or preparation.status != "bound": + failure_class = ( + "action_capability" + if unbound_error is not None + or isinstance(preparation_error, ActionCapabilityError) + else ( + "preparation_error" + if preparation_error is not None + else _preparation_failure_class( + preparation, + scene_error=scene_error, + analysis=analysis, + request=normalized, + ) + ) + ) + failed_stage = ( + WorkflowStage.FINAL_INSPECTION + if inspection_error + else ( + WorkflowStage.UNBOUND_ACTION + if unbound_error is not None + else _failure_stage(failure_class, normalized) + ) + ) + if state.stages[failed_stage].value in { + "pending", + "running", + "succeeded", + }: + state = fail_stage( + state, + failed_stage, + reason=( + str(unbound_error) + if unbound_error is not None + else ( + str(preparation_error) + if preparation_error is not None + else ( + str(scene_error) + if scene_error is not None + else ( + preparation.status + if preparation is not None + else failure_class + ) + ) + ) + ), + ) + return self._publish( + transaction, + staging, + normalized, + workflow_cfg, + planning_cfg, + execution_cfg, + run_metadata, + state, + attempts, + status=( + "input_conflict" + if failure_class == "input_conflict" + else "failed" + ), + failure_class=failure_class, + ) + + final_candidate_id = preparation.selected_candidate_id + if not isinstance(final_candidate_id, str) or not final_candidate_id: + raise ValueError( + "A bound preparation must select one non-empty candidate ID." + ) + selected_attempt = attempts[-1] + final_unbound = getattr(preparation, "unbound_action_plan", None) + if ( + final_unbound is None + and final_candidate_id != unbound_plan["candidate_id"] + ): + final_candidate = next( + item + for item in candidate_set["candidates"] + if item["candidate_id"] == final_candidate_id + ) + final_unbound = self.action_agent.draft(final_candidate) + elif final_unbound is None: + final_unbound = unbound_plan + if str(final_unbound.get("candidate_id")) != final_candidate_id: + raise ValueError( + "Final UnboundActionPlan candidate does not match preparation." + ) + selected_attempt["final_unbound_action_plan"] = deepcopy( + dict(final_unbound) + ) + selected_attempt["unbound_transition"] = { + "initial_candidate_id": str(unbound_plan["candidate_id"]), + "initial_hash": canonical_hash(unbound_plan), + "final_candidate_id": final_candidate_id, + "final_hash": canonical_hash(final_unbound), + "changed": final_unbound != unbound_plan, + } + _write_json( + preparation.output_dir.parent / "final_unbound_action_plan.json", + final_unbound, + ) + _write_json( + preparation.output_dir.parent / "attempt.json", + selected_attempt, + ) + + for stage in ( + WorkflowStage.FINAL_BINDING, + WorkflowStage.STATIC_FEASIBILITY, + WorkflowStage.GROUNDED_ACTION, + ): + state = start_stage(state, stage) + state = complete_stage(state, stage) + if not execute: + final_root = staging / "final" + final_bundle = final_root / "bundle" + final_root.mkdir() + shutil.copytree(preparation.output_dir, final_bundle) + selected_attempt["status"] = "prepared" + _write_json( + preparation.output_dir.parent / "attempt.json", + selected_attempt, + ) + _write_json( + final_root / "selection.json", + { + "scene_attempt": selected_attempt["scene_attempt"], + "candidate_id": final_candidate_id, + "action_attempt": None, + "execution_report": None, + }, + ) + return self._publish( + transaction, + staging, + normalized, + workflow_cfg, + planning_cfg, + execution_cfg, + run_metadata, + state, + attempts, + status="prepared", + failure_class=None, + final_bundle=final_bundle, + ) + state = start_stage(state, WorkflowStage.EXECUTION) + + successful_report: Mapping[str, Any] | None = None + successful_action_root: Path | None = None + success_terms = _bundle_success_terms(preparation.output_dir) + for action_index in range(1, workflow_cfg.max_action_attempts + 1): + action_seed = int(base_seed) + action_index - 1 + action_root = ( + preparation.output_dir.parent + / "action_attempts" + / f"action_{action_index:04d}" + ) + action_record: dict[str, Any] = { + "action_attempt": action_index, + "seed": action_seed, + "status": "running", + "successful_environments": 0, + "required_successes": execution_cfg.required_successes, + "error": None, + } + try: + report = self.action_executor( + preparation.output_dir, + action_root, + seed=action_seed, + num_envs=execution_cfg.num_envs, + dataset_saving=bool(dataset_saving), + ) + successes = _environment_successes( + report, + required_semantic_steps=success_terms, + ) + if len(successes) != execution_cfg.num_envs: + raise ValueError( + "Execution report environment count does not match " + "TaskEngineExecutionCfg.num_envs." + ) + action_record["successful_environments"] = sum(successes) + accepted = ( + str(report.get("status")) not in {"rejected", "aborted"} + and sum(successes) >= execution_cfg.required_successes + ) + action_record["status"] = "succeeded" if accepted else "failed" + _write_json(action_root / "task_engine_attempt.json", action_record) + selected_attempt["action_attempts"].append(deepcopy(action_record)) + if accepted: + successful_report = deepcopy(dict(report)) + successful_action_root = action_root + break + except Exception as exc: + action_record["status"] = "failed" + action_record["error"] = _error_record(exc) + action_root.mkdir(parents=True, exist_ok=True) + _write_json(action_root / "task_engine_attempt.json", action_record) + selected_attempt["action_attempts"].append(deepcopy(action_record)) + _write_json( + preparation.output_dir.parent / "attempt.json", selected_attempt + ) + + if successful_report is None: + selected_attempt["status"] = "execution_failed" + _write_json( + preparation.output_dir.parent / "attempt.json", selected_attempt + ) + state = fail_stage( + state, + WorkflowStage.EXECUTION, + reason="All bounded Action Engine execution attempts failed.", + ) + return self._publish( + transaction, + staging, + normalized, + workflow_cfg, + planning_cfg, + execution_cfg, + run_metadata, + state, + attempts, + status="failed", + failure_class="action_execution", + ) + + selected_attempt["status"] = "succeeded" + _write_json( + preparation.output_dir.parent / "attempt.json", selected_attempt + ) + state = complete_stage(state, WorkflowStage.EXECUTION) + final_root = staging / "final" + final_bundle = final_root / "bundle" + final_root.mkdir() + shutil.copytree(preparation.output_dir, final_bundle) + _write_json( + final_root / "selection.json", + { + "scene_attempt": selected_attempt["scene_attempt"], + "candidate_id": final_candidate_id, + "action_attempt": int(successful_action_root.name.split("_")[-1]), + "execution_report": successful_report, + "success_spec_steps": list(success_terms), + }, + ) + return self._publish( + transaction, + staging, + normalized, + workflow_cfg, + planning_cfg, + execution_cfg, + run_metadata, + state, + attempts, + status="succeeded", + failure_class=None, + final_bundle=final_bundle, + ) + + def _draft_with_fallback( + self, + candidate_set: Mapping[str, Any], + selection: CandidateSelection, + ) -> tuple[Mapping[str, Any], list[dict[str, Any]]]: + selected_id = selection.selected_candidate_id + resolved_ids = { + str(item["candidate_id"]) + for item in selection.binding_report["candidates"] + if item["status"] == "resolved" + } + ordered = [selected_id] + [ + str(item["candidate_id"]) + for item in candidate_set["candidates"] + if item["candidate_id"] != selected_id + and item["candidate_id"] in resolved_ids + ] + failures = [] + for candidate_id in ordered: + candidate = next( + item + for item in candidate_set["candidates"] + if item["candidate_id"] == candidate_id + ) + try: + return self.action_agent.draft(candidate), failures + except ActionCapabilityError: + raise + except (TypeError, ValueError) as exc: + failures.append( + { + "candidate_id": candidate_id, + "stage": "unbound_action", + "draft": deepcopy(candidate["draft"]), + "error": _error_record(exc), + } + ) + raise ValueError( + "No selected task candidate can be represented by Action Engine." + ) + + @staticmethod + def _publish( + transaction: ArtifactTransaction, + staging: Path, + request: Mapping[str, Any], + workflow_cfg: TaskEngineWorkflowCfg, + planning_cfg: TaskEnginePlanningCfg, + execution_cfg: TaskEngineExecutionCfg, + run_metadata: Mapping[str, Any], + state: TaskEngineState, + attempts: Sequence[Mapping[str, Any]], + *, + status: str, + failure_class: str | None, + final_bundle: Path | None = None, + ) -> TaskEngineRunResult: + state_path = staging / "workflow_state.json" + manifest_path = staging / "run_manifest.json" + _write_json(state_path, state.to_dict()) + _write_json( + manifest_path, + { + "schema_version": TASK_ENGINE_RUN_MANIFEST_SCHEMA, + "run_id": run_metadata["run_id"], + "created_at": run_metadata["created_at"], + "output_root": Path(request["output_dir"]).parent.as_posix(), + "run_dir": Path(request["output_dir"]).as_posix(), + "status": status, + "failure_class": failure_class, + "request": deepcopy(dict(request)), + "configuration": { + "workflow": { + "max_parallel_workers": workflow_cfg.max_parallel_workers, + "max_scene_attempts": workflow_cfg.max_scene_attempts, + "max_action_attempts": workflow_cfg.max_action_attempts, + }, + "planning": { + "candidate_count": planning_cfg.candidate_count, + "planning_mode": planning_cfg.planning_mode, + "max_episodes": planning_cfg.max_episodes, + "max_episode_steps": planning_cfg.max_episode_steps, + }, + "execution": { + "num_envs": execution_cfg.num_envs, + "success_policy": execution_cfg.success_policy, + "min_successful_envs": execution_cfg.min_successful_envs, + "dataset_saving": bool(run_metadata["dataset_saving"]), + }, + }, + "attempts": deepcopy(list(attempts)), + "final_bundle": ( + None if final_bundle is None else final_bundle.as_posix() + ), + }, + ) + published = transaction.commit() + return TaskEngineRunResult( + status=status, + output_dir=published, + manifest_path=published / manifest_path.name, + state_path=published / state_path.name, + final_bundle=( + None if final_bundle is None else published / "final" / "bundle" + ), + failure_class=failure_class, + ) + + +def _complete_materialized_scene( + state: TaskEngineState, + *, + has_edit: bool, +) -> TaskEngineState: + if has_edit and state.stages[WorkflowStage.SCENE_EDIT].value == "running": + state = complete_stage(state, WorkflowStage.SCENE_EDIT) + state = start_stage(state, WorkflowStage.SCENE_FINALIZATION) + if state.stages[WorkflowStage.SCENE_FINALIZATION].value == "running": + state = complete_stage(state, WorkflowStage.SCENE_FINALIZATION) + return state + + +def _scene_remediable( + preparation: PreparationResult, + *, + analysis: SceneAnalysis, + request: Mapping[str, Any], +) -> bool: + if preparation.status != "infeasible": + return False + report = preparation.feasibility_report + if not isinstance(report, Mapping) or report.get("remediation_class") != ( + "scene_remediable" + ): + return False + if analysis.input_kind == "image": + return True + return request["scene_edit_prompt"] is not None + + +def _is_scene_remediable_error(error: Exception) -> bool: + """Return whether one typed Scene failure may create a new attempt.""" + return isinstance(error, (SceneRemediableError, SceneServiceError)) + + +def _preparation_failure_class( + preparation: PreparationResult | None, + *, + scene_error: Exception | None, + analysis: SceneAnalysis, + request: Mapping[str, Any], +) -> str: + if scene_error is not None: + return "scene_materialization" + if preparation is None: + return "scene_materialization" + if preparation.status == "planning_failed": + return "action_capability" + if preparation.status in {"ambiguous", "unsatisfied"}: + return "input_conflict" + if preparation.status == "infeasible": + report = preparation.feasibility_report + remediation = ( + str(report.get("remediation_class")) + if isinstance(report, Mapping) + else "terminal" + ) + if remediation == "action_capability": + return "action_capability" + if remediation == "input_conflict": + return "input_conflict" + if remediation != "scene_remediable": + return "terminal_feasibility" + if ( + analysis.input_kind == "gym_project" + and request["scene_edit_prompt"] is None + ): + return "read_only_scene_infeasible" + return "scene_infeasible" + return "preparation" + + +def _failure_stage( + failure_class: str, + request: Mapping[str, Any], +) -> WorkflowStage: + if failure_class == "action_capability": + return WorkflowStage.GROUNDED_ACTION + if failure_class == "preparation_error": + return WorkflowStage.FINAL_BINDING + if failure_class == "input_conflict": + return WorkflowStage.FINAL_BINDING + if failure_class in { + "scene_infeasible", + "read_only_scene_infeasible", + "terminal_feasibility", + }: + return WorkflowStage.STATIC_FEASIBILITY + if failure_class == "scene_materialization": + return ( + WorkflowStage.SCENE_EDIT + if request["scene_edit_prompt"] is not None + else WorkflowStage.SCENE_FINALIZATION + ) + return WorkflowStage.GROUNDED_ACTION + + +def _environment_successes( + report: Mapping[str, Any], + *, + required_semantic_steps: Sequence[str] = (), +) -> list[bool]: + environments = report.get("environments") + if not isinstance(environments, Sequence) or isinstance(environments, (str, bytes)): + raise ValueError("Execution report environments must be a sequence.") + values = [] + for item in environments: + if not isinstance(item, Mapping) or not isinstance(item.get("success"), bool): + raise ValueError("Every execution environment requires boolean success.") + success = bool(item["success"]) + if required_semantic_steps: + semantics = item.get("semantic_success") + if not isinstance(semantics, Mapping): + success = False + else: + success = success and all( + semantics.get(step_id) is True + for step_id in required_semantic_steps + ) + values.append(success) + if not values: + raise ValueError("Execution report must contain at least one environment.") + return values + + +def _bundle_success_terms(bundle: Path) -> tuple[str, ...]: + path = bundle / "grounded_task_plan.json" + if not path.is_file(): + return () + try: + value = _read_json(path) + success_spec = value.get("success_spec") + terms = success_spec.get("terms") if isinstance(success_spec, Mapping) else None + strict = isinstance(value.get("schema_version"), str) + if not isinstance(terms, Sequence) or isinstance(terms, (str, bytes)): + if strict: + raise ValueError("GroundedTaskPlan has no valid SuccessSpec terms.") + return () + result = tuple( + str(item["step_id"]) + for item in terms + if isinstance(item, Mapping) and isinstance(item.get("step_id"), str) + ) + if len(result) != len(terms) or (strict and not result): + if strict: + raise ValueError("GroundedTaskPlan SuccessSpec terms are invalid.") + return () + return result + except OSError: + return () + + +def _highest_vote_candidate(candidate_set: Mapping[str, Any]) -> Mapping[str, Any]: + candidates = candidate_set.get("candidates") + if not isinstance(candidates, Sequence) or isinstance(candidates, (str, bytes)): + raise TypeError("TaskCandidateSet.candidates must be a sequence.") + values = [item for item in candidates if isinstance(item, Mapping)] + if not values: + raise ValueError("TaskCandidateSet requires at least one candidate.") + return max( + values, + key=lambda item: ( + int(item.get("vote_count", 0)), + str(item.get("candidate_id", "")), + ), + ) + + +def _copy_trajectory_record(report: Mapping[str, Any], output_root: Path) -> str | None: + raw = report.get("record_dir") + if not isinstance(raw, str) or not raw: + return None + source = Path(raw).expanduser().resolve() + if not source.is_dir(): + return None + destination = output_root / "trajectory" + if source == destination or destination in source.parents: + return source.as_posix() + shutil.copytree(source, destination) + return destination.as_posix() + + +def _revision_record(revision: SceneRevision) -> dict[str, Any]: + return { + "source": revision.source.as_posix(), + "output_root": ( + None if revision.output_root is None else revision.output_root.as_posix() + ), + "revision_id": revision.revision_id, + "seed": revision.seed, + "edit_plan": deepcopy(revision.edit_plan), + "source_fingerprint": ( + None + if revision.source_fingerprint is None + else revision.source_fingerprint.to_dict() + ), + } + + +def _error_record(error: Exception) -> dict[str, str]: + return { + "type": type(error).__name__, + "failure_type": ( + "scene_remediable" if _is_scene_remediable_error(error) else "terminal" + ), + "message": str(error), + } + + +def _read_json(path: Path) -> dict[str, Any]: + value = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, dict): + raise ValueError(f"JSON artifact must contain an object: {path}") + return value + + +def _write_json(path: Path, value: 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/task_engine/workflow_contracts.py b/embodichain/gen_sim/task_engine/workflow_contracts.py new file mode 100644 index 000000000..a6b61b04b --- /dev/null +++ b/embodichain/gen_sim/task_engine/workflow_contracts.py @@ -0,0 +1,179 @@ +# ---------------------------------------------------------------------------- +# 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 inputs for Task Engine cross-engine workflows.""" + +from __future__ import annotations + +from collections.abc import Mapping +from copy import deepcopy +import json +from pathlib import Path +from typing import Any, Final, Literal, TypeAlias + +__all__ = [ + "TASK_RUN_REQUEST_SCHEMA", + "SceneInputKind", + "TaskRunRequest", + "scene_input_kind", + "validate_scene_history_root", + "validate_scene_output_separation", + "validate_task_run_request", +] + +TASK_RUN_REQUEST_SCHEMA: Final = "embodichain.task-engine-run-request/v1" +TaskRunRequest: TypeAlias = dict[str, Any] +SceneInputKind = Literal["image", "gym_project"] + +_REQUEST_KEYS = frozenset( + { + "schema_version", + "task_id", + "task_instruction", + "image_path", + "gym_project", + "scene_edit_prompt", + "output_dir", + } +) + + +def validate_task_run_request(value: Mapping[str, Any]) -> TaskRunRequest: + """Validate and detach one Task Engine run request. + + Version 1 deliberately has no ``scene_generation_prompt``. Image workflows + use the image-only Scene Engine generation behavior and may apply one + optional edit after that initial scene has been generated. + """ + if not isinstance(value, Mapping): + raise TypeError("TaskRunRequest must be a mapping.") + result = deepcopy(dict(value)) + if set(result) != _REQUEST_KEYS: + missing = sorted(_REQUEST_KEYS - set(result)) + extra = sorted(set(result) - _REQUEST_KEYS) + raise ValueError( + f"TaskRunRequest fields differ; missing={missing}, extra={extra}." + ) + if result.get("schema_version") != TASK_RUN_REQUEST_SCHEMA: + raise ValueError( + "TaskRunRequest.schema_version must be " f"{TASK_RUN_REQUEST_SCHEMA!r}." + ) + result["task_id"] = _nonempty(result.get("task_id"), "task_id") + result["task_instruction"] = _nonempty( + result.get("task_instruction"), "task_instruction" + ) + result["output_dir"] = _path(result.get("output_dir"), "output_dir") + + image_path = _optional_path(result.get("image_path"), "image_path") + gym_project = _optional_path(result.get("gym_project"), "gym_project") + if (image_path is None) == (gym_project is None): + raise ValueError( + "TaskRunRequest requires exactly one of image_path or gym_project." + ) + result["image_path"] = image_path + result["gym_project"] = gym_project + if gym_project is not None: + validate_scene_output_separation(gym_project, result["output_dir"]) + + edit_prompt = result.get("scene_edit_prompt") + if edit_prompt is not None: + edit_prompt = _nonempty(edit_prompt, "scene_edit_prompt") + result["scene_edit_prompt"] = edit_prompt + _json_safe(result) + return result + + +def scene_input_kind(request: Mapping[str, Any]) -> SceneInputKind: + """Return the selected scene input kind after validating ``request``.""" + normalized = validate_task_run_request(request) + return "image" if normalized["image_path"] is not None else "gym_project" + + +def validate_scene_output_separation( + gym_project: str | Path, + output_dir: str | Path, +) -> None: + """Reject output paths that could replace or modify a read-only source. + + Args: + gym_project: Existing Gym project directory or configuration path. + output_dir: Transactional output directory for the Task Engine run. + + Raises: + ValueError: If either path contains the other or both paths are equal. + """ + source = Path(gym_project).expanduser().resolve() + output = Path(output_dir).expanduser().resolve() + if source == output or source in output.parents or output in source.parents: + raise ValueError( + "Task Engine output_dir and source Gym project must not overlap." + ) + + +def validate_scene_history_root( + gym_project: str | Path, + output_root: str | Path, +) -> None: + """Protect a source project before reserving a history-directory child. + + A prior run may live below the same history root because every new run is + published to a distinct timestamped child. The inverse remains unsafe: + creating the history root at or below the source project would write a + reservation and output artifacts into the read-only source tree. + + Args: + gym_project: Existing Gym project directory or configuration path. + output_root: Parent directory under which a new run will be reserved. + + Raises: + ValueError: If the history root is equal to or contained by the source + project boundary. + """ + source = Path(gym_project).expanduser().resolve() + protected_root = source.parent if source.is_file() else source + history_root = Path(output_root).expanduser().resolve() + if protected_root == history_root or protected_root in history_root.parents: + raise ValueError( + "Task Engine output_root must not be inside the read-only source " + "Gym project." + ) + + +def _path(value: Any, field_name: str) -> str: + text = _nonempty(value, field_name) + return Path(text).expanduser().resolve().as_posix() + + +def _optional_path(value: Any, field_name: str) -> str | None: + if value is None: + return None + return _path(value, field_name) + + +def _nonempty(value: Any, field_name: str) -> str: + if not isinstance(value, str): + raise TypeError(f"TaskRunRequest.{field_name} must be a string.") + result = value.strip() + if not result: + raise ValueError(f"TaskRunRequest.{field_name} must not be empty.") + return result + + +def _json_safe(value: Any) -> None: + try: + json.dumps(value, ensure_ascii=False, allow_nan=False) + except (TypeError, ValueError) as exc: + raise ValueError("TaskRunRequest must contain strict JSON data.") from exc diff --git a/embodichain/lab/sim/atomic_actions/affordance.py b/embodichain/lab/sim/atomic_actions/affordance.py index ad6b7c8af..c21332e82 100644 --- a/embodichain/lab/sim/atomic_actions/affordance.py +++ b/embodichain/lab/sim/atomic_actions/affordance.py @@ -16,6 +16,8 @@ from __future__ import annotations +from collections.abc import Callable + import torch from dataclasses import dataclass, field from typing import Any, TYPE_CHECKING @@ -118,16 +120,27 @@ def get_valid_grasp_poses( [0, 0, -1], dtype=torch.float32 ), object_part: str = "center", + grasp_cost_fn: ( + Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor] | None + ) = None, ) -> list[tuple[torch.Tensor, torch.Tensor]]: if self._generator is None: self._init_generator() approach_direction = self._resolve_approach_direction(approach_direction) results = [] for i, obj_pose in enumerate(obj_poses): + pose_cost_fn = None + if grasp_cost_fn is not None: + pose_cost_fn = lambda grasp_poses, costs: grasp_cost_fn( + obj_pose, + grasp_poses, + costs, + ) is_success, grasp_poses, _, costs = self._generator.get_valid_grasp_poses( object_pose=obj_pose, approach_direction=approach_direction, object_part=object_part, + pose_cost_fn=pose_cost_fn, ) if grasp_poses.shape == (4, 4): grasp_poses = grasp_poses.unsqueeze(0) diff --git a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py index ffbd145bc..993dd714b 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/pick_up.py +++ b/embodichain/lab/sim/atomic_actions/primitives/pick_up.py @@ -73,6 +73,33 @@ translate_pose_world, ) +_UPRIGHT_SIDE_GRASP_MAX_AXIS_ALIGNMENT = 0.65 +_UPRIGHT_SIDE_GRASP_MIN_AXIS_FRACTION = 0.35 +_UPRIGHT_SIDE_GRASP_MAX_AXIS_FRACTION = 0.75 +_UPRIGHT_SIDE_GRASP_HEIGHT_COST_WEIGHT = 2.0 + + +def _upright_yaw_pose_variants( + target_pose: torch.Tensor, + sample_count: int, +) -> torch.Tensor: + """Return object poses with evenly sampled world-Z yaw rotations.""" + signed_steps = [0] + for step in range(1, (sample_count + 1) // 2): + signed_steps.extend((step, -step)) + if sample_count % 2 == 0: + signed_steps.append(sample_count // 2) + angles = target_pose.new_tensor(signed_steps) * (2.0 * math.pi / sample_count) + yaw = target_pose.new_zeros((sample_count, 3, 3)) + yaw[:, 0, 0] = torch.cos(angles) + yaw[:, 0, 1] = -torch.sin(angles) + yaw[:, 1, 0] = torch.sin(angles) + yaw[:, 1, 1] = torch.cos(angles) + yaw[:, 2, 2] = 1.0 + variants = target_pose[:, None].repeat(1, sample_count, 1, 1) + variants[:, :, :3, :3] = torch.matmul(yaw[None], target_pose[:, None, :3, :3]) + return variants + @dataclass(frozen=True, slots=True, eq=False) class GraspGoal(ObjectActionGoal): @@ -119,6 +146,9 @@ class PickUpOptions(ActionOptions): downstream_object_target_poses: tuple[torch.Tensor, ...] = () """Future object poses that must be reachable with the selected grasp.""" + upright_yaw_samples: int = 1 + """Equivalent world-yaw samples for semantically upright downstream targets.""" + obj_upright_direction: torch.Tensor | None = None """Optional object local direction used to choose the upright grasp rotation.""" @@ -134,6 +164,8 @@ def __post_init__(self) -> None: raise ValueError("lift_height must be non-negative.") if self.pre_grasp_distance < 0.0: raise ValueError("pre_grasp_distance must be non-negative.") + if self.upright_yaw_samples < 1: + raise ValueError("upright_yaw_samples must be positive.") if self.approach_direction.shape != (3,): raise ValueError("approach_direction must have shape (3,).") if not torch.isfinite(self.approach_direction).all(): @@ -416,10 +448,22 @@ def _resolve_grasp_pose( options: PickUpOptions, approach_direction: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: + grasp_cost_fn = None + if options.rotate_upright is not None: + grasp_cost_fn = lambda object_pose, grasp_poses, costs: ( + self._upright_grasp_costs( + semantics, + object_pose, + grasp_poses, + costs, + options, + ) + ) grasp_poses_result = semantics.affordance.get_valid_grasp_poses( obj_poses=object_pose, approach_direction=approach_direction, object_part=options.pick_object_part, + grasp_cost_fn=grasp_cost_fn, ) num_envs = object_pose.shape[0] n_max_pose = max(r[0].shape[0] for r in grasp_poses_result) @@ -480,6 +524,11 @@ def _select_feasible_grasp_variants( object_poses, options, ) + upright_compatible = self._upright_grasp_compatibility_mask( + grasp_variants, + object_poses, + options, + ) pre_grasp_variants = grasp_variants.clone() pre_grasp_z = pre_grasp_variants[..., :3, 2] @@ -504,7 +553,11 @@ def _select_feasible_grasp_variants( grasp_variants, options, approach_direction ) pickup_success = ( - alignment_success & pre_grasp_success & grasp_success & lift_success + upright_compatible + & alignment_success + & pre_grasp_success + & grasp_success + & lift_success ) downstream_success_counts: list[list[int]] = [] object_to_eef_variants = torch.matmul( @@ -527,17 +580,35 @@ def _select_feasible_grasp_variants( f"(4, 4) or ({num_envs}, 4, 4), but got " f"{object_target_pose.shape}." ) - downstream_eef_variants = torch.matmul( - object_target_pose[:, None, None], object_to_eef_variants - ) - downstream_success, downstream_seed = self._compute_batch_candidate_ik( - downstream_eef_variants, downstream_seed, manipulator + object_target_variants = _upright_yaw_pose_variants( + object_target_pose, + options.upright_yaw_samples, ) + downstream_success = torch.zeros_like(pickup_success) + selected_qpos = downstream_seed + for yaw_target in object_target_variants.unbind(dim=1): + downstream_eef_variants = torch.matmul( + yaw_target[:, None, None], object_to_eef_variants + ) + yaw_success, yaw_qpos = self._compute_batch_candidate_ik( + downstream_eef_variants, + downstream_seed, + manipulator, + ) + newly_solved = ~downstream_success & yaw_success + selected_qpos = torch.where( + newly_solved[..., None], yaw_qpos, selected_qpos + ) + downstream_success |= yaw_success + if bool((pickup_success & downstream_success).any(dim=(1, 2)).all()): + break + downstream_seed = selected_qpos pickup_success &= downstream_success downstream_success_counts.append(pickup_success.sum(dim=(1, 2)).tolist()) if not pickup_success.any(dim=(1, 2)).all(): logger.log_warning( "PickUp found no candidate with a feasible vertical pickup path: " + f"upright_compatible={upright_compatible.sum(dim=(1, 2)).tolist()}, " f"aligned={alignment_success.sum(dim=(1, 2)).tolist()}, " f"pre_grasp={pre_grasp_success.sum(dim=(1, 2)).tolist()}, " f"grasp={(pre_grasp_success & grasp_success).sum(dim=(1, 2)).tolist()}, " @@ -620,14 +691,14 @@ def _upright_adjusted_grasp_poses( if options.rotate_upright is None: return grasp_xpos - if options.obj_upright_direction is None: - upright_direction = torch.tensor( - [0, 0, 1], dtype=torch.float32, device=self.device - ) - else: - upright_direction = options.obj_upright_direction.to( - device=self.device, dtype=torch.float32 - ) + upright_direction = self._normalized_obj_upright_direction(options).to( + device=grasp_xpos.device, + dtype=grasp_xpos.dtype, + ) + object_pose = object_pose.to( + device=grasp_xpos.device, + dtype=grasp_xpos.dtype, + ) obj_upright = torch.matmul(object_pose[:, :3, :3], upright_direction) adjusted_grasp_xpos = grasp_xpos.clone() grasp_ry = adjusted_grasp_xpos[..., :3, 1] @@ -646,5 +717,115 @@ def _upright_adjusted_grasp_poses( ) return adjusted_grasp_xpos + def _upright_grasp_compatibility_mask( + self, + grasp_xpos: torch.Tensor, + object_poses: torch.Tensor, + options: PickUpOptions, + ) -> torch.Tensor: + """Reject upright grasps that clamp the object's support and top faces.""" + shape = grasp_xpos.shape[:3] + if options.rotate_upright is None: + return torch.ones(shape, dtype=torch.bool, device=grasp_xpos.device) + local_upright = self._normalized_obj_upright_direction(options).to( + device=grasp_xpos.device, + dtype=grasp_xpos.dtype, + ) + object_poses = object_poses.to( + device=grasp_xpos.device, + dtype=grasp_xpos.dtype, + ) + world_upright = torch.matmul(object_poses[:, :3, :3], local_upright) + closing_axes = torch.nn.functional.normalize( + grasp_xpos[..., :3, 0], + dim=-1, + ) + axis_alignment = torch.abs( + torch.sum(closing_axes * world_upright[:, None, None, :], dim=-1) + ) + return axis_alignment <= _UPRIGHT_SIDE_GRASP_MAX_AXIS_ALIGNMENT + + def _upright_grasp_costs( + self, + semantics: ObjectSemantics, + object_pose: torch.Tensor, + grasp_poses: torch.Tensor, + costs: torch.Tensor, + options: PickUpOptions, + ) -> torch.Tensor: + """Rank side grasps before generator top-k truncation.""" + local_upright = self._normalized_obj_upright_direction(options).to( + device=grasp_poses.device, + dtype=grasp_poses.dtype, + ) + object_pose = object_pose.to( + device=grasp_poses.device, + dtype=grasp_poses.dtype, + ) + world_upright = torch.matmul(object_pose[:3, :3], local_upright) + closing_axes = torch.nn.functional.normalize( + grasp_poses[:, :3, 0], + dim=-1, + ) + axis_alignment = torch.abs( + torch.sum(closing_axes * world_upright[None, :], dim=-1) + ) + adjusted = torch.where( + axis_alignment <= _UPRIGHT_SIDE_GRASP_MAX_AXIS_ALIGNMENT, + costs, + torch.full_like(costs, torch.inf), + ) + + vertices = semantics.geometry.get("mesh_vertices") + if vertices is None: + return adjusted + vertices = torch.as_tensor( + vertices, + dtype=grasp_poses.dtype, + device=grasp_poses.device, + ) + if vertices.ndim != 2 or vertices.shape[-1] != 3 or vertices.numel() == 0: + return adjusted + vertex_axis_positions = torch.matmul(vertices, local_upright) + axis_min = vertex_axis_positions.min() + axis_extent = vertex_axis_positions.max() - axis_min + if float(axis_extent) <= 1.0e-6: + return adjusted + + relative_centers = grasp_poses[:, :3, 3] - object_pose[None, :3, 3] + center_axis_positions = torch.sum( + relative_centers * world_upright[None, :], + dim=-1, + ) + center_fractions = (center_axis_positions - axis_min) / axis_extent + interval = ( + _UPRIGHT_SIDE_GRASP_MAX_AXIS_FRACTION + - _UPRIGHT_SIDE_GRASP_MIN_AXIS_FRACTION + ) + height_penalty = ( + torch.clamp( + _UPRIGHT_SIDE_GRASP_MIN_AXIS_FRACTION - center_fractions, + min=0.0, + ) + + torch.clamp( + center_fractions - _UPRIGHT_SIDE_GRASP_MAX_AXIS_FRACTION, + min=0.0, + ) + ) / interval + return adjusted + _UPRIGHT_SIDE_GRASP_HEIGHT_COST_WEIGHT * height_penalty + + def _normalized_obj_upright_direction( + self, + options: PickUpOptions, + ) -> torch.Tensor: + direction = options.obj_upright_direction + if direction is None: + direction = torch.tensor([0, 0, 1], dtype=torch.float32) + direction = direction.to(device=self.device, dtype=torch.float32) + norm = torch.linalg.vector_norm(direction) + if norm <= 1.0e-6: + logger.log_error("obj_upright_direction must be non-zero.", ValueError) + return direction / norm + __all__ = ["GraspGoal", "PickUp", "PickUpOptions"] diff --git a/embodichain/toolkits/graspkit/pg_grasp/antipodal_generator.py b/embodichain/toolkits/graspkit/pg_grasp/antipodal_generator.py index abc0466ef..c53d92613 100644 --- a/embodichain/toolkits/graspkit/pg_grasp/antipodal_generator.py +++ b/embodichain/toolkits/graspkit/pg_grasp/antipodal_generator.py @@ -14,8 +14,11 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import os import argparse +from collections.abc import Callable import open3d as o3d import time import torch @@ -613,6 +616,9 @@ def get_valid_grasp_poses( approach_direction: torch.Tensor, object_part: str = "center", visualize_collision: bool = False, + pose_cost_fn: ( + Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None + ) = None, ): if self._hit_point_pairs is None: logger.log_warning( @@ -659,6 +665,7 @@ def get_valid_grasp_poses( approach_direction=approach_direction, mesh_vert_transformed=mesh_vert_transformed, visualize_collision=visualize_collision, + pose_cost_fn=pose_cost_fn, ) def get_dual_arm_valid_grasp_poses( @@ -762,6 +769,9 @@ def _filter_valid_grasp_poses( mesh_vert_transformed: torch.Tensor, object_pose: torch.Tensor, visualize_collision: bool = False, + pose_cost_fn: ( + Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None + ) = None, ): grasp_x = F.normalize(hit_points_ - origin_points_, dim=-1) cos_angle = torch.clamp((grasp_x * approach_direction).sum(dim=-1), -1.0, 1.0) @@ -850,6 +860,17 @@ def _filter_valid_grasp_poses( center_cost = center_distance / center_distance.max() length_cost = 1 - valid_open_lengths / valid_open_lengths.max() total_cost = 0.2 * angle_cost + 0.2 * length_cost + 0.6 * center_cost + if pose_cost_fn is not None: + adjusted_cost = pose_cost_fn(valid_grasp_poses, total_cost) + if adjusted_cost.shape != total_cost.shape: + logger.log_error( + "pose_cost_fn must preserve the grasp cost shape.", + ValueError, + ) + total_cost = adjusted_cost.to( + device=total_cost.device, + dtype=total_cost.dtype, + ) n_valid = valid_grasp_poses.shape[0] if n_valid == 0: diff --git a/setup.py b/setup.py index ac131d484..06f687a3a 100644 --- a/setup.py +++ b/setup.py @@ -133,7 +133,9 @@ def main(): package_dir=get_package_dir(), package_data={ "embodichain": ["VERSION"], + "embodichain.gen_sim.action_engine.config": ["*.yaml"], "embodichain.gen_sim.simready_pipeline.configs": ["*.json"], + "embodichain.gen_sim.task_engine": ["*.yaml"], "embodichain_tasks.configs": ["**/*.json", "**/*.yaml", "**/*.yml"], }, cmdclass=cmdclass, diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 000000000..3dbe08328 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1,21 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""EmbodiChain test package.""" + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/tests/benchmark/gen_sim/test_e1_e2_benchmark.py b/tests/benchmark/gen_sim/test_e1_e2_benchmark.py new file mode 100644 index 000000000..7d72f7c67 --- /dev/null +++ b/tests/benchmark/gen_sim/test_e1_e2_benchmark.py @@ -0,0 +1,36 @@ +# ---------------------------------------------------------------------------- +# 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 pathlib import Path + +from embodichain.gen_sim.action_engine.evaluation.e1_e2_scene_action import ( + run_benchmark, +) + + +def test_e1_e2_contract_benchmark_is_reproducibly_executable(tmp_path: Path) -> None: + results, report = run_benchmark(iterations=2, output_dir=tmp_path) + + assert tuple(item.scenario for item in results) == ("E1", "E2") + assert all(item.success_rate == 1.0 for item in results) + assert all(item.feasibility_status == "runtime_probe" for item in results) + assert all(item.action_count >= 3 for item in results) + markdown = report.read_text(encoding="utf-8") + assert markdown.count("## Time & Memory") == 1 + assert markdown.count("## Success & Other Metrics") == 1 + assert markdown.count("## Leaderboard") == 1 diff --git a/tests/gen_sim/__init__.py b/tests/gen_sim/__init__.py new file mode 100644 index 000000000..cdeead7b0 --- /dev/null +++ b/tests/gen_sim/__init__.py @@ -0,0 +1,21 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Generative simulation tests.""" + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/tests/gen_sim/action_engine/__init__.py b/tests/gen_sim/action_engine/__init__.py new file mode 100644 index 000000000..e2bb4c0aa --- /dev/null +++ b/tests/gen_sim/action_engine/__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 tests.""" diff --git a/tests/gen_sim/action_engine/capabilities/__init__.py b/tests/gen_sim/action_engine/capabilities/__init__.py new file mode 100644 index 000000000..046cb429b --- /dev/null +++ b/tests/gen_sim/action_engine/capabilities/__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 capability tests.""" diff --git a/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py b/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py new file mode 100644 index 000000000..7049a79c4 --- /dev/null +++ b/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py @@ -0,0 +1,210 @@ +# ---------------------------------------------------------------------------- +# 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 dataclass +from types import SimpleNamespace + +import torch + +from embodichain.gen_sim.action_engine.capabilities import ( + AtomicCapability, + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.runtime.actions import AtomicActionAdapter +from embodichain.gen_sim.action_engine.runtime.grounding import ActionGrounder +from embodichain.gen_sim.action_engine.runtime.loader import load_execution_program +from embodichain.gen_sim.action_engine.runtime.models import GroundedAction +from embodichain.gen_sim.action_engine.runtime.state import ExecutionState +from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph + +from ..task_fixtures import make_task_spec +from embodichain.gen_sim.action_engine.planning.linker import link_seed_graph +from embodichain.lab.sim.atomic_actions import ( + ActionOptions, + ActionPlan, + EndEffectorPoseGoal, + PlannerDiagnostics, + TimedTrajectory, +) + + +@dataclass(frozen=True, slots=True) +class _TestOptions(ActionOptions): + marker: str = "test" + + +class _TestAction: + skill_id = "test_retreat" + end_effector_roles: tuple[str, ...] = () + + +class _TestEngine: + def plan(self, invocation, context): + assert isinstance(invocation.skill_options, _TestOptions) + positions = context.robot.qpos[:, None, :] + return ActionPlan( + skill_id=invocation.skill_id, + plan_success=torch.ones(context.batch_size, dtype=torch.bool), + trajectory=TimedTrajectory.from_positions( + positions, + env_ids=context.env_ids, + control_dt=invocation.motion_policy.control_dt, + ), + recovery_policy=invocation.recovery_policy, + planned_scene_version=context.scene.version, + planned_collision_world_revision=(0,) * context.batch_size, + diagnostics=PlannerDiagnostics(backend="test"), + ) + + +class _Robot: + dof = 2 + uid = "test_robot" + control_parts = {"left_arm": [0], "right_arm": [1]} + + def get_qpos(self): + return torch.zeros((1, 2)) + + def get_joint_ids(self, *, name: str): + return self.control_parts.get(name, []) + + +class _Entity: + def get_local_pose(self, *, to_matrix: bool): + assert to_matrix + return torch.eye(4).unsqueeze(0) + + +class _Sim: + def get_rigid_object(self, _uid: str): + return _Entity() + + +def test_new_descriptor_reuses_loader_and_adapter_without_dispatch_changes() -> None: + registry = build_atomic_capability_registry() + calls = [] + + def target_hook(**kwargs): + calls.append("target") + pose = kwargs["object_pose"].clone() + return GroundedAction( + action_class="TestRetreat", + arm=kwargs["arm"], + control="arm", + target=EndEffectorPoseGoal(xpos=pose), + cfg=kwargs["policy"], + object_pose=pose, + target_object_pose=pose, + motion_policy=kwargs["policy"], + ) + + def config_hook(**_kwargs): + calls.append("config") + return _TestOptions() + + registry.register( + AtomicCapability( + "TestRetreat", + _TestAction, + _TestOptions, + frozenset({"policy_pose"}), + frozenset({"arm"}), + "single_arm", + "preserve", + "eef_pose", + motion_base="MoveEndEffector", + target_materializer_hook=target_hook, + config_materializer_hook=config_hook, + contract_resolver_hook=registry.get( + "MoveEndEffector" + ).contract_resolver_hook, + ) + ) + task, requirements = make_task_spec("E1") + bindings = { + item["role_id"]: f"uid_{item['role_id']}" for item in requirements["objects"] + } + graph = instantiate_seed_graph(task, bindings) + graph = deepcopy(graph) + graph["capability_catalog_hash"] = registry.catalog_hash() + cleanup = next( + node for node in graph["nodes"] if node["atomic_action"] == "MoveEndEffector" + ) + cleanup["atomic_action"] = "TestRetreat" + cleanup.pop("contract") + for group in graph["task_groups"]: + group.pop("contract") + graph["metadata"].pop("action_contract_linker") + graph = link_seed_graph(graph, registry=registry) + + program = load_execution_program(graph, registry=registry) + assert any( + action["atomic_action_class"] == "TestRetreat" + for edge in program.edges + for action in edge.actions + ) + + env = SimpleNamespace( + num_envs=1, + device=torch.device("cpu"), + robot=_Robot(), + sim=_Sim(), + agent_robot_profile="dual_ur10", + get_agent_arm_control_part=lambda is_left: ( + "left_arm" if is_left else "right_arm" + ), + get_agent_eef_control_part=lambda _is_left: None, + ) + adapter = AtomicActionAdapter( + env, + grasp_policy={}, + capability_registry=registry, + ) + adapter._atomic_engine = _TestEngine() + step = next( + step + for step in program.semantic_steps + if any( + action["atomic_action_class"] == "TestRetreat" + for edge_id in step.edge_ids + for action in next( + edge for edge in program.edges if edge.id == edge_id + ).actions + ) + ) + action = next( + action + for edge in program.edges + for action in edge.actions + if action["atomic_action_class"] == "TestRetreat" + ) + grounder = ActionGrounder( + program, + env, + lambda _uid: None, + capability_registry=registry, + ) + state = ExecutionState(last_qpos=torch.zeros((1, 2))) + grounded = grounder.ground(action, step, arm="left_arm", state=state) + outcome = adapter.plan( + grounded, + state, + ) + assert outcome.success.tolist() == [True] + assert calls == ["target", "config"] 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/compiler/__init__.py b/tests/gen_sim/action_engine/compiler/__init__.py new file mode 100644 index 000000000..e7977347e --- /dev/null +++ b/tests/gen_sim/action_engine/compiler/__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 compiler tests.""" diff --git a/tests/gen_sim/action_engine/compiler/test_compiler.py b/tests/gen_sim/action_engine/compiler/test_compiler.py new file mode 100644 index 000000000..6144fcc3f --- /dev/null +++ b/tests/gen_sim/action_engine/compiler/test_compiler.py @@ -0,0 +1,572 @@ +# ---------------------------------------------------------------------------- +# 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 collections.abc import Mapping +from typing import Any + +import pytest + +from embodichain.gen_sim.action_engine.compiler import compile_task_agent +from embodichain.gen_sim.action_engine.domain import ( + EXECUTION_PROGRAM_SCHEMA, + TASK_AGENT_SCHEMA, +) + + +def _program(step: Mapping[str, Any]) -> dict[str, Any]: + return { + "schema_version": TASK_AGENT_SCHEMA, + "task": "operator_demo", + "goal": "Exercise one semantic operator.", + "semantic_steps": [dict(step)], + } + + +def test_place_relative_carries_payloads_through_single_arm_action_bindings() -> None: + execution = compile_task_agent( + _program( + { + "id": "s01_place_carrier", + "operator": "place_relative", + "object": "paper_cup", + "goal": { + "reference_object": "popcorn_bucket", + "relation": "on", + "payloads": [{"object": "glue_stick", "slot": "center"}], + }, + } + ) + ) + + step = execution["semantic_steps"][0] + assert step["goal"]["payloads"] == [{"object": "glue_stick", "slot": "center"}] + carrying_actions = [ + action + for edge in execution["edges"] + for action in edge["actions"] + if action["atomic_action_class"] in {"PickUp", "MoveHeldObject", "Place"} + ] + assert carrying_actions + assert all( + action["target_binding"]["payloads"] == step["goal"]["payloads"] + for action in carrying_actions + ) + + +@pytest.mark.parametrize( + ("step", "expected_action"), + [ + ( + { + "id": "s01_line", + "operator": "arrange_line", + "objects": ["can_a", "can_b"], + "goal": {"axis": "world_y", "anchor": "table_center"}, + }, + "MoveHeldObject", + ), + ( + { + "id": "s01_stack", + "operator": "build_stack", + "objects": ["block_a", "block_b"], + "goal": {"stack_mode": "on_top", "anchor": "table_center"}, + }, + "PickUp", + ), + ( + { + "id": "s01_place", + "operator": "place_relative", + "object": "cup", + "goal": {"reference_object": "tray", "relation": "on"}, + }, + "Place", + ), + ( + { + "id": "s01_hover", + "operator": "hold_hover", + "object": "cup", + "goal": {}, + }, + "MoveJoints", + ), + ( + { + "id": "s01_transport", + "operator": "coordinated_transport", + "object": "tray", + "goal": {"direction": "front", "terminal_behavior": "place"}, + }, + "CoordinatedPickment", + ), + ( + { + "id": "s01_orient", + "operator": "orient_object", + "object": "cup", + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + }, + }, + "MoveHeldObject", + ), + ( + { + "id": "s01_press", + "operator": "press", + "object": "button", + "goal": {"terminal_state": "activated"}, + }, + "Press", + ), + ( + { + "id": "s01_coordinated_place", + "operator": "coordinated_place", + "object": "cup", + "goal": {"support_object": "tray", "relation": "on"}, + }, + "CoordinatedPlacement", + ), + ], +) +def test_every_builtin_operator_compiles( + step: Mapping[str, Any], + expected_action: str, +) -> None: + execution = compile_task_agent(_program(step)) + action_classes = { + action["atomic_action_class"] + for edge in execution["edges"] + for action in edge["actions"] + } + + assert execution["schema_version"] == EXECUTION_PROGRAM_SCHEMA + assert expected_action in action_classes + assert execution["nodes"][0]["id"] == execution["start"] + assert execution["goal"] in {node["id"] for node in execution["nodes"]} + + +def test_collective_operator_expands_and_composes_with_press() -> None: + program = { + "schema_version": TASK_AGENT_SCHEMA, + "task": "arrange_then_press", + "goal": "Arrange both cans, then press the button.", + "semantic_steps": [ + { + "id": "s01_line", + "operator": "arrange_line", + "objects": ["can_a", "can_b"], + "goal": {"axis": "world_y", "anchor": "table_center"}, + "depends_on": [], + }, + { + "id": "s02_press", + "operator": "press", + "object": "button", + "goal": {}, + "depends_on": ["s01_line"], + }, + ], + } + + execution = compile_task_agent(program) + steps = {step["id"]: step for step in execution["semantic_steps"]} + + assert list(steps) == ["s01_line__01", "s01_line__02", "s02_press"] + assert steps["s02_press"]["depends_on"] == [ + "s01_line__01", + "s01_line__02", + ] + assert execution["edges"][-1]["depends_on"] == [ + steps["s01_line__01"]["edge_ids"][-1], + steps["s01_line__02"]["edge_ids"][-1], + ] + assert "route" not in repr(execution) + + +def test_coordinated_place_picks_both_objects_before_placement() -> None: + execution = compile_task_agent( + _program( + { + "id": "s01_coordinated_place", + "operator": "coordinated_place", + "object": "cup", + "goal": {"support_object": "tray", "relation": "on"}, + } + ) + ) + step = execution["semantic_steps"][0] + first_edge, placement_edge = [ + next(edge for edge in execution["edges"] if edge["id"] == edge_id) + for edge_id in step["edge_ids"] + ] + + assert [action["atomic_action_class"] for action in first_edge["actions"]] == [ + "PickUp", + "PickUp", + ] + assert [action["actor"] for action in first_edge["actions"]] == [ + {"mode": "required", "arm": "left_arm"}, + {"mode": "required", "arm": "right_arm"}, + ] + assert [action["target_binding"]["object"] for action in first_edge["actions"]] == [ + "cup", + "tray", + ] + assert [action["motion_policy"] for action in first_edge["actions"]] == [ + {"modifiers": []}, + {"modifiers": []}, + ] + assert placement_edge["actions"][0]["atomic_action_class"] == ( + "CoordinatedPlacement" + ) + assert placement_edge["depends_on"] == [first_edge["id"]] + + +def test_independent_required_arms_create_allocation_group() -> None: + program = { + "schema_version": TASK_AGENT_SCHEMA, + "task": "parallel_place", + "goal": "Place two objects with opposite arms.", + "semantic_steps": [ + { + "id": "s01_left", + "operator": "place_relative", + "object": "left_object", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": {"reference_object": "left_tray", "relation": "on"}, + "depends_on": [], + }, + { + "id": "s02_right", + "operator": "place_relative", + "object": "right_object", + "actor": {"mode": "required", "arm": "right_arm"}, + "goal": {"reference_object": "right_tray", "relation": "on"}, + "depends_on": [], + }, + ], + } + + execution = compile_task_agent(program) + + assert execution["allocation_groups"] == [ + { + "id": "g01_distinct_arms", + "semantic_step_ids": ["s01_left", "s02_right"], + "arm_constraint": "distinct_arms", + "execution_policy": "parallel_if_feasible", + "parallel_action_classes": ["PickUp"], + "workspace_policy": "shared_target_serial", + } + ] + + +def test_orient_object_composes_upright_motion_modifier() -> None: + execution = compile_task_agent( + { + "schema_version": TASK_AGENT_SCHEMA, + "task": "upright", + "goal": "Stand the can upright.", + "semantic_steps": [ + { + "id": "s01_orient", + "operator": "orient_object", + "object": "can", + "actor": {"mode": "auto"}, + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + }, + "depends_on": [], + } + ], + "allocation_groups": [], + } + ) + + assert [edge["actions"][0]["motion_policy"] for edge in execution["edges"]] == [ + {"modifiers": [{"type": "orientation", "mode": "upright"}]}, + {"modifiers": [{"type": "orientation", "mode": "upright"}]}, + {"modifiers": [{"type": "orientation", "mode": "upright"}]}, + {"modifiers": [{"type": "orientation", "mode": "upright"}]}, + {"modifiers": [{"type": "orientation", "mode": "upright"}]}, + {"modifiers": []}, + ] + move_phases = [ + edge["actions"][0]["target_binding"]["phase"] + for edge in execution["edges"] + if edge["actions"][0]["atomic_action_class"] == "MoveHeldObject" + ] + assert move_phases == ["staging", "final"] + assert execution["semantic_steps"][0]["goal"] == { + "relation": "none", + "reference_state": "live", + "orientation_goal": "upright", + "orientation_axis": "none", + "position_anchor": "initial_xy", + "support_object": "table", + "upright_local_axis": "auto", + } + + +def test_auto_pickups_require_shared_explicit_allocation_group() -> None: + program = { + "schema_version": TASK_AGENT_SCHEMA, + "task": "dual_arm_basket", + "goal": "Use both arms to place the cube and cup in the basket.", + "semantic_steps": [ + { + "id": "s01_cube", + "operator": "place_relative", + "object": "cube", + "actor": { + "mode": "auto", + "allocation_group": "dual_arms_1", + }, + "goal": {"reference_object": "basket", "relation": "inside"}, + "depends_on": [], + }, + { + "id": "s02_cup", + "operator": "place_relative", + "object": "cup", + "actor": { + "mode": "auto", + "allocation_group": "dual_arms_1", + }, + "goal": {"reference_object": "basket", "relation": "inside"}, + "depends_on": [], + }, + ], + } + + execution = compile_task_agent(program) + pickup_edges = [ + next( + edge + for edge in execution["edges"] + if edge["semantic_step_id"] == step_id + and edge["actions"][0]["atomic_action_class"] == "PickUp" + ) + for step_id in ("s01_cube", "s02_cup") + ] + transport_edges = [ + next( + edge + for edge in execution["edges"] + if edge["semantic_step_id"] == step_id + and edge["actions"][0]["atomic_action_class"] == "MoveHeldObject" + ) + for step_id in ("s01_cube", "s02_cup") + ] + + assert execution["allocation_groups"] == [ + { + "id": "g01_distinct_arms", + "semantic_step_ids": ["s01_cube", "s02_cup"], + "arm_constraint": "distinct_arms", + "execution_policy": "parallel_if_feasible", + "parallel_action_classes": ["PickUp"], + "workspace_policy": "shared_target_serial", + } + ] + assert all("workspace:basket" not in edge["resources"] for edge in pickup_edges) + assert all("workspace:basket" in edge["resources"] for edge in transport_edges) + + for step in program["semantic_steps"]: + step["actor"].pop("allocation_group") + assert compile_task_agent(program)["allocation_groups"] == [] + + +def test_unrelated_dependent_is_allowed_while_hold_reserves_arm() -> None: + program = { + "schema_version": TASK_AGENT_SCHEMA, + "task": "hold_then_press", + "goal": "Hold the cube and then press the button.", + "semantic_steps": [ + { + "id": "s01_hold", + "operator": "hold_hover", + "object": "cube", + "actor": {"mode": "auto"}, + "goal": {}, + "depends_on": [], + }, + { + "id": "s02_press", + "operator": "press", + "object": "button", + "actor": {"mode": "auto"}, + "goal": {}, + "depends_on": ["s01_hold"], + }, + ], + } + + execution = compile_task_agent(program) + steps = {step["id"]: step for step in execution["semantic_steps"]} + + assert steps["s01_hold"]["postcondition"] == { + "type": "object_held", + "object": "cube", + } + assert steps["s02_press"]["depends_on"] == ["s01_hold"] + + +def test_hold_may_follow_an_ancestor_that_previously_used_the_same_object() -> None: + program = { + "schema_version": TASK_AGENT_SCHEMA, + "task": "place_then_hold", + "goal": "Place the cube, then pick it up and keep holding it.", + "semantic_steps": [ + { + "id": "s01_place", + "operator": "place_relative", + "object": "cube", + "actor": {"mode": "auto"}, + "goal": {"reference_object": "tray", "relation": "on"}, + "depends_on": [], + }, + { + "id": "s02_hold", + "operator": "hold_hover", + "object": "cube", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": {}, + "depends_on": ["s01_place"], + }, + ], + } + + execution = compile_task_agent(program) + + assert execution["semantic_steps"][-1]["postcondition"]["type"] == "object_held" + + +@pytest.mark.parametrize( + ("operator", "actor", "goal"), + [ + ("press", {"mode": "required", "arm": "left_arm"}, {}), + ( + "coordinated_transport", + {"mode": "coordinated", "arms": ["left_arm", "right_arm"]}, + {"direction": "none", "terminal_behavior": "hold"}, + ), + ], +) +def test_required_hold_rejects_later_steps_that_need_its_arm( + operator: str, + actor: dict, + goal: dict, +) -> None: + program = { + "schema_version": TASK_AGENT_SCHEMA, + "task": "occupied_arm", + "goal": "Keep holding the cube, then operate the button.", + "semantic_steps": [ + { + "id": "s01_hold", + "operator": "hold_hover", + "object": "cube", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": {}, + "depends_on": [], + }, + { + "id": "s02_other", + "operator": operator, + "object": "button", + "actor": actor, + "goal": goal, + "depends_on": ["s01_hold"], + }, + ], + } + + with pytest.raises(ValueError, match="reserves arm 'left_arm'"): + compile_task_agent(program) + + +def test_held_object_cannot_be_reused_by_an_independent_step() -> None: + program = { + "schema_version": TASK_AGENT_SCHEMA, + "task": "conflicting_object_ownership", + "goal": "Hold and place the same cube.", + "semantic_steps": [ + { + "id": "s01_hold", + "operator": "hold_hover", + "object": "cube", + "actor": {"mode": "auto"}, + "goal": {}, + "depends_on": [], + }, + { + "id": "s02_place", + "operator": "place_relative", + "object": "cube", + "actor": {"mode": "auto"}, + "goal": {"reference_object": "tray", "relation": "on"}, + "depends_on": [], + }, + ], + } + + with pytest.raises(ValueError, match="reserves object 'cube'"): + compile_task_agent(program) + + +def test_unknown_operator_is_rejected_before_graph_construction() -> None: + with pytest.raises(ValueError, match="Unknown semantic operator"): + compile_task_agent( + _program( + { + "id": "s01_unknown", + "operator": "teleport", + "object": "cube", + "goal": {}, + } + ) + ) + + +def test_coordinated_transport_rejects_unknown_direction() -> None: + with pytest.raises(ValueError, match="direction"): + compile_task_agent( + _program( + { + "id": "s01_transport", + "operator": "coordinated_transport", + "object": "shared_box", + "actor": { + "mode": "coordinated", + "arms": ["left_arm", "right_arm"], + }, + "goal": { + "direction": "somewhere_vague", + "terminal_behavior": "hold", + }, + "depends_on": [], + } + ) + ) diff --git a/tests/gen_sim/action_engine/compiler/test_v2.py b/tests/gen_sim/action_engine/compiler/test_v2.py new file mode 100644 index 000000000..6409a6a7a --- /dev/null +++ b/tests/gen_sim/action_engine/compiler/test_v2.py @@ -0,0 +1,172 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import pytest + +from embodichain.gen_sim.action_engine.compiler import ( + compile_task_agent, + compile_task_agent_v2, + seed_graph_to_execution_program, +) +from embodichain.gen_sim.action_engine.domain import TASK_AGENT_SCHEMA +from embodichain.gen_sim.action_engine.protocol import SEED_GRAPH_SCHEMA + + +def _task_agent() -> dict: + return { + "schema_version": TASK_AGENT_SCHEMA, + "task": "place-cup", + "goal": "Place the cup in the tray.", + "semantic_steps": [ + { + "id": "s01_place_cup", + "operator": "place_relative", + "object": "cup", + "actor": {"mode": "auto"}, + "goal": { + "relation": "inside", + "reference_object": "tray", + "reference_state": "live", + "slot": "auto", + "orientation_goal": "preserve", + "orientation_axis": "none", + }, + "depends_on": [], + } + ], + "allocation_groups": [], + } + + +def test_v2_compiler_preserves_mature_atomic_action_topology() -> None: + known = {"cup", "tray"} + legacy = compile_task_agent(_task_agent(), known_objects=known) + seed = compile_task_agent_v2(_task_agent(), known_objects=known) + materialized = seed_graph_to_execution_program(seed, known_objects=known) + + legacy_actions = [ + action["atomic_action_class"] + for edge in legacy["edges"] + for action in edge["actions"] + ] + seed_actions = [node["atomic_action"] for node in seed["nodes"]] + materialized_actions = [ + action["atomic_action_class"] + for edge in materialized["edges"] + for action in edge["actions"] + ] + assert seed["schema_version"] == SEED_GRAPH_SCHEMA + assert seed_actions == legacy_actions + assert materialized_actions == legacy_actions + assert seed["task_groups"][0]["task_type"] == "E1" + + +@pytest.mark.parametrize( + ("operator", "objects", "goal", "actor"), + [ + ( + "orient_object", + ["can"], + { + "orientation_goal": "upright", + "orientation_axis": "long_axis", + "position_anchor": "initial_xy", + "support_object": "table", + "upright_local_axis": "long_axis", + }, + {"mode": "auto"}, + ), + ( + "coordinated_transport", + ["tray"], + { + "direction": "up", + "terminal_behavior": "hold", + "orientation_goal": "preserve", + "orientation_axis": "none", + }, + {"mode": "coordinated", "arms": ["left_arm", "right_arm"]}, + ), + ( + "build_stack", + ["cube_a", "cube_b"], + { + "anchor": "table_center", + "stack_mode": "on_top", + "orientation_goal": "preserve", + "orientation_axis": "none", + }, + {"mode": "auto"}, + ), + ( + "arrange_line", + ["cube_a", "cube_b"], + { + "anchor": "table_center", + "axis": "world_y", + "order_by": "explicit", + "order_constraint": "ordered", + "order_direction": "given", + "orientation_goal": "preserve", + "orientation_axis": "none", + "participation": "auto", + }, + {"mode": "auto"}, + ), + ], +) +def test_v2_preserves_all_current_task_recipe_topologies( + operator: str, + objects: list[str], + goal: dict, + actor: dict, +) -> None: + step = { + "id": "task_01", + "operator": operator, + "actor": actor, + "goal": goal, + "depends_on": [], + } + if operator in {"build_stack", "arrange_line"}: + step["objects"] = objects + else: + step["object"] = objects[0] + task = { + "schema_version": TASK_AGENT_SCHEMA, + "task": f"regression-{operator}", + "goal": f"Regression task for {operator}.", + "semantic_steps": [step], + "allocation_groups": [], + } + known = {*objects, "table"} + legacy = compile_task_agent(task, known_objects=known) + seed = compile_task_agent_v2(task, known_objects=known) + rematerialized = seed_graph_to_execution_program(seed, known_objects=known) + + def signature(program: dict) -> dict[str, list[list[str]]]: + edges = {edge["id"]: edge for edge in program["edges"]} + return { + step["id"]: [ + [action["atomic_action_class"] for action in edges[edge_id]["actions"]] + for edge_id in step["edge_ids"] + ] + for step in program["semantic_steps"] + } + + assert signature(rematerialized) == signature(legacy) diff --git a/tests/gen_sim/action_engine/config/__init__.py b/tests/gen_sim/action_engine/config/__init__.py new file mode 100644 index 000000000..355d915ff --- /dev/null +++ b/tests/gen_sim/action_engine/config/__init__.py @@ -0,0 +1,17 @@ +# ---------------------------------------------------------------------------- +# 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 diff --git a/tests/gen_sim/action_engine/config/test_runtime_policy.py b/tests/gen_sim/action_engine/config/test_runtime_policy.py new file mode 100644 index 000000000..119ec5c71 --- /dev/null +++ b/tests/gen_sim/action_engine/config/test_runtime_policy.py @@ -0,0 +1,307 @@ +# ---------------------------------------------------------------------------- +# 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 hashlib +import json + +import pytest + +from embodichain.gen_sim.action_engine.config import ( + ArmSelectionPolicyCfg, + RuntimePolicyCfg, + default_runtime_policy, + generation_defaults, + resolve_agent_runtime_policy, + runtime_policy_hash, +) +from embodichain.lab.sim.atomic_actions.primitives.place import PlaceOptions + + +def test_default_runtime_policy_preserves_current_arm_selection_behavior() -> None: + policy = default_runtime_policy("dual_ur10") + + assert policy.arm_selection.as_mapping() == { + "crossing_deadband_ratio": 0.08, + "allow_cross_side_fallback": False, + "pickup_crossing_weight": 1.0, + "placement_crossing_weight": 1.5, + "motion_cost_scale": pytest.approx(3.141592653589793), + "fallback_workspace_half_width": 0.5, + "orient_object_preferred_arm_deadband": 0.02, + } + + +def test_defaults_cover_current_execution_and_generation_policy() -> None: + runtime = default_runtime_policy("dual_ur10") + generation = generation_defaults() + + assert runtime.execution == { + "max_transitions": 1000, + "semantic_step_settle_steps": 10, + "max_retries_per_action": 2, + "max_graph_revisions": 8, + "max_recovery_actions": 12, + "support_stability_samples": 3, + "support_stability_interval_steps": 5, + "support_linear_velocity_tolerance": pytest.approx(0.02), + "support_angular_velocity_tolerance": pytest.approx(0.2), + } + assert runtime.planner == { + "backend": "curobo", + "single_arm_strategy": "motion_gen", + "coordinated_strategy": "ik_interp", + "fallback_strategy": "ik_interp", + "allow_fallback": True, + "dynamic_collision": False, + "static_obstacle_uids": [], + "dynamic_obstacle_uids": [], + "curobo": { + "log_level": "error", + "obstacle_representation": "cuboid", + "multi_env": False, + "use_cuda_graph": True, + "preserve_plan_samples": False, + "max_attempts": 5, + "collision_activation_distance": pytest.approx(0.01), + }, + } + assert runtime.grounding["arrangement"]["row_search_radius"] == 0.25 + assert runtime.grasp["antipodal_n_sample"] == 10000 + assert runtime.motion_modifiers["orientation"]["upright"]["MoveHeldObject"][ + "surface_clearance" + ] == pytest.approx(0.05) + assert runtime.motion_modifiers["handover_role"]["transfer"]["PickUp"] == { + "sample_interval": 80, + "hand_interp_steps": 5, + "pick_object_part": "top", + } + assert runtime.motion_defaults["HandOver"]["receive_pick_object_part"] == "bottom" + assert runtime.motion_defaults["CoordinatedPickment"][ + "middle_empty_ratio" + ] == pytest.approx(0.4) + assert runtime.motion_defaults["CoordinatedPickment"][ + "is_filter_ground_collision" + ] is False + assert runtime.predicate_fallbacks["upright_max_tilt"] == pytest.approx( + 0.2617993877991494 + ) + assert generation["physics"]["rigid_object"]["mass"] == pytest.approx(0.1) + assert generation["environment"]["arm_aim_yaw_offset"] == { + "left": pytest.approx(0.0), + "right": pytest.approx(0.0), + } + assert generation["environment"]["recording"] == { + "enabled": True, + "resolution": [640, 360], + "interval_step": 5, + } + assert generation["scene"]["object_length_sample_points"] == 5000 + assert generation["dataset"]["control_frequency"] == 25 + assert generation["randomization"]["table_height_delta_range"] == [ + [-0.05], + [0.05], + ] + + +def test_place_defaults_fit_the_mainline_motion_sample_budget() -> None: + place = default_runtime_policy("dual_ur10").motion_defaults["Place"] + sample_count = int(place["sample_interval"]) + hand_steps = PlaceOptions().hand_interp_steps + motion_steps = sample_count - hand_steps + down_steps = int(round(motion_steps) * 0.6) + back_steps = motion_steps - down_steps + cartesian_count = int(place["cartesian_waypoint_count"]) + + assert 1 + 2 * cartesian_count <= down_steps + assert 1 + cartesian_count <= back_steps + + +def test_default_runtime_policy_returns_detached_profile_snapshots() -> None: + first = default_runtime_policy("dual_ur10") + second = default_runtime_policy("dual_ur10") + franka = default_runtime_policy("dual_franka") + + first.arm_selection.pickup_crossing_weight = 9.0 + first.motion_defaults["PickUp"]["lift_height"] = 9.0 + + assert second.arm_selection.pickup_crossing_weight == 1.0 + assert second.motion_defaults["PickUp"]["lift_height"] == 0.30 + assert franka.arm_selection.pickup_crossing_weight == 1.0 + assert franka.motion_defaults["MoveEndEffector"]["retreat_height"] == 0.10 + + +def test_generation_defaults_return_detached_values() -> None: + first = generation_defaults() + second = generation_defaults() + + first["physics"]["rigid_object"]["mass"] = 9.0 + + assert second["physics"]["rigid_object"]["mass"] == pytest.approx(0.1) + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("crossing_deadband_ratio", 1.0, "crossing_deadband_ratio"), + ("pickup_crossing_weight", -0.1, "pickup_crossing_weight"), + ("placement_crossing_weight", -0.1, "placement_crossing_weight"), + ("motion_cost_scale", 0.0, "motion_cost_scale"), + ("fallback_workspace_half_width", 0.0, "fallback_workspace_half_width"), + ], +) +def test_arm_selection_policy_rejects_invalid_values( + field: str, + value: float, + message: str, +) -> None: + values = default_runtime_policy("dual_ur10").arm_selection.as_mapping() + values[field] = value + + with pytest.raises(ValueError, match=message): + ArmSelectionPolicyCfg.from_mapping(values) + + +def test_arm_selection_policy_requires_boolean_cross_side_fallback() -> None: + values = default_runtime_policy("dual_ur10").arm_selection.as_mapping() + values["allow_cross_side_fallback"] = "false" + + with pytest.raises(TypeError, match="allow_cross_side_fallback"): + ArmSelectionPolicyCfg.from_mapping(values) + + +def test_arm_selection_policy_loads_old_snapshot_without_fallback_field() -> None: + snapshot = default_runtime_policy("dual_ur10").as_mapping() + snapshot["arm_selection"].pop("allow_cross_side_fallback") + payload = json.dumps( + snapshot, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ).encode("utf-8") + + policy = resolve_agent_runtime_policy( + { + "robot_profile": "dual_ur10", + "runtime_policy": snapshot, + "runtime_policy_hash": hashlib.sha256(payload).hexdigest(), + } + ) + + assert policy.arm_selection.allow_cross_side_fallback is False + + +def test_agent_policy_snapshot_is_hash_verified_and_legacy_config_falls_back() -> None: + policy = default_runtime_policy("dual_ur5") + snapshot = policy.as_mapping() + config = { + "robot_profile": "dual_ur5", + "runtime_policy": snapshot, + "runtime_policy_hash": runtime_policy_hash(policy), + } + + resolved = resolve_agent_runtime_policy(config) + assert resolved.as_mapping() == snapshot + + tampered = deepcopy(config) + tampered["runtime_policy"]["motion_defaults"]["PickUp"]["lift_height"] = 8.0 + with pytest.raises(ValueError, match="hash does not match"): + resolve_agent_runtime_policy(tampered) + + legacy = resolve_agent_runtime_policy({"robot_profile": "dual_ur5"}) + assert legacy.as_mapping() == snapshot + + +def test_narrow_v1_policy_snapshot_is_migrated_to_complete_runtime_policy() -> None: + snapshot = { + "schema_version": "action_engine_runtime_policy_v1", + "arm_selection": { + "crossing_deadband_ratio": 0.08, + "pickup_crossing_weight": 2.0, + "placement_crossing_weight": 1.5, + "motion_cost_scale": 3.141592653589793, + "fallback_workspace_half_width": 0.5, + }, + } + payload = json.dumps( + snapshot, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ).encode("utf-8") + + resolved = resolve_agent_runtime_policy( + { + "robot_profile": "dual_ur10", + "runtime_policy": snapshot, + "runtime_policy_hash": hashlib.sha256(payload).hexdigest(), + } + ) + + assert resolved.arm_selection.pickup_crossing_weight == 2.0 + assert resolved.motion_defaults["PickUp"]["lift_height"] == 0.30 + + +def test_v3_policy_snapshot_is_migrated_with_default_planner_policy() -> None: + expected = default_runtime_policy("dual_ur10") + snapshot = expected.as_mapping() + snapshot.pop("planner") + snapshot["schema_version"] = "action_engine_runtime_policy_v3" + payload = json.dumps( + snapshot, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ).encode("utf-8") + + resolved = resolve_agent_runtime_policy( + { + "robot_profile": "dual_ur10", + "runtime_policy": snapshot, + "runtime_policy_hash": hashlib.sha256(payload).hexdigest(), + } + ) + + assert resolved.schema_version == "action_engine_runtime_policy_v6" + assert resolved.planner == expected.planner + + +def test_curobo_policy_rejects_coordinated_motion_generation() -> None: + snapshot = default_runtime_policy("dual_ur10").as_mapping() + snapshot["planner"]["coordinated_strategy"] = "motion_gen" + + with pytest.raises(ValueError, match="coordinated_strategy"): + RuntimePolicyCfg.from_mapping(snapshot) + + +@pytest.mark.parametrize( + ("patch", "message"), + [ + ({"fallback_strategy": "motion_gen"}, "fallback_strategy"), + ({"backend": "toppra", "dynamic_collision": True}, "dynamic_collision"), + ], +) +def test_planner_policy_rejects_unsupported_combinations( + patch: dict[str, object], + message: str, +) -> None: + snapshot = default_runtime_policy("dual_ur10").as_mapping() + snapshot["planner"].update(patch) + + with pytest.raises(ValueError, match=message): + RuntimePolicyCfg.from_mapping(snapshot) diff --git a/tests/gen_sim/action_engine/domain/__init__.py b/tests/gen_sim/action_engine/domain/__init__.py new file mode 100644 index 000000000..c8e03f284 --- /dev/null +++ b/tests/gen_sim/action_engine/domain/__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 domain tests.""" diff --git a/tests/gen_sim/action_engine/domain/test_programs.py b/tests/gen_sim/action_engine/domain/test_programs.py new file mode 100644 index 000000000..8fdf8074b --- /dev/null +++ b/tests/gen_sim/action_engine/domain/test_programs.py @@ -0,0 +1,150 @@ +# ---------------------------------------------------------------------------- +# 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.compiler import compile_task_agent +from embodichain.gen_sim.action_engine.domain import ( + TASK_AGENT_SCHEMA, + execution_program_hash, + validate_execution_program, + validate_task_agent, +) + + +def _task_agent() -> dict: + return { + "schema_version": TASK_AGENT_SCHEMA, + "task": "place_demo", + "goal": "Place the cup on the tray.", + "semantic_steps": [ + { + "id": "s01_place", + "operator": "place_relative", + "object": "cup", + "actor": {"mode": "auto"}, + "goal": {"reference_object": "tray", "relation": "on"}, + "depends_on": [], + } + ], + } + + +def test_task_validation_is_detached_and_adds_unambiguous_defaults() -> None: + source = _task_agent() + del source["semantic_steps"][0]["actor"] + del source["semantic_steps"][0]["depends_on"] + + validated = validate_task_agent(source) + validated["semantic_steps"][0]["goal"]["relation"] = "inside" + + assert source["semantic_steps"][0]["goal"]["relation"] == "on" + assert validated["semantic_steps"][0]["actor"] == {"mode": "auto"} + assert validated["semantic_steps"][0]["depends_on"] == [] + + +@pytest.mark.parametrize( + "actor", + [ + {"mode": "auto", "allocation_group": "dual_arms_1"}, + { + "mode": "required", + "arm": "left_arm", + "allocation_group": "dual_arms_1", + }, + ], +) +def test_single_arm_allocation_group_is_validated_and_preserved(actor: dict) -> None: + source = _task_agent() + source["semantic_steps"][0]["actor"] = actor + + validated = validate_task_agent(source) + execution = compile_task_agent(validated) + + assert validated["semantic_steps"][0]["actor"] == actor + assert execution["semantic_steps"][0]["actor"] == actor + assert all( + action["actor"] == actor + for edge in execution["edges"] + for action in edge["actions"] + ) + + +def test_allocation_group_must_be_nonempty_and_single_arm_only() -> None: + source = _task_agent() + source["semantic_steps"][0]["actor"]["allocation_group"] = " " + with pytest.raises(ValueError, match="allocation_group"): + validate_task_agent(source) + + source["semantic_steps"][0]["actor"] = { + "mode": "coordinated", + "arms": ["left_arm", "right_arm"], + "allocation_group": "dual_arms_1", + } + with pytest.raises(ValueError, match="unknown fields"): + validate_task_agent(source) + + +def test_task_validation_rejects_cycles_and_grounded_values() -> None: + cyclic = _task_agent() + cyclic["semantic_steps"].extend( + [ + { + "id": "s02", + "operator": "press", + "object": "button", + "depends_on": ["s03"], + }, + { + "id": "s03", + "operator": "press", + "object": "button", + "depends_on": ["s02"], + }, + ] + ) + with pytest.raises(ValueError, match="cycle"): + validate_task_agent(cyclic) + + grounded = _task_agent() + grounded["semantic_steps"][0]["goal"]["target_pose"] = [0.0, 0.0, 0.0] + with pytest.raises(ValueError, match="grounded runtime data"): + validate_task_agent(grounded) + + +def test_execution_hash_is_stable_and_validation_is_strict() -> None: + execution = compile_task_agent(_task_agent()) + reordered = {key: execution[key] for key in reversed(list(execution))} + + assert execution_program_hash(execution) == execution_program_hash(reordered) + assert len(execution_program_hash(execution)) == 64 + + broken = deepcopy(execution) + broken["edges"][0]["target_binding"] = {} + with pytest.raises(ValueError, match="unknown fields"): + validate_execution_program(broken) + + +def test_execution_validation_rejects_unowned_edges() -> None: + execution = compile_task_agent(_task_agent()) + execution["semantic_steps"][0]["edge_ids"].pop() + + with pytest.raises(ValueError, match="unowned edges"): + validate_execution_program(execution) diff --git a/tests/gen_sim/action_engine/domain/test_task_contracts.py b/tests/gen_sim/action_engine/domain/test_task_contracts.py new file mode 100644 index 000000000..765e05aa8 --- /dev/null +++ b/tests/gen_sim/action_engine/domain/test_task_contracts.py @@ -0,0 +1,71 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import pytest + +from embodichain.gen_sim.action_engine.domain import ( + RELATIONS, + TASK_CONTRACTS, + TASK_TYPES, + TERMINAL_BEHAVIORS, + TRANSPORT_DIRECTIONS, + normalize_placement_relation, + task_contract, + task_success_type, +) + + +def test_task_contract_catalog_covers_the_canonical_protocol() -> None: + assert set(TASK_CONTRACTS) == set(TASK_TYPES) + assert all(contract.core_actions for contract in TASK_CONTRACTS.values()) + assert {contract.source_structure for contract in TASK_CONTRACTS.values()} == { + "articulation", + "rigid_object", + } + assert task_contract("E2").success_type == "object_upright" + assert task_contract("E5").scene_affordances == { + "dual_graspable", + "rigid", + } + + +def test_e5_success_depends_only_on_terminal_behavior() -> None: + assert task_success_type("E5", {"terminal_behavior": "hold"}) == ( + "held_by_both_grippers" + ) + assert task_success_type("E5", {"terminal_behavior": "place"}) == "semantic_goal" + with pytest.raises(ValueError, match="terminal_behavior"): + task_success_type("E5", {"terminal_behavior": "none"}) + + +def test_symbolic_transport_values_are_language_neutral_protocol_enums() -> None: + assert {"on", "inside", "behind", "left_of"} <= RELATIONS + assert {"none", "up", "left", "world_y"} <= TRANSPORT_DIRECTIONS + assert TERMINAL_BEHAVIORS == {"none", "hold", "place"} + + +@pytest.mark.parametrize("relation", ["above", "on_top", "on_top_of"]) +def test_released_hover_and_legacy_support_relations_normalize_to_on( + relation: str, +) -> None: + assert normalize_placement_relation(relation) == "on" + + +def test_placement_relation_normalization_rejects_non_spatial_semantics() -> None: + with pytest.raises(ValueError, match="Unsupported placement relation"): + normalize_placement_relation("visual_slot") diff --git a/tests/gen_sim/action_engine/domain/test_v2.py b/tests/gen_sim/action_engine/domain/test_v2.py new file mode 100644 index 000000000..79604756b --- /dev/null +++ b/tests/gen_sim/action_engine/domain/test_v2.py @@ -0,0 +1,281 @@ +# ---------------------------------------------------------------------------- +# 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.capabilities import ( + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.domain import ( + motion_policy, + public_task_spec, + seed_graph_hash, + validate_public_task_spec, + validate_scene_requirements, + validate_seed_graph, + validate_task_spec, +) +from embodichain.gen_sim.action_engine.protocol import ( + SCENE_REQUIREMENTS_SCHEMA, + SEED_GRAPH_SCHEMA, + TASK_SPEC_SCHEMA, +) +from embodichain.gen_sim.action_engine.planning.linker import link_seed_graph + + +def _seed_graph() -> dict: + registry = build_atomic_capability_registry() + draft = { + "schema_version": SEED_GRAPH_SCHEMA, + "task_id": "place-cup", + "instruction": "Place the cup in the tray.", + "level": "L1", + "reasoning_type": "none", + "planner_route": "offline", + "nodes": [ + { + "id": "pick_cup", + "atomic_action": "PickUp", + "object_uid": "cup", + "actor": {"mode": "auto"}, + "control": "arm", + "target_binding": {"kind": "object", "object": "cup"}, + "depends_on": [], + "task_instance_id": "e1_001", + "task_type": "E1", + "role": "primary", + "precondition": {"type": "object_not_fallen", "object": "cup"}, + "postcondition": {"type": "object_held", "object": "cup"}, + "motion_policy": motion_policy(), + }, + { + "id": "place_cup", + "atomic_action": "Place", + "object_uid": "cup", + "actor": {"mode": "auto"}, + "control": "arm", + "target_binding": {"kind": "current_held_pose"}, + "depends_on": ["pick_cup"], + "task_instance_id": "e1_001", + "task_type": "E1", + "role": "primary", + "precondition": {"type": "object_held", "object": "cup"}, + "postcondition": { + "type": "object_in_container", + "object": "cup", + "container": "tray", + }, + "motion_policy": motion_policy(), + }, + ], + "task_groups": [ + { + "id": "e1_001", + "task_type": "E1", + "role": "primary", + "operator": "place_relative", + "object_uid": "cup", + "actor": {"mode": "auto"}, + "goal": {"relation": "inside", "reference_object": "tray"}, + "depends_on": [], + "node_ids": ["pick_cup", "place_cup"], + "success": { + "type": "object_in_container", + "object": "cup", + "container": "tray", + }, + } + ], + "success": { + "type": "object_in_container", + "object": "cup", + "container": "tray", + }, + "capability_catalog_hash": registry.catalog_hash(), + "metadata": {}, + } + return link_seed_graph( + draft, + registry=registry, + task_order=["e1_001"], + known_objects={"cup", "tray"}, + ) + + +def test_seed_graph_validates_direct_atomic_action_nodes() -> None: + registry = build_atomic_capability_registry() + graph = validate_seed_graph( + _seed_graph(), + known_objects={"cup", "tray"}, + known_actions=registry.names(), + executable_actions=registry.executable_names(), + require_executable=True, + ) + assert [node["atomic_action"] for node in graph["nodes"]] == ["PickUp", "Place"] + assert graph["task_groups"][0]["node_ids"] == ["pick_cup", "place_cup"] + + +def test_seed_graph_rejects_grounded_motion_and_cycles() -> None: + grounded = _seed_graph() + grounded["nodes"][0]["target_binding"]["target_pose"] = [0.0, 0.0, 0.0] + with pytest.raises(ValueError, match="grounded motion data"): + validate_seed_graph(grounded) + + cyclic = _seed_graph() + cyclic["nodes"][0]["depends_on"] = ["place_cup"] + with pytest.raises(ValueError, match="dependency cycle"): + validate_seed_graph(cyclic) + + +def test_seed_graph_rejects_unknown_uids_illegal_groups_and_resource_conflicts() -> ( + None +): + with pytest.raises(ValueError, match="unknown object"): + validate_seed_graph(_seed_graph(), known_objects={"tray"}) + + illegal_group = _seed_graph() + illegal_group["task_groups"][0]["task_type"] = "E9" + for node in illegal_group["nodes"]: + node["task_type"] = "E9" + with pytest.raises(ValueError, match="core actions"): + validate_seed_graph(illegal_group) + + conflicting = _seed_graph() + conflicting["nodes"][1]["depends_on"] = [] + conflicting["task_groups"][0]["contract"]["entry_node_ids"] = [ + "pick_cup", + "place_cup", + ] + conflicting["task_groups"][0]["contract"]["terminal_node_ids"] = [ + "pick_cup", + "place_cup", + ] + with pytest.raises(ValueError, match="resource conflicts"): + validate_seed_graph(conflicting) + + +def test_seed_graph_hash_is_order_stable_and_detached() -> None: + graph = _seed_graph() + original = deepcopy(graph) + first = seed_graph_hash(graph) + second = seed_graph_hash({key: graph[key] for key in reversed(graph)}) + assert first == second + assert graph == original + + +def test_task_spec_enforces_reasoning_level_and_repetition_shape() -> None: + spec = { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "upright-cans", + "level": "L2", + "instruction": "Stand both cans upright.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "e2_1", + "task_type": "E2", + "params": {"object_role": "can_1"}, + "depends_on": [], + }, + { + "id": "e2_2", + "task_type": "E2", + "params": {"object_role": "can_2"}, + "depends_on": [], + }, + ], + "success": {"type": "all_upright"}, + "oracle": {"object_roles": ["can_1", "can_2"]}, + "metadata": {}, + } + assert validate_task_spec(spec)["level"] == "L2" + spec["level"] = "L4" + with pytest.raises(ValueError, match="non-'none'"): + validate_task_spec(spec) + + +def test_public_l4_task_hides_oracle_and_reference_instances() -> None: + spec = { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "complete-mouth", + "level": "L4", + "instruction": "Complete the missing mouth.", + "reasoning_type": "visual_semantics", + "task_instances": [ + { + "id": "hidden_e1", + "task_type": "E1", + "params": {"object_role": "mouth", "target_role": "face"}, + "depends_on": [], + "role": "primary", + } + ], + "success": {"type": "visual_part_complete"}, + "oracle": {"missing_part": "mouth"}, + "metadata": {}, + } + + public = public_task_spec(spec) + + assert "oracle" not in public + assert "task_instances" not in public + assert validate_public_task_spec(public) == public + + +def test_scene_requirements_validate_task_first_handoff() -> None: + requirements = validate_scene_requirements( + { + "schema_version": SCENE_REQUIREMENTS_SCHEMA, + "task_id": "upright-cans", + "objects": [ + { + "role_id": "can", + "category": "can", + "count": 2, + "affordances": ["graspable", "orientable"], + "initial_state": {"orientation": "fallen"}, + "attributes": {}, + } + ], + "cameras": [{"role": "overview", "requires_rgb": True}], + "spatial_constraints": [], + "distractor_count": 0, + "metadata": {}, + } + ) + assert requirements["objects"][0]["count"] == 2 + + +def test_planning_only_capability_is_rejected_before_execution() -> None: + registry = build_atomic_capability_registry() + graph = _seed_graph() + graph["nodes"][0]["atomic_action"] = "Pour" + graph["nodes"][0]["target_binding"] = { + "kind": "pour_goal", + "object": "cup", + "reference_object": "tray", + } + with pytest.raises(ValueError, match="planning-only"): + validate_seed_graph( + graph, + known_actions=registry.names(), + executable_actions=registry.executable_names(), + require_executable=True, + ) 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/generation/test_generation.py b/tests/gen_sim/action_engine/generation/test_generation.py new file mode 100644 index 000000000..cc1cfd91f --- /dev/null +++ b/tests/gen_sim/action_engine/generation/test_generation.py @@ -0,0 +1,1533 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Focused tests for the independent Action Engine generation boundary.""" + +from __future__ import annotations + +import inspect +import json +from pathlib import Path +import sys +from types import ModuleType + +import numpy as np +import pytest + +from embodichain.gen_sim.action_engine.protocol import ( + SCENE_REQUIREMENTS_SCHEMA, + TASK_SPEC_SCHEMA, +) +from embodichain.gen_sim.action_engine.cli import ( + generate_action_agent_config as cli_module, +) +from embodichain.gen_sim.action_engine.cli.generate_action_agent_config import ( + build_parser, +) +from embodichain.gen_sim.action_engine.generation.artifacts import ( + artifact_paths, + write_generation_artifacts, +) +from embodichain.gen_sim.action_engine.generation import ( + config_builder as config_builder_module, +) +from embodichain.gen_sim.action_engine.generation.config_builder import ( + build_agent_config, + build_fast_gym_config, +) +from embodichain.gen_sim.action_engine.generation.generator import ( + _add_ab_camera_requirements, + _scene_requirements_from_bindings, + _task_spec_role_bindings, + generate_action_engine_config, +) +from embodichain.gen_sim.action_engine.generation.source_scene import ( + prepare_scene, + resolve_gym_config_path, + resolve_source_scene, +) +from embodichain.gen_sim.action_engine.tasks import GroundedTaskSpec + + +@pytest.fixture +def gym_export(tmp_path: Path) -> Path: + export = tmp_path / "gym_export" + assets = export / "mesh_assets" + assets.mkdir(parents=True) + (assets / "table.glb").write_bytes(b"not-a-real-glb") + (assets / "can.glb").write_bytes(b"not-a-real-glb") + state = export / "scene_state" + state.mkdir() + (state / "result.json").write_text("{}\n", encoding="utf-8") + + config = { + "id": "Prompt2Scene-test-v0", + "env": {"events": {}, "observations": {}, "dataset": {}}, + "robot": {}, + "sensor": [], + "light": {}, + "background": [ + { + "uid": "table_0", + "description": "A white table.", + "shape": {"shape_type": "Mesh", "fpath": "mesh_assets/table.glb"}, + "init_pos": [0.0, 0.0, 0.0], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + } + ], + "rigid_object": [ + { + "uid": "interact_can_0", + "description": "A red soda can.", + "shape": { + "shape_type": "Mesh", + "fpath": "mesh_assets/can.glb", + "acd_method": "coacd", + "max_convex_hull_num": 32, + }, + "attrs": {"mass": 0.01}, + "init_pos": [1.0, 2.0, 0.7], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + "max_convex_hull_num": 32, + } + ], + } + (export / "gym_config.json").write_text( + json.dumps(config), + encoding="utf-8", + ) + return export + + +@pytest.fixture +def scene_export(tmp_path: Path) -> Path: + export = tmp_path / "scene_export" + assets = export / "mesh_assets" + assets.mkdir(parents=True) + (assets / "table.glb").write_bytes(b"not-a-real-glb") + (assets / "bottle_001.glb").write_bytes(b"not-a-real-glb") + (assets / "bottle_002.glb").write_bytes(b"not-a-real-glb") + + config = { + "format": "embodichain.scene-export/v1", + "scene_id": "scene-export-test", + "background": [ + { + "uid": "table", + "description": "A white table.", + "shape": {"shape_type": "Mesh", "fpath": "mesh_assets/table.glb"}, + "init_pos": [0.0, 0.0, 0.0], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + } + ], + "rigid_object": [ + { + "uid": uid, + "name": f"Bottle {index}", + "description": f"Bottle instance {index}.", + "shape": { + "shape_type": "Mesh", + "fpath": f"mesh_assets/{uid}.glb", + }, + "init_pos": [float(index), float(index + 1), 0.7], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + } + for index, uid in enumerate(("bottle_001", "bottle_002"), start=1) + ], + } + (export / "scene_config.json").write_text( + json.dumps(config), + encoding="utf-8", + ) + return export + + +def _existing_v2_task_spec(task_id: str = "direct_task") -> dict[str, object]: + return { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": task_id, + "level": "L1", + "instruction": "test-instruction", + "reasoning_type": "none", + "task_instances": [ + { + "id": "task_01", + "task_type": "E2", + "params": {"object_role": "object_01"}, + "depends_on": [], + "role": "primary", + } + ], + "success": {"type": "semantic_goal", "task_instance_id": "task_01"}, + "oracle": {}, + "metadata": {"role_bindings": {"object_01": "interact_can"}}, + } + + +def test_prepare_scene_normalizes_uid_paths_and_prompt2scene_transform( + gym_export: Path, +) -> None: + scene = prepare_scene(gym_export) + + assert scene.uid_map == { + "table_0": "table", + "interact_can_0": "interact_can", + } + assert scene.z_rotation_degrees == -90.0 + assert scene.rigid_objects[0]["init_pos"] == [2.0, -1.0, 0.7] + assert scene.rigid_objects[0]["max_convex_hull_num"] == 16 + assert scene.rigid_objects[0]["acd_method"] == "vhacd" + assert scene.rigid_objects[0]["shape"]["acd_method"] == "vhacd" + assert scene.rigid_objects[0]["shape"]["max_convex_hull_num"] == 16 + mesh_path = Path(scene.rigid_objects[0]["shape"]["fpath"]) + assert mesh_path.is_absolute() + assert mesh_path.is_file() + assert scene.planner_objects[1]["source_uid"] == "interact_can_0" + assert scene.planner_objects[1]["uid"] == "interact_can" + + +def test_prepare_scene_supports_scene_export_v1(scene_export: Path) -> None: + scene = prepare_scene(scene_export.parent) + + assert scene.source_config_path == scene_export / "scene_config.json" + assert scene.uid_map == { + "table": "table", + "bottle_001": "bottle_001", + "bottle_002": "bottle_002", + } + assert scene.planner_objects[1]["name"] == "Bottle 1" + assert scene.z_rotation_degrees == -90.0 + assert scene.rigid_objects[0]["init_pos"] == [2.0, -1.0, 0.7] + assert all( + Path(config["shape"]["fpath"]).is_file() + for config in (*scene.background, *scene.rigid_objects) + ) + + +def test_prepare_scene_requires_exactly_one_background(gym_export: Path) -> None: + source_path = gym_export / "gym_config.json" + source = json.loads(source_path.read_text(encoding="utf-8")) + source["background"].append( + { + **source["background"][0], + "uid": "floor_0", + "description": "A floor beneath the work surface.", + } + ) + source_path.write_text(json.dumps(source), encoding="utf-8") + + with pytest.raises(ValueError, match="exactly one background"): + prepare_scene(gym_export) + + +def test_prepare_scene_does_not_treat_physics_attrs_as_semantics( + gym_export: Path, +) -> None: + scene = prepare_scene(gym_export) + rigid_object = next( + item for item in scene.planner_objects if item["role"] == "rigid_object" + ) + + assert rigid_object["attributes"] == {} + + +@pytest.mark.parametrize( + "companion_relative_path", + ( + Path("gym_export/scene_config.json"), + Path("scene_export/scene_config.json"), + ), +) +def test_source_scene_resolution_prefers_gym_config_in_mixed_export( + tmp_path: Path, + companion_relative_path: Path, +) -> None: + gym_export = tmp_path / "gym_export" + gym_export.mkdir(parents=True) + gym_config = gym_export / "gym_config.json" + gym_config.write_text("{}", encoding="utf-8") + companion = tmp_path / companion_relative_path + companion.parent.mkdir(parents=True, exist_ok=True) + companion.write_text( + json.dumps({"format": "embodichain.scene-export/v1"}), encoding="utf-8" + ) + + resolved = resolve_source_scene(tmp_path) + + assert resolved.path == gym_config + assert resolved.source_format == "legacy_gym_config" + assert resolved.is_prompt2scene is True + assert resolve_gym_config_path(tmp_path) == resolved.path + + +def test_explicit_scene_export_config_overrides_mixed_layout( + gym_export: Path, + scene_export: Path, +) -> None: + resolved = resolve_source_scene(scene_export / "scene_config.json") + + assert resolved.path == scene_export / "scene_config.json" + assert resolved.source_format == "embodichain.scene-export/v1" + assert resolved.is_prompt2scene is True + + +def test_scene_export_config_rejects_unknown_format(scene_export: Path) -> None: + config_path = scene_export / "scene_config.json" + config = json.loads(config_path.read_text(encoding="utf-8")) + config["format"] = "embodichain.scene-export/v2" + config_path.write_text(json.dumps(config), encoding="utf-8") + + with pytest.raises(ValueError, match="unsupported format"): + resolve_source_scene(config_path) + + +def test_fast_gym_config_has_runnable_franka_contract(gym_export: Path) -> None: + scene = prepare_scene(gym_export) + config = build_fast_gym_config( + scene, + task_name="line_task", + task_description="Arrange the can.", + robot_profile="franka", + execution_program_hash="a" * 64, + max_episodes=1, + max_episode_steps=2000, + randomize_scene=True, + ) + + assert config["id"] == "ActionEngine-v1" + assert config["robot"]["uid"] == "DualFrankaPanda" + assert config["robot"]["init_pos"][2] == pytest.approx(0.35) + assert config["sensor"][0]["uid"] == "cam_high" + assert config["env"]["extensions"]["agent_robot_profile"] == "dual_franka" + assert config["env"]["extensions"]["agent_static_obstacle_uids"] == ["table"] + assert config["env"]["extensions"]["agent_dynamic_obstacle_uids"] == [ + "interact_can" + ] + assert "agent_grasp_runtime_defaults" not in config["env"]["extensions"] + assert config["env"]["extensions"]["agent_arm_slots"] == { + "left": {"arm": "left_arm", "eef": "left_eef"}, + "right": {"arm": "right_arm", "eef": "right_eef"}, + } + assert config["env"]["extensions"]["arm_aim_yaw_offset"] == { + "left": pytest.approx(0.0), + "right": pytest.approx(0.0), + } + assert config["env"]["extensions"]["action_engine"]["seed_task_graph"] == ( + "seed_task_graph.json" + ) + assert ( + config["env"]["extensions"]["action_engine"]["defaults_schema_version"] + == "action_engine_defaults_v1" + ) + registry = config["env"]["events"]["register_info_to_env"]["params"]["registry"] + assert [entry["entity_cfg"]["uid"] for entry in registry] == ["interact_can"] + assert "randomize_interact_can_pose" in config["env"]["events"] + assert "randomize_table_height" in config["env"]["events"] + recorder = config["env"]["events"]["record_camera"] + assert recorder["interval_step"] == 5 + assert recorder["params"]["resolution"] == [640, 360] + assert recorder["params"]["intrinsics"] == pytest.approx( + [280.0, 280.0, 320.0, 180.0] + ) + object_length = config["env"]["events"]["prepare_extra_attr"]["params"]["attrs"][0] + assert object_length["func_kwargs"]["sample_points"] == 5000 + assert ( + config["env"]["dataset"]["lerobot"]["params"]["robot_meta"]["control_freq"] + == 25 + ) + assert config["env"]["observations"]["norm_robot_eef_joint"]["params"][ + "joint_ids" + ] == list(range(14, 26)) + + +def test_offline_recording_can_be_disabled_but_ab_keeps_audience_recorder( + gym_export: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + defaults = dict(config_builder_module._GENERATION_DEFAULTS) + environment = dict(defaults["environment"]) + recording = dict(environment["recording"]) + recording["enabled"] = False + environment["recording"] = recording + defaults["environment"] = environment + monkeypatch.setattr(config_builder_module, "_GENERATION_DEFAULTS", defaults) + scene = prepare_scene(gym_export) + + offline = build_fast_gym_config( + scene, + task_name="offline_task", + task_description="Offline recording policy.", + robot_profile="franka", + execution_program_hash="a" * 64, + max_episodes=1, + max_episode_steps=100, + ) + ab = build_fast_gym_config( + scene, + task_name="ab_task", + task_description="A/B recording policy.", + robot_profile="franka", + execution_program_hash="b" * 64, + max_episodes=1, + max_episode_steps=100, + planning_mode="ab", + seed_task_graph_path="offline/seed_task_graph.json", + ) + + assert "record_camera" not in offline["env"]["events"] + assert ab["env"]["events"]["record_camera"]["params"]["name"] == ( + "record_cam_audience_view" + ) + assert ab["env"]["events"]["record_camera"]["interval_step"] == 5 + + +@pytest.mark.parametrize( + ("override", "message"), + [ + ({"enabled": "yes"}, "enabled must be a boolean"), + ({"resolution": [640]}, "resolution must contain two positive integers"), + ({"interval_step": 0}, "interval_step must be positive"), + ], +) +def test_recording_policy_rejects_invalid_generation_defaults( + gym_export: Path, + monkeypatch: pytest.MonkeyPatch, + override: dict[str, object], + message: str, +) -> None: + defaults = dict(config_builder_module._GENERATION_DEFAULTS) + environment = dict(defaults["environment"]) + recording = dict(environment["recording"]) + recording.update(override) + environment["recording"] = recording + defaults["environment"] = environment + monkeypatch.setattr(config_builder_module, "_GENERATION_DEFAULTS", defaults) + + with pytest.raises(ValueError, match=message): + build_fast_gym_config( + prepare_scene(gym_export), + task_name="invalid_recording", + task_description="Invalid recording policy.", + robot_profile="franka", + execution_program_hash="c" * 64, + max_episodes=1, + max_episode_steps=100, + ) + + +def test_fast_gym_config_preserves_unicode_instruction_and_uses_task_name_label( + gym_export: Path, +) -> None: + scene = prepare_scene(gym_export) + task_name = "task1000" + task_description = "unicode-λ-instruction" + + config = build_fast_gym_config( + scene, + task_name=task_name, + task_description=task_description, + robot_profile="franka", + execution_program_hash="a" * 64, + max_episodes=1, + max_episode_steps=2000, + ) + + params = config["env"]["dataset"]["lerobot"]["params"] + assert params["instruction"]["lang"] == task_description + assert params["extra"]["task_name"] == task_name + assert params["extra"]["task_description"] == task_name + + +def test_ab_config_uses_offline_branch_and_four_vlm_cameras( + gym_export: Path, + tmp_path: Path, +) -> None: + scene = prepare_scene(gym_export) + graph_path = "offline/seed_task_graph.json" + config = build_fast_gym_config( + scene, + task_name="ab_task", + task_description="test-instruction", + robot_profile="ur10", + execution_program_hash="d" * 64, + max_episodes=1, + max_episode_steps=100, + planning_mode="ab", + seed_task_graph_path=graph_path, + ) + agent = build_agent_config( + task_name="ab_task", + robot_profile="ur10", + execution_program_hash="d" * 64, + source_config_path=scene.source_config_path, + uid_map=scene.uid_map, + planning_mode="ab", + seed_task_graph_path=graph_path, + vlm_model="mimo-vlm", + vlm_camera_uids=[ + "vlm_front", + "vlm_left", + "vlm_rear", + "vlm_right", + ], + ) + paths = artifact_paths(tmp_path, planning_mode="ab") + + assert paths.seed_task_graph == tmp_path.resolve() / graph_path + assert config["env"]["extensions"]["action_engine"]["planning_mode"] == "ab" + assert config["env"]["extensions"]["action_engine"]["seed_task_graph"] == ( + graph_path + ) + vlm_sensors = [ + sensor for sensor in config["sensor"] if sensor["uid"].startswith("vlm_") + ] + assert [sensor["uid"] for sensor in vlm_sensors] == [ + "vlm_front", + "vlm_left", + "vlm_rear", + "vlm_right", + ] + assert all( + sensor["enable_color"] and sensor["enable_depth"] for sensor in vlm_sensors + ) + assert agent["planning_mode"] == "ab" + assert agent["offline_seed_task_graph"] == graph_path + assert agent["vlm_model"] == "mimo-vlm" + assert agent["vlm_camera_uids"] == [ + "vlm_front", + "vlm_left", + "vlm_rear", + "vlm_right", + ] + assert agent["online_planning"] == { + "vlm_model": "mimo-vlm", + "camera_uids": ["vlm_front", "vlm_left", "vlm_rear", "vlm_right"], + } + + +def test_ab_builders_default_to_the_offline_graph_path(gym_export: Path) -> None: + scene = prepare_scene(gym_export) + config = build_fast_gym_config( + scene, + task_name="ab_default_path", + task_description="A/B path smoke test.", + robot_profile="ur10", + execution_program_hash="e" * 64, + max_episodes=1, + max_episode_steps=10, + planning_mode="ab", + ) + agent = build_agent_config( + task_name="ab_default_path", + robot_profile="ur10", + execution_program_hash="e" * 64, + source_config_path=scene.source_config_path, + uid_map=scene.uid_map, + planning_mode="ab", + ) + + expected = "offline/seed_task_graph.json" + assert config["env"]["extensions"]["action_engine"]["seed_task_graph"] == expected + assert agent["seed_task_graph"] == expected + assert agent["online_planning"]["camera_uids"] == [ + "vlm_front", + "vlm_left", + "vlm_rear", + "vlm_right", + ] + + +def test_ab_scene_requirements_declare_four_vlm_views() -> None: + requirements = { + "schema_version": "action_engine_scene_requirements_v2", + "task_id": "ab", + "objects": [ + { + "role_id": "object_01", + "category": "can", + "count": 1, + "affordances": ["graspable"], + "initial_state": {}, + "attributes": {}, + } + ], + "cameras": [], + "spatial_constraints": [], + "distractor_count": 0, + "metadata": {}, + } + output = _add_ab_camera_requirements(requirements) + assert [item["uid"] for item in output["cameras"]] == [ + "vlm_front", + "vlm_left", + "vlm_rear", + "vlm_right", + ] + assert all(item["modalities"] == ["rgb", "depth"] for item in output["cameras"]) + + +def test_ab_builder_rejects_noncanonical_vlm_camera_ids(gym_export: Path) -> None: + scene = prepare_scene(gym_export) + with pytest.raises(ValueError, match="canonical"): + build_agent_config( + task_name="ab_invalid_cameras", + robot_profile="ur10", + execution_program_hash="f" * 64, + source_config_path=scene.source_config_path, + uid_map=scene.uid_map, + planning_mode="ab", + vlm_camera_uids=["front", "left", "rear", "right"], + ) + + +@pytest.mark.parametrize( + ("profile", "robot_uid", "solver_type"), + [ + ("dual_ur3", "DualUR3", "ur3"), + ("dual_ur5", "DualUR5", "ur5"), + ("dual_ur10", "DualUR10", "ur10"), + ("dual_franka", "DualFrankaPanda", None), + ], +) +def test_fast_gym_config_supports_all_robot_profiles( + gym_export: Path, + profile: str, + robot_uid: str, + solver_type: str | None, +) -> None: + expected_tcp = [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.2], + [0.0, 0.0, 0.0, 1.0], + ] + expected_hand_mount = [ + [0.0, 1.0, 0.0, 0.0], + [-1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0], + ] + scene = prepare_scene(gym_export) + config = build_fast_gym_config( + scene, + task_name="profile_task", + task_description="Profile smoke test.", + robot_profile=profile, + execution_program_hash="b" * 64, + max_episodes=1, + max_episode_steps=20, + ) + + assert config["robot"]["uid"] == robot_uid + assert config["env"]["extensions"]["agent_robot_profile"] == profile + for arm in ("left_arm", "right_arm"): + assert config["robot"]["solver_cfg"][arm]["tcp"] == expected_tcp + components = { + component["component_type"]: component + for component in config["robot"]["urdf_cfg"]["components"] + } + for hand in ("left_hand", "right_hand"): + assert components[hand]["transform"] == expected_hand_mount + if solver_type is not None: + assert config["robot"]["solver_cfg"]["left_arm"]["ur_type"] == solver_type + + +@pytest.mark.parametrize( + ( + "profile", + "expected_position_xy", + "expected_rotation", + "expected_world_x", + ), + [ + ("ur10", [2.0, 0.0], [0.0, 0.0, 0.0], 0.9), + ("franka", [-0.7, 0.0], [0.0, 0.0, 180.0], 0.55), + ], +) +def test_dual_robot_profiles_use_identity_mounts_and_same_side_arm_names( + gym_export: Path, + profile: str, + expected_position_xy: list[float], + expected_rotation: list[float], + expected_world_x: float, +) -> None: + scene = prepare_scene(gym_export) + config = build_fast_gym_config( + scene, + task_name="dual_ur_frame_task", + task_description="Verify the Dual-UR world frame.", + robot_profile=profile, + execution_program_hash="c" * 64, + max_episodes=1, + max_episode_steps=20, + ) + + robot = config["robot"] + robot_yaw = np.deg2rad(float(robot["init_rot"][2])) + robot_rotation = np.array( + [ + [np.cos(robot_yaw), -np.sin(robot_yaw), 0.0], + [np.sin(robot_yaw), np.cos(robot_yaw), 0.0], + [0.0, 0.0, 1.0], + ] + ) + robot_position = np.asarray(robot["init_pos"], dtype=np.float64) + components = { + component["component_type"]: np.asarray( + component["transform"], dtype=np.float64 + ) + for component in robot["urdf_cfg"]["components"] + if component["component_type"] in {"left_arm", "right_arm"} + } + world_transforms = {} + for side, component in components.items(): + world = np.eye(4) + world[:3, :3] = robot_rotation @ component[:3, :3] + world[:3, 3] = robot_position + robot_rotation @ component[:3, 3] + world_transforms[side] = world + + assert robot["init_pos"][:2] == pytest.approx(expected_position_xy) + assert robot["init_rot"] == pytest.approx(expected_rotation) + assert world_transforms["left_arm"][:3, 3] == pytest.approx( + [expected_world_x, -0.3, world_transforms["left_arm"][2, 3]] + ) + assert world_transforms["right_arm"][:3, 3] == pytest.approx( + [expected_world_x, 0.3, world_transforms["right_arm"][2, 3]] + ) + np.testing.assert_allclose(components["left_arm"][:3, :3], np.eye(3), atol=1.0e-12) + np.testing.assert_allclose(components["right_arm"][:3, :3], np.eye(3), atol=1.0e-12) + np.testing.assert_allclose( + world_transforms["left_arm"][:3, :3], robot_rotation, atol=1.0e-12 + ) + np.testing.assert_allclose( + world_transforms["right_arm"][:3, :3], robot_rotation, atol=1.0e-12 + ) + + +def test_fast_gym_config_keeps_scene_deterministic_by_default( + gym_export: Path, +) -> None: + scene = prepare_scene(gym_export) + config = build_fast_gym_config( + scene, + task_name="deterministic_task", + task_description="Keep the source scene fixed.", + robot_profile="ur10", + execution_program_hash="d" * 64, + max_episodes=1, + max_episode_steps=20, + ) + + events = config["env"]["events"] + assert "randomize_interact_can_pose" not in events + assert "randomize_table_height" not in events + + +@pytest.mark.parametrize( + ("alias", "canonical"), + [ + ("franka", "dual_franka"), + ("ur5", "dual_ur5"), + ("ur10", "dual_ur10"), + ], +) +def test_required_cli_robot_aliases_build_runnable_profiles( + gym_export: Path, + alias: str, + canonical: str, +) -> None: + scene = prepare_scene(gym_export) + config = build_fast_gym_config( + scene, + task_name="profile_alias_task", + task_description="Profile alias smoke test.", + robot_profile=alias, + execution_program_hash="c" * 64, + max_episodes=1, + max_episode_steps=20, + ) + + assert config["env"]["extensions"]["agent_robot_profile"] == canonical + + +def test_source_scene_scale_policies_are_deterministic(gym_export: Path) -> None: + preserved = prepare_scene(gym_export) + multiplied = prepare_scene( + gym_export, + body_scale_policy="multiply", + body_scale=(2.0, 3.0, 4.0), + ) + absolute = prepare_scene( + gym_export, + body_scale_policy="absolute", + body_scale=(2.0, 3.0, 4.0), + ) + + assert preserved.body_scale_policy == "preserve" + assert multiplied.rigid_objects[0]["body_scale"] == [2.0, 3.0, 4.0] + assert absolute.rigid_objects[0]["body_scale"] == [2.0, 3.0, 4.0] + assert multiplied.asset_hashes == absolute.asset_hashes + + +def test_artifact_writer_refuses_implicit_overwrite(tmp_path: Path) -> None: + payload = {"value": 1} + paths = write_generation_artifacts( + tmp_path, + gym_config=payload, + agent_config=payload, + task_spec=payload, + scene_requirements=payload, + seed_task_graph=payload, + seed_task_graph_png=b"\x89PNG\r\n\x1a\nold", + overwrite=False, + ) + assert json.loads(paths.gym_config.read_text(encoding="utf-8")) == payload + assert paths.seed_task_graph_png.read_bytes().startswith(b"\x89PNG") + + # A leftover PNG participates in the same preflight as every JSON artifact. + for path in ( + paths.gym_config, + paths.agent_config, + paths.task_spec, + paths.scene_requirements, + paths.execution_program, + ): + path.unlink() + with pytest.raises(FileExistsError, match="--overwrite"): + write_generation_artifacts( + tmp_path, + gym_config=payload, + agent_config=payload, + task_spec=payload, + scene_requirements=payload, + seed_task_graph=payload, + seed_task_graph_png=b"\x89PNG\r\n\x1a\nnew", + overwrite=False, + ) + + replaced = write_generation_artifacts( + tmp_path, + gym_config=payload, + agent_config=payload, + task_spec=payload, + scene_requirements=payload, + seed_task_graph=payload, + seed_task_graph_png=b"\x89PNG\r\n\x1a\nnew", + overwrite=True, + ) + assert replaced.seed_task_graph_png.read_bytes().endswith(b"new") + + +def test_artifact_writer_creates_ab_branch_directory(tmp_path: Path) -> None: + payload = {"value": "ab"} + paths = write_generation_artifacts( + tmp_path, + gym_config=payload, + agent_config=payload, + task_spec=payload, + scene_requirements=payload, + seed_task_graph=payload, + seed_task_graph_png=b"\x89PNG\r\n\x1a\nab", + overwrite=False, + planning_mode="ab", + ) + + assert paths.seed_task_graph.parent == tmp_path / "offline" + assert json.loads(paths.seed_task_graph.read_text(encoding="utf-8")) == payload + + +def test_generation_calls_interpreter_recipe_and_renderer_once( + gym_export: Path, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from embodichain.gen_sim.action_engine import tasks + from embodichain.gen_sim.action_engine.generation import generator + + planner_call: dict[str, object] = {} + recipe_calls: list[tuple[object, object]] = [] + rendered: dict[str, object] = {} + published: dict[str, object] = {} + + def fake_interpret_and_ground(**kwargs): + planner_call.update(kwargs) + task_spec = _existing_v2_task_spec(str(kwargs["task_name"])) + task_spec["instruction"] = str(kwargs["task_description"]) + bindings = {"object_01": "interact_can"} + requirements = _scene_requirements_from_bindings( + str(kwargs["task_name"]), + kwargs["scene_objects"], + bindings, + ) + return GroundedTaskSpec(task_spec, requirements, bindings) + + monkeypatch.setattr( + tasks, + "interpret_and_ground_task_spec", + fake_interpret_and_ground, + ) + real_recipe = tasks.instantiate_seed_graph + + def capture_recipe(task_spec, role_bindings): + recipe_calls.append((task_spec, role_bindings)) + return real_recipe(task_spec, role_bindings) + + monkeypatch.setattr(tasks, "instantiate_seed_graph", capture_recipe) + renderer_module = ModuleType( + "embodichain.gen_sim.action_engine.graph_visualization" + ) + + def fake_renderer(program): + rendered["program"] = program + return b"\x89PNG\r\n\x1a\nseed" + + renderer_module.render_seed_task_graph_png = fake_renderer + monkeypatch.setitem(sys.modules, renderer_module.__name__, renderer_module) + real_writer = generator.write_generation_artifacts + + def capture_writer(*args, **kwargs): + published["program"] = kwargs["seed_task_graph"] + return real_writer(*args, **kwargs) + + monkeypatch.setattr(generator, "write_generation_artifacts", capture_writer) + output_dir = tmp_path / "configs" + paths = generate_action_engine_config( + gym_export, + output_dir, + task_name="line_task", + task_description="test-instruction", + robot_profile="franka", + ) + + assert planner_call["task_name"] == "line_task" + assert planner_call["task_description"] == "test-instruction" + assert planner_call["robot_profile"] == "franka" + assert len(recipe_calls) == 1 + planner_objects = planner_call["scene_objects"] + assert isinstance(planner_objects, list) + assert {obj["uid"] for obj in planner_objects} == {"table", "interact_can"} + assert {path.name for path in output_dir.iterdir()} == { + "fast_gym_config.json", + "agent_config.json", + "task_spec.json", + "scene_requirements.json", + "seed_task_graph.json", + "seed_task_graph.png", + } + assert paths.seed_task_graph_png.read_bytes() == b"\x89PNG\r\n\x1a\nseed" + assert rendered["program"] is published["program"] + + agent_config = json.loads(paths.agent_config.read_text(encoding="utf-8")) + assert agent_config["schema_version"] == "action_engine_config_v2" + assert agent_config["task_spec"] == "task_spec.json" + assert agent_config["scene_requirements"] == "scene_requirements.json" + assert agent_config["seed_task_graph"] == "seed_task_graph.json" + assert len(agent_config["seed_task_graph_hash"]) == 64 + assert agent_config["runtime_policy"]["schema_version"] == ( + "action_engine_runtime_policy_v6" + ) + assert agent_config["runtime_policy"]["planner"]["dynamic_collision"] is True + assert agent_config["runtime_policy"]["planner"]["static_obstacle_uids"] == [ + "table" + ] + assert agent_config["runtime_policy"]["planner"]["dynamic_obstacle_uids"] == [ + "interact_can" + ] + assert len(agent_config["runtime_policy_hash"]) == 64 + assert "png" not in json.dumps(agent_config).lower() + + from embodichain.gen_sim.action_engine.runtime import ( + load_agent_execution_program, + ) + + regenerated = load_agent_execution_program( + agent_config, + agent_config_path=paths.agent_config, + regenerate=True, + ) + assert regenerated.task == "line_task" + assert regenerated.seed_graph is not None + + +def test_existing_v2_task_spec_bypasses_text_planner_and_derives_scene_requirements( + gym_export: Path, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from embodichain.gen_sim.action_engine import tasks + + renderer_module = ModuleType( + "embodichain.gen_sim.action_engine.graph_visualization" + ) + renderer_module.render_seed_task_graph_png = lambda _program: b"direct-task-png" + monkeypatch.setitem(sys.modules, renderer_module.__name__, renderer_module) + + def unexpected_text_planner(**_kwargs): + raise AssertionError("an existing TaskSpec must not invoke text planning") + + monkeypatch.setattr( + tasks, "interpret_and_ground_task_spec", unexpected_text_planner + ) + input_path = tmp_path / "task_spec.json" + input_path.write_text( + json.dumps(_existing_v2_task_spec()), + encoding="utf-8", + ) + + paths = generate_action_engine_config( + gym_export, + tmp_path / "generated", + task_name="direct_task", + task_spec=input_path, + robot_profile="ur10", + ) + + persisted_task = json.loads(paths.task_spec.read_text(encoding="utf-8")) + persisted_requirements = json.loads( + paths.scene_requirements.read_text(encoding="utf-8") + ) + assert persisted_task["metadata"]["role_bindings"] == {"object_01": "interact_can"} + assert [item["role_id"] for item in persisted_requirements["objects"]] == [ + "object_01" + ] + assert persisted_requirements["metadata"]["source"] == ("task_spec_role_bindings") + gym_config = json.loads(paths.gym_config.read_text(encoding="utf-8")) + assert ( + gym_config["env"]["dataset"]["lerobot"]["params"]["instruction"]["lang"] + == "test-instruction" + ) + + +def test_existing_v2_task_spec_uses_validated_scene_requirements_sidecar( + gym_export: Path, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + renderer_module = ModuleType( + "embodichain.gen_sim.action_engine.graph_visualization" + ) + renderer_module.render_seed_task_graph_png = lambda _program: b"sidecar-png" + monkeypatch.setitem(sys.modules, renderer_module.__name__, renderer_module) + + input_dir = tmp_path / "task-first" + input_dir.mkdir() + task = _existing_v2_task_spec("sidecar_task") + requirements = { + "schema_version": SCENE_REQUIREMENTS_SCHEMA, + "task_id": "sidecar_task", + "objects": [ + { + "role_id": "object_01", + "category": "can", + "count": 1, + "affordances": ["graspable", "orientable"], + "initial_state": {"orientation": "fallen"}, + "attributes": {"color": "red"}, + } + ], + "cameras": [], + "spatial_constraints": [], + "distractor_count": 0, + "metadata": {"task_first": True}, + } + (input_dir / "task_spec.json").write_text( + json.dumps(task), + encoding="utf-8", + ) + (input_dir / "scene_requirements.json").write_text( + json.dumps(requirements), + encoding="utf-8", + ) + + paths = generate_action_engine_config( + gym_export, + tmp_path / "generated-sidecar", + task_name="sidecar_task", + task_spec=input_dir / "task_spec.json", + robot_profile="ur10", + ) + + assert json.loads(paths.scene_requirements.read_text(encoding="utf-8")) == ( + requirements + ) + + +def test_task_factory_style_sidecar_binds_roles_without_text_llm( + gym_export: Path, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + renderer_module = ModuleType( + "embodichain.gen_sim.action_engine.graph_visualization" + ) + renderer_module.render_seed_task_graph_png = lambda _program: b"task-first-png" + monkeypatch.setitem(sys.modules, renderer_module.__name__, renderer_module) + source_path = gym_export / "gym_config.json" + source = json.loads(source_path.read_text(encoding="utf-8")) + source["rigid_object"][0]["category"] = "can" + source["rigid_object"][0]["attributes"] = {"color": "red"} + source["rigid_object"][0]["affordances"] = ["graspable", "orientable"] + source["rigid_object"][0]["initial_state"] = {"orientation": "fallen"} + source_path.write_text(json.dumps(source), encoding="utf-8") + + input_dir = tmp_path / "task-first-unbound" + input_dir.mkdir() + task = _existing_v2_task_spec("task_first_unbound") + task["metadata"] = {"fixture": "abstract-task"} + requirements = { + "schema_version": SCENE_REQUIREMENTS_SCHEMA, + "task_id": "task_first_unbound", + "objects": [ + { + "role_id": "object_01", + "category": "can", + "count": 1, + "affordances": ["graspable", "orientable"], + "initial_state": {"orientation": "fallen"}, + "attributes": {"color": "red"}, + } + ], + "cameras": [], + "spatial_constraints": [], + "distractor_count": 0, + "metadata": {"task_first": True}, + } + (input_dir / "task_spec.json").write_text(json.dumps(task), encoding="utf-8") + (input_dir / "scene_requirements.json").write_text( + json.dumps(requirements), encoding="utf-8" + ) + + paths = generate_action_engine_config( + gym_export, + tmp_path / "generated-unbound-sidecar", + task_name="task_first_unbound", + task_spec=input_dir / "task_spec.json", + robot_profile="ur10", + ) + + task_artifact = json.loads(paths.task_spec.read_text(encoding="utf-8")) + assert task_artifact["metadata"]["role_bindings"] == {"object_01": "interact_can"} + + +def test_task_spec_input_rejects_natural_language_conflict( + gym_export: Path, + tmp_path: Path, +) -> None: + task = _existing_v2_task_spec() + with pytest.raises(ValueError, match="task_spec cannot be combined"): + generate_action_engine_config( + gym_export, + tmp_path / "conflict-description", + task_name="direct_task", + task_description="do something", + task_spec=task, + robot_profile="ur10", + ) + + +def test_task_spec_role_binding_accepts_legacy_oracle_and_rejects_conflicts() -> None: + task = _existing_v2_task_spec() + task["metadata"] = {} + task["oracle"] = {"role_bindings": {"object_01": "interact_can"}} + assert _task_spec_role_bindings(task, ["table", "interact_can"]) == { + "object_01": "interact_can" + } + + task["metadata"] = {"role_bindings": {"object_01": "table"}} + with pytest.raises(ValueError, match="Conflicting role_bindings"): + _task_spec_role_bindings(task, ["table", "interact_can"]) + + +def test_task_spec_role_binding_merges_non_overlapping_handoffs() -> None: + task = _existing_v2_task_spec() + task["task_instances"][0]["params"]["target_role"] = "object_02" + task["metadata"] = {"role_bindings": {"object_01": "interact_can"}} + task["oracle"] = {"role_bindings": {"object_02": "interact_target"}} + + assert _task_spec_role_bindings( + task, + ["table", "interact_can", "interact_target"], + ) == {"object_01": "interact_can", "object_02": "interact_target"} + + +def test_task_factory_sidecar_requires_static_affordance_and_state_evidence() -> None: + task = _existing_v2_task_spec("missing-static-evidence") + task["metadata"] = {} + requirements = { + "objects": [ + { + "role_id": "object_01", + "category": "can", + "count": 1, + "affordances": ["graspable", "orientable"], + "initial_state": {"orientation": "fallen"}, + "attributes": {"color": "red"}, + } + ] + } + scene = [ + { + "runtime_uid": "interact_can", + "description": "A red soda can.", + "init_pos": [0.0, 0.0, 0.7], + } + ] + + with pytest.raises(ValueError, match="requires one unambiguous scene match"): + _task_spec_role_bindings( + task, + ["interact_can"], + scene_requirements=requirements, + scene_objects=scene, + robot_profile="ur10", + ) + + +@pytest.mark.parametrize( + ("scene_metadata", "required_attributes"), + ( + ({}, {}), + ({"category": "can"}, {"color": "red"}), + ), +) +def test_task_factory_sidecar_does_not_infer_semantics_from_description( + scene_metadata: dict, + required_attributes: dict, +) -> None: + task = _existing_v2_task_spec("no-text-evidence") + task["metadata"] = {} + requirements = { + "objects": [ + { + "role_id": "object_01", + "category": "can", + "count": 1, + "affordances": [], + "initial_state": {}, + "attributes": required_attributes, + } + ] + } + scene = [ + { + "runtime_uid": "mystery_object", + "role": "rigid_object", + "description": "A red soda can.", + "init_pos": [0.0, 0.0, 0.7], + **scene_metadata, + } + ] + + with pytest.raises(ValueError, match="requires one unambiguous scene match"): + _task_spec_role_bindings( + task, + ["mystery_object"], + scene_requirements=requirements, + scene_objects=scene, + robot_profile="ur10", + ) + + +def test_ab_generation_writes_shared_and_offline_branch_artifacts( + gym_export: Path, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + renderer_module = ModuleType( + "embodichain.gen_sim.action_engine.graph_visualization" + ) + renderer_module.render_seed_task_graph_png = lambda _program: b"ab-seed-png" + monkeypatch.setitem(sys.modules, renderer_module.__name__, renderer_module) + + output_dir = tmp_path / "ab-config" + paths = generate_action_engine_config( + gym_export, + output_dir, + task_name="ab_task", + task_spec=_existing_v2_task_spec("ab_task"), + robot_profile="ur10", + planning_mode="ab", + vlm_model="mimo-vlm", + ) + + assert paths.seed_task_graph == output_dir / "offline/seed_task_graph.json" + assert paths.seed_task_graph_png == output_dir / "offline/seed_task_graph.png" + assert not (output_dir / "seed_task_graph.json").exists() + agent_config = json.loads(paths.agent_config.read_text(encoding="utf-8")) + assert agent_config["planning_mode"] == "ab" + assert agent_config["offline_seed_task_graph"] == "offline/seed_task_graph.json" + assert agent_config["online_planning"]["vlm_model"] == "mimo-vlm" + scene_requirements = json.loads( + paths.scene_requirements.read_text(encoding="utf-8") + ) + assert [camera["uid"] for camera in scene_requirements["cameras"]] == [ + "vlm_front", + "vlm_left", + "vlm_rear", + "vlm_right", + ] + gym_config = json.loads(paths.gym_config.read_text(encoding="utf-8")) + assert [ + sensor["uid"] + for sensor in gym_config["sensor"] + if sensor["uid"].startswith("vlm_") + ] == ["vlm_front", "vlm_left", "vlm_rear", "vlm_right"] + + +def test_invalid_explicit_task_fails_before_output_asset_materialization( + gym_export: Path, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from embodichain.gen_sim.action_engine import tasks + from embodichain.gen_sim.action_engine.generation import generator + + normalized = False + recipe_called = False + writer_called = False + + def reject_task(**_kwargs): + raise ValueError("object selector is ambiguous") + + def record_normalization(*_args, **_kwargs): + nonlocal normalized + normalized = True + raise AssertionError("normalization must not run after planning failure") + + def unexpected_recipe(*_args, **_kwargs): + nonlocal recipe_called + recipe_called = True + raise AssertionError("recipe must not run after interpretation failure") + + def unexpected_writer(*_args, **_kwargs): + nonlocal writer_called + writer_called = True + raise AssertionError("writer must not run after interpretation failure") + + monkeypatch.setattr(tasks, "interpret_and_ground_task_spec", reject_task) + monkeypatch.setattr(tasks, "instantiate_seed_graph", unexpected_recipe) + monkeypatch.setattr(generator, "normalize_scene_assets", record_normalization) + monkeypatch.setattr(generator, "write_generation_artifacts", unexpected_writer) + output_dir = tmp_path / "invalid" + + with pytest.raises(ValueError, match="ambiguous"): + generate_action_engine_config( + gym_export, + output_dir, + task_name="invalid_task", + task_description="test-instruction", + robot_profile="franka", + ) + + assert normalized is False + assert recipe_called is False + assert writer_called is False + assert not output_dir.exists() + + +def test_agent_config_uses_relative_program_paths(gym_export: Path) -> None: + scene = prepare_scene(gym_export) + config = build_agent_config( + task_name="line_task", + robot_profile="franka", + execution_program_hash="b" * 64, + source_config_path=scene.source_config_path, + uid_map=scene.uid_map, + ) + assert config["task_spec"] == "task_spec.json" + assert config["scene_requirements"] == "scene_requirements.json" + assert config["seed_task_graph"] == "seed_task_graph.json" + assert config["runtime_policy"]["arm_selection"]["pickup_crossing_weight"] == 1.0 + assert config["runtime_policy"]["motion_defaults"]["PickUp"][ + "lift_height" + ] == pytest.approx(0.30) + assert config["runtime_policy"]["grasp"]["max_open_length"] == pytest.approx(0.15) + assert len(config["runtime_policy_hash"]) == 64 + + +def test_agent_config_anchors_absolute_motion_heights_to_tabletop( + gym_export: Path, +) -> None: + scene = prepare_scene(gym_export) + config = build_agent_config( + task_name="high_table_task", + robot_profile="franka", + execution_program_hash="c" * 64, + source_config_path=scene.source_config_path, + uid_map=scene.uid_map, + table_top_z=1.05, + ) + + policy = config["runtime_policy"] + assert policy["motion_defaults"]["MoveEndEffector"][ + "maximum_eef_height" + ] == pytest.approx(1.45) + assert policy["grounding"]["handover"]["maximum_eef_height"] == pytest.approx(1.85) + + +def test_documented_cli_accepts_franka_profile() -> None: + args = build_parser().parse_args( + [ + "--gym_project", + "gym_export", + "--output_dir", + "configs/task4_2", + "--task_name", + "task4_2", + "--task_description", + "Arrange the cans in a line.", + "--robot-profile", + "franka", + "--overwrite", + ] + ) + assert args.robot_profile == "franka" + assert args.overwrite is True + + +def test_generation_cli_defaults_to_mature_robot_without_scene_randomization() -> None: + args = build_parser().parse_args( + [ + "--gym_project", + "gym_export", + "--output_dir", + "configs/task2_3", + "--task_name", + "task2_3", + "--task_description", + "Upright both objects.", + ] + ) + + assert args.robot_profile == "ur10" + assert args.randomize_scene is False + assert args.planning_mode == "offline" + assert not hasattr(args, "instruction_parser") + assert not hasattr(args, "task_agent") + + +def test_generation_cli_accepts_ab_models() -> None: + args = build_parser().parse_args( + [ + "--gym_project", + "gym_export", + "--output_dir", + "configs/ab", + "--task_name", + "ab", + "--task_description", + "test-instruction", + "--planning-mode", + "ab", + "--llm-model", + "text-model", + "--vlm-model", + "vision-model", + ] + ) + + assert args.planning_mode == "ab" + assert args.llm_model == "text-model" + assert args.vlm_model == "vision-model" + + +def test_generation_cli_accepts_existing_task_spec_without_description() -> None: + args = build_parser().parse_args( + [ + "--gym_project", + "gym_export", + "--output_dir", + "configs/direct", + "--task_name", + "direct_task", + "--task-spec", + "tasks/direct_task/task_spec.json", + ] + ) + + assert args.task_spec == "tasks/direct_task/task_spec.json" + assert cli_module._resolve_task_description(args) == "" + + +def test_generation_cli_reports_seed_png_path( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + paths = artifact_paths(tmp_path) + monkeypatch.setattr( + cli_module, + "generate_action_engine_config", + lambda *_args, **_kwargs: paths, + ) + monkeypatch.setattr( + sys, + "argv", + [ + "generate_action_agent_config", + "--gym_project", + "gym_export", + "--output_dir", + str(tmp_path), + "--task_name", + "task4_2", + "--task_description", + "Arrange cans.", + ], + ) + + cli_module.cli() + + assert ( + f"Generated Seed graph PNG: {paths.seed_task_graph_png}" + in capsys.readouterr().out + ) + + +@pytest.mark.parametrize( + "removed_args", + [ + ["--instruction-parser", "llm"], + ["--instruction_parser", "llm"], + ["--task-agent", "task-agent.json"], + ["--task_agent", "task-agent.json"], + ], +) +def test_generation_cli_rejects_removed_arguments(removed_args: list[str]) -> None: + base_args = [ + "--gym-project", + "gym_export", + "--output-dir", + "configs/task", + "--task-name", + "task", + "--task-description", + "Upright the can.", + ] + + with pytest.raises(SystemExit, match="2"): + build_parser().parse_args([*base_args, *removed_args]) + + +def test_removed_python_parameters_are_absent() -> None: + parameters = inspect.signature(generate_action_engine_config).parameters + + assert "instruction_parser" not in parameters + assert "task_agent" not in parameters diff --git a/tests/gen_sim/action_engine/planning/__init__.py b/tests/gen_sim/action_engine/planning/__init__.py new file mode 100644 index 000000000..de3758a5b --- /dev/null +++ b/tests/gen_sim/action_engine/planning/__init__.py @@ -0,0 +1,21 @@ +# ---------------------------------------------------------------------------- +# 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 Engine planning tests.""" + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/tests/gen_sim/action_engine/planning/test_linker.py b/tests/gen_sim/action_engine/planning/test_linker.py new file mode 100644 index 000000000..104438d5b --- /dev/null +++ b/tests/gen_sim/action_engine/planning/test_linker.py @@ -0,0 +1,407 @@ +# ---------------------------------------------------------------------------- +# 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.domain import seed_graph_hash +from embodichain.gen_sim.action_engine.planning.linker import ( + link_seed_graph, + link_task_dependencies, +) +from embodichain.gen_sim.action_engine.protocol import ( + SEED_GRAPH_SCHEMA, + TASK_SPEC_SCHEMA, +) +from embodichain.gen_sim.action_engine.runtime.loader import load_execution_program +from embodichain.gen_sim.action_engine.tasks.recipes import instantiate_seed_graph + + +def _handover_task() -> dict: + return { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "handover_then_place", + "level": "L3", + "instruction": "Stand both cans, hand over the purple can, then place it.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "task_01", + "task_type": "E2", + "params": { + "object_role": "purple", + "required_arm": "right_arm", + }, + "depends_on": [], + "role": "primary", + }, + { + "id": "task_02", + "task_type": "E2", + "params": { + "object_role": "orange", + "required_arm": "left_arm", + }, + "depends_on": [], + "role": "primary", + }, + { + "id": "task_03", + "task_type": "E4", + "params": { + "object_role": "purple", + "transfer_arm": "right_arm", + "receive_arm": "left_arm", + }, + "depends_on": ["task_02"], + "role": "primary", + }, + { + "id": "task_04", + "task_type": "E1", + "params": { + "object_role": "purple", + "target_role": "orange", + "relation": "left_of", + "required_arm": "left_arm", + }, + "depends_on": ["task_03"], + "role": "primary", + }, + ], + "success": {"type": "all_complete"}, + "oracle": {}, + "metadata": {}, + } + + +def _handover_graph() -> dict: + return instantiate_seed_graph( + _handover_task(), + {"purple": "purple_can", "orange": "orange_can"}, + ) + + +def _unlink_for_rebuild(graph: dict) -> None: + graph["metadata"].pop("action_contract_linker", None) + for group in graph["task_groups"]: + group.pop("contract", None) + + +def test_task_linker_preserves_parallel_arms_and_waits_for_both_before_handover() -> ( + None +): + linked = link_task_dependencies( + _handover_task(), + {"purple": "purple_can", "orange": "orange_can"}, + ) + by_id = {item["id"]: item for item in linked["task_instances"]} + + assert by_id["task_01"]["depends_on"] == [] + assert by_id["task_02"]["depends_on"] == [] + assert by_id["task_03"]["depends_on"] == ["task_02", "task_01"] + + +def test_resource_dependency_provenance_is_persisted_in_seed_graph() -> None: + task = _handover_task() + task["task_instances"] = task["task_instances"][:3] + handover = task["task_instances"][2] + handover["params"] = { + "object_role": "orange", + "transfer_arm": "left_arm", + "receive_arm": "right_arm", + } + + graph = instantiate_seed_graph( + task, + {"purple": "purple_can", "orange": "orange_can"}, + ) + + provenance = graph["metadata"]["action_contract_task_linker"] + assert provenance["linked_dependencies"] == [ + { + "from": "task_01", + "to": "task_03", + "reason": "resource", + "detail": "arm:right_arm", + } + ] + + +def test_same_object_e2_handover_gets_direct_causal_edge_through_a_chain() -> None: + task = _handover_task() + task["task_instances"][1]["depends_on"] = ["task_01"] + linked = link_task_dependencies( + task, + {"purple": "purple_can", "orange": "orange_can"}, + ) + handover = next( + item for item in linked["task_instances"] if item["id"] == "task_03" + ) + + assert handover["depends_on"] == ["task_02", "task_01"] + + +def test_handover_ownership_flows_through_home_terminal_barrier() -> None: + graph = _handover_graph() + groups = {group["id"]: group for group in graph["task_groups"]} + nodes = {node["id"]: node for node in graph["nodes"]} + handover_group = groups["task_03"] + terminal_id = handover_group["contract"]["terminal_node_ids"][0] + terminal = nodes[terminal_id] + receiver_entry = nodes[groups["task_04"]["contract"]["entry_node_ids"][0]] + handover = next( + node + for node in graph["nodes"] + if node["task_instance_id"] == "task_03" and node["atomic_action"] == "HandOver" + ) + + assert terminal["atomic_action"] == "MoveJoints" + assert terminal["contract"]["completion"] == "terminal_barrier" + assert terminal["contract"]["failure_policy"] == "best_effort" + retreat = next( + node + for node in graph["nodes"] + if node["task_instance_id"] == "task_03" + and node["atomic_action"] == "MoveEndEffector" + ) + assert retreat["contract"]["failure_policy"] == "safety_required" + assert terminal_id in receiver_entry["depends_on"] + assert { + (effect["op"], effect["atom"]["predicate"], effect["atom"].get("arm")) + for effect in handover["contract"]["effects"] + } >= { + ("delete", "object_held", "right_arm"), + ("add", "object_held", "left_arm"), + } + + +def test_linker_is_idempotent_and_hash_stable() -> None: + graph = _handover_graph() + relinked = link_seed_graph( + graph, + task_order=["task_01", "task_02", "task_03", "task_04"], + known_objects={"purple_can", "orange_can", "table"}, + ) + + assert relinked == graph + assert seed_graph_hash(relinked) == seed_graph_hash(graph) + + +def test_linker_rejects_missing_cleanup_wrong_holder_and_duplicate_pickup() -> None: + missing_cleanup = deepcopy(_handover_graph()) + home = next( + node + for node in missing_cleanup["nodes"] + if node["task_instance_id"] == "task_03" + and node["atomic_action"] == "MoveJoints" + ) + missing_cleanup["nodes"].remove(home) + next(group for group in missing_cleanup["task_groups"] if group["id"] == "task_03")[ + "node_ids" + ].remove(home["id"]) + for node in missing_cleanup["nodes"]: + node["depends_on"] = [ + dependency for dependency in node["depends_on"] if dependency != home["id"] + ] + _unlink_for_rebuild(missing_cleanup) + with pytest.raises(ValueError, match="terminal barrier"): + link_seed_graph(missing_cleanup) + + wrong_holder = deepcopy(_handover_graph()) + staging = next( + node + for node in wrong_holder["nodes"] + if node["task_instance_id"] == "task_03" + and node["atomic_action"] == "MoveHeldObject" + ) + staging["actor"] = {"mode": "required", "arm": "left_arm"} + staging.pop("contract") + _unlink_for_rebuild(wrong_holder) + with pytest.raises(ValueError, match="no producer|unavailable state"): + link_seed_graph(wrong_holder) + + duplicate_pickup = deepcopy(_handover_graph()) + pickup = next( + node + for node in duplicate_pickup["nodes"] + if node["task_instance_id"] == "task_01" and node["atomic_action"] == "PickUp" + ) + staging = next( + node + for node in duplicate_pickup["nodes"] + if node["task_instance_id"] == "task_01" + and node["atomic_action"] == "MoveHeldObject" + ) + repeated = deepcopy(pickup) + repeated["id"] = "task_01__duplicate_pickup" + repeated["depends_on"] = [pickup["id"]] + repeated.pop("contract") + staging["depends_on"] = [repeated["id"]] + group = next( + group for group in duplicate_pickup["task_groups"] if group["id"] == "task_01" + ) + pickup_index = group["node_ids"].index(pickup["id"]) + group["node_ids"].insert(pickup_index + 1, repeated["id"]) + duplicate_pickup["nodes"].insert( + duplicate_pickup["nodes"].index(pickup) + 1, repeated + ) + _unlink_for_rebuild(duplicate_pickup) + with pytest.raises(ValueError, match="requires unavailable state"): + link_seed_graph(duplicate_pickup) + + +def test_unavailable_arm_reports_current_holder_and_requested_object() -> None: + task = _handover_task() + placement = task["task_instances"][3] + placement["params"].update( + { + "object_role": "orange", + "target_role": "purple", + "required_arm": "left_arm", + } + ) + + with pytest.raises( + ValueError, + match=( + "left_arm.*currently holds 'purple_can'.*" "primary object is 'orange_can'" + ), + ): + instantiate_seed_graph( + task, + {"purple": "purple_can", "orange": "orange_can"}, + ) + + +def test_readers_remain_parallel_and_writer_waits_for_both() -> None: + task = { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "read_write", + "level": "L3", + "instruction": "Inspect a shared target, then manipulate it.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "read_left", + "task_type": "E1", + "params": { + "object_role": "a", + "target_role": "target", + "required_arm": "left_arm", + }, + "depends_on": [], + "role": "primary", + }, + { + "id": "read_right", + "task_type": "E1", + "params": { + "object_role": "b", + "target_role": "target", + "required_arm": "right_arm", + }, + "depends_on": [], + "role": "primary", + }, + { + "id": "write_target", + "task_type": "E2", + "params": { + "object_role": "target", + "required_arm": "left_arm", + }, + "depends_on": [], + "role": "primary", + }, + ], + "success": {}, + "oracle": {}, + "metadata": {}, + } + linked = link_task_dependencies( + task, + {"a": "object_a", "b": "object_b", "target": "shared_target"}, + ) + by_id = {item["id"]: item for item in linked["task_instances"]} + + assert by_id["read_left"]["depends_on"] == [] + assert by_id["read_right"]["depends_on"] == [] + assert by_id["write_target"]["depends_on"] == ["read_left", "read_right"] + + +def test_explicit_distinct_arm_allocation_keeps_auto_groups_parallel() -> None: + task = { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "allocated_auto", + "level": "L2", + "instruction": "Stand both objects upright in parallel.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "first", + "task_type": "E2", + "params": {"object_role": "first_object"}, + "depends_on": [], + "role": "primary", + }, + { + "id": "second", + "task_type": "E2", + "params": {"object_role": "second_object"}, + "depends_on": [], + "role": "primary", + }, + ], + "success": {}, + "oracle": {}, + "metadata": { + "allocation_groups": [ + { + "id": "distinct_pair", + "task_instance_ids": ["first", "second"], + "arm_constraint": "distinct_arms", + } + ] + }, + } + bindings = {"first_object": "first_uid", "second_object": "second_uid"} + linked = link_task_dependencies(task, bindings) + graph = instantiate_seed_graph(linked, bindings) + + assert all(not item["depends_on"] for item in linked["task_instances"]) + assert all(not group["depends_on"] for group in graph["task_groups"]) + + +def test_v2_and_resolver_mismatch_require_regeneration() -> None: + with pytest.raises( + ValueError, match="lacks persisted Action Contracts.*regenerate" + ): + load_execution_program({"schema_version": "action_engine_seed_graph_v2"}) + + graph = _handover_graph() + graph["nodes"][0]["contract"]["claims"][0]["access"] = "shared_read" + with pytest.raises( + ValueError, match="does not match the current capability resolver" + ): + load_execution_program(graph) + + +def test_seed_graph_schema_is_v3() -> None: + assert _handover_graph()["schema_version"] == SEED_GRAPH_SCHEMA diff --git a/tests/gen_sim/action_engine/planning/test_online_v2.py b/tests/gen_sim/action_engine/planning/test_online_v2.py new file mode 100644 index 000000000..290f667bd --- /dev/null +++ b/tests/gen_sim/action_engine/planning/test_online_v2.py @@ -0,0 +1,480 @@ +# ---------------------------------------------------------------------------- +# 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 threading import Barrier + +import pytest +import torch + +import embodichain.gen_sim.action_engine.planning.online as online_module +import embodichain.gen_sim.action_engine.planning.planner as planner_module +import embodichain.gen_sim.action_engine.planning.vision as vision_module +from embodichain.gen_sim.action_engine.domain import public_task_spec +from embodichain.gen_sim.action_engine.planning import ( + CameraObservation, + SceneObservation, + analyze_visual_scene, + fuse_seed_graphs, + plan_candidates_parallel, + plan_online_seed_graph, + select_seed_graph, + validate_visual_facts, +) +from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph + +from ..task_fixtures import make_task_level + + +def _task(level: str, *, reasoning: str | None = None): + task, requirements = make_task_level(level, reasoning=reasoning) + bindings = { + item["role_id"]: f"uid_{item['role_id']}" for item in requirements["objects"] + } + return task, requirements, bindings + + +def test_online_planner_sees_public_task_and_returns_complete_seed_graph() -> None: + task, _, bindings = _task("L4", reasoning="visual_semantics") + offline = instantiate_seed_graph(task, bindings) + body = {key: deepcopy(offline[key]) for key in ("nodes", "task_groups", "success")} + for node in body["nodes"]: + node.pop("contract") + for group in body["task_groups"]: + group.pop("contract") + visual_move = next( + node for node in body["nodes"] if node["atomic_action"] == "MoveHeldObject" + ) + visual_move["target_binding"] = { + "kind": "visual_constraint", + "camera_uid": "front", + "normalized_keypoint": [0.2, 0.3], + } + camera = CameraObservation( + "front", + torch.zeros((8, 8, 3), dtype=torch.uint8), + None, + None, + None, + ) + observation = SceneObservation( + (camera,), + tuple({"uid": uid} for uid in bindings.values()), + ) + uid = next(iter(bindings.values())) + facts = { + "entities": [ + { + "uid": uid, + "camera_uid": "front", + "bbox": [0.1, 0.2, 0.3, 0.4], + "keypoints": {"center": [0.2, 0.3]}, + "confidence": 0.9, + } + ], + "relations": [], + "task_predicates": [], + "confidence": 0.9, + } + prompts = [] + + def caller(**kwargs): + prompts.append(kwargs["prompt"]) + return body + + graph, observed_facts = plan_online_seed_graph( + public_task_spec(task), + observation, + visual_facts=facts, + graph_caller=caller, + ) + + assert graph["planner_route"] == "online" + assert observed_facts == facts + assert "oracle" not in prompts[0] + assert '"task_instances"' not in prompts[0] + assert '"E4"' in prompts[0] + assert "Transfer one held object" in prompts[0] + assert graph["metadata"]["oracle_exposed"] is False + assert any( + node["target_binding"]["kind"] == "visual_constraint" for node in graph["nodes"] + ) + + +def test_offline_and_online_candidates_plan_concurrently_with_isolated_views() -> None: + task, _, bindings = _task("L1") + offline = instantiate_seed_graph(task, bindings) + online = deepcopy(offline) + online["planner_route"] = "online" + barrier = Barrier(2) + views = {} + + def offline_planner(*, task_spec): + views["offline"] = task_spec + barrier.wait(timeout=2.0) + return offline + + def online_planner(*, task_spec): + views["online"] = task_spec + barrier.wait(timeout=2.0) + return online + + pair = plan_candidates_parallel( + task, + offline_planner=offline_planner, + online_planner=online_planner, + ) + + assert "oracle" in views["offline"] + assert "oracle" not in views["online"] + assert pair.offline["planner_route"] == "offline" + assert pair.online["planner_route"] == "online" + + +def test_visual_facts_reject_unknown_uid_and_out_of_range_keypoint() -> None: + value = { + "entities": [ + { + "uid": "unknown", + "camera_uid": "front", + "bbox": [0.0, 0.0, 1.2, 1.0], + "keypoints": {}, + "confidence": 1.0, + } + ], + "relations": [], + "task_predicates": [], + "confidence": 1.0, + } + with pytest.raises(ValueError, match="unknown UID"): + validate_visual_facts(value, known_uids={"known"}, camera_uids={"front"}) + + +def test_visual_facts_reject_visible_entity_without_image_evidence() -> None: + value = { + "entities": [ + { + "uid": "known", + "camera_uid": "front", + "visible": True, + "confidence": 0.9, + } + ], + "relations": [], + "task_predicates": [], + "confidence": 0.9, + } + + with pytest.raises(ValueError, match="bbox or keypoint"): + validate_visual_facts(value, known_uids={"known"}, camera_uids={"front"}) + + +def test_visual_facts_reject_non_numeric_image_coordinates() -> None: + value = { + "entities": [ + { + "uid": "known", + "camera_uid": "front", + "bbox": ["0.1", 0.2, 0.3, 0.4], + "confidence": 0.9, + } + ], + "relations": [], + "task_predicates": [], + "confidence": 0.9, + } + + with pytest.raises(ValueError, match="must be numeric"): + validate_visual_facts(value, known_uids={"known"}, camera_uids={"front"}) + + +def test_visual_fact_caller_receives_rgb_depth_and_calibration_evidence() -> None: + task, _, bindings = _task("L1") + uid = next(iter(bindings.values())) + observation = SceneObservation( + ( + CameraObservation( + "front", + torch.zeros((4, 5, 3), dtype=torch.uint8), + torch.linspace(0.0, 1.0, 20, dtype=torch.float32).reshape(4, 5), + torch.eye(3), + torch.eye(4), + ), + ), + ({"uid": uid},), + ) + captured = {} + + def caller(**kwargs): + captured.update(kwargs) + return { + "entities": [ + { + "uid": uid, + "camera_uid": "front", + "bbox": [0.1, 0.2, 0.3, 0.4], + "confidence": 0.9, + } + ], + "relations": [], + "task_predicates": [], + "confidence": 0.9, + } + + facts = analyze_visual_scene(observation, task, caller=caller) + + assert facts["entities"][0]["uid"] == uid + assert len(captured["images"]) == 2 + assert '"depth_image_index": 1' in captured["prompt"] + assert '"intrinsics": [[1.0, 0.0, 0.0]' in captured["prompt"] + assert captured["schema"]["properties"]["task_predicates"]["maxItems"] == 0 + + +def test_visual_task_predicates_are_limited_to_the_current_task() -> None: + task, _, bindings = _task("L4", reasoning="visual_semantics") + uid = next(iter(bindings.values())) + observation = SceneObservation( + ( + CameraObservation( + "front", + torch.zeros((4, 5, 3), dtype=torch.uint8), + None, + None, + None, + ), + ), + ({"uid": uid},), + ) + captured = {} + + def caller(**kwargs): + captured.update(kwargs) + return { + "entities": [ + { + "uid": uid, + "camera_uid": "front", + "bbox": [0.1, 0.2, 0.3, 0.4], + "confidence": 0.9, + } + ], + "relations": [], + "task_predicates": [{"type": "mouth_completed", "confidence": 0.9}], + "confidence": 0.9, + } + + facts = analyze_visual_scene(observation, task, caller=caller) + + predicate_type = captured["schema"]["properties"]["task_predicates"]["items"][ + "properties" + ]["type"] + assert predicate_type["enum"] == ["mouth_completed"] + assert facts["task_predicates"][0]["type"] == "mouth_completed" + + +def test_visual_facts_reject_unrequested_task_predicate() -> None: + value = { + "entities": [], + "relations": [], + "task_predicates": [{"type": "mouth_completed", "confidence": 0.9}], + "confidence": 0.9, + } + + with pytest.raises(ValueError, match="task_predicates.*must be one of"): + validate_visual_facts( + value, + known_uids={"known"}, + camera_uids={"front"}, + ) + + +def test_production_online_graph_caller_receives_reset_time_multiview_evidence( + monkeypatch: pytest.MonkeyPatch, +) -> None: + observation = SceneObservation( + ( + CameraObservation( + "front", + torch.zeros((4, 5, 3), dtype=torch.uint8), + torch.zeros((4, 5), dtype=torch.float32), + torch.eye(3), + torch.eye(4), + ), + ), + ({"uid": "known"},), + ) + captured = {} + + def caller(**kwargs): + captured.update(kwargs) + return {"nodes": [], "task_groups": [], "success": {}} + + monkeypatch.setattr(vision_module, "_default_structured_caller", caller) + monkeypatch.setattr(vision_module, "_vlm_model", lambda model: f"resolved:{model}") + + result = online_module._default_graph_caller( + prompt="plan", + schema={"type": "object"}, + model="mimo", + observation=observation, + ) + + assert result == {"nodes": [], "task_groups": [], "success": {}} + assert captured["model"] == "resolved:mimo" + assert len(captured["images"]) == 2 + + +def test_default_vision_caller_disables_custom_socket_transport( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import langchain_openai + + captured = {} + + class FakeRunnable: + def invoke(self, _messages): + return {"facts": []} + + class FakeChatOpenAI: + def __init__(self, **kwargs): + captured.update(kwargs) + + def with_structured_output(self, _schema, **_kwargs): + return FakeRunnable() + + monkeypatch.setattr(langchain_openai, "ChatOpenAI", FakeChatOpenAI) + monkeypatch.setattr( + planner_module, + "_load_llm_settings", + lambda *, model: { + "api_key": "test-key", + "model": model or "test-model", + "base_url": "https://example.test/v1", + "default_query": {}, + }, + ) + + vision_module._default_structured_caller( + prompt="inspect", + images=(), + schema={"type": "object"}, + model="test-model", + ) + + assert captured["http_socket_options"] == () + + +def test_visual_facts_reject_unstructured_entity_fields() -> None: + value = { + "entities": [ + { + "uid": "known", + "camera_uid": "front", + "bbox": [0.1, 0.2, 0.3, 0.4], + "semantic_label": "can", + "confidence": 0.9, + } + ], + "relations": [], + "task_predicates": [], + "confidence": 0.9, + } + + with pytest.raises(ValueError, match="unsupported fields"): + validate_visual_facts(value, known_uids={"known"}, camera_uids={"front"}) + + +def test_visual_facts_reject_noncanonical_relation_type() -> None: + value = { + "entities": [], + "relations": [ + {"type": "obstructs", "uids": ["box", "sign"], "confidence": 0.9} + ], + "task_predicates": [], + "confidence": 0.9, + } + + with pytest.raises(ValueError, match="relation type"): + validate_visual_facts( + value, + known_uids={"box", "sign"}, + camera_uids={"front"}, + ) + + +def test_visual_facts_require_ordered_relation_participants() -> None: + value = { + "entities": [], + "relations": [{"type": "occludes", "uids": ["box"], "confidence": 0.9}], + "task_predicates": [], + "confidence": 0.9, + } + + with pytest.raises(ValueError, match="exactly 2 UIDs"): + validate_visual_facts( + value, + known_uids={"box"}, + camera_uids={"front"}, + ) + + +def test_selection_prefers_exact_offline_and_l4_online() -> None: + task, _, bindings = _task("L1") + offline = instantiate_seed_graph(task, bindings) + online = deepcopy(offline) + online["planner_route"] = "online" + selected, evaluations = select_seed_graph( + offline, + online, + task, + known_objects=set(bindings.values()) | {"table"}, + exact_template_match=True, + ) + assert selected["metadata"]["selected_from"] == "offline" + assert evaluations["offline"].score > evaluations["online"].score + + l4, _, l4_bindings = _task("L4", reasoning="logic") + l4_offline = instantiate_seed_graph(l4, l4_bindings) + l4_online = deepcopy(l4_offline) + l4_online["planner_route"] = "online" + selected, _ = select_seed_graph( + l4_offline, + l4_online, + l4, + known_objects=set(l4_bindings.values()) | {"table"}, + visual_confidence=0.95, + ) + assert selected["metadata"]["selected_from"] == "online" + + +def test_fusion_keeps_whole_task_groups() -> None: + task, _, bindings = _task("L2") + offline = instantiate_seed_graph(task, bindings) + online = deepcopy(offline) + online["planner_route"] = "online" + routes = { + group["id"]: ("offline" if index % 2 == 0 else "online") + for index, group in enumerate(offline["task_groups"]) + } + fused = fuse_seed_graphs(offline, online, routes) + + assert fused["planner_route"] == "fused" + assert all( + all(node_id.startswith(routes[group["id"]]) for node_id in group["node_ids"]) + for group in fused["task_groups"] + ) diff --git a/tests/gen_sim/action_engine/planning/test_planner.py b/tests/gen_sim/action_engine/planning/test_planner.py new file mode 100644 index 000000000..c089ccef8 --- /dev/null +++ b/tests/gen_sim/action_engine/planning/test_planner.py @@ -0,0 +1,730 @@ +# ---------------------------------------------------------------------------- +# 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 inspect +import json +from pathlib import Path +from typing import Any + +import pytest + +from embodichain.gen_sim.action_engine.domain import TASK_AGENT_SCHEMA +from embodichain.gen_sim.action_engine.planning import plan_task +from embodichain.gen_sim.action_engine.planning import planner as planner_module + + +def _scene() -> list[dict[str, Any]]: + return [ + { + "uid": "table", + "runtime_uid": "table", + "source_uid": "table", + "role": "background", + "description": "A table.", + }, + *[ + { + "uid": f"interact_soda_can_{index}_0", + "runtime_uid": f"interact_soda_can_{index}", + "source_uid": f"interact_soda_can_{index}_0", + "role": "rigid_object", + "description": "An aluminum soda can.", + } + for index in range(5) + ], + ] + + +def _dual_arm_scene() -> list[dict[str, Any]]: + return [ + { + "uid": uid, + "runtime_uid": uid, + "source_uid": uid, + "role": "rigid_object", + "description": description, + } + for uid, description in ( + ("cube", "A cube on the left side of the table."), + ("cup", "A paper cup on the right side of the table."), + ("basket", "A basket near the center of the table."), + ) + ] + + +def _stack_scene() -> list[dict[str, Any]]: + return [ + { + "uid": uid, + "runtime_uid": uid, + "source_uid": f"{uid}_0", + "role": role, + "description": description, + } + for uid, role, description in ( + ("table", "background", "A table."), + ("paper_cup", "rigid_object", "A paper cup."), + ("popcorn_bucket", "rigid_object", "A popcorn bucket."), + ("earbuds_case", "rigid_object", "A blue earbuds case."), + ) + ] + + +def test_injected_planner_returns_only_semantics_and_resolves_aliases() -> None: + observed: dict[str, Any] = {} + + def caller(*, prompt: str, model: str | None) -> dict[str, Any]: + observed.update(prompt=prompt, model=model) + return { + "semantic_steps": [ + { + "id": "s01_place", + "operator": "place_relative", + "object": "interact_soda_can_0_0", + "goal": {"reference_object": "table", "relation": "on"}, + }, + { + "id": "s02_orient", + "operator": "orient_object", + "object": "interact_soda_can_1_0", + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + }, + }, + ] + } + + program = plan_task( + task_name="injected", + task_description="Place one object and then orient another.", + scene_objects=_scene(), + model="test-model", + llm_caller=caller, + ) + + assert program["schema_version"] == TASK_AGENT_SCHEMA + assert program["semantic_steps"][0]["object"] == "interact_soda_can_0" + assert program["semantic_steps"][1]["depends_on"] == ["s01_place"] + assert "Do not select a task route" in observed["prompt"] + assert observed["model"] == "test-model" + + +def test_planner_repairs_a_non_visible_skill_once() -> None: + calls = 0 + + def caller(**_kwargs: Any) -> dict[str, Any]: + nonlocal calls + calls += 1 + if calls == 1: + return { + "semantic_steps": [ + { + "id": "s01_hold", + "operator": "hold_hover", + "object": "cube", + "goal": {}, + } + ] + } + return { + "semantic_steps": [ + { + "id": "s01_place_cube", + "operator": "place_relative", + "object": "cube", + "goal": { + "reference_object": "basket", + "relation": "inside", + "orientation_goal": "preserve", + "orientation_axis": "none", + }, + "depends_on": [], + }, + { + "id": "s02_place_cup", + "operator": "place_relative", + "object": "cup", + "goal": { + "reference_object": "basket", + "relation": "inside", + }, + "depends_on": [], + }, + ], + "allocation_groups": [ + { + "id": "dual_arms_1", + "semantic_step_ids": ["s01_place_cube", "s02_place_cup"], + "arm_constraint": "distinct_arms", + } + ], + } + + program = plan_task( + task_name="dual_arm_basket", + task_description="test-instruction", + scene_objects=_dual_arm_scene(), + llm_caller=caller, + ) + + assert [ + (step["id"], step["operator"], step["object"], step["depends_on"]) + for step in program["semantic_steps"] + ] == [ + ("s01_place_cube", "place_relative", "cube", []), + ("s02_place_cup", "place_relative", "cup", []), + ] + assert calls == 2 + assert program["allocation_groups"][0]["arm_constraint"] == "distinct_arms" + + +def test_planner_repairs_build_stack_singular_object_contract() -> None: + prompts: list[str] = [] + + def caller(*, prompt: str, **_kwargs: Any) -> dict[str, Any]: + prompts.append(prompt) + if len(prompts) == 1: + return { + "semantic_steps": [ + { + "id": "s01_build_stack", + "operator": "build_stack", + "object": "paper_cup", + "goal": { + "anchor": "popcorn_bucket", + "stack_mode": "on_top", + }, + "depends_on": [], + } + ], + "allocation_groups": [], + } + return { + "semantic_steps": [ + { + "id": "s01_build_stack", + "operator": "build_stack", + "objects": ["paper_cup", "earbuds_case"], + "goal": { + "anchor": "popcorn_bucket", + "stack_mode": "on_top", + "orientation_goal": "preserve", + "orientation_axis": "none", + }, + "depends_on": [], + } + ], + "allocation_groups": [], + } + + program = plan_task( + task_name="task3_2", + task_description="test-instruction", + scene_objects=_stack_scene(), + llm_caller=caller, + ) + + assert len(prompts) == 2 + assert "build_stack requires an 'objects' list" in prompts[1] + assert program["semantic_steps"][0]["objects"] == [ + "paper_cup", + "earbuds_case", + ] + assert program["semantic_steps"][0]["goal"]["anchor"] == "popcorn_bucket" + + +def test_planner_rejects_a_non_visible_skill_after_one_repair() -> None: + def caller(**_kwargs: Any) -> dict[str, Any]: + return { + "semantic_steps": [ + { + "id": "s01_hold", + "operator": "hold_hover", + "object": "cube", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": {}, + "depends_on": [], + }, + { + "id": "s02_place", + "operator": "place_relative", + "object": "cube", + "actor": {"mode": "required", "arm": "right_arm"}, + "goal": { + "reference_object": "basket", + "relation": "inside", + }, + "depends_on": ["s01_hold"], + }, + ] + } + + with pytest.raises(ValueError, match="after one repair"): + plan_task( + task_name="conflicting_arms", + task_description="Hold the cube with the left arm, then place it " + "with the right arm.", + scene_objects=_dual_arm_scene(), + llm_caller=caller, + ) + + +def test_spatial_two_sided_phrase_does_not_invent_arm_constraint() -> None: + def caller(**_kwargs: Any) -> dict[str, Any]: + return { + "semantic_steps": [ + { + "id": "s01_left", + "operator": "orient_object", + "object": "cube", + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + }, + "depends_on": [], + }, + { + "id": "s02_right", + "operator": "orient_object", + "object": "cup", + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + }, + "depends_on": [], + }, + ], + "allocation_groups": [], + } + + program = plan_task( + task_name="two_sided_upright", + task_description="test-instruction", + scene_objects=_dual_arm_scene(), + llm_caller=caller, + ) + + assert program["allocation_groups"] == [] + + +def test_planner_does_not_infer_arm_group_from_instruction_text() -> None: + def caller(**_kwargs: Any) -> dict[str, Any]: + return { + "semantic_steps": [ + { + "id": step_id, + "operator": "orient_object", + "object": object_uid, + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + }, + "depends_on": [], + } + for step_id, object_uid in (("s01", "cube"), ("s02", "cup")) + ], + "allocation_groups": [], + } + + program = plan_task( + task_name="explicit_both_arms", + task_description="test-instruction", + scene_objects=_dual_arm_scene(), + llm_caller=caller, + ) + + assert program["allocation_groups"] == [] + + +def test_planner_does_not_expose_internal_operator_contracts() -> None: + def caller(**_kwargs: Any) -> dict[str, Any]: + return { + "semantic_steps": [ + { + "id": "s01_hold", + "operator": "hold_hover", + "object": "cube", + "actor": {"mode": "auto"}, + "goal": {}, + "depends_on": [], + }, + { + "id": "s02_place", + "operator": "place_relative", + "object": "cube", + "actor": {"mode": "auto"}, + "goal": { + "reference_object": "basket", + "relation": "inside", + }, + "depends_on": ["s01_hold"], + }, + ] + } + + with pytest.raises(ValueError, match="after one repair"): + plan_task( + task_name="nondefault_hover", + task_description="Hold the cube in a special pose, then place it.", + scene_objects=_dual_arm_scene(), + llm_caller=caller, + ) + + +def test_plan_task_has_no_rule_fallback_parameter() -> None: + assert "deterministic_fallback" not in inspect.signature(plan_task).parameters + + +def test_arrange_line_preserves_structured_orientation_output() -> None: + def caller(**_kwargs: Any) -> dict[str, Any]: + return { + "semantic_steps": [ + { + "id": "s01_line", + "operator": "arrange_line", + "objects": [ + "interact_soda_can_0", + "interact_soda_can_1", + ], + "goal": { + "axis": "world_y", + "order_constraint": "free", + "orientation_goal": "upright", + "orientation_axis": "long_axis", + }, + "depends_on": [], + } + ], + "allocation_groups": [], + } + + program = plan_task( + task_name="neutral_line", + task_description="test-instruction", + scene_objects=_scene(), + llm_caller=caller, + ) + + goal = program["semantic_steps"][0]["goal"] + assert goal["orientation_goal"] == "upright" + assert goal["orientation_axis"] == "long_axis" + + +def test_arrange_line_preserves_structured_axis_output() -> None: + def caller(**_kwargs: Any) -> dict[str, Any]: + return { + "semantic_steps": [ + { + "id": "s01_line", + "operator": "arrange_line", + "objects": [ + "interact_soda_can_0", + "interact_soda_can_1", + ], + "goal": { + "anchor": "table_center", + "axis": "world_x", + "order_constraint": "free", + }, + "depends_on": [], + } + ], + "allocation_groups": [], + } + + program = plan_task( + task_name="ambiguous_line_axis", + task_description="test-instruction", + scene_objects=_scene(), + llm_caller=caller, + ) + + assert program["semantic_steps"][0]["goal"]["axis"] == "world_x" + + +def test_instruction_text_does_not_override_structured_axis_output() -> None: + def caller(**_kwargs: Any) -> dict[str, Any]: + return { + "semantic_steps": [ + { + "id": "s01_line", + "operator": "arrange_line", + "objects": [ + "interact_soda_can_0", + "interact_soda_can_1", + ], + "goal": { + "anchor": "table_center", + "axis": "world_y", + "order_constraint": "free", + }, + "depends_on": [], + } + ], + "allocation_groups": [], + } + + program = plan_task( + task_name="front_to_back_line_axis", + task_description="test-instruction", + scene_objects=_scene(), + llm_caller=caller, + ) + + assert program["semantic_steps"][0]["goal"]["axis"] == "world_y" + + +def test_arrange_line_preserves_explicit_orientation_request() -> None: + def caller(**_kwargs: Any) -> dict[str, Any]: + return { + "semantic_steps": [ + { + "id": "s01_line", + "operator": "arrange_line", + "objects": [ + "interact_soda_can_0", + "interact_soda_can_1", + ], + "goal": { + "axis": "world_y", + "order_constraint": "free", + "orientation_goal": "upright", + "orientation_axis": "none", + }, + "depends_on": [], + } + ], + "allocation_groups": [], + } + + program = plan_task( + task_name="upright_line", + task_description="test-instruction", + scene_objects=_scene(), + llm_caller=caller, + ) + + assert program["semantic_steps"][0]["goal"]["orientation_goal"] == "upright" + + +def test_planner_rejects_route_or_graph_output() -> None: + def caller(**_kwargs: Any) -> dict[str, Any]: + return {"route": "arrangement_line", "semantic_steps": []} + + with pytest.raises(ValueError, match="only 'semantic_steps'"): + plan_task( + task_name="bad", + task_description="Arrange objects.", + scene_objects=_scene(), + llm_caller=caller, + ) + + +def test_llm_settings_read_gen_sim_dotenv( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + env_path = tmp_path / ".env" + env_path.write_text( + "\n".join( + ( + "# Local Action Engine credentials", + 'export OPENAI_API_KEY="dotenv-key"', + "OPENAI_BASE_URL=https://dotenv.example/v1/", + "OPENAI_MODEL=dotenv-model", + ) + ), + encoding="utf-8", + ) + config_path = tmp_path / "gen_config.json" + config_path.write_text( + json.dumps( + { + "llm": { + "openai_compatible": { + "api_key": "json-key", + "base_url": "https://json.example/v1", + "model": "json-model", + "default_query": {"api-version": "test"}, + } + } + } + ), + encoding="utf-8", + ) + monkeypatch.setattr(planner_module, "_GEN_SIM_ENV_PATH", env_path) + monkeypatch.setattr(planner_module, "_GEN_CONFIG_PATH", config_path) + for name in ( + "OPENAI_API_KEY", + "OPENAI_BASE_URL", + "OPENAI_API_BASE", + "OPENAI_MODEL", + "LLM_MODEL", + "LLM_URL", + ): + monkeypatch.delenv(name, raising=False) + + settings = planner_module._load_llm_settings(model=None) + + assert settings == { + "api_key": "dotenv-key", + "base_url": "https://dotenv.example/v1", + "model": "dotenv-model", + "default_query": {"api-version": "test"}, + } + + +def test_process_environment_and_model_argument_override_dotenv( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + env_path = tmp_path / ".env" + env_path.write_text( + "\n".join( + ( + "OPENAI_API_KEY=dotenv-key", + "OPENAI_BASE_URL=https://dotenv.example/v1", + "OPENAI_MODEL=dotenv-model", + ) + ), + encoding="utf-8", + ) + missing_config = tmp_path / "missing.json" + monkeypatch.setattr(planner_module, "_GEN_SIM_ENV_PATH", env_path) + monkeypatch.setattr(planner_module, "_GEN_CONFIG_PATH", missing_config) + monkeypatch.setenv("OPENAI_API_KEY", "shell-key") + monkeypatch.setenv("OPENAI_API_BASE", "https://shell.example/v1/") + monkeypatch.setenv("LLM_MODEL", "shell-model") + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_MODEL", raising=False) + + settings = planner_module._load_llm_settings(model="argument-model") + + assert settings["api_key"] == "shell-key" + assert settings["base_url"] == "https://shell.example/v1" + assert settings["model"] == "argument-model" + + +def test_partial_process_transport_does_not_mix_with_dotenv( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + env_path = tmp_path / ".env" + env_path.write_text( + "\n".join( + ( + "OPENAI_API_KEY=dotenv-key", + "OPENAI_BASE_URL=https://dotenv.example/v1", + "OPENAI_MODEL=dotenv-model", + ) + ), + encoding="utf-8", + ) + monkeypatch.setattr(planner_module, "_GEN_SIM_ENV_PATH", env_path) + monkeypatch.setattr( + planner_module, + "_GEN_CONFIG_PATH", + tmp_path / "missing.json", + ) + for name in ( + "OPENAI_API_KEY", + "OPENAI_BASE_URL", + "OPENAI_API_BASE", + "OPENAI_MODEL", + "LLM_MODEL", + "LLM_URL", + ): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("OPENAI_API_KEY", "unrelated-process-key") + + settings = planner_module._load_llm_settings(model=None) + + assert settings["api_key"] == "dotenv-key" + assert settings["base_url"] == "https://dotenv.example/v1" + assert settings["model"] == "dotenv-model" + + +def test_default_llm_caller_disables_custom_socket_transport( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import langchain_openai + + captured: dict[str, Any] = {} + + class FakeRunnable: + def invoke(self, _messages: Any) -> dict[str, list[Any]]: + return {"semantic_steps": [], "allocation_groups": []} + + class FakeChatOpenAI: + def __init__(self, **kwargs: Any) -> None: + captured.update(kwargs) + + def with_structured_output( + self, + _schema: dict[str, Any], + **_kwargs: Any, + ) -> FakeRunnable: + return FakeRunnable() + + monkeypatch.setattr(langchain_openai, "ChatOpenAI", FakeChatOpenAI) + monkeypatch.setattr( + planner_module, + "_load_llm_settings", + lambda *, model: { + "api_key": "test-key", + "model": model or "test-model", + "base_url": "https://example.test/v1", + "default_query": {}, + }, + ) + + planner_module._default_llm_caller(prompt="plan", model="test-model") + + assert captured["http_socket_options"] == () + + +def test_structured_output_transport_selects_json_mode_only_for_mimo() -> None: + calls: list[dict[str, Any]] = [] + + class FakeClient: + def with_structured_output(self, schema: dict[str, Any], **kwargs: Any) -> str: + calls.append({"schema": schema, "kwargs": kwargs}) + return "structured" + + schema = {"type": "object"} + client = FakeClient() + mimo = planner_module._structured_output_runnable( + client, + schema, + settings={ + "model": "mimo-v2.5", + "base_url": "https://token-plan-cn.xiaomimimo.com/v1", + }, + ) + generic = planner_module._structured_output_runnable( + client, + schema, + settings={"model": "gpt-test", "base_url": "https://example.test/v1"}, + ) + + assert mimo == generic == "structured" + assert [call["kwargs"] for call in calls] == [ + {"method": "json_mode"}, + {"method": "json_schema"}, + ] diff --git a/tests/gen_sim/action_engine/runtime/__init__.py b/tests/gen_sim/action_engine/runtime/__init__.py new file mode 100644 index 000000000..66ef3ba1a --- /dev/null +++ b/tests/gen_sim/action_engine/runtime/__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 + +"""Runtime contract tests for Action Engine.""" diff --git a/tests/gen_sim/action_engine/runtime/test_actions.py b/tests/gen_sim/action_engine/runtime/test_actions.py new file mode 100644 index 000000000..4e1bb3caf --- /dev/null +++ b/tests/gen_sim/action_engine/runtime/test_actions.py @@ -0,0 +1,787 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Focused contracts for the public atomic-action adapter.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +from embodichain.gen_sim.action_engine.runtime import actions +from embodichain.gen_sim.action_engine.runtime.actions import AtomicActionAdapter +from embodichain.gen_sim.action_engine.runtime.models import ( + ActionOutcome, + GroundedAction, +) +from embodichain.gen_sim.action_engine.runtime.state import ExecutionState +from embodichain.lab.sim.atomic_actions import ( + Affordance, + ActionPlan, + AntipodalAffordance, + CoordinatedPickGoal, + EndEffectorPoseGoal, + GraspGoal, + HeldObjectState, + JointPositionGoal, + ObjectSemantics, + PlannerDiagnostics, + RecoveryPolicy, + SceneSnapshot, + StateDelta, + TimedTrajectory, +) +from embodichain.lab.sim.planners import CuroboPlannerCfg +from embodichain.toolkits.graspkit.pg_grasp import GraspGeneratorCfg + + +class _MeshEntity: + def get_vertices(self, *, env_ids: list[int], scale: bool) -> torch.Tensor: + assert env_ids == [0] + assert scale + return torch.tensor( + [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + ], + dtype=torch.float32, + ) + + def get_triangles(self, *, env_ids: list[int]) -> torch.Tensor: + assert env_ids == [0] + return torch.tensor([[0, 1, 2]], dtype=torch.int64) + + +class _PoseEntity: + def __init__(self, pose: torch.Tensor) -> None: + self.pose = pose + + def get_local_pose(self, *, to_matrix: bool) -> torch.Tensor: + assert to_matrix + return self.pose.clone() + + +class _PlannerRobot: + uid = "test_robot" + dof = 8 + + _ids = { + "physical_left_arm": [0, 1], + "physical_left_eef": [2, 3], + "physical_right_arm": [4, 5], + "physical_right_eef": [6, 7], + } + + 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: + assert to_matrix + pose = torch.eye(4).repeat(2, 1, 1) + pose[:, 1, 3] = 0.3 if name == "physical_left_arm" else -0.3 + return pose + + +def _planner_env( + *, + table: Any | None = None, + rigid_objects: dict[str, Any] | None = None, +) -> SimpleNamespace: + entities = dict(rigid_objects or {}) + if table is not None: + entities["table"] = table + return SimpleNamespace( + num_envs=2, + device=torch.device("cpu"), + robot=_PlannerRobot(), + sim=SimpleNamespace(get_rigid_object=entities.get), + left_arm_joints=[0, 1], + left_eef_joints=[2, 3], + right_arm_joints=[4, 5], + right_eef_joints=[6, 7], + open_state=torch.zeros(2), + close_state=torch.ones(2), + get_agent_arm_control_part=lambda is_left: ( + "physical_left_arm" if is_left else "physical_right_arm" + ), + get_agent_eef_control_part=lambda is_left: ( + "physical_left_eef" if is_left else "physical_right_eef" + ), + ) + + +def test_semantics_prewarms_vhacd_cache_before_affordance( + monkeypatch: Any, +) -> None: + """The lazy shared checker must see V-HACD's pickle, never create CoACD.""" + events: list[str] = [] + observed: dict[str, Any] = {} + entity = _MeshEntity() + env = SimpleNamespace( + num_envs=1, + device=torch.device("cpu"), + sim=SimpleNamespace( + get_rigid_object=lambda uid: entity if uid == "cube" else None + ), + agent_grasp_runtime_defaults={"max_decomposition_hulls": 8}, + ) + + def fake_prepare(**kwargs: Any) -> SimpleNamespace: + events.append("cache") + observed.update(kwargs) + return SimpleNamespace(status="hit") + + def fake_affordance(**kwargs: Any) -> Affordance: + events.append("affordance") + observed["generator_cfg"] = kwargs["generator_cfg"] + observed["gripper_collision_cfg"] = kwargs["gripper_collision_cfg"] + return Affordance() + + monkeypatch.setattr( + actions, + "ensure_vhacd_grasp_collision_cache", + fake_prepare, + ) + monkeypatch.setattr(actions, "AntipodalAffordance", fake_affordance) + + adapter = AtomicActionAdapter(env) + first = adapter.semantics("cube") + second = adapter.semantics("cube") + + assert first is second + assert events == ["cache", "affordance"] + assert observed["max_decomposition_hulls"] == 8 + assert observed["mesh_vertices"].dtype == torch.float32 + assert observed["mesh_triangles"].dtype == torch.int64 + assert observed["generator_cfg"].n_deviated_approach_directions == 4 + assert observed["gripper_collision_cfg"] is not None + + +def test_planner_policy_uses_curobo_for_single_arm_and_ik_for_dual_arm() -> None: + adapter = AtomicActionAdapter(_planner_env()) + goal = JointPositionGoal(target=torch.zeros(2, 2)) + + single = adapter._invocation( + GroundedAction("MoveJoints", "left_arm", "arm", goal, {}), + adapter.capabilities.get("MoveJoints"), + ) + coordinated_goal = CoordinatedPickGoal( + semantics=ObjectSemantics( + label="tray", + geometry={}, + affordance=AntipodalAffordance(), + ), + object_target_pose=torch.eye(4), + object_initial_pose=torch.eye(4), + ) + coordinated = adapter._invocation( + GroundedAction( + "CoordinatedPickment", + "coordinated", + "arm", + coordinated_goal, + {}, + ), + adapter.capabilities.get("CoordinatedPickment"), + ) + hand = adapter._invocation( + GroundedAction("MoveJoints", "left_arm", "hand", goal, {}), + adapter.capabilities.get("MoveJoints"), + ) + + assert single.motion_policy.planner == "curobo" + assert single.motion_policy.strategy == "motion_gen" + assert coordinated.motion_policy.strategy == "ik_interp" + assert torch.allclose( + coordinated.skill_options.left_to_right_arm_direction, + torch.tensor([0.0, -1.0, 0.0]), + ) + assert hand.motion_policy.strategy == "ik_interp" + + +def test_coordinated_pickment_scopes_ground_filter_to_gensim_goal_copy() -> None: + adapter = AtomicActionAdapter(_planner_env()) + original_cfg = GraspGeneratorCfg(is_filter_ground_collision=True) + affordance = AntipodalAffordance(generator_cfg=original_cfg) + goal = CoordinatedPickGoal( + semantics=ObjectSemantics( + label="tray", + geometry={}, + affordance=affordance, + ), + object_target_pose=torch.eye(4), + object_initial_pose=torch.eye(4), + ) + grounded = GroundedAction( + "CoordinatedPickment", + "coordinated", + "arm", + goal, + { + "middle_empty_ratio": 0.7, + "is_filter_ground_collision": False, + }, + ) + + invocation = adapter._invocation( + grounded, + adapter.capabilities.get("CoordinatedPickment"), + ) + + scoped_affordance = invocation.goal.semantics.affordance + assert isinstance(scoped_affordance, AntipodalAffordance) + assert scoped_affordance is not affordance + assert affordance.generator_cfg is original_cfg + assert original_cfg.is_filter_ground_collision is True + assert scoped_affordance.generator_cfg is not original_cfg + assert scoped_affordance.generator_cfg.is_filter_ground_collision is False + assert invocation.skill_options.middle_empty_ratio == pytest.approx(0.7) + + +def test_retreat_uses_row_local_motion_planner_reachability_search( + monkeypatch: Any, +) -> None: + env = _planner_env() + adapter = AtomicActionAdapter(env) + reference = torch.eye(4).repeat(2, 1, 1) + reference[:, 2, 3] = 1.05 + requested = reference.clone() + requested[:, 2, 3] = 1.35 + height_thresholds = torch.tensor([1.24, 1.00]) + attempted_targets: list[torch.Tensor] = [] + + def plan(invocation: Any, _context: Any) -> ActionPlan: + target = invocation.goal.xpos.clone() + attempted_targets.append(target) + height_reachable = target[:, 2, 3] <= height_thresholds + baseward_reachable = target[:, 1, 3] < -0.05 + success = height_reachable | baseward_reachable + terminal = target[:, 2, 3, None].repeat(1, 8) + positions = torch.stack((torch.zeros_like(terminal), terminal), dim=1) + return ActionPlan( + skill_id="move_end_effector", + plan_success=success, + trajectory=TimedTrajectory.from_positions( + positions, + env_ids=torch.arange(2), + control_dt=0.01, + ), + recovery_policy=RecoveryPolicy(), + planned_scene_version=0, + planned_collision_world_revision=(0, 0), + diagnostics=PlannerDiagnostics(backend="fake"), + expected_effects=StateDelta(), + ) + + monkeypatch.setattr(adapter, "_engine", lambda: SimpleNamespace(plan=plan)) + grounded = GroundedAction( + "MoveEndEffector", + "right_arm", + "arm", + EndEffectorPoseGoal(xpos=requested), + { + "sample_interval": 10, + "retreat_height": 0.30, + "minimum_retreat_height": 0.05, + "retreat_distance": 0.10, + }, + motion_policy={ + "collision_safety": "required", + "retreat_reachability_search": True, + "retreat_reference_pose": reference, + "minimum_retreat_height": 0.05, + "retreat_distance": 0.10, + }, + ) + + outcome = adapter.plan( + grounded, + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + assert len(attempted_targets) > 1 + assert bool(outcome.success.all()) + selected_z = outcome.grounded.target.xpos[:, 2, 3] + assert selected_z.tolist() == pytest.approx([1.20, 1.35]) + assert outcome.grounded.target.xpos[:, 1, 3].tolist() == pytest.approx([0.0, -0.10]) + search = outcome.planner_trace["reachability_search"] + assert search["strategy"] == "bounded_motion_planner" + assert search["selected_target_z"].tolist() == pytest.approx([1.20, 1.35]) + assert len(search["attempts"]) == len(attempted_targets) + + +def test_curobo_generator_receives_generated_static_obstacles( + monkeypatch: Any, +) -> None: + table = object() + can = object() + captured: dict[str, Any] = {} + + def fake_motion_generator(*, cfg: Any) -> object: + captured["cfg"] = cfg + return object() + + monkeypatch.setattr(actions, "MotionGenerator", fake_motion_generator) + adapter = AtomicActionAdapter( + _planner_env(table=table, rigid_objects={"can": can}), + planner_policy={ + "dynamic_collision": True, + "dynamic_obstacle_uids": ["can"], + }, + ) + + generator = adapter._generator() + + assert generator is adapter._motion_generator + planner = captured["cfg"].planner_cfg + assert isinstance(planner, CuroboPlannerCfg) + assert planner.world.rigid_objects == {"table": table, "can": can} + assert planner.world.dynamic_obstacle_names == ["can"] + assert planner.world.obstacle_representation == "cuboid" + assert planner.world.collision_cache == {"cuboid": 8, "mesh": 2} + + +def test_curobo_generator_sizes_collision_cache_for_large_scene( + monkeypatch: Any, +) -> None: + rigid_objects = {f"object_{index:02d}": object() for index in range(13)} + captured: dict[str, Any] = {} + + def fake_motion_generator(*, cfg: Any) -> object: + captured["cfg"] = cfg + return object() + + monkeypatch.setattr(actions, "MotionGenerator", fake_motion_generator) + adapter = AtomicActionAdapter( + _planner_env(rigid_objects=rigid_objects), + planner_policy={ + "dynamic_collision": True, + "dynamic_obstacle_uids": list(rigid_objects), + }, + ) + + adapter._generator() + + planner = captured["cfg"].planner_cfg + assert planner.world.collision_cache == {"cuboid": 13, "mesh": 2} + + +def test_dynamic_scene_parks_contact_target_and_held_rows() -> None: + actual = torch.eye(4).repeat(2, 1, 1) + actual[:, 2, 3] = torch.tensor([0.7, 0.8]) + entities = {uid: _PoseEntity(actual.clone()) for uid in ("target", "held", "other")} + adapter = AtomicActionAdapter( + _planner_env(rigid_objects=entities), + planner_policy={ + "dynamic_collision": True, + "dynamic_obstacle_uids": list(entities), + }, + ) + held_semantics = ObjectSemantics( + label="held", + entity=entities["held"], + geometry={}, + affordance=Affordance(), + ) + held = HeldObjectState( + semantics=held_semantics, + object_to_eef=torch.eye(4).repeat(2, 1, 1), + grasp_xpos=torch.eye(4).repeat(2, 1, 1), + env_mask=torch.tensor([True, False]), + ) + state = ExecutionState( + last_qpos=torch.zeros(2, 8), + held_objects={"physical_left_arm": held}, + ) + grounded = GroundedAction( + "PickUp", + "right_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + object_uid="target", + ) + + scene = adapter._scene_snapshot(grounded, state) + + assert torch.equal( + scene.entities["target"].pose[:, 2, 3], + actual[:, 2, 3] + actions._COLLISION_PARKING_Z_OFFSET, + ) + assert scene.entities["held"].pose[0, 2, 3] == ( + actual[0, 2, 3] + actions._COLLISION_PARKING_Z_OFFSET + ) + assert scene.entities["held"].pose[1, 2, 3] == actual[1, 2, 3] + assert torch.equal(scene.entities["other"].pose, actual) + + +def test_released_object_returns_to_live_dynamic_collision_pose() -> None: + actual = torch.eye(4).repeat(2, 1, 1) + actual[:, 0, 3] = torch.tensor([0.2, 0.4]) + entity = _PoseEntity(actual) + adapter = AtomicActionAdapter( + _planner_env(rigid_objects={"released": entity}), + planner_policy={ + "dynamic_collision": True, + "dynamic_obstacle_uids": ["released"], + }, + ) + grounded = GroundedAction( + "MoveJoints", + "left_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + object_uid="released", + ) + + scene = adapter._scene_snapshot( + grounded, + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + assert torch.equal(scene.entities["released"].pose, actual) + + +def test_default_scene_provider_advances_only_after_material_change() -> None: + actual = torch.eye(4).repeat(2, 1, 1) + entity = _PoseEntity(actual.clone()) + adapter = AtomicActionAdapter( + _planner_env(rigid_objects={"can": entity}), + planner_policy={ + "dynamic_collision": True, + "dynamic_obstacle_uids": ["can"], + }, + ) + grounded = GroundedAction( + "MoveJoints", + "left_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + object_uid="can", + ) + state = ExecutionState(last_qpos=torch.zeros(2, 8)) + + first = adapter._scene_snapshot(grounded, state) + unchanged = adapter._scene_snapshot(grounded, state) + entity.pose[:, 0, 3] += 0.1 + changed = adapter._scene_snapshot(grounded, state) + + assert first.version == unchanged.version == 0 + assert changed.version == 1 + assert changed.collision_world_revisions(2) == (1, 1) + + +def test_external_scene_provider_is_used_by_planning_snapshot() -> None: + pose = torch.eye(4).repeat(2, 1, 1) + + class _Provider: + def snapshot(self, *, timestamp: float, env_ids: torch.Tensor) -> SceneSnapshot: + assert timestamp == 0.0 + assert torch.equal(env_ids, torch.tensor([0, 1])) + return SceneSnapshot( + timestamp=timestamp, + version=7, + entities={"can": actions.EntityState(pose)}, + ) + + adapter = AtomicActionAdapter( + _planner_env(), + scene_provider=_Provider(), + ) + grounded = GroundedAction( + "MoveJoints", + "left_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + object_uid="can", + ) + + scene = adapter._scene_snapshot( + grounded, + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + assert scene.version == 7 + assert torch.equal(scene.entities["can"].pose, pose) + + +def test_start_session_delegates_to_shared_atomic_engine(monkeypatch: Any) -> None: + adapter = AtomicActionAdapter(_planner_env()) + grounded = GroundedAction( + "MoveJoints", + "left_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + ) + state = ExecutionState(last_qpos=torch.zeros(2, 8)) + marker = object() + captured: dict[str, Any] = {} + + monkeypatch.setattr(adapter, "_planning_context", lambda *_args: "context") + monkeypatch.setattr(adapter, "_invocation", lambda *_args: "invocation") + + class _Engine: + def start(self, invocations: tuple[Any, ...], context: Any) -> object: + captured["invocations"] = invocations + captured["context"] = context + return marker + + monkeypatch.setattr(adapter, "_engine", lambda: _Engine()) + + result = adapter.start_session(grounded, state) + + assert result is marker + assert captured == {"invocations": ("invocation",), "context": "context"} + + +def test_retreat_parks_intentional_contact_objects() -> None: + actual = torch.eye(4).repeat(2, 1, 1) + entities = { + uid: _PoseEntity(actual.clone()) for uid in ("released", "container", "other") + } + adapter = AtomicActionAdapter( + _planner_env(rigid_objects=entities), + planner_policy={ + "dynamic_collision": True, + "dynamic_obstacle_uids": list(entities), + }, + ) + grounded = GroundedAction( + "MoveEndEffector", + "left_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + motion_policy={ + "collision_exclusion_uids": ["released", "container"], + }, + object_uid="released", + ) + + scene = adapter._scene_snapshot( + grounded, + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + parked_z = actual[:, 2, 3] + actions._COLLISION_PARKING_Z_OFFSET + assert torch.equal(scene.entities["released"].pose[:, 2, 3], parked_z) + assert torch.equal(scene.entities["container"].pose[:, 2, 3], parked_z) + assert torch.equal(scene.entities["other"].pose, actual) + + +def test_action_outcome_commits_state_delta_only_for_verified_rows() -> None: + semantics = ObjectSemantics( + label="cube", + entity=object(), + geometry={}, + affordance=Affordance(), + ) + held = HeldObjectState( + semantics=semantics, + object_to_eef=torch.eye(4).repeat(2, 1, 1), + grasp_xpos=torch.eye(4).repeat(2, 1, 1), + ) + prior = ExecutionState(last_qpos=torch.zeros(2, 3)) + trajectory = torch.stack( + (torch.zeros(2, 3), torch.ones(2, 3)), + dim=1, + ) + delta = StateDelta(held_object_updates={"physical_left_arm": held}) + projected = ExecutionState.from_task_state( + delta.apply(prior.to_task_state(), torch.ones(2, dtype=torch.bool)), + last_qpos=trajectory[:, -1], + ) + grounded = GroundedAction( + "PickUp", + "left_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + ) + outcome = ActionOutcome( + trajectory=trajectory, + success=torch.ones(2, dtype=torch.bool), + next_state=projected, + grounded=grounded, + prior_state=prior, + expected_effects=delta, + ) + + committed = outcome.state_after(torch.tensor([True, False])) + + assert torch.equal(committed.last_qpos[0], torch.ones(3)) + assert torch.equal(committed.last_qpos[1], torch.zeros(3)) + committed_held = committed.get_held_object("physical_left_arm") + assert committed_held is not None + assert torch.equal(committed_held.env_mask, torch.tensor([True, False])) + + +def test_fallback_rows_keep_the_fallback_plan_effects(monkeypatch: Any) -> None: + env = _planner_env() + adapter = AtomicActionAdapter(env) + semantics = ObjectSemantics( + label="cube", + entity=object(), + geometry={}, + affordance=Affordance(), + ) + + def held_at(x: float) -> HeldObjectState: + relation = torch.eye(4).repeat(2, 1, 1) + relation[:, 0, 3] = x + return HeldObjectState( + semantics=semantics, + object_to_eef=relation, + grasp_xpos=torch.eye(4).repeat(2, 1, 1), + ) + + def action_plan( + success: torch.Tensor, + terminal: float, + held: HeldObjectState, + ) -> ActionPlan: + positions = torch.full((2, 2, 8), terminal) + return ActionPlan( + skill_id="pick_up", + plan_success=success, + trajectory=TimedTrajectory.from_positions( + positions, + env_ids=torch.arange(2), + control_dt=0.01, + ), + recovery_policy=RecoveryPolicy(), + planned_scene_version=0, + planned_collision_world_revision=(0, 0), + diagnostics=PlannerDiagnostics(backend="fake"), + expected_effects=StateDelta( + held_object_updates={"physical_left_arm": held} + ), + ) + + plans = iter( + ( + action_plan(torch.tensor([True, False]), 1.0, held_at(1.0)), + action_plan(torch.tensor([True, True]), 2.0, held_at(2.0)), + ) + ) + strategies: list[str] = [] + + def plan(invocation: Any, _context: Any) -> ActionPlan: + strategies.append(invocation.motion_policy.strategy) + return next(plans) + + monkeypatch.setattr( + adapter, + "_engine", + lambda: SimpleNamespace(plan=plan), + ) + grounded = GroundedAction( + "PickUp", + "left_arm", + "arm", + GraspGoal(semantics=semantics), + {}, + ) + + outcome = adapter.plan( + grounded, + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + assert strategies == ["motion_gen", "ik_interp"] + assert torch.equal(outcome.success, torch.tensor([True, True])) + assert torch.equal(outcome.next_state.last_qpos[0], torch.ones(8)) + assert torch.equal(outcome.next_state.last_qpos[1], torch.full((8,), 2.0)) + held = outcome.next_state.get_held_object("physical_left_arm") + assert held is not None + assert held.object_to_eef[0, 0, 3] == 1.0 + assert held.object_to_eef[1, 0, 3] == 2.0 + assert torch.equal( + outcome.planner_trace["primary_success"], torch.tensor([True, False]) + ) + assert torch.equal( + outcome.planner_trace["fallback_attempted"], torch.tensor([False, True]) + ) + assert torch.equal( + outcome.planner_trace["fallback_used"], torch.tensor([False, True]) + ) + + +def test_collision_required_cleanup_does_not_use_unsafe_fallback( + monkeypatch: Any, +) -> None: + pose = torch.eye(4).repeat(2, 1, 1) + adapter = AtomicActionAdapter( + _planner_env(rigid_objects={"released": _PoseEntity(pose)}), + planner_policy={ + "dynamic_collision": True, + "dynamic_obstacle_uids": ["released"], + }, + ) + failed_plan = ActionPlan( + skill_id="move_joints", + plan_success=torch.tensor([False, False]), + trajectory=TimedTrajectory.from_positions( + torch.zeros(2, 2, 8), + env_ids=torch.arange(2), + control_dt=0.01, + ), + recovery_policy=RecoveryPolicy(), + planned_scene_version=1, + planned_collision_world_revision=(1, 1), + diagnostics=PlannerDiagnostics(backend="fake"), + expected_effects=StateDelta(), + ) + strategies: list[str] = [] + + def plan(invocation: Any, _context: Any) -> ActionPlan: + strategies.append(invocation.motion_policy.strategy) + assert invocation.motion_policy.dynamic_collision_mode.value == "required" + return failed_plan + + monkeypatch.setattr(adapter, "_engine", lambda: SimpleNamespace(plan=plan)) + grounded = GroundedAction( + "MoveJoints", + "left_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + motion_policy={"collision_safety": "required"}, + object_uid="released", + ) + + outcome = adapter.plan( + grounded, + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + assert strategies == ["motion_gen"] + assert not bool(outcome.success.any()) + assert outcome.planner_trace["fallback_allowed"] is False + assert not bool(outcome.planner_trace["fallback_attempted"].any()) + assert not bool(outcome.planner_trace["fallback_used"].any()) + assert outcome.planner_trace["collision_obstacle_positions"]["released"].shape == ( + 2, + 3, + ) diff --git a/tests/gen_sim/action_engine/runtime/test_atomic_compat.py b/tests/gen_sim/action_engine/runtime/test_atomic_compat.py new file mode 100644 index 000000000..e6c2d7573 --- /dev/null +++ b/tests/gen_sim/action_engine/runtime/test_atomic_compat.py @@ -0,0 +1,87 @@ +# ---------------------------------------------------------------------------- +# 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 types import SimpleNamespace + +import pytest +import torch + +from embodichain.gen_sim.action_engine.runtime.actions import AtomicActionAdapter +from embodichain.gen_sim.action_engine.runtime.atomic_compat import ( + ExactTargetMoveHeldObject, + ExactTargetMoveHeldObjectOptions, +) +from embodichain.lab.sim.atomic_actions import MoveHeldObject, MoveHeldObjectOptions + + +def test_exact_target_transport_only_disables_rotation_when_requested( + monkeypatch: pytest.MonkeyPatch, +) -> None: + applied = [] + result = object() + + def fake_apply(self, move_eef_xpos, end_arm_xpos) -> None: + del self, move_eef_xpos, end_arm_xpos + applied.append(True) + + def fake_plan(self, request, context): + del request, context + self._apply_automatic_transport_rotation(torch.eye(4), torch.eye(4)) + return result + + monkeypatch.setattr( + MoveHeldObject, + "_apply_automatic_transport_rotation", + fake_apply, + ) + monkeypatch.setattr(MoveHeldObject, "_plan", fake_plan) + action = ExactTargetMoveHeldObject() + disabled_request = SimpleNamespace( + skill_options=ExactTargetMoveHeldObjectOptions( + allow_automatic_transport_rotation=False, + ) + ) + enabled_request = SimpleNamespace( + skill_options=ExactTargetMoveHeldObjectOptions(), + ) + + assert action._plan(disabled_request, object()) is result + assert not applied + assert action._plan(enabled_request, object()) is result + assert applied == [True] + + +@pytest.mark.parametrize( + ("yaw_samples", "expected"), + [(1, True), (8, False)], +) +def test_semantic_transport_config_scopes_rotation_override( + yaw_samples: int, + expected: bool, +) -> None: + adapter = AtomicActionAdapter.__new__(AtomicActionAdapter) + action = SimpleNamespace(cfg={"upright_yaw_samples": yaw_samples}) + capability = SimpleNamespace( + config_type=MoveHeldObjectOptions, + target_materializer="semantic_held_object", + ) + + options = adapter._build_single_arm_config(action, capability) + + assert isinstance(options, ExactTargetMoveHeldObjectOptions) + assert options.allow_automatic_transport_rotation is expected diff --git a/tests/gen_sim/action_engine/runtime/test_grasp_collision_cache.py b/tests/gen_sim/action_engine/runtime/test_grasp_collision_cache.py new file mode 100644 index 000000000..a08771054 --- /dev/null +++ b/tests/gen_sim/action_engine/runtime/test_grasp_collision_cache.py @@ -0,0 +1,354 @@ +# ---------------------------------------------------------------------------- +# 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 hashlib +import json +import os +from pathlib import Path +import pickle +from typing import Callable + +import numpy as np +import pytest +import torch + +from embodichain.gen_sim.action_engine.runtime import grasp_collision_cache +from embodichain.gen_sim.action_engine.runtime.grasp_collision_cache import ( + GraspCollisionCacheError, + ensure_vhacd_grasp_collision_cache, + grasp_collision_cache_path, +) + + +def _tetrahedron() -> tuple[torch.Tensor, torch.Tensor]: + vertices = torch.tensor( + [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + ], + dtype=torch.float32, + ) + triangles = torch.tensor( + [ + [0, 2, 1], + [0, 1, 3], + [0, 3, 2], + [1, 2, 3], + ], + dtype=torch.int64, + ) + return vertices, triangles + + +def _plane_equations() -> list[tuple[np.ndarray, np.ndarray]]: + return [ + ( + np.asarray( + [ + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + ], + dtype=np.float32, + ), + np.asarray([-1.0, -1.0, -1.0], dtype=np.float32), + ), + ( + np.asarray([[1.0, 1.0, 1.0]], dtype=np.float32), + np.asarray([-1.0], dtype=np.float32), + ), + ] + + +def _install_fake_decomposer( + monkeypatch: pytest.MonkeyPatch, +) -> list[tuple[tuple[int, ...], tuple[int, ...], int]]: + calls: list[tuple[tuple[int, ...], tuple[int, ...], int]] = [] + + def fake_decompose( + vertices: np.ndarray, + triangles: np.ndarray, + max_decomposition_hulls: int, + ) -> list[tuple[np.ndarray, np.ndarray]]: + calls.append( + ( + tuple(vertices.shape), + tuple(triangles.shape), + max_decomposition_hulls, + ) + ) + return _plane_equations() + + monkeypatch.setattr( + grasp_collision_cache, + "_compute_vhacd_plane_equations", + fake_decompose, + ) + return calls + + +def test_cache_key_and_payload_match_main_checker_contract( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + vertices, triangles = _tetrahedron() + _install_fake_decomposer(monkeypatch) + + result = ensure_vhacd_grasp_collision_cache( + mesh_vertices=vertices, + mesh_triangles=triangles, + max_decomposition_hulls=16, + cache_dir=tmp_path, + ) + + expected_hash = hashlib.md5( + vertices.numpy().tobytes() + triangles.numpy().tobytes() + ).hexdigest() + assert result.cache_path == tmp_path / f"{expected_hash}_16.pkl" + with result.cache_path.open("rb") as cache_file: + payload = pickle.load(cache_file) + assert set(payload) == {"plane_equations", "plane_equation_counts"} + assert payload["plane_equations"].shape == (2, 3, 4) + assert payload["plane_equations"].dtype == torch.float32 + assert payload["plane_equation_counts"].tolist() == [3, 1] + assert payload["plane_equation_counts"].dtype == torch.int32 + + +def test_main_checker_loads_prepared_cache_without_running_coacd( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + import embodichain.lab.sim + from embodichain.toolkits.graspkit.pg_grasp import collision_checker + + vertices, triangles = _tetrahedron() + _install_fake_decomposer(monkeypatch) + result = ensure_vhacd_grasp_collision_cache( + mesh_vertices=vertices, + mesh_triangles=triangles, + max_decomposition_hulls=16, + cache_dir=tmp_path, + ) + + def fail_coacd(*args: object, **kwargs: object) -> None: + raise AssertionError("The prepared V-HACD cache must bypass CoACD.") + + monkeypatch.setattr(embodichain.lab.sim, "CONVEX_DECOMP_DIR", tmp_path) + monkeypatch.setattr(collision_checker, "convex_decomposition_coacd", fail_coacd) + checker = collision_checker.ConvexCollisionChecker( + vertices, + triangles, + max_decomposition_hulls=16, + ) + + assert checker.cache_path == result.cache_path.as_posix() + assert checker.plane_equations["plane_equation_counts"].tolist() == [3, 1] + + +def test_matching_vhacd_metadata_returns_cache_hit( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + vertices, triangles = _tetrahedron() + calls = _install_fake_decomposer(monkeypatch) + kwargs = { + "mesh_vertices": vertices, + "mesh_triangles": triangles, + "max_decomposition_hulls": 16, + "cache_dir": tmp_path, + } + + first = ensure_vhacd_grasp_collision_cache(**kwargs) + second = ensure_vhacd_grasp_collision_cache(**kwargs) + + assert first.status == "generated" + assert second.status == "hit" + assert len(calls) == 1 + + +def test_non_vhacd_metadata_forces_cache_replacement( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + vertices, triangles = _tetrahedron() + calls = _install_fake_decomposer(monkeypatch) + kwargs = { + "mesh_vertices": vertices, + "mesh_triangles": triangles, + "max_decomposition_hulls": 16, + "cache_dir": tmp_path, + } + first = ensure_vhacd_grasp_collision_cache(**kwargs) + metadata = json.loads(first.metadata_path.read_text(encoding="utf-8")) + metadata["backend"] = "coacd" + first.metadata_path.write_text(json.dumps(metadata), encoding="utf-8") + + replaced = ensure_vhacd_grasp_collision_cache(**kwargs) + + assert replaced.status == "replaced" + assert len(calls) == 2 + repaired = json.loads(replaced.metadata_path.read_text(encoding="utf-8")) + assert repaired["backend"] == "vhacd" + + +def test_modified_cache_fails_checksum_and_is_rebuilt( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + vertices, triangles = _tetrahedron() + calls = _install_fake_decomposer(monkeypatch) + kwargs = { + "mesh_vertices": vertices, + "mesh_triangles": triangles, + "max_decomposition_hulls": 16, + "cache_dir": tmp_path, + } + first = ensure_vhacd_grasp_collision_cache(**kwargs) + first.cache_path.write_bytes(b"not a valid collision cache") + + replaced = ensure_vhacd_grasp_collision_cache(**kwargs) + + assert replaced.status == "replaced" + assert len(calls) == 2 + with replaced.cache_path.open("rb") as cache_file: + assert "plane_equations" in pickle.load(cache_file) + + +def test_cache_and_metadata_are_published_by_atomic_replace( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + vertices, triangles = _tetrahedron() + _install_fake_decomposer(monkeypatch) + replacements: list[tuple[Path, Path]] = [] + real_replace: Callable[[os.PathLike[str], os.PathLike[str]], None] = os.replace + + def recording_replace( + source: os.PathLike[str], + destination: os.PathLike[str], + ) -> None: + replacements.append((Path(source), Path(destination))) + real_replace(source, destination) + + monkeypatch.setattr(grasp_collision_cache.os, "replace", recording_replace) + + result = ensure_vhacd_grasp_collision_cache( + mesh_vertices=vertices, + mesh_triangles=triangles, + max_decomposition_hulls=16, + cache_dir=tmp_path, + ) + + assert [destination for _, destination in replacements] == [ + result.cache_path, + result.metadata_path, + ] + assert all( + source.parent == destination.parent for source, destination in replacements + ) + assert all(not source.exists() for source, _ in replacements) + + +@pytest.mark.parametrize( + ("vertices", "triangles", "message"), + [ + ( + torch.empty((0, 3), dtype=torch.float32), + torch.tensor([[0, 1, 2]], dtype=torch.int64), + "mesh_vertices", + ), + ( + torch.zeros((3, 3), dtype=torch.float32), + torch.tensor([[0, 1]], dtype=torch.int64), + "mesh_triangles", + ), + ( + torch.tensor([[0.0, 0.0, 0.0], [1.0, float("nan"), 0.0], [0.0, 1.0, 0.0]]), + torch.tensor([[0, 1, 2]], dtype=torch.int64), + "finite", + ), + ( + torch.zeros((3, 3), dtype=torch.float32), + torch.tensor([[0, 1, 3]], dtype=torch.int64), + "indices", + ), + ], +) +def test_invalid_mesh_is_rejected_before_decomposition( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + vertices: torch.Tensor, + triangles: torch.Tensor, + message: str, +) -> None: + calls = _install_fake_decomposer(monkeypatch) + + with pytest.raises(ValueError, match=message): + ensure_vhacd_grasp_collision_cache( + mesh_vertices=vertices, + mesh_triangles=triangles, + max_decomposition_hulls=16, + cache_dir=tmp_path, + ) + + assert calls == [] + + +@pytest.mark.parametrize("max_decomposition_hulls", [True, 0, -1, 1.5]) +def test_invalid_hull_limit_is_rejected( + tmp_path: Path, + max_decomposition_hulls: object, +) -> None: + vertices, triangles = _tetrahedron() + + with pytest.raises((TypeError, ValueError), match="max_decomposition_hulls"): + ensure_vhacd_grasp_collision_cache( + mesh_vertices=vertices, + mesh_triangles=triangles, + max_decomposition_hulls=max_decomposition_hulls, # type: ignore[arg-type] + cache_dir=tmp_path, + ) + + +def test_symlinked_cache_path_is_refused( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + vertices, triangles = _tetrahedron() + _install_fake_decomposer(monkeypatch) + cache_path = grasp_collision_cache_path( + vertices, + triangles, + 16, + cache_dir=tmp_path, + ) + victim = tmp_path / "victim.pkl" + victim.write_bytes(b"do not overwrite") + cache_path.symlink_to(victim) + + with pytest.raises(GraspCollisionCacheError, match="symlink"): + ensure_vhacd_grasp_collision_cache( + mesh_vertices=vertices, + mesh_triangles=triangles, + max_decomposition_hulls=16, + cache_dir=tmp_path, + ) + + assert victim.read_bytes() == b"do not overwrite" 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..10268150c --- /dev/null +++ b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py @@ -0,0 +1,5853 @@ +# ---------------------------------------------------------------------------- +# 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 ( + 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) + binding = adapter._binding( + grounded, + adapter.capabilities.get("CoordinatedPlacement"), + ) + assert binding.manipulators == { + "placing": "physical_left_arm", + "support": "physical_right_arm", + } + assert binding.end_effectors == { + "placing": "physical_left_eef", + "support": "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/task_fixtures.py b/tests/gen_sim/action_engine/task_fixtures.py new file mode 100644 index 000000000..c845cdfab --- /dev/null +++ b/tests/gen_sim/action_engine/task_fixtures.py @@ -0,0 +1,229 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Language-neutral structured fixtures for Action Engine tests.""" + +from __future__ import annotations + +from typing import Any + +from embodichain.gen_sim.action_engine.domain import ( + task_success_type, + validate_scene_requirements, + validate_task_spec, +) +from embodichain.gen_sim.action_engine.protocol import ( + SCENE_REQUIREMENTS_SCHEMA, + TASK_SPEC_SCHEMA, +) + +__all__ = ["make_task_level", "make_task_spec"] + +_OBJECT_FIXTURES = { + "E1": ("can", ["graspable", "placeable"], {}), + "E2": ("can", ["graspable", "orientable"], {"orientation": "fallen"}), + "E3": ("container", ["graspable", "pourable"], {"held_by": "left_arm"}), + "E4": ("cup", ["graspable", "handover"], {}), + "E5": ("tray", ["dual_graspable", "rigid"], {}), + "E6": ("drawer", ["articulated", "pullable"], {"joint_state": "closed"}), + "E7": ("drawer", ["articulated", "pushable"], {"joint_state": "open"}), + "E8": ("knob", ["turnable"], {}), + "E9": ("button", ["pressable"], {"activation": "inactive"}), +} + + +def make_task_spec( + task_type: str = "E1", + *, + task_id: str | None = None, +) -> tuple[dict[str, Any], dict[str, Any]]: + """Build one validated L1 TaskSpec and matching scene requirements.""" + if task_type not in _OBJECT_FIXTURES: + raise ValueError(f"Unsupported fixture task type {task_type!r}.") + category, affordances, initial_state = _OBJECT_FIXTURES[task_type] + object_role = "object_01" + params: dict[str, Any] = {"object_role": object_role} + objects = [ + { + "role_id": object_role, + "category": category, + "count": 1, + "affordances": affordances, + "initial_state": initial_state, + "attributes": {}, + } + ] + if task_type in {"E1", "E3"}: + target_role = "target_01" + params.update({"target_role": target_role, "relation": "inside"}) + if task_type == "E3": + params["source_role"] = params.pop("object_role") + objects.append( + { + "role_id": target_role, + "category": "container", + "count": 1, + "affordances": ["container", "support_surface"], + "initial_state": {}, + "attributes": {}, + } + ) + elif task_type == "E2": + params.update( + { + "orientation_goal": "upright", + "support_role": "table", + "upright_local_axis": "long_axis", + } + ) + elif task_type == "E4": + params.update( + { + "transfer_arm": "left_arm", + "receive_arm": "right_arm", + "orientation_goal": "none", + } + ) + elif task_type == "E5": + params.update({"direction": "up", "terminal_behavior": "hold"}) + elif task_type == "E6": + params["target_state"] = "open" + elif task_type == "E7": + params["target_state"] = "closed" + elif task_type == "E8": + params["target_setting"] = 2 + elif task_type == "E9": + params["target_state"] = "activated" + + effective_id = task_id or f"fixture-{task_type.lower()}" + task = validate_task_spec( + { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": effective_id, + "level": "L1", + "instruction": "test-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": {"fixture": True}, + } + ) + requirements = validate_scene_requirements( + { + "schema_version": SCENE_REQUIREMENTS_SCHEMA, + "task_id": effective_id, + "objects": objects, + "cameras": [], + "spatial_constraints": [ + {"type": "reachable", "roles": "all_interaction_objects"} + ], + "distractor_count": 0, + "metadata": {"fixture": True}, + } + ) + return task, requirements + + +def make_task_level( + level: str, + *, + reasoning: str | None = None, +) -> tuple[dict[str, Any], dict[str, Any]]: + """Build a validated fixture for one public TaskSpec level.""" + if level == "L1": + return make_task_spec("E1") + first, requirements = make_task_spec("E1", task_id=f"fixture-{level.lower()}") + if level == "L2": + second = { + "id": "task_02", + "task_type": "E1", + "params": { + "object_role": "object_02", + "target_role": "target_02", + "relation": "inside", + }, + "depends_on": ["task_01"], + "role": "primary", + } + first["level"] = "L2" + first["task_instances"].append(second) + first["success"] = { + "op": "all", + "terms": [ + {"type": "semantic_goal", "task_instance_id": "task_01"}, + {"type": "semantic_goal", "task_instance_id": "task_02"}, + ], + } + requirements["objects"].extend( + [ + { + "role_id": "object_02", + "category": "can", + "count": 1, + "affordances": ["graspable", "placeable"], + "initial_state": {}, + "attributes": {}, + }, + { + "role_id": "target_02", + "category": "container", + "count": 1, + "affordances": ["container", "support_surface"], + "initial_state": {}, + "attributes": {}, + }, + ] + ) + return validate_task_spec(first), validate_scene_requirements(requirements) + if level == "L4": + first["level"] = "L4" + first["reasoning_type"] = reasoning or "visual_semantics" + first["success"] = { + "visual_semantics": { + "type": "visual_relation", + "relation": "mouth_completed", + }, + "pattern": { + "type": "visual_relation", + "relation": "pattern_completed", + }, + "logic": {"type": "sum_equals", "value": 5}, + "memory": {"type": "original_order_restored"}, + "common_sense": {"type": "functional_place_setting"}, + "constraint": {"type": "stable_unobstructed"}, + }[first["reasoning_type"]] + first["oracle"] = {"fixture": True} + requirements["cameras"] = [ + { + "role": "reasoning_view", + "modalities": ["rgb", "depth"], + "coverage": "all_interaction_objects", + } + ] + return validate_task_spec(first), validate_scene_requirements(requirements) + raise ValueError(f"Unsupported fixture task level {level!r}.") diff --git a/tests/gen_sim/action_engine/tasks/__init__.py b/tests/gen_sim/action_engine/tasks/__init__.py new file mode 100644 index 000000000..d9480994f --- /dev/null +++ b/tests/gen_sim/action_engine/tasks/__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 task generation tests.""" diff --git a/tests/gen_sim/action_engine/tasks/test_factory.py b/tests/gen_sim/action_engine/tasks/test_factory.py new file mode 100644 index 000000000..bfe65cba9 --- /dev/null +++ b/tests/gen_sim/action_engine/tasks/test_factory.py @@ -0,0 +1,486 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Task-first generation, instantiation, and scene hand-off contracts.""" + +from __future__ import annotations + +from copy import deepcopy + +from embodichain.gen_sim.action_engine.tasks import ( + ground_instruction_draft, + instantiate_seed_graph, +) + + +def _selector( + kind: str = "none", + *, + reference: str = "", + step_id: str = "", + quantifier: str = "one", +) -> dict: + return { + "kind": kind, + "step_id": step_id, + "reference": reference, + "quantifier": quantifier, + "count": 0, + } + + +def _intent_step( + step_id: str, + task_type: str, + object_selector: dict, + **updates, +) -> dict: + step = { + "id": step_id, + "task_type": task_type, + "object": object_selector, + "target": _selector(), + "relation": "none", + "required_arm": "auto", + "transfer_arm": "none", + "receive_arm": "none", + "orientation_goal": "upright" if task_type == "E2" else "none", + "target_state": "none", + "target_setting": 0, + "layout": "none", + "axis": "none", + "direction": "none", + "terminal_behavior": "none", + "depends_on": [], + } + step.update(updates) + return step + + +def _ground_draft( + task_id: str, + instruction: str, + scene_objects: list[dict], + steps: list[dict], + bindings: dict[str, list[str]], +): + return ground_instruction_draft( + task_id, + instruction, + {"steps": steps}, + scene_objects, + robot_profile="ur10", + reference_bindings=bindings, + ) + + +def test_orient_then_handover_releases_then_reacquires_with_role_side_pickup() -> None: + task = { + "schema_version": "action_engine_task_spec_v2", + "task_id": "orient_then_handover", + "level": "L3", + "instruction": "test-instruction-orient-handover", + "reasoning_type": "none", + "task_instances": [ + { + "id": "orient", + "task_type": "E2", + "params": { + "object_role": "can", + "required_arm": "left_arm", + }, + "depends_on": [], + "role": "primary", + }, + { + "id": "handover", + "task_type": "E4", + "params": { + "object_role": "can", + "transfer_arm": "right_arm", + "receive_arm": "left_arm", + }, + "depends_on": ["orient"], + "role": "primary", + }, + ], + "success": {"type": "handover_complete"}, + "oracle": {}, + "metadata": {}, + } + + graph = instantiate_seed_graph(task, {"can": "interact_can"}) + orient_nodes = [ + node for node in graph["nodes"] if node["task_instance_id"] == "orient" + ] + handover_nodes = [ + node for node in graph["nodes"] if node["task_instance_id"] == "handover" + ] + orient = next(group for group in graph["task_groups"] if group["id"] == "orient") + + assert [node["atomic_action"] for node in orient_nodes] == [ + "PickUp", + "MoveHeldObject", + "MoveHeldObject", + "Place", + "MoveEndEffector", + "MoveJoints", + ] + assert orient_nodes[0]["motion_policy"] == { + "modifiers": [{"type": "orientation", "mode": "upright"}] + } + assert [node["atomic_action"] for node in handover_nodes] == [ + "PickUp", + "MoveHeldObject", + "HandOver", + "MoveEndEffector", + "MoveJoints", + ] + assert handover_nodes[0]["motion_policy"] == { + "modifiers": [{"type": "handover_role", "mode": "transfer"}] + } + assert handover_nodes[0]["depends_on"] == [orient["node_ids"][-1]] + assert orient["actor"] == {"mode": "required", "arm": "left_arm"} + assert handover_nodes[0]["actor"] == {"mode": "required", "arm": "right_arm"} + assert orient_nodes[-1]["contract"]["completion"] == "terminal_barrier" + assert orient_nodes[-2]["contract"]["failure_policy"] == "safety_required" + assert orient_nodes[-1]["contract"]["failure_policy"] == "best_effort" + assert not any( + effect["atom"]["predicate"] == "arm_home" + for effect in orient["contract"]["exit_effects"] + ) + assert any( + requirement["predicate"] == "object_free" + for requirement in handover_nodes[0]["contract"]["requires"] + ) + + +def test_handover_to_place_uses_receiver_hold_without_repickup() -> None: + task = { + "schema_version": "action_engine_task_spec_v2", + "task_id": "handover_then_place", + "level": "L3", + "instruction": "test-instruction-handover-place", + "reasoning_type": "none", + "task_instances": [ + { + "id": "task_01", + "task_type": "E4", + "params": { + "object_role": "yellow_can", + "transfer_arm": "left_arm", + "receive_arm": "right_arm", + "orientation_goal": "preserve", + }, + "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": {}, + } + + graph = instantiate_seed_graph( + task, + { + "yellow_can": "interact_yellow_can", + "purple_can": "interact_purple_can", + }, + ) + + handover = next(group for group in graph["task_groups"] if group["id"] == "task_01") + placement = next( + group for group in graph["task_groups"] if group["id"] == "task_02" + ) + placement_nodes = [ + node for node in graph["nodes"] if node["task_instance_id"] == "task_02" + ] + + assert [ + node["atomic_action"] + for node in graph["nodes"] + if node["task_instance_id"] == "task_01" + ] == [ + "PickUp", + "MoveHeldObject", + "HandOver", + "MoveEndEffector", + "MoveJoints", + ] + assert handover["actor"] == {"mode": "required", "arm": "left_arm"} + assert graph["nodes"][0]["motion_policy"] == { + "modifiers": [{"type": "handover_role", "mode": "transfer"}] + } + assert graph["nodes"][1]["target_binding"]["kind"] == "handover_staging" + assert graph["nodes"][2]["motion_policy"] == {"modifiers": []} + handover_retreat = graph["nodes"][3] + assert handover_retreat["actor"] == {"mode": "required", "arm": "left_arm"} + assert handover_retreat["target_binding"] == { + "kind": "policy_pose", + "source": "handover", + "operation": "retreat", + } + assert handover_retreat["motion_policy"] == {"modifiers": []} + assert handover_retreat["depends_on"] == [graph["nodes"][2]["id"]] + handover_home = graph["nodes"][4] + assert handover_home["actor"] == {"mode": "required", "arm": "left_arm"} + assert handover_home["target_binding"] == { + "kind": "joint_state", + "source": "initial", + "operation": "handover_home", + } + assert handover_home["motion_policy"] == {"modifiers": []} + assert handover_home["depends_on"] == [handover_retreat["id"]] + assert [node["atomic_action"] for node in placement_nodes] == [ + "MoveHeldObject", + "MoveHeldObject", + "Place", + "MoveEndEffector", + "MoveJoints", + ] + assert placement["actor"] == {"mode": "required", "arm": "right_arm"} + assert placement_nodes[0]["precondition"] == { + "type": "object_held", + "object": "interact_yellow_can", + "arm": "right_arm", + } + assert placement_nodes[0]["depends_on"] == [handover["node_ids"][-1]] + + +def test_structured_draft_grounds_handover_then_receiver_placement() -> None: + scene = [ + { + "runtime_uid": "table", + "role": "background", + "description": "A table.", + "init_pos": [0.0, 0.0, 0.0], + }, + { + "runtime_uid": "interact_yellow_can", + "role": "rigid_object", + "description": "A yellow soda can.", + "init_pos": [0.0, -0.25, 0.75], + }, + { + "runtime_uid": "interact_purple_can", + "role": "rigid_object", + "description": "A purple soda can.", + "init_pos": [0.0, 0.25, 0.75], + }, + ] + + planned = _ground_draft( + "handover_then_place", + "test-instruction-handover-place", + scene, + [ + _intent_step( + "handover", + "E4", + _selector("scene_ref", reference="object-alpha"), + transfer_arm="left_arm", + receive_arm="right_arm", + ), + _intent_step( + "place", + "E1", + _selector("step_result", step_id="handover"), + target=_selector("scene_ref", reference="object-beta"), + relation="right_of", + required_arm="right_arm", + ), + ], + { + "handover.object": ["interact_yellow_can"], + "place.target": ["interact_purple_can"], + }, + ) + graph = instantiate_seed_graph(planned.task_spec, planned.role_bindings) + + assert planned.task_spec["level"] == "L3" + assert [item["task_type"] for item in planned.task_spec["task_instances"]] == [ + "E4", + "E1", + ] + assert planned.role_bindings == { + "object_01": "interact_yellow_can", + "object_02": "interact_purple_can", + } + assert [node["atomic_action"] for node in graph["nodes"]] == [ + "PickUp", + "MoveHeldObject", + "HandOver", + "MoveEndEffector", + "MoveJoints", + "MoveHeldObject", + "MoveHeldObject", + "Place", + "MoveEndEffector", + "MoveJoints", + ] + assert graph["task_groups"][1]["goal"]["relation"] == "right_of" + assert graph["task_groups"][1]["goal"]["relation_frame"] == "robot" + + +def test_seed_graph_adds_missing_same_object_e2_handover_dependency() -> None: + scene = [ + { + "runtime_uid": "purple_can", + "role": "rigid_object", + "description": "A purple soda can.", + "init_pos": [0.0, -0.25, 0.7], + }, + { + "runtime_uid": "orange_can", + "role": "rigid_object", + "description": "An orange soda can.", + "init_pos": [0.0, 0.2, 0.7], + }, + ] + planned = _ground_draft( + "missing_same_object_edge", + "test-instruction-multi-step", + scene, + [ + _intent_step( + "orient_purple", + "E2", + _selector("scene_ref", reference="object-alpha"), + required_arm="right_arm", + ), + _intent_step( + "orient_orange", + "E2", + _selector("scene_ref", reference="object-beta"), + required_arm="left_arm", + depends_on=["orient_purple"], + ), + _intent_step( + "handover_purple", + "E4", + _selector("step_result", step_id="orient_purple"), + transfer_arm="right_arm", + receive_arm="left_arm", + depends_on=["orient_orange"], + ), + _intent_step( + "place_purple", + "E1", + _selector("step_result", step_id="handover_purple"), + target=_selector("scene_ref", reference="object-beta"), + relation="left_of", + required_arm="left_arm", + ), + ], + { + "orient_purple.object": ["purple_can"], + "orient_orange.object": ["orange_can"], + "place_purple.target": ["orange_can"], + }, + ) + underconstrained = deepcopy(planned.task_spec) + underconstrained["task_instances"][2]["depends_on"] = ["task_02"] + + graph = instantiate_seed_graph(underconstrained, planned.role_bindings) + handover = next(group for group in graph["task_groups"] if group["id"] == "task_03") + purple = next(group for group in graph["task_groups"] if group["id"] == "task_01") + staging = next( + node + for node in graph["nodes"] + if node["task_instance_id"] == "task_03" + and node["atomic_action"] == "MoveHeldObject" + ) + assert handover["depends_on"] == ["task_02", "task_01"] + assert [ + node["atomic_action"] + for node in graph["nodes"] + if node["task_instance_id"] == "task_03" + ] == ["PickUp", "MoveHeldObject", "HandOver", "MoveEndEffector", "MoveJoints"] + pickup = next( + node + for node in graph["nodes"] + if node["task_instance_id"] == "task_03" and node["atomic_action"] == "PickUp" + ) + assert purple["node_ids"][-1] in pickup["depends_on"] + assert staging["depends_on"] == [pickup["id"]] + + +def test_structured_draft_treats_table_as_support_in_generic_line_task() -> None: + scene = [ + { + "runtime_uid": "table", + "role": "background", + "description": "A table.", + "init_pos": [0.0, 0.0, 0.0], + }, + { + "runtime_uid": "interact_red_can", + "role": "rigid_object", + "description": "A red soda can.", + "init_pos": [0.0, -0.25, 0.75], + }, + { + "runtime_uid": "interact_blue_cup", + "role": "rigid_object", + "description": "A blue cup.", + "init_pos": [0.0, 0.25, 0.75], + }, + ] + + planned = _ground_draft( + "arrange_line", + "test-instruction-line", + scene, + [ + _intent_step( + "line", + "E1", + _selector( + "scene_ref", + reference="object-set", + quantifier="all", + ), + layout="line", + ) + ], + {"line.object": ["interact_red_can", "interact_blue_cup"]}, + ) + graph = instantiate_seed_graph(planned.task_spec, planned.role_bindings) + + assert planned.task_spec["level"] == "L2" + assert set(planned.role_bindings.values()) == { + "interact_red_can", + "interact_blue_cup", + } + assert all(group["operator"] == "arrange_line" for group in graph["task_groups"]) diff --git a/tests/gen_sim/action_engine/tasks/test_grounding.py b/tests/gen_sim/action_engine/tasks/test_grounding.py new file mode 100644 index 000000000..e927c5fb6 --- /dev/null +++ b/tests/gen_sim/action_engine/tasks/test_grounding.py @@ -0,0 +1,339 @@ +# ---------------------------------------------------------------------------- +# 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 + +import pytest + +from embodichain.gen_sim.action_engine.tasks.grounding import ( + ground_scene_references, +) +from embodichain.gen_sim.action_engine.tasks.assembly import SceneInventory + + +def _selector( + reference: str, + *, + quantifier: str = "one", + count: int = 0, +) -> dict: + return { + "kind": "scene_ref", + "step_id": "", + "reference": reference, + "quantifier": quantifier, + "count": count, + } + + +def _scene() -> list[dict]: + return [ + { + "runtime_uid": "table", + "uid": "table", + "role": "background", + "category": "dining_table", + "name": "work table", + "description": "A rectangular work table.", + "init_pos": [0.0, 0.0, 0.0], + }, + { + "runtime_uid": "cutting_board", + "uid": "cutting_board", + "role": "rigid_object", + "category": "cutting_board", + "name": "wood board", + "description": "A large rectangular wooden cutting board.", + "attributes": { + "size": "large", + "geometry": {"position": [0.0, 0.2, 0.7], "note": "flat"}, + }, + "initial_state": {"orientation": "fallen"}, + "init_pos": [0.0, 0.2, 0.7], + }, + { + "runtime_uid": "salt_shaker", + "uid": "salt_shaker", + "role": "rigid_object", + "category": "salt_shaker", + "description": "A small glass salt shaker.", + "affordances": ["graspable"], + "init_pos": [0.0, -0.2, 0.7], + }, + ] + + +def _intent( + *, + object_selector: dict | None = None, + target_selector: dict | None = None, +) -> dict: + return { + "steps": [ + { + "id": "move", + "task_type": "E1", + "object": object_selector or _selector("object-alpha"), + "target": target_selector or _selector("target-alpha"), + "relation": "on", + } + ] + } + + +def _binding( + reference_id: str, + uids: list[str], + *, + status: str = "resolved", + confidence: float = 1.0, + **extra: object, +) -> dict: + return { + "reference_id": reference_id, + "status": status, + "uids": uids, + "confidence": confidence, + **extra, + } + + +def _run(intent: dict, caller) -> object: + scene = _scene() + return ground_scene_references( + instruction="test-instruction", + intent=intent, + inventory=SceneInventory(scene, robot_profile="franka"), + scene_objects=scene, + model="test-model", + caller=caller, + ) + + +def test_grounding_prompt_preserves_open_semantics_and_redacts_geometry() -> None: + captured: dict[str, object] = {} + + def caller(**kwargs): + captured.update(kwargs) + return { + "bindings": [ + _binding("move.object", ["cutting_board"]), + _binding("move.target", ["table"]), + ] + } + + result = _run(_intent(), caller) + + prompt = str(captured["prompt"]) + assert result.bindings == { + "move.object": ("cutting_board",), + "move.target": ("table",), + } + assert '"category": "cutting_board"' in prompt + assert '"category": "salt_shaker"' in prompt + assert '"name": "wood board"' in prompt + assert '"orientation": "fallen"' in prompt + assert '"size": "large"' in prompt + prompt_inventory = json.loads(prompt.split("Redacted scene inventory:\n", 1)[1]) + side_by_uid = {item["uid"]: item["side"] for item in prompt_inventory} + assert side_by_uid["cutting_board"] == "right" + assert side_by_uid["salt_shaker"] == "left" + assert '"position"' not in prompt + assert '"init_pos"' not in prompt + + +@pytest.mark.parametrize("robot_profile", ["ur5", "ur10", "franka"]) +def test_scene_inventory_uses_the_shared_final_world_lateral_axis( + robot_profile: str, +) -> None: + inventory = SceneInventory(_scene(), robot_profile=robot_profile) + + assert inventory.left_score(inventory.by_uid["salt_shaker"]) > 0.0 + assert inventory.left_score(inventory.by_uid["cutting_board"]) < 0.0 + + +def test_grounding_repairs_one_invalid_uid_in_the_same_batch() -> None: + responses = [ + { + "bindings": [ + _binding("move.object", ["invented"]), + _binding("move.target", ["table"]), + ] + }, + { + "bindings": [ + _binding("move.object", ["cutting_board"]), + _binding("move.target", ["table"]), + ] + }, + ] + prompts: list[str] = [] + + def caller(**kwargs): + prompts.append(kwargs["prompt"]) + return deepcopy(responses[len(prompts) - 1]) + + result = _run(_intent(), caller) + + assert result.attempts == 2 + assert "previous grounding JSON failed" in prompts[1] + assert result.bindings["move.object"] == ("cutting_board",) + + +@pytest.mark.parametrize( + "response,error", + [ + ( + { + "bindings": [ + _binding( + "move.object", + ["cutting_board"], + status="ambiguous", + ), + _binding("move.target", ["table"]), + ] + }, + "was not resolved", + ), + ( + { + "bindings": [ + _binding( + "move.object", + [], + status="not_found", + confidence=0.0, + ), + _binding("move.target", ["table"]), + ] + }, + "was not resolved", + ), + ( + { + "bindings": [ + _binding("move.object", ["cutting_board"], confidence=0.49), + _binding("move.target", ["table"]), + ] + }, + "confidence is below", + ), + ( + { + "bindings": [ + _binding("move.object", ["cutting_board", "cutting_board"]), + _binding("move.target", ["table"]), + ] + }, + "duplicate UIDs", + ), + ( + { + "bindings": [ + _binding("move.object", ["cutting_board"]), + _binding("move.object", ["salt_shaker"]), + _binding("move.target", ["table"]), + ] + }, + "Duplicate grounding binding", + ), + ( + { + "bindings": [ + _binding("move.object", ["cutting_board", "salt_shaker"]), + _binding("move.target", ["table"]), + ] + }, + "quantifier=one requires exactly one UID", + ), + ( + {"bindings": [_binding("move.object", ["cutting_board"])]}, + "omitted requests", + ), + ( + { + "bindings": [ + _binding("move.object", ["table"]), + _binding("move.target", ["cutting_board"]), + ] + }, + "candidate range", + ), + ( + { + "bindings": [ + _binding("move.object", ["cutting_board"]), + _binding("move.target", ["cutting_board"]), + ] + }, + "same UID", + ), + ( + { + "bindings": [ + _binding( + "move.object", + ["cutting_board"], + affordances=["graspable"], + ), + _binding("move.target", ["table"]), + ] + }, + "unsupported", + ), + ], +) +def test_grounding_fails_closed_after_one_repair(response: dict, error: str) -> None: + with pytest.raises(ValueError, match=f"after one repair.*{error}"): + _run(_intent(), lambda **_kwargs: deepcopy(response)) + + +def test_grounding_enforces_count_and_accepts_an_open_world_set() -> None: + intent = _intent( + object_selector=_selector("object-set", quantifier="count", count=2) + ) + response = { + "bindings": [ + _binding("move.object", ["cutting_board", "salt_shaker"]), + _binding("move.target", ["table"]), + ] + } + + result = _run(intent, lambda **_kwargs: response) + assert result.bindings["move.object"] == ("cutting_board", "salt_shaker") + + invalid = deepcopy(response) + invalid["bindings"][0]["uids"] = ["cutting_board"] + with pytest.raises(ValueError, match="requires exactly 2 UIDs"): + _run(intent, lambda **_kwargs: invalid) + + +def test_grounding_accepts_a_nonempty_all_binding() -> None: + intent = _intent(object_selector=_selector("object-set", quantifier="all")) + response = { + "bindings": [ + _binding("move.object", ["cutting_board", "salt_shaker"]), + _binding("move.target", ["table"]), + ] + } + + result = _run(intent, lambda **_kwargs: response) + + assert result.bindings["move.object"] == ("cutting_board", "salt_shaker") diff --git a/tests/gen_sim/action_engine/tasks/test_interpretation.py b/tests/gen_sim/action_engine/tasks/test_interpretation.py new file mode 100644 index 000000000..a9462b2f9 --- /dev/null +++ b/tests/gen_sim/action_engine/tasks/test_interpretation.py @@ -0,0 +1,1885 @@ +# ---------------------------------------------------------------------------- +# 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 + +import embodichain.gen_sim.action_engine.tasks.interpretation as interpretation_module +import embodichain.gen_sim.task_engine.interpretation as task_interpretation_module +from embodichain.gen_sim.action_engine.tasks.assembly import SceneInventory +from embodichain.gen_sim.action_engine.tasks import ( + INSTRUCTION_INTENT_SCHEMA, + instantiate_seed_graph, + interpret_and_ground_task_spec, + validate_instruction_intent, +) + + +def _selector(kind: str = "none", **values): + legacy_kind = kind + if kind == "selector": + kind = "scene_ref" + reference = values.pop("reference", "") + if legacy_kind == "selector": + uid = str(values.pop("uid", "")).strip() + legacy_terms = [ + str(values.pop(field, "")).strip() + for field in ("side", "color", "category") + ] + reference = reference or uid + if not reference: + reference = " ".join( + term for term in legacy_terms if term not in {"", "none"} + ) + result = { + "kind": kind, + "step_id": "", + "reference": reference, + "quantifier": "one", + "count": 0, + } + result.update(values) + return result + + +def _grounding(**bindings): + return { + "bindings": [ + { + "reference_id": reference_id, + "status": "resolved", + "uids": [uid] if isinstance(uid, str) else list(uid), + "confidence": 1.0, + } + for reference_id, uid in bindings.items() + ] + } + + +def _grounding_caller(**bindings): + response = _grounding(**bindings) + return lambda **_kwargs: deepcopy(response) + + +def _step(step_id: str, task_type: str, object_selector: dict, **values): + result = { + "id": step_id, + "task_type": task_type, + "object": object_selector, + "target": _selector(), + "relation": "none", + "required_arm": "auto", + "transfer_arm": "none", + "receive_arm": "none", + "orientation_goal": "upright" if task_type == "E2" else "none", + "target_state": "none", + "target_setting": 0, + "layout": "none", + "axis": "none", + "direction": "none", + "terminal_behavior": "none", + "depends_on": [], + } + result.update(values) + return result + + +def _scene(): + return [ + { + "runtime_uid": "purple_can", + "uid": "purple_can", + "role": "rigid_object", + "description": "A purple soda can.", + "init_pos": [0.0, -0.25, 0.7], + }, + { + "runtime_uid": "orange_can", + "uid": "orange_can", + "role": "rigid_object", + "description": "An orange soda can.", + "init_pos": [0.0, 0.2, 0.7], + }, + ] + + +def _scene_with_table(): + return [ + { + "runtime_uid": "table", + "uid": "table", + "role": "background", + "description": "A table.", + "init_pos": [0.0, 0.0, 0.0], + }, + *_scene(), + ] + + +def _scene_export_style_scene(): + return [ + { + "runtime_uid": "table", + "uid": "table", + "role": "background", + "description": "A light grey dining table.", + "init_pos": [0.0, 0.0, 0.0], + }, + { + "runtime_uid": "carrot_001", + "uid": "carrot_001", + "role": "rigid_object", + "category": "carrot", + "description": ( + "A single orange carrot with a green top located at the top left " + "of the table." + ), + "init_pos": [0.28, 0.47, 1.06], + }, + { + "runtime_uid": "cutting_board_001", + "uid": "cutting_board_001", + "role": "rigid_object", + "category": "cutting_board", + "description": ( + "A rectangular cutting board located in the upper middle-left " + "area of the table." + ), + "init_pos": [0.14, 0.21, 1.07], + }, + { + "runtime_uid": "peeler_001", + "uid": "peeler_001", + "role": "rigid_object", + "category": "vegetable_peeler", + "description": "A black-handled vegetable peeler.", + "init_pos": [-0.13, -0.61, 1.07], + }, + ] + + +def _payload_scene(): + return [ + { + "runtime_uid": "glue_stick", + "uid": "glue_stick", + "role": "object", + "category": "glue_stick", + "description": "A solid glue stick.", + "init_pos": [0.0, -0.2, 0.7], + }, + { + "runtime_uid": "paper_cup", + "uid": "paper_cup", + "role": "object", + "category": "cup", + "description": "A paper cup.", + "init_pos": [0.0, 0.0, 0.7], + }, + { + "runtime_uid": "popcorn_bucket", + "uid": "popcorn_bucket", + "role": "object", + "category": "bucket", + "description": "A popcorn bucket.", + "init_pos": [0.0, 0.25, 0.7], + }, + ] + + +def _handover_intent(): + return { + "steps": [ + _step( + "orient", + "E2", + _selector("scene_ref", reference="object-alpha"), + required_arm="right_arm", + ), + _step( + "handover", + "E4", + _selector("step_result", step_id="orient"), + transfer_arm="right_arm", + receive_arm="left_arm", + depends_on=["orient"], + ), + _step( + "place", + "E1", + _selector("step_result", step_id="handover"), + target=_selector("scene_ref", reference="object-beta"), + relation="left_of", + required_arm="left_arm", + depends_on=["handover"], + ), + ] + } + + +def _two_object_handover_intent_with_missing_place_target(): + return { + "steps": [ + _step( + "orient_purple", + "E2", + _selector("scene_ref", reference="object-alpha"), + required_arm="right_arm", + ), + _step( + "orient_orange", + "E2", + _selector("scene_ref", reference="object-beta"), + required_arm="left_arm", + depends_on=["orient_purple"], + ), + _step( + "handover", + "E4", + _selector("step_result", step_id="orient_purple"), + transfer_arm="right_arm", + receive_arm="left_arm", + depends_on=["orient_orange"], + ), + _step( + "place", + "E1", + _selector("step_result", step_id="handover"), + relation="left_of", + required_arm="left_arm", + depends_on=["handover"], + ), + ] + } + + +def _two_object_handover_intent(): + intent = _two_object_handover_intent_with_missing_place_target() + intent["steps"][3]["target"] = _selector( + "scene_ref", + reference="object-beta", + ) + return intent + + +def test_llm_intent_handles_handover_pronoun_and_elliptical_place() -> None: + calls = [] + + def caller(**kwargs): + calls.append(kwargs) + return _handover_intent() + + grounded = interpret_and_ground_task_spec( + "handover_task", + "instruction-marker", + _scene(), + robot_profile="ur10", + model="test-model", + caller=caller, + grounding_caller=_grounding_caller( + **{ + "orient.object": "purple_can", + "place.target": "orange_can", + } + ), + ) + graph = instantiate_seed_graph(grounded.task_spec, grounded.role_bindings) + + assert [item["task_type"] for item in grounded.task_spec["task_instances"]] == [ + "E2", + "E4", + "E1", + ] + assert grounded.role_bindings == { + "object_01": "purple_can", + "object_02": "orange_can", + } + placement_actions = [ + node["atomic_action"] for node in graph["nodes"] if node["task_type"] == "E1" + ] + orient_actions = [ + node["atomic_action"] for node in graph["nodes"] if node["task_type"] == "E2" + ] + handover_nodes = [node for node in graph["nodes"] if node["task_type"] == "E4"] + assert orient_actions == [ + "PickUp", + "MoveHeldObject", + "MoveHeldObject", + "Place", + "MoveEndEffector", + "MoveJoints", + ] + assert [node["atomic_action"] for node in handover_nodes] == [ + "PickUp", + "MoveHeldObject", + "HandOver", + "MoveEndEffector", + "MoveJoints", + ] + assert handover_nodes[0]["motion_policy"] == { + "modifiers": [{"type": "handover_role", "mode": "transfer"}] + } + assert any( + requirement["predicate"] == "object_free" + for requirement in handover_nodes[0]["contract"]["requires"] + ) + assert placement_actions == [ + "MoveHeldObject", + "MoveHeldObject", + "Place", + "MoveEndEffector", + "MoveJoints", + ] + out_of_order = deepcopy(grounded.task_spec) + out_of_order["task_instances"] = list(reversed(out_of_order["task_instances"])) + reordered_graph = instantiate_seed_graph( + out_of_order, + grounded.role_bindings, + ) + assert [group["task_type"] for group in reordered_graph["task_groups"]] == [ + "E2", + "E4", + "E1", + ] + assert "instruction-marker" in calls[0]["prompt"] + assert calls[0]["model"] == "test-model" + + +def test_six_step_repeated_objects_preserve_two_handover_continuations() -> None: + intent = { + "steps": [ + _step( + "orient_purple", + "E2", + _selector("scene_ref", reference="object-alpha"), + required_arm="right_arm", + ), + _step( + "orient_orange", + "E2", + _selector("scene_ref", reference="object-beta"), + required_arm="left_arm", + ), + _step( + "handover_orange", + "E4", + _selector("step_result", step_id="orient_orange"), + transfer_arm="left_arm", + receive_arm="right_arm", + depends_on=["orient_orange"], + ), + _step( + "place_orange", + "E1", + _selector("step_result", step_id="handover_orange"), + target=_selector("scene_ref", reference="target-gamma"), + relation="on", + required_arm="right_arm", + depends_on=["handover_orange"], + ), + _step( + "handover_purple", + "E4", + _selector("step_result", step_id="orient_purple"), + transfer_arm="right_arm", + receive_arm="left_arm", + # The model may preserve only the object-lineage dependency. + # Stable lowering must not let this step leapfrog an earlier + # placement that releases its transfer arm. + depends_on=["orient_purple"], + ), + _step( + "place_purple", + "E1", + _selector("step_result", step_id="handover_purple"), + target=_selector("scene_ref", reference="object-beta"), + relation="on", + required_arm="left_arm", + depends_on=["handover_purple"], + ), + ] + } + scene = [ + *_scene(), + { + "runtime_uid": "notebook", + "uid": "notebook", + "role": "rigid_object", + "description": "A spiral notebook.", + "init_pos": [0.2, 0.0, 0.7], + }, + ] + + grounded = interpret_and_ground_task_spec( + "two_handover_task", + "test-instruction-multi-step", + scene, + robot_profile="franka", + model="test-model", + caller=lambda **_kwargs: deepcopy(intent), + grounding_caller=_grounding_caller( + **{ + "orient_purple.object": "purple_can", + "orient_orange.object": "orange_can", + "place_orange.target": "notebook", + "place_purple.target": "orange_can", + } + ), + ) + assert [ + instance["task_type"] for instance in grounded.task_spec["task_instances"] + ] == ["E2", "E2", "E4", "E1", "E4", "E1"] + graph = instantiate_seed_graph(grounded.task_spec, grounded.role_bindings) + groups = {group["id"]: group for group in graph["task_groups"]} + actions_by_group = { + group_id: [ + node["atomic_action"] + for node in graph["nodes"] + if node["task_instance_id"] == group_id + ] + for group_id in groups + } + + assert actions_by_group["task_04"][0] == "MoveHeldObject" + assert "PickUp" not in actions_by_group["task_04"] + assert actions_by_group["task_05"][0] == "PickUp" + assert actions_by_group["task_06"][0] == "MoveHeldObject" + assert "PickUp" not in actions_by_group["task_06"] + assert groups["task_04"]["contract"]["entry_requires"] == [ + { + "predicate": "object_held", + "object_uid": "orange_can", + "arm": "right_arm", + } + ] + assert groups["task_06"]["contract"]["entry_requires"] == [ + { + "predicate": "object_held", + "object_uid": "purple_can", + "arm": "left_arm", + } + ] + + +def test_e5_relative_transport_emits_pickment_and_optional_release() -> None: + scene = [ + { + "runtime_uid": "table", + "uid": "table", + "role": "background", + "category": "table", + "description": "table", + "init_pos": [0.0, 0.0, 0.0], + }, + { + "runtime_uid": "plastic_tray", + "uid": "plastic_tray", + "role": "object", + "category": "tray", + "description": "plastic tray", + "init_pos": [0.0, 0.0, 0.7], + }, + { + "runtime_uid": "banana_left", + "uid": "banana_left", + "role": "object", + "category": "banana", + "description": "left banana", + "init_pos": [0.0, 0.25, 0.7], + }, + ] + intent = { + "steps": [ + _step( + "move_tray", + "E5", + _selector("selector", uid="plastic_tray"), + target=_selector("selector", uid="banana_left"), + relation="behind", + direction="none", + terminal_behavior="hold", + ) + ] + } + grounded = interpret_and_ground_task_spec( + "dual_tray", + "test-instruction-relative-transport", + scene, + robot_profile="franka", + model="test-model", + caller=lambda **_kwargs: intent, + grounding_caller=_grounding_caller( + **{ + "move_tray.object": "plastic_tray", + "move_tray.target": "banana_left", + } + ), + ) + + graph = instantiate_seed_graph(grounded.task_spec, grounded.role_bindings) + assert [node["atomic_action"] for node in graph["nodes"]] == ["CoordinatedPickment"] + assert graph["task_groups"][0]["operator"] == "coordinated_transport" + assert graph["task_groups"][0]["goal"] == { + "direction": "none", + "terminal_behavior": "hold", + "orientation_goal": "none", + "orientation_axis": "none", + "relation_frame": "robot", + "reference_object": "banana_left", + "reference_state": "live", + "relation": "behind", + } + + released_spec = deepcopy(grounded.task_spec) + released_spec["task_instances"][0]["params"]["terminal_behavior"] = "place" + released = instantiate_seed_graph(released_spec, grounded.role_bindings) + assert [node["atomic_action"] for node in released["nodes"]] == [ + "CoordinatedPickment", + "MoveJoints", + "MoveJoints", + ] + release_nodes = released["nodes"][1:] + assert all( + node["depends_on"] == [released["nodes"][0]["id"]] for node in release_nodes + ) + assert {node["actor"]["arm"] for node in release_nodes} == { + "left_arm", + "right_arm", + } + assert {node["control"] for node in release_nodes} == {"hand"} + assert len({node["sync_group"] for node in release_nodes}) == 1 + assert all(node["precondition"] == {} for node in release_nodes) + assert { + node["target_binding"]["coordinated_release_role"] for node in release_nodes + } == {"participant", "commit"} + contracts = { + node["target_binding"]["coordinated_release_role"]: node["contract"] + for node in release_nodes + } + coordinated_hold = { + "predicate": "object_coordinated_held", + "object_uid": "plastic_tray", + } + assert contracts["participant"]["requires"] == [coordinated_hold] + assert contracts["participant"]["effects"] == [] + assert contracts["commit"]["requires"] == [coordinated_hold] + assert { + ( + effect["op"], + effect["atom"]["predicate"], + effect["atom"].get("arm"), + ) + for effect in contracts["commit"]["effects"] + } == { + ("delete", "object_coordinated_held", None), + ("add", "object_free", None), + ("add", "arm_free", "left_arm"), + ("add", "arm_free", "right_arm"), + } + from embodichain.gen_sim.action_engine.runtime import load_execution_program + + program = load_execution_program(released) + assert [ + action["atomic_action_class"] + for edge in program.edges + for action in edge.actions + ] == ["CoordinatedPickment", "MoveJoints", "MoveJoints"] + assert len(program.edges[-1].actions) == 2 + + in_place_spec = deepcopy(released_spec) + in_place_params = in_place_spec["task_instances"][0]["params"] + in_place_params.pop("target_role") + in_place_params.update({"direction": "none", "relation": "none"}) + in_place = instantiate_seed_graph(in_place_spec, grounded.role_bindings) + assert [node["atomic_action"] for node in in_place["nodes"]] == [ + "CoordinatedPickment", + "MoveJoints", + "MoveJoints", + ] + assert "reference_object" not in in_place["task_groups"][0]["goal"] + + +def test_e5_accepts_generic_rigid_object_without_exported_affordances() -> None: + scene = [ + { + "runtime_uid": "interact_wooden_block", + "uid": "interact_wooden_block", + "role": "rigid_object", + "description": "A long rectangular wooden block.", + "init_pos": [0.0, 0.0, 0.7], + } + ] + intent = { + "steps": [ + _step( + "move_block", + "E5", + _selector("selector", uid="interact_wooden_block"), + direction="left", + terminal_behavior="hold", + ) + ] + } + + grounded = interpret_and_ground_task_spec( + "dual_block", + "test-instruction-directional-transport", + scene, + robot_profile="franka", + model="test-model", + caller=lambda **_kwargs: intent, + grounding_caller=_grounding_caller( + **{"move_block.object": "interact_wooden_block"} + ), + ) + + instance = grounded.task_spec["task_instances"][0] + assert instance["task_type"] == "E5" + assert grounded.role_bindings[instance["params"]["object_role"]] == ( + "interact_wooden_block" + ) + assert instance["params"]["direction"] == "left" + + +def test_task1_2_open_reference_generates_coordinated_pick_move_and_release() -> None: + scene = [ + { + "runtime_uid": "table", + "uid": "table", + "role": "background", + "description": "A white table.", + "init_pos": [0.0, 0.0, 0.0], + }, + { + "runtime_uid": "interact_apple", + "uid": "interact_apple", + "role": "rigid_object", + "description": "A red apple.", + "init_pos": [0.0, -0.2, 0.7], + }, + { + "runtime_uid": "interact_wooden_tray", + "uid": "interact_wooden_tray", + "role": "rigid_object", + "description": "A long rectangular wooden tray.", + "init_pos": [0.0, 0.0, 0.7], + }, + { + "runtime_uid": "interact_rubiks_cube", + "uid": "interact_rubiks_cube", + "role": "rigid_object", + "description": "A Rubik's cube.", + "init_pos": [0.0, 0.2, 0.7], + }, + ] + intent = { + "steps": [ + _step( + "move_block", + "E5", + _selector("scene_ref", reference="object-alpha"), + direction="left", + terminal_behavior="place", + ) + ] + } + + grounded = interpret_and_ground_task_spec( + "task1_2", + "test-instruction-directional-place", + scene, + robot_profile="franka", + model="test-model", + caller=lambda **_kwargs: intent, + grounding_caller=_grounding_caller( + **{"move_block.object": "interact_wooden_tray"} + ), + ) + graph = instantiate_seed_graph(grounded.task_spec, grounded.role_bindings) + + instance = grounded.task_spec["task_instances"][0] + assert instance["params"]["direction"] == "left" + assert instance["params"]["terminal_behavior"] == "place" + assert grounded.task_spec["success"]["terms"] == [ + {"type": "semantic_goal", "task_instance_id": instance["id"]} + ] + assert [node["atomic_action"] for node in graph["nodes"]] == [ + "CoordinatedPickment", + "MoveJoints", + "MoveJoints", + ] + assert grounded.scene_requirements["objects"][0]["category"] == "rigid_object" + assert grounded.task_spec["metadata"]["instruction_call_count"] == 1 + assert grounded.task_spec["metadata"]["scene_grounding_call_count"] == 1 + + +def test_e5_pick_and_hold_defaults_missing_direction_to_up() -> None: + scene = [ + { + "runtime_uid": "table", + "uid": "table", + "role": "background", + "description": "A wooden table.", + "init_pos": [0.0, 0.0, 0.0], + }, + { + "runtime_uid": "wooden_tray", + "uid": "wooden_tray", + "role": "rigid_object", + "description": "A shallow round wooden serving tray.", + "init_pos": [0.0, 0.0, 0.7], + }, + ] + intent = { + "steps": [ + _step( + "lift_tray", + "E5", + _selector("scene_ref", reference="object-alpha"), + required_arm="none", + direction="none", + terminal_behavior="hold", + ) + ] + } + + grounded = interpret_and_ground_task_spec( + "lift_tray", + "test-instruction-hold", + scene, + robot_profile="franka", + model="test-model", + caller=lambda **_kwargs: deepcopy(intent), + grounding_caller=_grounding_caller(**{"lift_tray.object": "wooden_tray"}), + ) + graph = instantiate_seed_graph(grounded.task_spec, grounded.role_bindings) + + instance = grounded.task_spec["task_instances"][0] + assert instance["params"]["direction"] == "up" + assert instance["params"]["terminal_behavior"] == "hold" + assert [node["atomic_action"] for node in graph["nodes"]] == ["CoordinatedPickment"] + assert grounded.task_spec["success"]["terms"] == [ + {"type": "held_by_both_grippers", "task_instance_id": instance["id"]} + ] + assert grounded.task_spec["metadata"]["instruction_call_count"] == 1 + assert grounded.task_spec["metadata"]["instruction_intent_normalizations"] == [ + { + "path": "steps[0].direction", + "from": "none", + "to": "up", + "reason": "e5_hold_defaults_to_lift", + } + ] + + +@pytest.mark.parametrize( + ("scene_update", "error"), + ( + ({"affordances": ["rigid"]}, "missing affordances.*dual_graspable"), + ({"role": "articulation"}, "requires .*rigid.object structure"), + ), +) +def test_e5_rejects_explicitly_incompatible_scene_evidence( + scene_update: dict, + error: str, +) -> None: + scene_object = { + "runtime_uid": "candidate", + "uid": "candidate", + "role": "rigid_object", + "description": "A candidate object.", + "init_pos": [0.0, 0.0, 0.7], + **scene_update, + } + intent = { + "steps": [ + _step( + "move_candidate", + "E5", + _selector("selector", uid="candidate"), + direction="left", + terminal_behavior="hold", + ) + ] + } + + with pytest.raises(ValueError, match=error): + interpret_and_ground_task_spec( + "invalid_dual_object", + "test-instruction-missing-object", + [scene_object], + robot_profile="franka", + model="test-model", + caller=lambda **_kwargs: intent, + grounding_caller=_grounding_caller( + **{"move_candidate.object": "candidate"} + ), + ) + + +@pytest.mark.parametrize( + ("scene_update", "should_succeed", "error"), + ( + ({"role": "articulation"}, True, ""), + ( + {"role": "articulation", "affordances": ["articulated"]}, + False, + "missing affordances.*pullable", + ), + ({"role": "rigid_object"}, False, "requires articulation structure"), + ), +) +def test_articulated_task_uses_structural_and_explicit_affordance_evidence( + scene_update: dict, + should_succeed: bool, + error: str, +) -> None: + scene_object = { + "runtime_uid": "cabinet_part", + "uid": "cabinet_part", + "description": "A cabinet moving part.", + "init_pos": [0.0, 0.0, 0.7], + **scene_update, + } + intent = { + "steps": [ + _step( + "open_part", + "E6", + _selector("scene_ref", reference="object-alpha"), + target_state="open", + ) + ] + } + + invoke = lambda: interpret_and_ground_task_spec( + "open_part", + "test-instruction-articulation", + [scene_object], + robot_profile="franka", + model="test-model", + caller=lambda **_kwargs: intent, + grounding_caller=_grounding_caller(**{"open_part.object": "cabinet_part"}), + ) + if should_succeed: + assert invoke().task_spec["task_instances"][0]["task_type"] == "E6" + else: + with pytest.raises(ValueError, match=error): + invoke() + + +def test_open_container_target_is_allowed_until_runtime_when_metadata_is_unknown() -> ( + None +): + scene = [ + { + "runtime_uid": "source_pitcher", + "uid": "source_pitcher", + "role": "rigid_object", + "category": "ceramic_pitcher", + "description": "A ceramic pitcher with water.", + "init_pos": [0.0, -0.2, 0.7], + }, + { + "runtime_uid": "custom_receiver", + "uid": "custom_receiver", + "role": "rigid_object", + "category": "handmade_vessel", + "description": "A handmade receiving vessel.", + "init_pos": [0.0, 0.2, 0.7], + }, + ] + intent = { + "steps": [ + _step( + "pour", + "E3", + _selector("scene_ref", reference="object-alpha"), + target=_selector("scene_ref", reference="target-alpha"), + relation="above", + ) + ] + } + + grounded = interpret_and_ground_task_spec( + "open_container", + "test-instruction-pour", + scene, + robot_profile="franka", + model="test-model", + caller=lambda **_kwargs: intent, + grounding_caller=_grounding_caller( + **{ + "pour.object": "source_pitcher", + "pour.target": "custom_receiver", + } + ), + ) + assert grounded.task_spec["task_instances"][0]["task_type"] == "E3" + + explicit = deepcopy(scene) + explicit[1]["affordances"] = ["support_surface"] + with pytest.raises(ValueError, match="none support containment"): + interpret_and_ground_task_spec( + "explicit_non_container", + "test-instruction-pour", + explicit, + robot_profile="franka", + model="test-model", + caller=lambda **_kwargs: intent, + grounding_caller=_grounding_caller( + **{ + "pour.object": "source_pitcher", + "pour.target": "custom_receiver", + } + ), + ) + + +def test_open_scene_reference_is_not_limited_by_fixed_selector_fields() -> None: + intent = { + "steps": [ + _step( + "orient", + "E2", + _selector("scene_ref", reference="object-alpha"), + ) + ] + } + grounded = interpret_and_ground_task_spec( + "open_reference", + "test-instruction-orient", + _scene(), + robot_profile="ur10", + caller=lambda **_kwargs: intent, + grounding_caller=_grounding_caller(**{"orient.object": "purple_can"}), + ) + assert grounded.role_bindings == {"object_01": "purple_can"} + + +def test_intent_rejects_atomic_actions_coordinates_and_extra_fields() -> None: + intent = _handover_intent() + intent["steps"][0]["atomic_action"] = "PickUp" + with pytest.raises(ValueError, match="forbidden fields"): + validate_instruction_intent(intent) + + intent = _handover_intent() + intent["steps"][0]["object"]["target_pose"] = [0.0, 0.0, 0.0] + with pytest.raises(ValueError, match="forbidden fields"): + validate_instruction_intent(intent) + + +def test_invalid_intent_gets_one_repair_attempt() -> None: + responses = [{"steps": []}, _handover_intent()] + prompts = [] + + def caller(**kwargs): + prompts.append(kwargs["prompt"]) + return deepcopy(responses[len(prompts) - 1]) + + grounded = interpret_and_ground_task_spec( + "repair", + "test-instruction-repair", + _scene(), + robot_profile="ur10", + caller=caller, + grounding_caller=_grounding_caller( + **{ + "orient.object": "purple_can", + "place.target": "orange_can", + } + ), + ) + assert len(prompts) == 2 + assert "previous JSON was invalid" in prompts[1] + assert grounded.task_spec["metadata"]["instruction_call_count"] == 2 + + +def test_interpreter_normalizes_registry_inapplicable_e4_required_arm() -> None: + intent = _handover_intent() + intent["steps"][1]["required_arm"] = "right_arm" + + with pytest.raises(ValueError, match="uses transfer_arm/receive_arm"): + validate_instruction_intent(intent) + + grounded = interpret_and_ground_task_spec( + "normalized_handover", + "test-instruction-handover-repair", + _scene(), + robot_profile="ur10", + caller=lambda **_kwargs: deepcopy(intent), + grounding_caller=_grounding_caller( + **{ + "orient.object": "purple_can", + "place.target": "orange_can", + } + ), + ) + + assert grounded.task_spec["metadata"]["instruction_call_count"] == 1 + assert grounded.task_spec["metadata"]["instruction_intent_normalizations"] == [ + { + "path": "steps[1].required_arm", + "from": "right_arm", + "to": "none", + "reason": "inapplicable_for_E4", + } + ] + + +def test_interpreter_resolves_same_arm_handover_from_step_result_ownership() -> None: + intent = { + "steps": [ + _step( + "orient_coke", + "E2", + _selector("scene_ref", reference="object-alpha"), + required_arm="right_arm", + ), + _step( + "orient_sprite", + "E2", + _selector("scene_ref", reference="object-beta"), + required_arm="left_arm", + ), + _step( + "handover_sprite", + "E4", + _selector("step_result", step_id="orient_sprite"), + required_arm="none", + transfer_arm="left_arm", + receive_arm="left_arm", + depends_on=["orient_sprite"], + ), + _step( + "place_sprite", + "E1", + _selector("step_result", step_id="handover_sprite"), + target=_selector("step_result", step_id="orient_coke"), + relation="on", + required_arm="right_arm", + depends_on=["orient_coke", "handover_sprite"], + ), + ] + } + + with pytest.raises(ValueError, match="transfer and receive arms must differ"): + validate_instruction_intent(intent) + + result = task_interpretation_module.interpret_instruction_draft( + "test-instruction-invalid-handover", + model="test-model", + caller=lambda **_kwargs: deepcopy(intent), + ) + + handover = result.intent["steps"][2] + assert (handover["transfer_arm"], handover["receive_arm"]) == ( + "left_arm", + "right_arm", + ) + assert result.attempts == 1 + assert result.normalizations == ( + { + "path": "steps[2].receive_arm", + "from": "left_arm", + "to": "right_arm", + "reason": "handover_arm_continuity", + }, + ) + + +def test_interpreter_repairs_direct_reference_handover_from_later_arm_semantics() -> ( + None +): + invalid_intent = { + "steps": [ + _step( + "orient_sprite", + "E2", + _selector("scene_ref", reference="object-beta"), + required_arm="left_arm", + ), + _step( + "handover_sprite", + "E4", + _selector("scene_ref", reference="object-beta"), + required_arm="none", + transfer_arm="left_arm", + receive_arm="left_arm", + depends_on=["orient_sprite"], + ), + _step( + "place_sprite", + "E1", + _selector("scene_ref", reference="object-beta"), + target=_selector("scene_ref", reference="object-alpha"), + relation="on", + required_arm="right_arm", + depends_on=["handover_sprite"], + ), + ] + } + repaired_intent = deepcopy(invalid_intent) + repaired_intent["steps"][1]["receive_arm"] = "right_arm" + prompts: list[str] = [] + + def caller(**kwargs): + prompts.append(kwargs["prompt"]) + return deepcopy(invalid_intent if len(prompts) == 1 else repaired_intent) + + result = task_interpretation_module.interpret_instruction_draft( + "test-instruction-same-arm-handover", + model="test-model", + caller=caller, + ) + + handover = result.intent["steps"][1] + assert (handover["transfer_arm"], handover["receive_arm"]) == ( + "left_arm", + "right_arm", + ) + assert result.attempts == 2 + assert result.normalizations == () + assert "Same-arm handover repair rule" in prompts[1] + + +def test_interpreter_does_not_merge_repeated_scene_reference_identity() -> None: + intent = { + "steps": [ + _step( + "orient_first_can", + "E2", + _selector("scene_ref", reference="object-token"), + required_arm="left_arm", + ), + _step( + "handover_second_can", + "E4", + _selector("scene_ref", reference="object-token"), + required_arm="none", + transfer_arm="left_arm", + receive_arm="left_arm", + depends_on=["orient_first_can"], + ), + _step( + "place_first_can", + "E1", + _selector("scene_ref", reference="object-token"), + target=_selector("scene_ref", reference="target-alpha"), + relation="on", + required_arm="right_arm", + depends_on=["handover_second_can"], + ), + ] + } + calls = 0 + + def caller(**_kwargs): + nonlocal calls + calls += 1 + return deepcopy(intent) + + with pytest.raises(ValueError, match="after one repair.*arms must differ"): + task_interpretation_module.interpret_instruction_draft( + "test-instruction-repeated-reference", + model="test-model", + caller=caller, + ) + + assert calls == 2 + + +def test_interpreter_does_not_guess_an_unconstrained_same_arm_handover() -> None: + intent = { + "steps": [ + _step( + "handover", + "E4", + _selector("scene_ref", reference="object-beta"), + transfer_arm="left_arm", + receive_arm="left_arm", + ) + ] + } + calls = 0 + + def caller(**_kwargs): + nonlocal calls + calls += 1 + return deepcopy(intent) + + with pytest.raises(ValueError, match="after one repair.*arms must differ"): + task_interpretation_module.interpret_instruction_draft( + "test-instruction-invalid-same-arm", + model="test-model", + caller=caller, + ) + + assert calls == 2 + + +def test_invalid_step_result_gets_repair_with_selector_rules() -> None: + """A malformed cross-step selector should reach the structured repair call.""" + invalid_intent = _handover_intent() + invalid_intent["steps"][1]["object"]["reference"] = "object-alpha" + responses = [invalid_intent, _handover_intent()] + prompts: list[str] = [] + + def caller(**kwargs): + prompts.append(kwargs["prompt"]) + return deepcopy(responses[len(prompts) - 1]) + + grounded = interpret_and_ground_task_spec( + "repair_step_result", + "test-instruction-step-result-repair", + _scene(), + robot_profile="ur10", + caller=caller, + grounding_caller=_grounding_caller( + **{ + "orient.object": "purple_can", + "place.target": "orange_can", + } + ), + ) + + assert len(prompts) == 2 + repair_prompt = prompts[1] + for term in ("step_result", "step_id", "reference"): + assert term in repair_prompt + assert "none" in repair_prompt + assert grounded.task_spec["metadata"]["instruction_call_count"] == 2 + + +def test_repeated_missing_e1_target_fails_without_local_guessing() -> None: + invalid_intent = _two_object_handover_intent_with_missing_place_target() + prompts: list[str] = [] + + def caller(**kwargs): + prompts.append(kwargs["prompt"]) + return deepcopy(invalid_intent) + + with pytest.raises(ValueError, match="after one repair"): + interpret_and_ground_task_spec( + "missing_target", + "test-instruction-missing-target", + _scene(), + robot_profile="ur10", + caller=caller, + ) + + assert len(prompts) == 2 + assert "Missing-target repair rule" in prompts[1] + + +def test_missing_target_completion_rejects_other_semantic_disagreement() -> None: + invalid_intent = _two_object_handover_intent_with_missing_place_target() + invalid_intent["steps"][-1]["required_arm"] = "right_arm" + + with pytest.raises(ValueError, match="after one repair"): + interpret_and_ground_task_spec( + "unsafe_target_completion", + "test-instruction-missing-target", + _scene(), + robot_profile="ur10", + caller=lambda **_kwargs: deepcopy(invalid_intent), + ) + + +def test_second_invalid_intent_fails_without_rule_fallback() -> None: + with pytest.raises(ValueError, match="after one repair"): + interpret_and_ground_task_spec( + "invalid", + "test-instruction-invalid", + _scene(), + robot_profile="ur10", + caller=lambda **_kwargs: {"steps": []}, + ) + + +def test_intent_infers_pronoun_dependency_from_canonical_symbols() -> None: + intent = { + "steps": [ + _step( + "handover", + "E4", + _selector("scene_ref", reference="object-alpha"), + transfer_arm="right_arm", + receive_arm="left_arm", + ), + _step( + "place", + "E1", + _selector("step_result", step_id="handover"), + target=_selector("scene_ref", reference="object-beta"), + relation="left_of", + required_arm="left_arm", + depends_on=["handover"], + ), + ] + } + + grounded = interpret_and_ground_task_spec( + "implicit_dependency", + "test-instruction-pronoun-dependency", + _scene(), + robot_profile="ur10", + caller=lambda **_kwargs: intent, + grounding_caller=_grounding_caller( + **{ + "handover.object": "purple_can", + "place.target": "orange_can", + } + ), + ) + + instances = grounded.task_spec["task_instances"] + assert [item["task_type"] for item in instances] == ["E4", "E1"] + assert instances[1]["depends_on"] == [instances[0]["id"]] + assert instances[1]["params"]["relation"] == "left_of" + assert instances[1]["params"]["required_arm"] == "left_arm" + + +def test_scene_grounding_rejects_unknown_uid() -> None: + intent = { + "steps": [ + _step( + "orient", + "E2", + _selector("scene_ref", reference="object-alpha"), + ) + ] + } + with pytest.raises(ValueError, match="after one repair.*unknown UIDs"): + interpret_and_ground_task_spec( + "unknown_uid", + "test-instruction-unknown-uid", + _scene(), + robot_profile="ur10", + caller=lambda **_kwargs: intent, + grounding_caller=_grounding_caller(**{"orient.object": "invented_uid"}), + ) + + +def test_instruction_intent_rejects_legacy_selector_protocol() -> None: + intent = _handover_intent() + intent["steps"][0]["object"] = { + "kind": "selector", + "step_id": "", + "uid": "purple_can", + "category": "can", + "color": "purple", + "side": "none", + "quantifier": "one", + "count": 0, + } + with pytest.raises(ValueError, match="requires exactly fields"): + validate_instruction_intent(intent) + + +def test_step_result_must_reference_a_preceding_step() -> None: + intent = { + "steps": [ + _step( + "handover", + "E4", + _selector("step_result", step_id="orient"), + transfer_arm="right_arm", + receive_arm="left_arm", + ), + _step( + "orient", + "E2", + _selector("scene_ref", reference="object-alpha"), + ), + ] + } + with pytest.raises(ValueError, match="preceding step"): + validate_instruction_intent(intent) + + +def test_step_result_selector_rejects_object_constraints() -> None: + intent = _handover_intent() + intent["steps"][1]["object"]["reference"] = "object-alpha" + + with pytest.raises(ValueError, match="may identify only a prior step_id"): + validate_instruction_intent(intent) + + +def test_instruction_intent_rejects_non_e_specific_parameters() -> None: + invalid_e9 = _step( + "press", + "E9", + _selector("selector", category="button"), + target_state="activated", + orientation_goal="upright", + ) + with pytest.raises(ValueError, match="orientation_goal"): + validate_instruction_intent({"steps": [invalid_e9]}) + + invalid_line = _step( + "line", + "E1", + _selector("selector", category="can", quantifier="all"), + layout="line", + relation="on", + ) + with pytest.raises(ValueError, match="line arrangement cannot carry a relation"): + validate_instruction_intent({"steps": [invalid_line]}) + + +def test_implicit_e1_relation_requires_an_unambiguous_support_target() -> None: + intent = { + "steps": [ + _step( + "place", + "E1", + _selector("selector", category="can", color="purple"), + target=_selector("selector", category="can", color="orange"), + ) + ] + } + + with pytest.raises(ValueError, match="omitted relation"): + interpret_and_ground_task_spec( + "ambiguous_implicit_place", + "test-instruction-implicit-relation", + _scene(), + robot_profile="ur10", + caller=lambda **_kwargs: intent, + grounding_caller=_grounding_caller( + **{ + "place.object": "purple_can", + "place.target": "orange_can", + } + ), + ) + + +def test_instruction_and_grounding_prompts_keep_their_boundaries() -> None: + captured: dict[str, dict] = {} + intent = { + "steps": [ + _step( + "place", + "E1", + _selector("selector", category="can", color="purple"), + target=_selector("selector", uid="table", category="table"), + relation="on", + ) + ] + } + + def caller(**kwargs): + captured["intent"] = kwargs + return intent + + def grounding_caller(**kwargs): + captured["grounding"] = kwargs + return _grounding(**{"place.object": "purple_can", "place.target": "table"}) + + grounded = interpret_and_ground_task_spec( + "onto_table", + "Put the purple can on the table.", + _scene_with_table(), + robot_profile="ur10", + caller=caller, + grounding_caller=grounding_caller, + ) + + assert '"uid": "table"' not in captured["intent"]["prompt"] + assert '"uid": "table"' in captured["grounding"]["prompt"] + assert '"core_actions"' not in captured["intent"]["prompt"] + assert captured["intent"]["schema"] == INSTRUCTION_INTENT_SCHEMA + assert grounded.role_bindings["object_02"] == "table" + + +def test_instruction_intent_schema_declares_every_required_selector_field() -> None: + selector_schema = INSTRUCTION_INTENT_SCHEMA["properties"]["steps"]["items"][ + "properties" + ]["object"] + + assert set(selector_schema["required"]) == set(selector_schema["properties"]) + assert "quantifier" in selector_schema["properties"] + + +def test_grounding_prompt_redacts_nested_scene_geometry() -> None: + scene = _scene() + scene[0]["attributes"] = { + "label": "purple", + "geometry": {"position": [0.0, 0.0, 0.7], "note": "can"}, + } + captured: dict[str, str] = {} + + def grounding_caller(**kwargs): + captured["prompt"] = kwargs["prompt"] + return _grounding(**{"orient.object": "purple_can"}) + + interpret_and_ground_task_spec( + "redacted_inventory", + "test-instruction-grounding-redaction", + scene, + robot_profile="ur10", + caller=lambda **_kwargs: { + "steps": [ + _step( + "orient", + "E2", + _selector("scene_ref", reference="object-alpha"), + ) + ] + }, + grounding_caller=grounding_caller, + ) + assert '"position"' not in captured["prompt"] + assert '"label": "purple"' in captured["prompt"] + + +def test_default_llm_parser_requires_the_documented_model_configuration( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("ACTION_ENGINE_LLM_MODEL", raising=False) + monkeypatch.delenv("OPENAI_MODEL", raising=False) + monkeypatch.setattr(task_interpretation_module, "_load_local_env", lambda: {}) + + with pytest.raises(ValueError, match="text LLM model is required"): + interpret_and_ground_task_spec( + "missing_model", + "test-instruction-model-config", + _scene(), + robot_profile="ur10", + ) + + +def test_injected_caller_skips_production_model_resolution( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def unexpected_model_resolution(_explicit: str | None) -> str | None: + raise AssertionError( + "injected callers must not resolve production model config" + ) + + monkeypatch.setattr( + task_interpretation_module, + "_instruction_model", + unexpected_model_resolution, + ) + + grounded = interpret_and_ground_task_spec( + "injected_caller", + "test-instruction-injected-caller", + _scene(), + robot_profile="ur10", + caller=lambda **_kwargs: { + "steps": [ + _step( + "orient", + "E2", + _selector("scene_ref", reference="object-alpha"), + ) + ] + }, + grounding_caller=_grounding_caller(**{"orient.object": "purple_can"}), + ) + + assert grounded.task_spec["metadata"]["instruction_model"] == "injected_caller" + + +def test_mimo_instruction_caller_uses_json_mode_and_disables_thinking( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """MiMo-compatible endpoints must not use the lossy JSON-schema route.""" + import langchain_openai + + calls: list[dict] = [] + responses = [ + { + "steps": [ + { + "id": "orient", + "task_type": "E2", + "object": _selector("scene_ref", reference="object-alpha"), + } + ] + }, + _handover_intent(), + _grounding( + **{ + "orient.object": "purple_can", + "place.target": "orange_can", + } + ), + ] + + class FakeRunnable: + def invoke(self, messages): + calls[-1]["messages"] = messages + return deepcopy(responses.pop(0)) + + class FakeChatOpenAI: + def __init__(self, **kwargs): + calls.append({"kwargs": kwargs}) + + def with_structured_output(self, schema, **kwargs): + calls[-1]["schema"] = schema + calls[-1]["structured_kwargs"] = kwargs + return FakeRunnable() + + monkeypatch.setattr(langchain_openai, "ChatOpenAI", FakeChatOpenAI) + monkeypatch.setattr( + task_interpretation_module, + "_load_llm_settings", + lambda *, model: { + "api_key": "test-key", + "model": model or "mimo-v2.5", + "base_url": "https://token-plan-cn.xiaomimimo.com/v1", + "default_query": {}, + }, + ) + + grounded = interpret_and_ground_task_spec( + "mimo_repair", + "test-instruction-json-mode", + _scene(), + robot_profile="ur10", + model="mimo-v2.5", + ) + + assert [item["task_type"] for item in grounded.task_spec["task_instances"]] == [ + "E2", + "E4", + "E1", + ] + assert len(calls) == 3 + for call in calls: + assert call["structured_kwargs"] == {"method": "json_mode"} + assert call["kwargs"]["http_socket_options"] == () + assert call["kwargs"]["max_completion_tokens"] == 4096 + assert call["kwargs"]["extra_body"] == {"thinking": {"type": "disabled"}} + repair_messages = calls[1]["messages"] + assert "previous JSON was invalid" in repair_messages[1].content + + +def test_instruction_prompt_contains_a_complete_shape_example() -> None: + prompt = interpretation_module._instruction_prompt("instruction-marker") + selector_rules = interpretation_module._instruction_selector_rules() + assert '"target_setting": 0' in prompt + assert '"depends_on": []' in prompt + assert "every step has all 16 step keys" in prompt + assert "step_result" in prompt + assert "open scene_ref.reference" in prompt + assert "Do not classify it or emit a scene UID" in prompt + assert "example object A" in prompt + assert "stale-object-reference" not in prompt + assert "step_result" in selector_rules + assert "step_id" in selector_rules + assert "reference" in selector_rules + + +def test_scene_export_spatial_descriptions_do_not_create_false_supports() -> None: + index = SceneInventory(_scene_export_style_scene(), robot_profile="franka") + + assert [entity.uid for entity in index.support] == ["table"] + assert {entity.uid for entity in index.movable} == { + "carrot_001", + "cutting_board_001", + "peeler_001", + } + + +def test_scene_export_exact_uids_ground_pick_and_place() -> None: + intent = { + "steps": [ + _step( + "step_1", + "E1", + _selector("selector", uid="carrot_001"), + target=_selector("selector", uid="cutting_board_001"), + relation="on", + required_arm="left_arm", + ) + ] + } + + grounded = interpret_and_ground_task_spec( + "scene_export_pick_place", + "test-instruction-exact-uids", + _scene_export_style_scene(), + robot_profile="franka", + model="test-model", + caller=lambda **_kwargs: intent, + grounding_caller=_grounding_caller( + **{ + "step_1.object": "carrot_001", + "step_1.target": "cutting_board_001", + } + ), + ) + + assert set(grounded.role_bindings.values()) == { + "carrot_001", + "cutting_board_001", + } + assert grounded.task_spec["task_instances"][0]["params"]["required_arm"] == ( + "left_arm" + ) + assert {item["category"] for item in grounded.scene_requirements["objects"]} == { + "carrot", + "cutting_board", + } + + +def test_multi_object_handover_keeps_both_order_and_holder_dependencies() -> None: + intent = _two_object_handover_intent() + grounded = interpret_and_ground_task_spec( + "multi_object_handover", + "test-instruction-multi-object-handover", + _scene(), + robot_profile="ur10", + caller=lambda **_kwargs: deepcopy(intent), + grounding_caller=_grounding_caller( + **{ + "orient_purple.object": "purple_can", + "orient_orange.object": "orange_can", + "place.target": "orange_can", + } + ), + ) + + instances = grounded.task_spec["task_instances"] + assert instances[2]["depends_on"] == ["task_02", "task_01"] + graph = instantiate_seed_graph(grounded.task_spec, grounded.role_bindings) + handover = next(group for group in graph["task_groups"] if group["id"] == "task_03") + handover_nodes = [ + node for node in graph["nodes"] if node["task_instance_id"] == "task_03" + ] + assert handover["depends_on"] == ["task_02", "task_01"] + assert [node["atomic_action"] for node in handover_nodes] == [ + "PickUp", + "MoveHeldObject", + "HandOver", + "MoveEndEffector", + "MoveJoints", + ] + purple = next(group for group in graph["task_groups"] if group["id"] == "task_01") + orange = next(group for group in graph["task_groups"] if group["id"] == "task_02") + assert handover_nodes[0]["depends_on"] == [ + orange["node_ids"][-1], + purple["node_ids"][-1], + ] + + +def test_single_arm_e1_propagates_direct_payload_into_goal_and_contracts() -> None: + intent = { + "steps": [ + _step( + "handover_glue", + "E4", + _selector("selector", uid="glue_stick"), + required_arm="left_arm", + transfer_arm="left_arm", + receive_arm="right_arm", + ), + _step( + "place_glue", + "E1", + _selector("step_result", step_id="handover_glue"), + target=_selector("selector", uid="paper_cup"), + relation="on", + required_arm="right_arm", + depends_on=["handover_glue"], + ), + _step( + "place_cup", + "E1", + _selector("selector", uid="paper_cup"), + target=_selector("selector", uid="popcorn_bucket"), + relation="on", + required_arm="right_arm", + depends_on=["place_glue"], + ), + ] + } + grounded = interpret_and_ground_task_spec( + "payload_chain", + "test-instruction-payload-propagation", + _payload_scene(), + robot_profile="ur10", + caller=lambda **_kwargs: deepcopy(intent), + grounding_caller=_grounding_caller( + **{ + "handover_glue.object": "glue_stick", + "place_glue.target": "paper_cup", + "place_cup.object": "paper_cup", + "place_cup.target": "popcorn_bucket", + } + ), + ) + + graph = instantiate_seed_graph(grounded.task_spec, grounded.role_bindings) + carrier_group = next( + group for group in graph["task_groups"] if group["id"] == "task_03" + ) + assert carrier_group["goal"]["payloads"] == [ + {"object": "glue_stick", "slot": "center"} + ] + carrier_nodes = [ + node + for node in graph["nodes"] + if node["task_instance_id"] == carrier_group["id"] + and node["atomic_action"] in {"PickUp", "MoveHeldObject", "Place"} + ] + assert carrier_nodes + for node in carrier_nodes: + assert node["target_binding"]["payloads"] == carrier_group["goal"]["payloads"] + assert any( + claim["resource"] == "object:glue_stick" and claim["access"] == "exclusive" + for claim in node["contract"]["claims"] + ) + + +def test_seed_graph_repairs_missing_e2_handover_lifecycle_edge() -> None: + intent = _two_object_handover_intent() + grounded = interpret_and_ground_task_spec( + "missing_lifecycle_edge", + "test-instruction-lifecycle-repair", + _scene(), + robot_profile="ur10", + caller=lambda **_kwargs: deepcopy(intent), + grounding_caller=_grounding_caller( + **{ + "orient_purple.object": "purple_can", + "orient_orange.object": "orange_can", + "place.target": "orange_can", + } + ), + ) + underconstrained = deepcopy(grounded.task_spec) + underconstrained["task_instances"][2]["depends_on"] = ["task_02"] + + graph = instantiate_seed_graph(underconstrained, grounded.role_bindings) + handover = next(group for group in graph["task_groups"] if group["id"] == "task_03") + purple = next(group for group in graph["task_groups"] if group["id"] == "task_01") + assert handover["depends_on"] == ["task_02", "task_01"] + pickup = next( + node + for node in graph["nodes"] + if node["task_instance_id"] == "task_03" and node["atomic_action"] == "PickUp" + ) + assert purple["node_ids"][-1] in pickup["depends_on"] diff --git a/tests/gen_sim/action_engine/tasks/test_language_decoupling.py b/tests/gen_sim/action_engine/tasks/test_language_decoupling.py new file mode 100644 index 000000000..b9f0768d3 --- /dev/null +++ b/tests/gen_sim/action_engine/tasks/test_language_decoupling.py @@ -0,0 +1,420 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Acceptance tests for the structured-LLM language boundary.""" + +from __future__ import annotations + +from copy import deepcopy +from pathlib import Path + +import pytest + +import embodichain.gen_sim.action_engine.tasks as action_engine_tasks +from embodichain.gen_sim.action_engine.tasks import ( + instantiate_seed_graph, + interpret_and_ground_task_spec, + validate_instruction_intent, +) +from embodichain.gen_sim.action_engine.tasks.assembly import SceneInventory + + +def _selector( + kind: str = "none", + *, + reference: str = "", + step_id: str = "", +) -> dict: + return { + "kind": kind, + "step_id": step_id, + "reference": reference, + "quantifier": "one", + "count": 0, + } + + +def _step(step_id: str, task_type: str, reference: str, **updates: object) -> dict: + step = { + "id": step_id, + "task_type": task_type, + "object": _selector("scene_ref", reference=reference), + "target": _selector(), + "relation": "none", + "required_arm": "none", + "transfer_arm": "none", + "receive_arm": "none", + "orientation_goal": "upright" if task_type == "E2" else "preserve", + "target_state": "none", + "target_setting": 0, + "layout": "none", + "axis": "none", + "direction": "none", + "terminal_behavior": "none", + "depends_on": [], + } + step.update(updates) + return step + + +def _binding(reference_id: str, *uids: str) -> dict: + return { + "reference_id": reference_id, + "status": "resolved", + "uids": list(uids), + "confidence": 1.0, + } + + +def _grounding_caller(*bindings: dict): + response = {"bindings": list(bindings)} + return lambda **_kwargs: deepcopy(response) + + +def _open_scene() -> list[dict]: + return [ + { + "runtime_uid": "work_surface", + "uid": "work_surface", + "role": "support_surface", + "category": "obsidian_dock", + "name": "the landing ledge", + "description": "A flat black ledge used as a work surface.", + "init_pos": [0.0, 0.0, 0.0], + }, + { + "runtime_uid": "aerogel_fixture_7", + "uid": "aerogel_fixture_7", + "role": "rigid_object", + "category": "aerogel_fixture", + "name": "translucent fixture", + "description": "A translucent rectangular fixture with a frosted edge.", + "init_pos": [0.0, 0.1, 0.7], + }, + { + "runtime_uid": "plantain_marker", + "uid": "plantain_marker", + "role": "rigid_object", + "category": "plantain_marker", + "description": "A curved yellow marker behind the fixture.", + "init_pos": [0.1, -0.2, 0.7], + }, + ] + + +def test_scene_inventory_preserves_open_category_labels() -> None: + scene = _open_scene() + scene[1]["category"] = "Prototype.Fixture/V2" + inventory = SceneInventory(scene, robot_profile="franka") + + assert inventory.by_uid["aerogel_fixture_7"].category == ("Prototype.Fixture/V2") + + +@pytest.mark.parametrize( + ("step", "invalid_field"), + [ + ( + _step( + "place", + "E1", + "object-alpha", + target=_selector("scene_ref", reference="target-alpha"), + relation="invalid-relation", + ), + "relation", + ), + ( + _step("orient", "E2", "object-alpha", required_arm="invalid-arm"), + "required_arm", + ), + ( + _step( + "orient", + "E2", + "object-alpha", + orientation_goal="invalid-orientation", + ), + "orientation_goal", + ), + ], +) +def test_llm_intent_rejects_natural_language_aliases( + step: dict, + invalid_field: str, +) -> None: + """Canonical protocol fields are not a second local language parser.""" + with pytest.raises(ValueError, match=invalid_field): + validate_instruction_intent({"steps": [step]}) + + +def test_noncanonical_llm_value_is_repaired_instead_of_locally_normalized() -> None: + invalid = { + "steps": [ + _step( + "place", + "E1", + "object-alpha", + target=_selector("scene_ref", reference="target-alpha"), + relation="invalid-relation", + ) + ] + } + valid = deepcopy(invalid) + valid["steps"][0]["relation"] = "left_of" + responses = [invalid, valid] + prompts: list[str] = [] + + def caller(**kwargs): + prompts.append(kwargs["prompt"]) + return deepcopy(responses[len(prompts) - 1]) + + grounded = interpret_and_ground_task_spec( + "strict_canonical_repair", + "test-instruction-invalid-relation", + _open_scene(), + robot_profile="franka", + model="test-model", + caller=caller, + grounding_caller=_grounding_caller( + _binding("place.object", "aerogel_fixture_7"), + _binding("place.target", "work_surface"), + ), + ) + + assert len(prompts) == 2 + assert "previous JSON was invalid" in prompts[1] + assert grounded.task_spec["metadata"]["instruction_call_count"] == 2 + assert grounded.task_spec["task_instances"][0]["params"]["relation"] == ("left_of") + assert "instruction_intent_normalizations" not in grounded.task_spec["metadata"] + + +def test_two_noncanonical_llm_responses_fail_without_grounding_or_rule_fallback() -> ( + None +): + invalid = { + "steps": [ + _step( + "place", + "E1", + "object-alpha", + target=_selector("scene_ref", reference="target-alpha"), + relation="invalid-relation", + ) + ] + } + grounding_called = False + + def unexpected_grounding(**_kwargs): + nonlocal grounding_called + grounding_called = True + raise AssertionError("invalid canonical intent must not reach grounding") + + with pytest.raises(ValueError, match="after one repair.*relation"): + interpret_and_ground_task_spec( + "strict_canonical_failure", + "test-instruction-invalid-relation", + _open_scene(), + robot_profile="franka", + model="test-model", + caller=lambda **_kwargs: deepcopy(invalid), + grounding_caller=unexpected_grounding, + ) + + assert grounding_called is False + + +def test_legacy_instruction_parser_modules_and_api_are_absent() -> None: + tasks_dir = Path(action_engine_tasks.__file__).resolve().parent + + assert not (tasks_dir / "deterministic.py").exists() + assert not (tasks_dir / "planning.py").exists() + assert not hasattr(action_engine_tasks, "plan_grounded_task_spec") + + +def test_production_sources_do_not_reference_legacy_instruction_parser() -> None: + action_engine_dir = Path(action_engine_tasks.__file__).resolve().parent.parent + forbidden = ( + "tasks.deterministic", + "tasks.planning", + "plan_grounded_task_spec", + "instruction_parser", + "deterministic_fallback", + ) + offenders: dict[str, list[str]] = {} + for path in action_engine_dir.rglob("*.py"): + source = path.read_text(encoding="utf-8") + matches = [term for term in forbidden if term in source] + if matches: + offenders[str(path.relative_to(action_engine_dir))] = matches + + assert offenders == {} + + +def test_llm_caller_exception_propagates_without_scene_grounding() -> None: + expected = RuntimeError("model unavailable") + grounding_called = False + + def fail_model(**_kwargs): + raise expected + + def unexpected_grounding(**_kwargs): + nonlocal grounding_called + grounding_called = True + raise AssertionError("failed interpretation must not reach grounding") + + with pytest.raises(RuntimeError) as caught: + interpret_and_ground_task_spec( + "model_failure", + "test-instruction-caller-error", + _open_scene(), + robot_profile="franka", + model="test-model", + caller=fail_model, + grounding_caller=unexpected_grounding, + ) + + assert caught.value is expected + assert grounding_called is False + + +def test_unfamiliar_wording_and_categories_flow_through_injected_llm_stages() -> None: + intent = { + "steps": [ + _step( + "relocate_fixture", + "E1", + "object-alpha", + target=_selector("scene_ref", reference="target-alpha"), + relation="on", + required_arm="auto", + ) + ] + } + + grounded = interpret_and_ground_task_spec( + "open_world_fixture", + "test-instruction-open-reference", + _open_scene(), + robot_profile="franka", + model="test-model", + caller=lambda **_kwargs: deepcopy(intent), + grounding_caller=_grounding_caller( + _binding("relocate_fixture.object", "aerogel_fixture_7"), + _binding("relocate_fixture.target", "work_surface"), + ), + ) + + assert set(grounded.role_bindings.values()) == { + "aerogel_fixture_7", + "work_surface", + } + assert {item["category"] for item in grounded.scene_requirements["objects"]} == { + "aerogel_fixture", + "obsidian_dock", + } + assert grounded.task_spec["metadata"]["instruction_call_count"] == 1 + assert grounded.task_spec["metadata"]["scene_grounding_call_count"] == 1 + + +@pytest.mark.parametrize( + ("name", "instruction", "step", "bindings", "actions", "success"), + [ + ( + "dual_lift", + "test-instruction-hold", + _step( + "lift_fixture", + "E5", + "object-alpha", + terminal_behavior="hold", + ), + [_binding("lift_fixture.object", "aerogel_fixture_7")], + ["CoordinatedPickment"], + "held_by_both_grippers", + ), + ( + "dual_move_place", + "test-instruction-directional-place", + _step( + "move_fixture", + "E5", + "object-alpha", + direction="left", + terminal_behavior="place", + ), + [_binding("move_fixture.object", "aerogel_fixture_7")], + ["CoordinatedPickment", "MoveJoints", "MoveJoints"], + "semantic_goal", + ), + ( + "dual_relative", + "test-instruction-relative-place", + _step( + "move_relative", + "E5", + "object-alpha", + target=_selector("scene_ref", reference="target-alpha"), + relation="behind", + terminal_behavior="hold", + ), + [ + _binding("move_relative.object", "aerogel_fixture_7"), + _binding("move_relative.target", "plantain_marker"), + ], + ["CoordinatedPickment"], + "held_by_both_grippers", + ), + ], +) +def test_e5_symbolic_intent_reaches_the_seed_graph( + name: str, + instruction: str, + step: dict, + bindings: list[dict], + actions: list[str], + success: str, +) -> None: + grounded = interpret_and_ground_task_spec( + name, + instruction, + _open_scene(), + robot_profile="franka", + model="test-model", + caller=lambda **_kwargs: {"steps": [deepcopy(step)]}, + grounding_caller=_grounding_caller(*bindings), + ) + instance = grounded.task_spec["task_instances"][0] + graph = instantiate_seed_graph(grounded.task_spec, grounded.role_bindings) + + assert [node["atomic_action"] for node in graph["nodes"]] == actions + assert grounded.task_spec["success"]["terms"] == [ + {"type": success, "task_instance_id": instance["id"]} + ] + assert instance["params"].get("direction") == ( + "up" if name == "dual_lift" else step["direction"] + ) + if name == "dual_relative": + assert graph["task_groups"][0]["goal"]["reference_object"] == ( + "plantain_marker" + ) + assert graph["task_groups"][0]["goal"]["relation"] == "behind" + if name == "dual_move_place": + release_nodes = graph["nodes"][1:] + assert {node["actor"]["arm"] for node in release_nodes} == { + "left_arm", + "right_arm", + } + assert len({node["sync_group"] for node in release_nodes}) == 1 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_graph_visualization.py b/tests/gen_sim/action_engine/test_graph_visualization.py new file mode 100644 index 000000000..efb1474ac --- /dev/null +++ b/tests/gen_sim/action_engine/test_graph_visualization.py @@ -0,0 +1,362 @@ +# ---------------------------------------------------------------------------- +# 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 io import BytesIO + +from PIL import Image, ImageStat +import pytest + +from embodichain.gen_sim.action_engine.compiler import ( + compile_task_agent, + compile_task_agent_v2, +) +from embodichain.gen_sim.action_engine.domain import ( + EXECUTION_PROGRAM_SCHEMA, + MOTION_POLICY_VERSION, + TASK_AGENT_SCHEMA, + validate_execution_program, +) +from embodichain.gen_sim.action_engine.graph_visualization import ( + _RuntimeOverlay, + _dag_levels, + _dag_positions, + _dependency_pairs, + _graph_data, + render_seed_task_graph_png, + render_task_graph_png, +) + +_PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n" + + +def _image(payload: bytes) -> Image.Image: + assert payload.startswith(_PNG_SIGNATURE) + image = Image.open(BytesIO(payload)).convert("RGB") + extrema = ImageStat.Stat(image).extrema + assert any(low != high for low, high in extrema) + return image + + +def _contains_color( + image: Image.Image, + color: str, + *, + minimum_pixels: int = 8, + tolerance: int = 4, +) -> bool: + target = tuple(bytes.fromhex(color.removeprefix("#"))) + matches = 0 + payload = image.tobytes() + for offset in range(0, len(payload), 3): + pixel = payload[offset : offset + 3] + if all( + abs(channel - expected) <= tolerance + for channel, expected in zip(pixel, target) + ): + matches += 1 + if matches >= minimum_pixels: + return True + return False + + +def _chain_program() -> dict[str, object]: + return compile_task_agent( + { + "schema_version": TASK_AGENT_SCHEMA, + "task": "unicode-λ-task", + "goal": "Pick up the cup and keep it hovering.", + "semantic_steps": [ + { + "id": "s01_hover", + "operator": "hold_hover", + "object": "cup", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": {}, + "depends_on": [], + } + ], + } + ) + + +def _action( + action_class: str, + arm: str | None, + target: str, +) -> dict[str, object]: + actor = {"mode": "auto"} if arm is None else {"mode": "required", "arm": arm} + return { + "atomic_action_class": action_class, + "actor": actor, + "control": "arm", + "target_binding": {"kind": "object", "object": target}, + "motion_policy": {"modifiers": []}, + } + + +def _fork_join_program() -> dict[str, object]: + program = { + "schema_version": EXECUTION_PROGRAM_SCHEMA, + "task": "fork_join_demo", + "goal_description": "Move two objects in parallel, then finish.", + "start": "v_start", + "goal": "v_goal", + "nodes": [ + {"id": "v_start", "semantic": "ready"}, + {"id": "v_left", "semantic": "left branch active"}, + {"id": "v_right", "semantic": "right branch active"}, + {"id": "v_join", "semantic": "branches complete"}, + {"id": "v_goal", "semantic": "task complete"}, + ], + "edges": [ + { + "id": "e_left_pick", + "source": "v_start", + "target": "v_left", + "semantic_step_id": "s_left", + "actions": [_action("PickUp", "left_arm", "left_object")], + "depends_on": [], + "resources": ["arm:left_arm"], + }, + { + "id": "e_right_pick", + "source": "v_start", + "target": "v_right", + "semantic_step_id": "s_right", + "actions": [_action("PickUp", "right_arm", "right_object")], + "depends_on": [], + "resources": ["arm:right_arm"], + }, + { + "id": "e_left_join", + "source": "v_left", + "target": "v_join", + "semantic_step_id": "s_left", + "actions": [_action("MoveHeldObject", "left_arm", "left_object")], + "depends_on": ["e_left_pick"], + "resources": ["arm:left_arm"], + }, + { + "id": "e_right_join", + "source": "v_right", + "target": "v_join", + "semantic_step_id": "s_right", + "actions": [_action("MoveHeldObject", "right_arm", "right_object")], + # Cross-branch dependency not implied by state continuity, so + # the renderer must draw a visible dashed dependency arrow. + "depends_on": ["e_right_pick", "e_left_pick"], + "resources": ["arm:right_arm"], + }, + { + "id": "e_finish", + "source": "v_join", + "target": "v_goal", + "semantic_step_id": "s_finish", + "actions": [_action("MoveJoints", None, "home")], + "depends_on": ["e_left_join", "e_right_join"], + "resources": ["arm:auto"], + }, + ], + "semantic_steps": [ + { + "id": "s_left", + "parent_step_id": "s_left", + "operator": "place_relative", + "object": "left_object", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": {"relation": "on"}, + "depends_on": [], + "postcondition": {"type": "semantic_goal"}, + "edge_ids": ["e_left_pick", "e_left_join"], + }, + { + "id": "s_right", + "parent_step_id": "s_right", + "operator": "place_relative", + "object": "right_object", + "actor": {"mode": "required", "arm": "right_arm"}, + "goal": {"relation": "on"}, + "depends_on": [], + "postcondition": {"type": "semantic_goal"}, + "edge_ids": ["e_right_pick", "e_right_join"], + }, + { + "id": "s_finish", + "parent_step_id": "s_finish", + "operator": "hold_hover", + "object": "home", + "actor": {"mode": "auto"}, + "goal": {}, + "depends_on": ["s_left", "s_right"], + "postcondition": {"type": "semantic_goal"}, + "edge_ids": ["e_finish"], + }, + ], + "allocation_groups": [ + { + "id": "g_parallel", + "semantic_step_ids": ["s_left", "s_right"], + "arm_constraint": "distinct_arms", + "execution_policy": "parallel_if_feasible", + "parallel_action_classes": ["PickUp"], + "workspace_policy": "shared_target_serial", + } + ], + "motion_policy_version": MOTION_POLICY_VERSION, + } + return validate_execution_program(program) + + +def test_seed_renderer_produces_a_compact_headless_png() -> None: + first = _image(render_seed_task_graph_png(_chain_program())) + second = _image(render_seed_task_graph_png(_chain_program())) + + assert first.size == second.size + assert first.width > first.height + assert first.height < 1_200 + + +def test_fork_join_layout_uses_actor_lanes_and_dependency_links() -> None: + program = _fork_join_program() + data = _graph_data(program, _RuntimeOverlay({}, {}, {})) + levels = _dag_levels(data.graph) + positions = _dag_positions( + data, + levels, + {"left": 2.6, "auto": 7.8, "right": 13.0}, + ) + + assert positions["v_left"][0] < 5.15 + assert positions["v_right"][0] > 10.45 + assert positions["v_start"][0] == pytest.approx(7.8) + assert positions["v_join"][0] == pytest.approx(7.8) + assert ("e_left_join", "e_finish") in _dependency_pairs(data) + assert ("e_right_join", "e_finish") in _dependency_pairs(data) + + image = _image(render_seed_task_graph_png(program)) + assert image.width > image.height + assert _contains_color(image, "#168A78") + assert _contains_color(image, "#D97706") + assert _contains_color(image, "#8A94A0") + + +def test_parallel_single_phase_edges_are_rendered_as_a_multigraph() -> None: + program = compile_task_agent( + { + "schema_version": TASK_AGENT_SCHEMA, + "task": "parallel_press", + "goal": "Press both independent buttons.", + "semantic_steps": [ + { + "id": "s_left", + "operator": "press", + "object": "left_button", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": {}, + "depends_on": [], + }, + { + "id": "s_right", + "operator": "press", + "object": "right_button", + "actor": {"mode": "required", "arm": "right_arm"}, + "goal": {}, + "depends_on": [], + }, + ], + } + ) + assert {(edge["source"], edge["target"]) for edge in program["edges"]} == { + ("v0_start", "v_goal") + } + + image = _image(render_seed_task_graph_png(program)) + + assert _contains_color(image, "#168A78") + assert _contains_color(image, "#D97706") + + +def test_runtime_renderer_overlays_observed_statuses() -> None: + program = _fork_join_program() + runtime = { + **program, + "runtime": { + "schema_version": "action_engine_runtime_record_v1", + "status": "failed", + "events": [ + { + "event": "edge", + "edge_id": "e_left_pick", + "arm": "left_arm", + "status": "executed", + }, + { + "event": "edge", + "edge_id": "e_right_pick", + "arm": "right_arm", + "status": "failed", + }, + ], + }, + } + + image = _image(render_task_graph_png(runtime)) + + assert _contains_color(image, "#25834B") + assert _contains_color(image, "#C43E3E") + + +def test_runtime_renderer_accepts_v2_seed_graph_envelope() -> None: + task_agent = { + "schema_version": TASK_AGENT_SCHEMA, + "task": "v2_runtime_overlay", + "goal": "Hold the cup.", + "semantic_steps": [ + { + "id": "hold", + "operator": "hold_hover", + "object": "cup", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": {}, + "depends_on": [], + } + ], + } + seed = compile_task_agent_v2(task_agent) + document = { + **seed, + "runtime": { + "schema_version": "action_engine_runtime_record_v2", + "status": "success", + "events": [], + }, + } + + image = _image(render_task_graph_png(document)) + + assert _contains_color(image, "#25834B") + + +def test_runtime_record_without_program_is_rejected() -> None: + with pytest.raises(ValueError, match="do not contain graph topology"): + render_task_graph_png( + { + "schema_version": "action_engine_runtime_record_v1", + "events": [], + } + ) diff --git a/tests/gen_sim/action_engine/test_motion_policy.py b/tests/gen_sim/action_engine/test_motion_policy.py new file mode 100644 index 000000000..aa253c6fb --- /dev/null +++ b/tests/gen_sim/action_engine/test_motion_policy.py @@ -0,0 +1,72 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import pytest + +from embodichain.gen_sim.action_engine.domain import motion_policy +from embodichain.gen_sim.action_engine.runtime.motion_policy import ( + resolve_motion_policy, +) + + +def test_upright_policy_matches_mature_runtime_across_robot_profiles() -> None: + upright = motion_policy(("orientation", "upright")) + franka = resolve_motion_policy("dual_franka", "PickUp", upright) + ur10 = resolve_motion_policy("dual_ur10", "PickUp", upright) + + assert franka["lift_height"] == pytest.approx(0.30) + assert ur10["lift_height"] == pytest.approx(0.30) + assert ur10["rotate_upright"] == pytest.approx(0.7853981633974483) + + +def test_policy_resolution_returns_detached_values() -> None: + upright = motion_policy(("orientation", "upright")) + first = resolve_motion_policy("ur10", "MoveHeldObject", upright) + first["surface_clearance"] = 123.0 + + second = resolve_motion_policy("dual_ur10", "MoveHeldObject", upright) + + assert second["surface_clearance"] == pytest.approx(0.05) + + +def test_unknown_action_base_is_rejected_instead_of_falling_back() -> None: + with pytest.raises(ValueError, match="Unknown Action Engine motion base"): + resolve_motion_policy("dual_franka", "TypoAction", motion_policy()) + + +def test_upright_and_handover_role_modifiers_compose_without_named_cross_product() -> ( + None +): + resolved = resolve_motion_policy( + "dual_franka", + "PickUp", + motion_policy( + ("orientation", "upright"), + ("handover_role", "transfer"), + ), + ) + + assert resolved["rotate_upright"] == pytest.approx(0.7853981633974483) + assert resolved["pick_object_part"] == "top" + assert "approach_direction_mode" not in resolved + assert resolved["sample_interval"] == 80 + + +def test_named_policy_strings_require_graph_regeneration() -> None: + with pytest.raises(ValueError, match="named string policies are no longer"): + resolve_motion_policy("dual_franka", "PickUp", "legacy_flat_name") diff --git a/tests/gen_sim/action_engine/test_orientation.py b/tests/gen_sim/action_engine/test_orientation.py new file mode 100644 index 000000000..57aa7ccd0 --- /dev/null +++ b/tests/gen_sim/action_engine/test_orientation.py @@ -0,0 +1,145 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import pytest + +from embodichain.gen_sim.action_engine.orientation import ( + AlignAxisConstraint, + MatchRotationConstraint, + compile_orientation_constraint, +) +from embodichain.gen_sim.action_engine.protocol import TASK_SPEC_SCHEMA +from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph + + +def test_unspecified_orientation_has_no_hard_constraint() -> None: + constraint = compile_orientation_constraint({}) + + assert constraint.terms == () + assert constraint.planning_preference == "minimize_rotation_from_current" + assert not constraint.requires_reference + + +def test_explicit_preserve_compiles_to_rotation_match() -> None: + constraint = compile_orientation_constraint({"orientation_goal": "preserve"}) + + assert constraint.terms == (MatchRotationConstraint(reference="step_start"),) + assert constraint.requires_reference + + +def test_upright_compiles_to_directed_axis_when_requested() -> None: + constraint = compile_orientation_constraint( + { + "orientation_goal": "upright", + "upright_local_axis": "z", + "orientation_directed": True, + } + ) + + assert constraint.terms == ( + AlignAxisConstraint( + local_axis="z", + target_axis="world_up", + directed=True, + ), + ) + + +def test_upright_rejects_non_boolean_directed_flag() -> None: + with pytest.raises(ValueError, match="orientation_directed must be a boolean"): + compile_orientation_constraint( + { + "orientation_goal": "upright", + "orientation_directed": "false", + } + ) + + +def test_legacy_long_axis_upright_remains_undirected() -> None: + constraint = compile_orientation_constraint( + { + "orientation_goal": "upright", + "upright_local_axis": "long_axis", + } + ) + + assert constraint.terms == ( + AlignAxisConstraint( + local_axis="long_axis", + target_axis="world_up", + directed=False, + ), + ) + + +def test_serialized_constraint_keeps_term_local_tolerance() -> None: + constraint = compile_orientation_constraint( + { + "orientation_constraint": { + "terms": [ + { + "type": "align_axis", + "local_axis": "z", + "target_axis": "world_up", + "directed": True, + "tolerance": 0.1, + "scope": "terminal", + } + ] + } + } + ) + + assert constraint.terms == ( + AlignAxisConstraint( + local_axis="z", + target_axis="world_up", + directed=True, + tolerance=0.1, + ), + ) + + +def test_new_placement_without_orientation_request_has_no_hard_constraint() -> None: + task = { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "place_can", + "level": "L1", + "instruction": "Place the can beside the box.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "place", + "task_type": "E1", + "params": { + "object_role": "can", + "target_role": "box", + "relation": "left_of", + }, + "depends_on": [], + "role": "primary", + } + ], + "success": {"type": "semantic_goal"}, + "oracle": {}, + "metadata": {}, + } + + graph = instantiate_seed_graph(task, {"can": "can", "box": "box"}) + + assert graph["task_groups"][0]["goal"]["orientation_goal"] == "none" 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) diff --git a/tests/gen_sim/scene_engine/test_clients.py b/tests/gen_sim/scene_engine/test_clients.py index 513f0c5c0..8ef856951 100644 --- a/tests/gen_sim/scene_engine/test_clients.py +++ b/tests/gen_sim/scene_engine/test_clients.py @@ -23,6 +23,7 @@ import pytest from embodichain.gen_sim.scene_engine.clients import geometry_generation +from embodichain.gen_sim.scene_engine.clients import image_generation from embodichain.gen_sim.scene_engine.clients import image_segmentation from embodichain.gen_sim.scene_engine.llms import load_config @@ -30,9 +31,16 @@ class _Response: """Minimal successful HTTP response used by client unit tests.""" - def __init__(self, payload: object, *, content: bytes = b"") -> None: + def __init__( + self, + payload: object, + *, + content: bytes = b"", + headers: dict[str, str] | None = None, + ) -> None: self._payload = payload self.content = content + self.headers = headers or {} def raise_for_status(self) -> None: return None @@ -79,7 +87,14 @@ def test_clients_load_their_required_dotenv_values( "SCENE_ENGINE_IMAGE_SEGMENTATION_TIMEOUT_S": "30", "SCENE_ENGINE_IMAGE_SEGMENTATION_MAX_ATTEMPTS": "2", "SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH": "/health", - "SCENE_ENGINE_IMAGE_SEGMENTATION_SINGLE_OBJECT_PATH": "/predict", + "SCENE_ENGINE_IMAGE_SEGMENTATION_BY_PROMPT_PATH": "/segment_by_prompt", + } + image_generation_values = { + "SCENE_ENGINE_IMAGE_GENERATION_BASE_URL": "http://image-generation/", + "SCENE_ENGINE_IMAGE_GENERATION_TIMEOUT_S": "120", + "SCENE_ENGINE_IMAGE_GENERATION_MAX_ATTEMPTS": "2", + "SCENE_ENGINE_IMAGE_GENERATION_HEALTH_PATH": "/health", + "SCENE_ENGINE_IMAGE_GENERATION_BY_PROMPT_PATH": "/generate_image_by_prompt", } llm_values = { "OPENAI_API_KEY": "test-key", @@ -96,18 +111,29 @@ def test_clients_load_their_required_dotenv_values( "read_scene_engine_env_values", lambda *_: segmentation_values, ) + monkeypatch.setattr( + image_generation, + "read_scene_engine_env_values", + lambda *_: image_generation_values, + ) monkeypatch.setattr( load_config, "read_scene_engine_env_values", lambda *_: llm_values ) geometry_client = geometry_generation.GeometryGenerationClient.from_dotenv() segmentation_client = image_segmentation.ImageSegmentationClient.from_dotenv() + image_generation_client = image_generation.ImageGenerationClient.from_dotenv() llm_client_config = load_config.load_llm_config() assert geometry_client._base_url == "http://geometry" assert geometry_client._generate_objects_path == "/objects" assert segmentation_client._base_url == "http://segment" - assert segmentation_client._segment_single_object_path == "/predict" + assert segmentation_client._segment_by_prompt_path == "/segment_by_prompt" + assert image_generation_client._base_url == "http://image-generation" + assert ( + image_generation_client._generate_image_by_prompt_path + == "/generate_image_by_prompt" + ) assert llm_client_config.default_query == {"api-version": "1"} assert llm_client_config.base_url == "http://llm/v1" @@ -146,15 +172,97 @@ def test_service_health_checks_use_the_configured_health_path() -> None: timeout_s=30, max_attempts=1, health_path="/health", - segment_single_object_path="/predict", + segment_by_prompt_path="/segment_by_prompt", session=segmentation_session, ) + image_generation_session = _Session(get_payload={"ok": True}) + image_generation_client = image_generation.ImageGenerationClient( + base_url="http://image-generation", + timeout_s=120, + max_attempts=1, + health_path="/health", + generate_image_by_prompt_path="/generate_image_by_prompt", + session=image_generation_session, + ) geometry_client.check_health() segmentation_client.check_health() + image_generation_client.check_health() assert geometry_session.get_calls == [("http://geometry/health", 10)] assert segmentation_session.get_calls == [("http://segment/health", 30)] + assert image_generation_session.get_calls == [ + ("http://image-generation/health", 10) + ] + + +def test_image_generation_client_posts_prompt_and_writes_png( + tmp_path: Path, +) -> None: + png_bytes = b"\x89PNG\r\n\x1a\nimage" + + class ImageGenerationSession(_Session): + def post(self, url: str, **kwargs: object) -> _Response: + self.post_call = {"url": url, **kwargs} + return _Response( + {}, + content=png_bytes, + headers={"content-type": "image/png", "x-generation-seed": "17"}, + ) + + session = ImageGenerationSession(get_payload={"ok": True}) + client = image_generation.ImageGenerationClient( + base_url="http://image-generation", + timeout_s=120, + max_attempts=1, + health_path="/health", + generate_image_by_prompt_path="/generate_image_by_prompt", + session=session, + ) + + output_path = client.generate_image_by_prompt( + prompt="a red mug on a wooden table", + output_path=tmp_path / "generated.png", + seed=17, + ) + + assert output_path.read_bytes() == png_bytes + assert session.post_call is not None + assert session.post_call["url"] == ( + "http://image-generation/generate_image_by_prompt" + ) + assert session.post_call["json"] == { + "prompt": "a red mug on a wooden table", + "seed": 17, + } + + +def test_image_generation_client_rejects_non_png_response(tmp_path: Path) -> None: + class ImageGenerationSession(_Session): + def post(self, url: str, **kwargs: object) -> _Response: + self.post_call = {"url": url, **kwargs} + return _Response( + {"ok": False, "error": "failed"}, + content=b'{"ok": false}', + headers={"content-type": "application/json"}, + ) + + session = ImageGenerationSession(get_payload={"ok": True}) + client = image_generation.ImageGenerationClient( + base_url="http://image-generation", + timeout_s=120, + max_attempts=1, + health_path="/health", + generate_image_by_prompt_path="/generate_image_by_prompt", + session=session, + ) + + with pytest.raises(RuntimeError, match="request failed after 1 attempts") as exc: + client.generate_image_by_prompt( + prompt="a red mug on a wooden table", + output_path=tmp_path / "generated.png", + ) + assert "response is not a PNG image" in str(exc.value.__cause__) def test_segmentation_client_posts_prompt_and_returns_rle_masks(tmp_path: Path) -> None: @@ -170,7 +278,7 @@ def test_segmentation_client_posts_prompt_and_returns_rle_masks(tmp_path: Path) timeout_s=30, max_attempts=1, health_path="/health", - segment_single_object_path="/predict", + segment_by_prompt_path="/segment_by_prompt", session=session, ) @@ -178,7 +286,7 @@ def test_segmentation_client_posts_prompt_and_returns_rle_masks(tmp_path: Path) rle_mask ] assert session.post_call is not None - assert session.post_call["url"] == "http://segment/predict" + assert session.post_call["url"] == "http://segment/segment_by_prompt" assert session.post_call["data"] == {"prompt": "table"} @@ -231,6 +339,7 @@ def post(self, url: str, **kwargs: object) -> _Response: { "ok": True, "result": { + "seed": 23, "objects": [ { "name": "cup", @@ -239,7 +348,7 @@ def post(self, url: str, **kwargs: object) -> _Response: "translation": [0, 0, 0], "scale": [1, 1, 1], } - ] + ], }, } ) @@ -266,9 +375,11 @@ def get(self, url: str, *, timeout: int) -> _Response: image_path=image_path, object_masks=[("cup", mask_path)], output_root=tmp_path / "output", + seed=23, ) assert objects[0]["mesh"] == "/results/cup.glb" assert session.post_call is not None assert session.post_call["url"] == "http://geometry/objects" + assert session.post_call["data"] == {"seed": "23"} assert (tmp_path / "output/cup.glb").read_bytes() == b"glTF-mesh" diff --git a/tests/gen_sim/scene_engine/test_pipeline_api.py b/tests/gen_sim/scene_engine/test_pipeline_api.py new file mode 100644 index 000000000..9b10361f7 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_pipeline_api.py @@ -0,0 +1,244 @@ +# ---------------------------------------------------------------------------- +# 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 pathlib import Path + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_edit_plan import SceneEditPlan +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, +) +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline import api + + +class _HealthyClient: + def __init__(self) -> None: + self.health_checks = 0 + + def check_health(self) -> None: + self.health_checks += 1 + + +def _materialization( + *, + scene: Scene, + scene_graph: SceneGraph, + output_root: Path, +) -> api.SceneMaterialization: + return api.SceneMaterialization( + scene=scene, + scene_graph=scene_graph, + output_root=output_root, + scene_config_path=output_root / "scene_export" / "scene_config.json", + ) + + +def _table_scene() -> tuple[Scene, SceneGraph]: + scene = Scene( + objects=[ + SceneObject( + id="table", + kind="table", + category="table", + name="table", + description="A work table.", + ) + ] + ) + graph = SceneGraph(nodes=[SceneGraphNode(object_id="table", parent_id=None)]) + return scene, graph + + +def test_analyze_image_persists_blueprint_and_artifact_hashes( + tmp_path: Path, + monkeypatch, +) -> None: + image_path = tmp_path / "input.png" + image_path.write_bytes(b"image") + scene, graph = _table_scene() + + def fake_understand_scene(**kwargs): + stage_root = Path(kwargs["output_root"]) / "scene_understanding" + stage_root.mkdir(parents=True) + (stage_root / "table-mask.png").write_bytes(b"mask") + return scene, graph + + monkeypatch.setattr(api, "understand_scene", fake_understand_scene) + segmentation = _HealthyClient() + package = api.analyze_image( + image_path, + tmp_path / "output", + vlm_client=object(), + image_segmentation_client=segmentation, + ) + document = json.loads(package.manifest_path.read_text(encoding="utf-8")) + + assert segmentation.health_checks == 1 + assert document["blueprint_id"] == package.blueprint_id + assert document["scene_graph"] == graph.to_dict() + assert document["artifacts"][0]["path"].endswith("table-mask.png") + assert len(document["artifacts"][0]["sha256"]) == 64 + + +def test_analyze_edit_persists_post_edit_blueprint( + tmp_path: Path, + monkeypatch, +) -> None: + scene, graph = _table_scene() + plan = SceneEditPlan(scene=scene, scene_graph=graph, operations=[]) + + class FakeImporter: + def __init__(self, *, output_root: Path) -> None: + self.output_root = output_root + + def import_scene_and_graph(self): + return scene, graph + + monkeypatch.setattr(api, "SceneExportImporter", FakeImporter) + monkeypatch.setattr( + api, + "understand_scene_edit", + lambda **_: (plan, graph), + ) + package = api.analyze_edit( + output_root=tmp_path, + edit_prompt="Keep the scene unchanged.", + vlm_client=object(), + ) + document = json.loads(package.manifest_path.read_text(encoding="utf-8")) + + assert document["blueprint_id"] == package.blueprint_id + assert document["scene_edit_plan"] == plan.to_dict() + assert document["updated_scene_graph"] == graph.to_dict() + + +def test_materialize_blueprint_does_not_mutate_audited_snapshot( + tmp_path: Path, + monkeypatch, +) -> None: + scene, graph = _table_scene() + manifest_path = tmp_path / "scene_blueprint.json" + manifest_path.write_text("audited blueprint\n", encoding="utf-8") + package = api.SceneBlueprintPackage( + blueprint_id="blueprint", + image_path=tmp_path / "input.png", + output_root=tmp_path, + manifest_path=manifest_path, + scene=scene, + scene_graph=graph, + ) + original_scene = deepcopy(scene.to_dict()) + original_graph = deepcopy(graph.to_dict()) + + def fake_generate_scene_and_refine(**kwargs): + assert kwargs["seed"] == 31 + assert kwargs["scene"] is not package.scene + assert kwargs["scene_graph"] is not package.scene_graph + kwargs["scene"].objects[0].name = "materialized table" + return kwargs["scene"] + + monkeypatch.setattr( + api, + "generate_scene_and_refine", + fake_generate_scene_and_refine, + ) + monkeypatch.setattr( + api, + "_export_materialization", + lambda *, scene, scene_graph, output_root: _materialization( + scene=scene, + scene_graph=scene_graph, + output_root=output_root, + ), + ) + + result = api.materialize_blueprint( + package, + vlm_client=object(), + geometry_generation_client=_HealthyClient(), + seed=31, + ) + + assert result.scene.objects[0].name == "materialized table" + assert package.scene.to_dict() == original_scene + assert package.scene_graph.to_dict() == original_graph + assert manifest_path.read_text(encoding="utf-8") == "audited blueprint\n" + + +def test_materialize_edit_does_not_mutate_audited_snapshot( + tmp_path: Path, + monkeypatch, +) -> None: + scene, graph = _table_scene() + plan = SceneEditPlan(scene=scene, scene_graph=graph, operations=[]) + manifest_path = tmp_path / "scene_edit_blueprint.json" + manifest_path.write_text("audited edit blueprint\n", encoding="utf-8") + package = api.SceneEditBlueprintPackage( + blueprint_id="edit-blueprint", + edit_prompt="Keep the scene unchanged.", + output_root=tmp_path, + manifest_path=manifest_path, + scene_edit_plan=plan, + updated_scene_graph=graph, + ) + original_plan = deepcopy(plan.to_dict()) + original_graph = deepcopy(graph.to_dict()) + + def fake_prepare_scene_edit_assets(**kwargs): + assert kwargs["seed"] == 32 + return [] + + monkeypatch.setattr( + api, "prepare_scene_edit_assets", fake_prepare_scene_edit_assets + ) + + def fake_edit_layout(**kwargs): + assert kwargs["scene_edit_plan"] is not package.scene_edit_plan + assert kwargs["updated_scene_graph"] is not package.updated_scene_graph + kwargs["scene"].objects[0].name = "edited table" + return kwargs["scene"] + + monkeypatch.setattr(api, "edit_layout", fake_edit_layout) + monkeypatch.setattr( + api, + "_export_materialization", + lambda *, scene, scene_graph, output_root: _materialization( + scene=scene, + scene_graph=scene_graph, + output_root=output_root, + ), + ) + clients = [_HealthyClient(), _HealthyClient(), _HealthyClient()] + + result = api.materialize_edit( + package, + vlm_client=object(), + image_generation_client=clients[0], + geometry_generation_client=clients[1], + image_segmentation_client=clients[2], + seed=32, + ) + + assert result.scene.objects[0].name == "edited table" + assert package.scene_edit_plan.to_dict() == original_plan + assert package.updated_scene_graph.to_dict() == original_graph + assert manifest_path.read_text(encoding="utf-8") == "audited edit blueprint\n" diff --git a/tests/gen_sim/scene_engine/test_scene_core_and_export.py b/tests/gen_sim/scene_engine/test_scene_core_and_export.py index 56c12af1c..20b40923f 100644 --- a/tests/gen_sim/scene_engine/test_scene_core_and_export.py +++ b/tests/gen_sim/scene_engine/test_scene_core_and_export.py @@ -24,11 +24,18 @@ import pytest from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, +) from embodichain.gen_sim.scene_engine.core.scene_object import ( ObjectPhysics, SceneObject, ) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_exporter import SceneExporter +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_importer import ( + SceneExportImporter, +) def _scene_object( @@ -60,6 +67,24 @@ def _physics(body_type: str) -> ObjectPhysics: ) +def _scene_graph(scene: Scene) -> SceneGraph: + if scene.table is None: + raise ValueError("Test scene must contain a table.") + return SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + *[ + SceneGraphNode( + object_id=asset.id, + parent_id="table", + parent_relation="on", + ) + for asset in scene.assets + ], + ] + ) + + def test_scene_returns_one_table_and_ordered_assets() -> None: table = _scene_object(object_id="table", kind="table") asset = _scene_object(object_id="cup", kind="asset") @@ -103,7 +128,9 @@ def test_object_physics_rejects_invalid_values( ) -def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> None: +def test_scene_export_rotates_the_complete_z_up_scene_by_default( + tmp_path: Path, +) -> None: table_glb = tmp_path / "table.glb" asset_glb = tmp_path / "cup.glb" table_glb.write_bytes(b"glTF-table") @@ -114,15 +141,26 @@ def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> No glb_path=table_glb, physics=_physics("kinematic"), ) + table.pos = [0.0, 0.0, 0.0] + table.support_contour_xy = [[-1.0, -0.5], [1.0, -0.5], [1.0, 0.5]] + table.support_optimization_rect_xy = [ + [-0.8, -0.3], + [0.8, -0.3], + [0.8, 0.3], + [-0.8, 0.3], + ] asset = _scene_object( object_id="cup", kind="asset", glb_path=asset_glb, physics=_physics("dynamic"), ) + asset.center_xy = [0.25, -0.5] + scene = Scene(objects=[table, asset]) export_path = SceneExporter( - scene=Scene(objects=[table, asset]), + scene=scene, + scene_graph=_scene_graph(scene), output_root=tmp_path / "output", ).export() exported = json.loads(export_path.read_text(encoding="utf-8")) @@ -133,10 +171,185 @@ def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> No assert (export_path.parent / "mesh_assets/cup/cup.glb").read_bytes() == b"glTF-cup" entry = exported["rigid_object"][0] assert entry["uid"] == "cup" + assert entry["category"] == "asset" + assert entry["name"] == "cup" assert entry["body_type"] == "dynamic" - assert entry["init_pos"] == [1.0, -3.0, 2.0] + assert np.allclose(entry["init_pos"], [-1.0, 3.0, 2.0]) assert entry["body_scale"] == [1.0, 2.0, 3.0] - assert np.allclose(entry["init_rot"], [0.0, 0.0, 0.0]) + assert np.allclose(entry["center_xy"], [-0.25, 0.5]) + assert np.allclose( + entry["init_rot"], + [0.0, 0.0, 180.0], + ) + exported_table = exported["background"][0] + assert np.allclose( + exported_table["support_contour_xy"], + [[1.0, 0.5], [-1.0, 0.5], [-1.0, -0.5]], + ) + assert np.allclose( + exported_table["support_optimization_rect_xy"], + [[0.8, 0.3], [-0.8, 0.3], [-0.8, -0.3], [0.8, -0.3]], + ) + exported_scene_json = json.loads((export_path.parent / "scene.json").read_text()) + exported_asset_json = exported_scene_json["objects"][1] + assert np.allclose(exported_asset_json["pos"], [-1.0, 2.0, -3.0]) + assert np.allclose(exported_asset_json["center_xy"], [-0.25, 0.5]) + assert json.loads((export_path.parent / "scene_graph.json").read_text()) == { + "nodes": [ + { + "object_id": "table", + "parent_id": None, + "parent_relation": None, + "table_region": None, + "orientation_state": None, + }, + { + "object_id": "cup", + "parent_id": "table", + "parent_relation": "on", + "table_region": None, + "orientation_state": None, + }, + ], + "relations": [], + } + + imported_scene, imported_graph = SceneExportImporter( + output_root=tmp_path / "output" + ).import_scene_and_graph() + assert [asset.id for asset in imported_scene.assets] == ["cup"] + assert imported_scene.assets[0].category == "asset" + assert imported_scene.assets[0].name == "cup" + assert np.allclose(imported_scene.assets[0].pos, [-1.0, 2.0, -3.0]) + assert np.allclose(imported_scene.assets[0].center_xy, [-0.25, 0.5]) + assert np.allclose( + imported_scene.table.support_contour_xy, + [[1.0, 0.5], [-1.0, 0.5], [-1.0, -0.5]], + ) + assert imported_graph.to_dict() == _scene_graph(scene).to_dict() + + +def test_scene_export_can_disable_the_default_z_up_rotation(tmp_path: Path) -> None: + table_glb = tmp_path / "table.glb" + asset_glb = tmp_path / "cup.glb" + table_glb.write_bytes(b"glTF-table") + asset_glb.write_bytes(b"glTF-cup") + table = _scene_object( + object_id="table", + kind="table", + glb_path=table_glb, + physics=_physics("kinematic"), + ) + table.pos = [0.0, 0.0, 0.0] + asset = _scene_object( + object_id="cup", + kind="asset", + glb_path=asset_glb, + physics=_physics("dynamic"), + ) + + scene = Scene(objects=[table, asset]) + export_path = SceneExporter( + scene=scene, + scene_graph=_scene_graph(scene), + output_root=tmp_path / "output", + rotate_z_up_180=False, + ).export() + exported = json.loads(export_path.read_text(encoding="utf-8")) + + assert exported["rigid_object"][0]["init_pos"] == [1.0, -3.0, 2.0] + assert np.allclose(exported["rigid_object"][0]["init_rot"], [0.0, 0.0, 0.0]) + + +def test_scene_graph_importer_restores_node_orientation_state() -> None: + imported_graph = SceneExportImporter._scene_graph_from_data( + { + "nodes": [ + { + "object_id": "table", + "parent_id": None, + "parent_relation": None, + "table_region": None, + "orientation_state": None, + }, + { + "object_id": "bottle_001", + "parent_id": "table", + "parent_relation": "on", + "table_region": None, + "orientation_state": "standing", + }, + ], + "relations": [], + } + ) + + assert imported_graph.node_by_id()["bottle_001"].orientation_state == "standing" + + +def test_scene_export_overwrites_an_existing_scene_export(tmp_path: Path) -> None: + table_glb = tmp_path / "table.glb" + cup_glb = tmp_path / "cup.glb" + banana_glb = tmp_path / "banana.glb" + table_glb.write_bytes(b"glTF-table") + cup_glb.write_bytes(b"glTF-cup") + banana_glb.write_bytes(b"glTF-banana") + output_root = tmp_path / "output" + + initial_table = _scene_object( + object_id="table", + kind="table", + glb_path=table_glb, + physics=_physics("kinematic"), + ) + initial_cup = _scene_object( + object_id="cup", + kind="asset", + glb_path=cup_glb, + physics=_physics("dynamic"), + ) + initial_scene = Scene(objects=[initial_table, initial_cup]) + SceneExporter( + scene=initial_scene, + scene_graph=_scene_graph(initial_scene), + output_root=output_root, + ).export() + + # The imported table mesh already occupies its final export location. + exported_table_glb = ( + output_root / "scene_export" / "mesh_assets" / "table" / "table.glb" + ) + updated_table = _scene_object( + object_id="table", + kind="table", + glb_path=exported_table_glb, + physics=_physics("kinematic"), + ) + banana = _scene_object( + object_id="banana", + kind="asset", + glb_path=banana_glb, + physics=_physics("dynamic"), + ) + updated_scene = Scene(objects=[updated_table, banana]) + SceneExporter( + scene=updated_scene, + scene_graph=_scene_graph(updated_scene), + output_root=output_root, + ).export() + + scene_export_root = output_root / "scene_export" + assert exported_table_glb.read_bytes() == b"glTF-table" + assert ( + scene_export_root / "mesh_assets" / "banana" / "banana.glb" + ).read_bytes() == b"glTF-banana" + assert not (scene_export_root / "mesh_assets" / "cup").exists() + assert ( + json.loads((scene_export_root / "scene.json").read_text(encoding="utf-8"))[ + "objects" + ][1]["id"] + == "banana" + ) def test_scene_export_requires_final_physics(tmp_path: Path) -> None: @@ -145,7 +358,13 @@ def test_scene_export_requires_final_physics(tmp_path: Path) -> None: table = _scene_object(object_id="table", kind="table", glb_path=glb_path) with pytest.raises(ValueError, match="no SimReady physics"): - SceneExporter(scene=Scene(objects=[table]), output_root=tmp_path).export() + SceneExporter( + scene=Scene(objects=[table]), + scene_graph=SceneGraph( + nodes=[SceneGraphNode(object_id="table", parent_id=None)] + ), + output_root=tmp_path, + ).export() def test_scene_export_rejects_backslash_in_object_id(tmp_path: Path) -> None: @@ -165,7 +384,9 @@ def test_scene_export_rejects_backslash_in_object_id(tmp_path: Path) -> None: ) with pytest.raises(ValueError, match="not safe for a GLB filename"): + scene = Scene(objects=[table, unsafe_asset]) SceneExporter( - scene=Scene(objects=[table, unsafe_asset]), + scene=scene, + scene_graph=_scene_graph(scene), output_root=tmp_path / "output", ).export() diff --git a/tests/gen_sim/scene_engine/test_scene_edit.py b/tests/gen_sim/scene_engine/test_scene_edit.py new file mode 100644 index 000000000..fd1a9e933 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_edit.py @@ -0,0 +1,135 @@ +# ---------------------------------------------------------------------------- +# 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 + +import pytest + +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_importer import ( + import_scene_from_output_root, +) + + +def _write_scene_export( + output_root: Path, + *, + include_table: bool = True, + include_asset_mesh: bool = True, +) -> None: + scene_export_root = output_root / "scene_export" + table_mesh_path = scene_export_root / "mesh_assets" / "table" / "table.glb" + asset_mesh_path = scene_export_root / "mesh_assets" / "cup" / "cup.glb" + table_mesh_path.parent.mkdir(parents=True) + asset_mesh_path.parent.mkdir(parents=True) + table_mesh_path.write_bytes(b"glTF-table") + if include_asset_mesh: + asset_mesh_path.write_bytes(b"glTF-cup") + + background = [] + if include_table: + background.append( + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "mesh_assets/table/table.glb", + }, + "attrs": {"mass": 1.0}, + "body_type": "kinematic", + "init_pos": [0.0, 0.0, 0.0], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + "max_convex_hull_num": 16, + } + ) + scene_config = { + "background": background, + "rigid_object": [ + { + "uid": "cup", + "shape": { + "shape_type": "Mesh", + "fpath": "mesh_assets/cup/cup.glb", + }, + "attrs": {"mass": 1.0}, + "body_type": "dynamic", + "init_pos": [1.0, -3.0, 2.0], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 2.0, 3.0], + "center_xy": [1.0, -3.0], + "max_convex_hull_num": 32, + } + ], + } + (scene_export_root / "scene_config.json").write_text( + json.dumps(scene_config), + encoding="utf-8", + ) + + +def test_import_scene_from_output_root_writes_y_up_scene_json(tmp_path: Path) -> None: + _write_scene_export(tmp_path) + (tmp_path / "scene_export" / "scene.json").write_text( + '{"old": true}', + encoding="utf-8", + ) + + scene = import_scene_from_output_root(tmp_path) + scene_json = json.loads( + (tmp_path / "scene_export" / "scene.json").read_text(encoding="utf-8") + ) + + assert scene.table is not None + assert scene.table.id == "table" + assert scene.assets[0].id == "cup" + assert scene.assets[0].pos == [1.0, 2.0, 3.0] + assert scene.assets[0].scale == [1.0, 2.0, 3.0] + assert scene.assets[0].center_xy == [1.0, -3.0] + assert scene.assets[0].simready_glb_path == str( + (tmp_path / "scene_export" / "mesh_assets" / "cup" / "cup.glb").resolve() + ) + assert scene_json["objects"][1]["id"] == "cup" + assert scene_json["objects"][1]["pos"] == [1.0, 2.0, 3.0] + + +def test_check_scene_export_for_edit_requires_export_directories( + tmp_path: Path, +) -> None: + with pytest.raises(ValueError, match="Output root"): + import_scene_from_output_root(tmp_path) + + (tmp_path / "scene_export").mkdir() + with pytest.raises(FileNotFoundError, match="mesh assets"): + import_scene_from_output_root(tmp_path) + + +def test_check_scene_export_for_edit_requires_table(tmp_path: Path) -> None: + _write_scene_export(tmp_path, include_table=False) + + with pytest.raises(ValueError, match="table"): + import_scene_from_output_root(tmp_path) + + +def test_check_scene_export_for_edit_requires_rigid_object_glb( + tmp_path: Path, +) -> None: + _write_scene_export(tmp_path, include_asset_mesh=False) + + with pytest.raises(FileNotFoundError, match="cup"): + import_scene_from_output_root(tmp_path) diff --git a/tests/gen_sim/scene_engine/test_scene_edit_plan.py b/tests/gen_sim/scene_engine/test_scene_edit_plan.py new file mode 100644 index 000000000..3b0ec57c3 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_edit_plan.py @@ -0,0 +1,526 @@ +# ---------------------------------------------------------------------------- +# 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 + +import pytest +from PIL import Image +import trimesh + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_edit_plan import ( + SceneEditOperation, + SceneEditPlan, +) +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, +) +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_understanding import ( + _apply_scene_edit_plan_to_scene_graph, + _build_updated_scene_graph, + _parse_scene_edit_operations, +) +from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_asset_preparation import ( + prepare_scene_edit_assets, +) + + +def _scene_and_graph() -> tuple[Scene, SceneGraph]: + scene = Scene( + objects=[ + SceneObject( + id="table", + kind="table", + category="table", + name="wooden table", + description="A wooden table.", + ), + SceneObject( + id="book_001", + kind="asset", + category="book", + name="blue book", + description="A blue book.", + ), + SceneObject( + id="orange_001", + kind="asset", + category="orange", + name="orange", + description="An orange.", + ), + ] + ) + scene_graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="book_001", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="orange_001", + parent_id="book_001", + parent_relation="on", + ), + ] + ) + return scene, scene_graph + + +def test_scene_edit_plan_accepts_add_without_a_position() -> None: + scene, scene_graph = _scene_and_graph() + + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="add", + object_id="cup_001", + category="cup", + name="green cup", + description="A small green ceramic cup.", + ) + ], + ) + + assert len(plan.operations) == 1 + assert plan.to_dict()["operations"] == [ + { + "op": "add", + "object_id": "cup_001", + "target_id": None, + "relation": None, + "table_region": None, + "category": "cup", + "name": "green cup", + "description": "A small green ceramic cup.", + "orientation_state": None, + } + ] + + +def test_scene_edit_plan_accepts_multiple_new_objects_with_the_same_category() -> None: + scene, scene_graph = _scene_and_graph() + + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="add", + object_id="orange_002", + category="orange", + name="small orange", + description="A small round orange with a textured peel.", + ), + SceneEditOperation( + op="add", + object_id="orange_003", + category="orange", + name="large orange", + description="A large round orange with a textured peel.", + ), + ], + ) + + assert [operation.category for operation in plan.operations] == ["orange", "orange"] + + +def test_scene_edit_parser_assigns_ids_to_same_category_adds_in_order() -> None: + scene, _ = _scene_and_graph() + draft = { + "operations": [ + { + "op": "add", + "object_id": None, + "target_id": None, + "relation": None, + "table_region": None, + "category": "orange", + "name": "small_orange", + "description": "A small round orange with a textured peel.", + "orientation_state": None, + }, + { + "op": "add", + "object_id": None, + "target_id": None, + "relation": None, + "table_region": None, + "category": "orange", + "name": "small_orange", + "description": "A small round orange with a textured peel.", + "orientation_state": None, + }, + ] + } + + operations = _parse_scene_edit_operations( + json.loads(json.dumps(draft)), scene=scene + ) + + assert [operation.object_id for operation in operations] == [ + "orange_002", + "orange_003", + ] + + +def test_scene_edit_plan_rejects_targets_outside_the_input_scene() -> None: + scene, scene_graph = _scene_and_graph() + + with pytest.raises(ValueError, match="existing scene objects"): + SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="add", + object_id="spoon_001", + target_id="new_orange_001", + relation="left_of", + category="spoon", + name="metal spoon", + description="A metal spoon.", + ) + ], + ) + + +def test_scene_edit_plan_requires_deleting_all_children_of_a_deleted_parent() -> None: + scene, scene_graph = _scene_and_graph() + + with pytest.raises(ValueError, match="all of its children"): + SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[SceneEditOperation(op="delete", object_id="book_001")], + ) + + +def test_scene_edit_plan_requires_a_position_for_move_operations() -> None: + scene, scene_graph = _scene_and_graph() + + with pytest.raises(ValueError, match="must specify target_id and relation"): + SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[SceneEditOperation(op="move", object_id="book_001")], + ) + + +def test_scene_edit_asset_preparation_skips_plans_without_adds( + tmp_path: Path, +) -> None: + scene, scene_graph = _scene_and_graph() + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="move", + object_id="book_001", + target_id="table", + relation="on", + ) + ], + ) + previous_asset_output = tmp_path / "scene_editing" / "asset_preparation" + previous_asset_output.mkdir(parents=True) + (previous_asset_output / "previous.txt").write_text("keep", encoding="utf-8") + + prepared_assets = prepare_scene_edit_assets( + scene_edit_plan=plan, + output_root=tmp_path, + image_generation_client=object(), # type: ignore[arg-type] + geometry_generation_client=object(), # type: ignore[arg-type] + image_segmentation_client=object(), # type: ignore[arg-type] + ) + + assert prepared_assets == [] + assert (previous_asset_output / "previous.txt").is_file() + + +def test_scene_edit_asset_preparation_generates_one_image_per_add( + tmp_path: Path, +) -> None: + class ImageGenerationClient: + def __init__(self) -> None: + self.requests: list[tuple[str, Path]] = [] + + def generate_image_by_prompt(self, *, prompt: str, output_path: Path) -> Path: + self.requests.append((prompt, output_path)) + Image.new("RGB", (6, 6), "white").save(output_path) + return output_path + + class ImageSegmentationClient: + def __init__(self) -> None: + self.requests: list[tuple[Path, str]] = [] + + def segment_single_object( + self, + *, + image_path: Path, + prompt: str, + ) -> list[dict[str, object]]: + self.requests.append((image_path, prompt)) + return [ + { + "size": [6, 6], + "counts": [7, 4, 2, 4, 2, 4, 2, 4, 7], + "starts_with": 1, + } + ] + + class GeometryGenerationClient: + def __init__(self) -> None: + self.requests: list[tuple[Path, list[tuple[str, Path]], Path]] = [] + + def generate_objects( + self, + *, + image_path: Path, + object_masks: list[tuple[str, Path]], + output_root: Path, + ) -> tuple[dict[str, object], list[dict[str, object]]]: + self.requests.append((image_path, object_masks, output_root)) + output_root.mkdir(parents=True, exist_ok=True) + for object_id, _ in object_masks: + trimesh.creation.box().export(output_root / f"{object_id}.glb") + return {}, [{"scale": [1.25, 1.5, 1.75]}] + + scene, scene_graph = _scene_and_graph() + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="add", + object_id="cup_001", + category="cup", + name="green cup", + description="A small green ceramic cup.", + ) + ], + ) + image_generation_client = ImageGenerationClient() + image_segmentation_client = ImageSegmentationClient() + geometry_generation_client = GeometryGenerationClient() + + prepared_assets = prepare_scene_edit_assets( + scene_edit_plan=plan, + output_root=tmp_path, + image_generation_client=image_generation_client, # type: ignore[arg-type] + geometry_generation_client=geometry_generation_client, # type: ignore[arg-type] + image_segmentation_client=image_segmentation_client, # type: ignore[arg-type] + ) + + expected_image_path = ( + tmp_path + / "scene_editing" + / "asset_preparation" + / "generated_images" + / "cup_001.png" + ) + assert [asset.id for asset in prepared_assets] == ["cup_001"] + assert prepared_assets[0].simready_glb_path is not None + assert prepared_assets[0].rot == [0.0, 0.0, 0.0] + assert prepared_assets[0].pos == [0.0, 0.0, 0.0] + assert prepared_assets[0].scale == [1.0, 1.0, 1.0] + assert image_generation_client.requests == [ + ("A small green ceramic cup.", expected_image_path) + ] + assert expected_image_path.is_file() + assert image_segmentation_client.requests == [ + (expected_image_path, "A small green ceramic cup.") + ] + generated_mask_path = ( + tmp_path + / "scene_editing" + / "asset_preparation" + / "generated_masks" + / "cup_001_mask.png" + ) + assert generated_mask_path.is_file() + with Image.open(generated_mask_path) as mask: + assert mask.getpixel((3, 3)) == 255 + assert mask.getpixel((0, 0)) == 0 + generated_glb_path = ( + tmp_path + / "scene_editing" + / "asset_preparation" + / "coarse_geometry" + / "cup_001.glb" + ) + assert geometry_generation_client.requests == [ + ( + expected_image_path, + [("cup_001", generated_mask_path)], + generated_glb_path.parent, + ) + ] + assert generated_glb_path.read_bytes().startswith(b"glTF") + + +def test_scene_edit_graph_builder_copies_the_pre_edit_graph() -> None: + scene, scene_graph = _scene_and_graph() + plan = SceneEditPlan(scene=scene, scene_graph=scene_graph) + + updated_scene_graph = _build_updated_scene_graph( + scene_graph=scene_graph, + scene_edit_plan=plan, + ) + + assert updated_scene_graph is not scene_graph + assert updated_scene_graph.nodes is not scene_graph.nodes + assert updated_scene_graph.relations is not scene_graph.relations + assert updated_scene_graph.to_dict() == scene_graph.to_dict() + + +def test_scene_edit_graph_builder_removes_deleted_nodes() -> None: + scene, scene_graph = _scene_and_graph() + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation(op="delete", object_id="orange_001"), + SceneEditOperation(op="delete", object_id="book_001"), + ], + ) + + updated_scene_graph = _build_updated_scene_graph( + scene_graph=scene_graph, + scene_edit_plan=plan, + ) + + assert set(updated_scene_graph.node_by_id()) == {"table"} + assert set(scene_graph.node_by_id()) == {"table", "book_001", "orange_001"} + + +def test_scene_edit_graph_builder_adds_unpositioned_objects_on_the_table() -> None: + scene, scene_graph = _scene_and_graph() + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="add", + object_id="cup_001", + category="cup", + name="green cup", + description="A small green ceramic cup.", + ) + ], + ) + + updated_scene_graph = _build_updated_scene_graph( + scene_graph=scene_graph, + scene_edit_plan=plan, + ) + + added_node = updated_scene_graph.node_by_id()["cup_001"] + assert added_node.parent_id == "table" + assert added_node.parent_relation == "on" + + +def test_scene_edit_graph_builder_updates_move_on_parent() -> None: + scene, scene_graph = _scene_and_graph() + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="move", + object_id="orange_001", + target_id="table", + relation="on", + ) + ], + ) + + updated_scene_graph = _build_updated_scene_graph( + scene_graph=scene_graph, + scene_edit_plan=plan, + ) + + assert updated_scene_graph.node_by_id()["orange_001"].parent_id == "table" + + +def test_scene_edit_graph_builder_adds_planar_relation_with_target_parent() -> None: + scene, scene_graph = _scene_and_graph() + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="add", + object_id="cup_001", + target_id="book_001", + relation="right_of", + category="cup", + name="green cup", + description="A small green ceramic cup.", + ) + ], + ) + + updated_scene_graph = _build_updated_scene_graph( + scene_graph=scene_graph, + scene_edit_plan=plan, + ) + + assert updated_scene_graph.node_by_id()["cup_001"].parent_id == "table" + assert any( + relation.source_id == "cup_001" + and relation.relation == "right_of" + and relation.target_id == "book_001" + for relation in updated_scene_graph.relations + ) + + +def test_scene_edit_plan_application_adds_new_nodes_before_relationship_updates() -> ( + None +): + scene, scene_graph = _scene_and_graph() + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="add", + object_id="cup_001", + target_id="book_001", + relation="right_of", + category="cup", + name="green cup", + description="A small green ceramic cup.", + ) + ], + ) + + _apply_scene_edit_plan_to_scene_graph( + scene_graph=scene_graph, + scene_edit_plan=plan, + ) + + assert scene_graph.node_by_id()["cup_001"].parent_id == "table" diff --git a/tests/gen_sim/scene_engine/test_scene_engine_config.py b/tests/gen_sim/scene_engine/test_scene_engine_config.py index 01914e144..3754210d9 100644 --- a/tests/gen_sim/scene_engine/test_scene_engine_config.py +++ b/tests/gen_sim/scene_engine/test_scene_engine_config.py @@ -87,6 +87,58 @@ def generate_scene(*, image_path: Path, output_root: Path) -> None: } +def test_scene_engine_cli_edits_existing_output_without_an_image( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + captured: dict[str, object] = {} + + def edit_scene(*, output_root: Path, edit_prompt: str) -> None: + captured["output_root"] = output_root + captured["edit_prompt"] = edit_prompt + + monkeypatch.setattr(start, "edit_scene", edit_scene) + output_root = tmp_path / "existing_output" + + start.cli_scene_engine(None, output_root, edit_prompt="move the cup right") + + assert captured == { + "output_root": output_root.resolve(), + "edit_prompt": "move the cup right", + } + + +def test_scene_engine_cli_generates_then_edits_when_both_inputs_exist( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + image_path = tmp_path / "scene.png" + image_path.write_bytes(b"png") + call_order: list[str] = [] + + def generate_scene(*, image_path: Path, output_root: Path) -> None: + call_order.append("generate") + + def edit_scene(*, output_root: Path, edit_prompt: str) -> None: + call_order.append("edit") + + monkeypatch.setattr(start, "generate_scene_from_image", generate_scene) + monkeypatch.setattr(start, "edit_scene", edit_scene) + + start.cli_scene_engine( + image_path, + tmp_path / "output", + edit_prompt="move the cup right", + ) + + assert call_order == ["generate", "edit"] + + +def test_scene_engine_cli_requires_an_image_or_edit_prompt(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="--image, --edit_prompt, or both"): + start.cli_scene_engine(None, tmp_path / "output") + + @pytest.mark.parametrize("image_name", ["missing.png", "scene.gif"]) def test_scene_engine_cli_rejects_invalid_image_inputs( tmp_path: Path, diff --git a/tests/gen_sim/scene_engine/test_scene_generation.py b/tests/gen_sim/scene_engine/test_scene_generation.py new file mode 100644 index 000000000..6a61a4d1b --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_generation.py @@ -0,0 +1,95 @@ +# ---------------------------------------------------------------------------- +# 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 numpy as np +from scipy.spatial.transform import Rotation + +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, +) +from embodichain.gen_sim.scene_engine.pipeline.generation.scene_generation import ( + _scene_graph_based_calibration, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( + layout_object_to_transform_matrix, + transform_matrix_to_layout_object, +) + + +def _y_up_layout_from_z_up_rotation( + object_id: str, + rotation_matrix: np.ndarray, +) -> dict[str, object]: + y_up_to_z_up_matrix = np.eye(4) + y_up_to_z_up_matrix[:3, :3] = np.array( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + z_up_to_y_up_matrix = np.linalg.inv(y_up_to_z_up_matrix) + z_up_transform = np.eye(4) + z_up_transform[:3, :3] = rotation_matrix + return transform_matrix_to_layout_object( + object_id, + z_up_to_y_up_matrix @ z_up_transform @ y_up_to_z_up_matrix, + ) + + +def _z_up_rotation_from_y_up_layout(layout: dict[str, object]) -> np.ndarray: + y_up_to_z_up_matrix = np.eye(4) + y_up_to_z_up_matrix[:3, :3] = np.array( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + return ( + y_up_to_z_up_matrix + @ layout_object_to_transform_matrix(layout) + @ np.linalg.inv(y_up_to_z_up_matrix) + )[:3, :3] + + +def test_scene_graph_calibration_makes_standing_asset_vertical() -> None: + scene_graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="bottle_001", + parent_id="table", + parent_relation="on", + orientation_state="standing", + ), + SceneGraphNode( + object_id="book_001", + parent_id="table", + parent_relation="on", + ), + ] + ) + lying_rotation = Rotation.from_euler("x", 90.0, degrees=True).as_matrix() + bottle_layout = _y_up_layout_from_z_up_rotation("bottle_001", lying_rotation) + book_layout = _y_up_layout_from_z_up_rotation("book_001", lying_rotation) + + calibrated_layouts = _scene_graph_based_calibration( + scene_graph=scene_graph, + assets_layout=[bottle_layout, book_layout], + ) + + bottle_axis = _z_up_rotation_from_y_up_layout(calibrated_layouts[0])[:, 2] + assert np.isclose(abs(bottle_axis[2]), 1.0) + assert np.allclose( + _z_up_rotation_from_y_up_layout(calibrated_layouts[1]), + lying_rotation, + ) diff --git a/tests/gen_sim/scene_engine/test_scene_graph.py b/tests/gen_sim/scene_engine/test_scene_graph.py new file mode 100644 index 000000000..d5b6adaf7 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_graph.py @@ -0,0 +1,313 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import pytest + +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, + SceneGraphRelation, +) + + +def test_scene_graph_accepts_layered_on_relations() -> None: + graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + table_region="center", + orientation_state="standing", + ), + SceneGraphNode( + object_id="cup", + parent_id="table", + parent_relation="on", + table_region="right_center", + ), + SceneGraphNode( + object_id="spoon", + parent_id="plate", + parent_relation="on", + ), + ], + relations=[ + SceneGraphRelation( + source_id="plate", + relation="left_of", + target_id="cup", + ), + ], + ) + + graph.validate() + assert graph.layer_by_id()["spoon"] == 2 + + +def test_scene_graph_rejects_planar_relations_without_common_parent() -> None: + with pytest.raises(ValueError, match="share one parent"): + SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="spoon", + parent_id="plate", + parent_relation="on", + ), + ], + relations=[ + SceneGraphRelation( + source_id="plate", + relation="right_of", + target_id="spoon", + ), + ], + ) + + +def test_scene_graph_rejects_conflicting_planar_relations() -> None: + with pytest.raises(ValueError, match="Conflicting planar relations"): + SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="cup", + parent_id="table", + parent_relation="on", + ), + ], + relations=[ + SceneGraphRelation( + source_id="plate", + relation="left_of", + target_id="cup", + ), + SceneGraphRelation( + source_id="cup", + relation="left_of", + target_id="plate", + ), + ], + ) + + +def test_scene_graph_requires_explicit_parent_relation() -> None: + with pytest.raises(ValueError, match="parent relation"): + SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + ), + ], + ) + + +def test_scene_graph_can_skip_validation_during_refresh() -> None: + graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + ), + ], + validate_on_refresh=False, + ) + + with pytest.raises(ValueError, match="parent relation"): + graph.validate() + + +def test_scene_graph_rejects_unsupported_parent_relation() -> None: + with pytest.raises(ValueError, match="must be on their parent"): + SceneGraphNode( + object_id="orange", + parent_id="box", + parent_relation="inside", + ) + + +def test_scene_graph_derives_layers_from_parent_links() -> None: + graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="spoon", + parent_id="plate", + parent_relation="on", + ), + ], + ) + + assert graph.layer_by_id() == { + "table": 0, + "plate": 1, + "spoon": 2, + } + + +def test_scene_graph_layer_by_id_requires_table_root() -> None: + graph = SceneGraph(nodes=[], validate_on_refresh=False) + + with pytest.raises(ValueError, match="table node"): + graph.layer_by_id() + + +def test_scene_graph_rejects_table_region_for_non_table_parent() -> None: + with pytest.raises(ValueError, match="only valid for objects on the table"): + SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="spoon", + parent_id="plate", + parent_relation="on", + table_region="center", + ), + ], + ) + + +def test_scene_graph_derives_support_and_inverse_planar_constraints() -> None: + graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="cup", + parent_id="table", + parent_relation="on", + ), + ], + relations=[ + SceneGraphRelation( + source_id="plate", + relation="left_of", + target_id="cup", + ), + SceneGraphRelation( + source_id="plate", + relation="left_of", + target_id="cup", + ), + ], + ) + + constraints = graph.derive_constraints() + + assert constraints == [ + {"source_id": "plate", "relation": "on", "target_id": "table"}, + {"source_id": "cup", "relation": "on", "target_id": "table"}, + {"source_id": "plate", "relation": "left_of", "target_id": "cup"}, + {"source_id": "cup", "relation": "right_of", "target_id": "plate"}, + ] + + +def test_scene_graph_materializes_inverse_planar_relations() -> None: + graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="cup", + parent_id="table", + parent_relation="on", + ), + ], + relations=[ + SceneGraphRelation( + source_id="plate", + relation="left_of", + target_id="cup", + ), + ], + ) + + assert [relation.to_dict() for relation in graph.relations] == [ + {"source_id": "plate", "relation": "left_of", "target_id": "cup"}, + {"source_id": "cup", "relation": "right_of", "target_id": "plate"}, + ] + + +def test_scene_graph_to_dict_serializes_graph_state() -> None: + graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + table_region="center", + orientation_state="standing", + ), + ], + ) + + graph_dict = graph.to_dict() + + assert graph_dict == { + "nodes": [ + { + "object_id": "table", + "parent_id": None, + "parent_relation": None, + "table_region": None, + "orientation_state": None, + }, + { + "object_id": "plate", + "parent_id": "table", + "parent_relation": "on", + "table_region": "center", + "orientation_state": "standing", + }, + ], + "relations": [], + } diff --git a/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py new file mode 100644 index 000000000..b1710ee80 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py @@ -0,0 +1,138 @@ +# ---------------------------------------------------------------------------- +# 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 pathlib import Path + +import numpy as np +import trimesh + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, +) +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_constructor import ( + SceneLayoutConstructor, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.table_surface_layout_optimizer import ( + _table_region_bounds, +) + + +def _asset( + *, + object_id: str, + glb_path: Path, + center_xy: list[float] | None = None, + pos: list[float] | None = None, +) -> SceneObject: + return SceneObject( + id=object_id, + kind="asset", + category=object_id, + name=object_id, + description=object_id, + simready_glb_path=str(glb_path), + rot=[0.0, 0.0, 0.0], + pos=pos or [0.0, 0.0, 0.0], + scale=[1.0, 1.0, 1.0], + center_xy=center_xy, + ) + + +def test_table_regions_put_front_at_larger_y() -> None: + table_bounds = np.asarray([[0.0, 0.0], [3.0, 3.0]]) + + assert np.allclose( + _table_region_bounds( + table_bounds=table_bounds, + table_region="back_center", + ), + [[1.0, 0.0], [2.0, 1.0]], + ) + assert np.allclose( + _table_region_bounds( + table_bounds=table_bounds, + table_region="front_center", + ), + [[1.0, 2.0], [2.0, 3.0]], + ) + + +def test_layout_constructor_places_new_child_on_parent_top( + tmp_path: Path, +) -> None: + book_glb = tmp_path / "book.glb" + cup_glb = tmp_path / "cup.glb" + # SimReady GLBs are y-up, so the book's short vertical axis is y. + trimesh.creation.box(extents=[1.0, 0.2, 1.0]).export(book_glb) + trimesh.creation.box(extents=[0.2, 0.2, 0.2]).export(cup_glb) + + table = SceneObject( + id="table", + kind="table", + category="table", + name="table", + description="table", + support_surface_z=0.0, + support_optimization_rect_xy=[ + [-2.0, -2.0], + [2.0, -2.0], + [2.0, 2.0], + [-2.0, 2.0], + ], + ) + book = _asset( + object_id="book_001", + glb_path=book_glb, + center_xy=[0.0, 0.0], + # This y-up position maps to a z-up center at z=0.52 m. + pos=[0.0, 0.52, 0.0], + ) + cup = _asset(object_id="cup_001", glb_path=cup_glb) + graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="book_001", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="cup_001", + parent_id="book_001", + parent_relation="on", + ), + ] + ) + + post_edit_scene = SceneLayoutConstructor( + formal_scene=Scene(objects=[table, book]), + goal_scene_graph=graph, + layout_variable_ids={"cup_001"}, + generated_scene_objects=[cup], + output_root=tmp_path, + ).construct() + + placed_cup = next( + asset for asset in post_edit_scene.assets if asset.id == "cup_001" + ) + assert placed_cup.center_xy == [0.0, 0.0] + # Book top is z=0.62 m and the cup is placed directly on it. + assert np.allclose(placed_cup.pos, [0.0, 0.72, 0.0]) diff --git a/tests/gen_sim/scene_engine/test_scene_understanding.py b/tests/gen_sim/scene_engine/test_scene_understanding.py index 1fdb51f7b..ba1cc3fae 100644 --- a/tests/gen_sim/scene_engine/test_scene_understanding.py +++ b/tests/gen_sim/scene_engine/test_scene_understanding.py @@ -20,10 +20,16 @@ import json from pathlib import Path +from PIL import Image, ImageDraw import pytest from embodichain.gen_sim.scene_engine.core.scene import Scene -from embodichain.gen_sim.scene_engine.pipeline import scene_understanding +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.generation import scene_understanding +from embodichain.gen_sim.scene_engine.pipeline.utils import image_segmentation_utils +from embodichain.gen_sim.scene_engine.pipeline.utils.image_segmentation_utils import ( + render_asset_mask_id_overlay, +) def _response(*, asset_name: str = "cup") -> str: @@ -55,11 +61,23 @@ def test_image_object_analysis_parses_code_fence_and_assigns_stable_ids() -> Non assert [asset.id for asset in scene.assets] == ["cup_001"] -def test_image_object_analysis_rejects_location_words_in_object_names() -> None: - with pytest.raises(ValueError, match="must not contain location"): - scene_understanding._parse_image_object_analysis_response( - _response(asset_name="left cup") - ) +def test_image_object_analysis_accepts_name_with_spatial_words() -> None: + scene = scene_understanding._parse_image_object_analysis_response( + _response(asset_name="left cup") + ) + + assert scene.assets[0].name == "left cup" + + +def test_image_object_analysis_accepts_description_with_structural_words() -> None: + response = json.loads(_response()) + response["assets"][0]["description"] = "A small ceramic cup with a lid on top." + + scene = scene_understanding._parse_image_object_analysis_response( + json.dumps(response) + ) + + assert scene.assets[0].description == "A small ceramic cup with a lid on top." def test_image_object_analysis_retries_then_updates_scene(tmp_path: Path) -> None: @@ -83,3 +101,226 @@ def complete(self, **_: object) -> str: assert scene.table is not None assert [asset.id for asset in scene.assets] == ["cup_001"] + + +def test_asset_mask_id_overlay_excludes_the_table_mask(tmp_path: Path) -> None: + image_path = tmp_path / "scene.png" + table_mask_path = tmp_path / "table_mask.png" + asset_mask_path = tmp_path / "bottle_mask.png" + output_path = tmp_path / "asset_masks_with_ids.png" + image_size = (512, 512) + Image.new("RGB", image_size, "black").save(image_path) + + table_mask = Image.new("L", image_size, 0) + ImageDraw.Draw(table_mask).rectangle((10, 10, 100, 100), fill=255) + table_mask.save(table_mask_path) + asset_mask = Image.new("L", image_size, 0) + ImageDraw.Draw(asset_mask).rectangle((380, 180, 450, 360), fill=255) + asset_mask.save(asset_mask_path) + + rendered_path = render_asset_mask_id_overlay( + image_path=image_path, + asset_masks=[("bottle_001", asset_mask_path)], + output_path=output_path, + ) + + with Image.open(rendered_path) as overlay: + assert overlay.getpixel((10, 10)) == (0, 0, 0) + assert overlay.getpixel((377, 180)) != (0, 0, 0) + + +def test_asset_mask_id_label_font_fits_the_mask_bbox() -> None: + image_size = (512, 512) + mask_bbox = (380, 180, 450, 360) + label = "bottle_001" + font = image_segmentation_utils._load_asset_id_label_font( + image_size=image_size, + mask_bbox=mask_bbox, + label=label, + ) + label_bounds = image_segmentation_utils._number_label_bounds( + draw=ImageDraw.Draw(Image.new("RGBA", image_size)), + label=label, + center=(0.0, 0.0), + font=font, + minimum_padding=2, + ) + + assert label_bounds[2] - label_bounds[0] <= round( + (mask_bbox[2] - mask_bbox[0]) * 0.9 + ) + + +def test_initial_scene_graph_places_every_asset_on_table(tmp_path: Path) -> None: + class VLM: + def complete(self, **_: object) -> str: + return json.dumps( + { + "orientation_states": [ + {"object_id": "cup_001", "orientation_state": "lying"}, + ] + } + ) + + scene = Scene( + objects=[ + SceneObject( + id="table", + kind="table", + category="table", + name="wooden table", + description="A wooden table.", + ), + SceneObject( + id="cup_001", + kind="asset", + category="cup", + name="blue cup", + description="A blue cup.", + ), + ], + ) + + overlay_path = tmp_path / "asset_masks_with_ids.png" + overlay_path.write_bytes(b"png") + scene_graph = scene_understanding._initialize_scene_graph_from_segmented_scene( + scene, + asset_mask_id_overlay_path=overlay_path, + vlm_client=VLM(), # type: ignore[arg-type] + ) + + assert scene_graph.to_dict() == { + "nodes": [ + { + "object_id": "table", + "parent_id": None, + "parent_relation": None, + "table_region": None, + "orientation_state": None, + }, + { + "object_id": "cup_001", + "parent_id": "table", + "parent_relation": "on", + "table_region": None, + "orientation_state": "lying", + }, + ], + "relations": [], + } + + +def test_scene_graph_initialization_uses_container_orientation_states( + tmp_path: Path, +) -> None: + class VLM: + def complete(self, **_: object) -> str: + return json.dumps( + { + "orientation_states": [ + { + "object_id": "bottle_001", + "orientation_state": "standing", + }, + {"object_id": "book_001", "orientation_state": "lying"}, + ] + } + ) + + overlay_path = tmp_path / "asset_masks_with_ids.png" + overlay_path.write_bytes(b"png") + scene = Scene( + objects=[ + SceneObject( + id="table", + kind="table", + category="table", + name="wooden table", + description="A wooden table.", + ), + SceneObject( + id="bottle_001", + kind="asset", + category="bottle", + name="blue bottle", + description="A blue bottle.", + ), + SceneObject( + id="book_001", + kind="asset", + category="book", + name="red book", + description="A red book.", + ), + ] + ) + + scene_graph = scene_understanding._initialize_scene_graph_from_segmented_scene( + scene, + asset_mask_id_overlay_path=overlay_path, + vlm_client=VLM(), # type: ignore[arg-type] + ) + + assert scene_graph.node_by_id()["bottle_001"].orientation_state == "standing" + assert scene_graph.node_by_id()["book_001"].orientation_state == "lying" + + +def test_scene_graph_initialization_requires_asset_mask_id_overlay( + tmp_path: Path, +) -> None: + scene = Scene( + objects=[ + SceneObject( + id="table", + kind="table", + category="table", + name="wooden table", + description="A wooden table.", + ) + ] + ) + + with pytest.raises(FileNotFoundError, match="Image input not found"): + scene_understanding._initialize_scene_graph_from_segmented_scene( + scene, + asset_mask_id_overlay_path=tmp_path / "missing.png", + vlm_client=object(), # type: ignore[arg-type] + ) + + +def test_scene_graph_initialization_info_lists_existing_object_ids() -> None: + scene = Scene( + objects=[ + SceneObject( + id="table", + kind="table", + category="table", + name="wooden table", + description="A wooden table.", + ), + SceneObject( + id="bottle_001", + kind="asset", + category="bottle", + name="blue bottle", + description="A blue bottle.", + center_xy=[0.2, -0.1], + ), + SceneObject( + id="book_001", + kind="asset", + category="book", + name="red book", + description="A red book.", + center_xy=[-0.1, 0.2], + ), + ] + ) + + simplified_scene_info = ( + scene_understanding._simplify_scene_info_for_graph_initialization(scene=scene) + ) + + assert simplified_scene_info == { + "asset_ids": ["bottle_001", "book_001"], + } diff --git a/tests/gen_sim/scene_engine/test_simready_processor_utils.py b/tests/gen_sim/scene_engine/test_simready_processor_utils.py new file mode 100644 index 000000000..51459b56b --- /dev/null +++ b/tests/gen_sim/scene_engine/test_simready_processor_utils.py @@ -0,0 +1,41 @@ +# ---------------------------------------------------------------------------- +# 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 pathlib import Path + +import pytest +import trimesh + +from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor_utils import ( + compute_uniform_xy_scale_for_target, +) + + +def test_uniform_scale_uses_the_z_up_tabletop_footprint(tmp_path: Path) -> None: + """Measure y-up GLBs against the VLM's z-up XY target footprint.""" + glb_path = tmp_path / "flat_fork.glb" + # In y-up, the thin vertical axis is y; in z-up it becomes the z axis. + trimesh.creation.box(extents=[2.0, 0.01, 0.5]).export(glb_path) + + scale = compute_uniform_xy_scale_for_target( + glb_path=glb_path, + target_xy_size_cm=[200.0, 50.0], + rotate_about_x=False, + ) + + assert scale == pytest.approx(1.0) diff --git a/tests/gen_sim/task_engine/__init__.py b/tests/gen_sim/task_engine/__init__.py new file mode 100644 index 000000000..b201491d8 --- /dev/null +++ b/tests/gen_sim/task_engine/__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. +# ---------------------------------------------------------------------------- + +"""Tests for Task Engine semantics and orchestration.""" + +from __future__ import annotations diff --git a/tests/gen_sim/task_engine/orchestration/__init__.py b/tests/gen_sim/task_engine/orchestration/__init__.py new file mode 100644 index 000000000..626fe57e2 --- /dev/null +++ b/tests/gen_sim/task_engine/orchestration/__init__.py @@ -0,0 +1,17 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Tests for Task Engine cross-engine orchestration.""" diff --git a/tests/gen_sim/task_engine/orchestration/test_architecture.py b/tests/gen_sim/task_engine/orchestration/test_architecture.py new file mode 100644 index 000000000..ad3593c0a --- /dev/null +++ b/tests/gen_sim/task_engine/orchestration/test_architecture.py @@ -0,0 +1,90 @@ +# ---------------------------------------------------------------------------- +# 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 ast +from pathlib import Path + +import embodichain.gen_sim as gen_sim_package +from embodichain.gen_sim.action_engine.agent import ActionAgent +from embodichain.gen_sim.task_engine import TaskAgent +from embodichain.gen_sim.task_engine import __main__ as task_engine_main +from embodichain.gen_sim.task_engine import cli as task_engine_cli +from embodichain.gen_sim.task_engine.orchestration.coordinator import ( + TaskEngineCoordinator, +) +from embodichain.gen_sim.task_engine.orchestration.scene_adapter import SceneAdapter + +_GEN_SIM_ROOT = Path(gen_sim_package.__file__).resolve().parent +_PURE_TASK_MODULES = ( + "agent.py", + "config.py", + "contracts.py", + "interpretation.py", + "ontology.py", + "state_machine.py", + "workflow_contracts.py", +) + + +def test_task_semantic_core_does_not_import_scene_action_or_orchestration() -> None: + forbidden = { + "embodichain.gen_sim.action_engine", + "embodichain.gen_sim.scene_engine", + "embodichain.gen_sim.task_engine.orchestration", + "embodichain.gen_sim.task_engine.scene", + } + offenders: list[str] = [] + for filename in _PURE_TASK_MODULES: + path = _GEN_SIM_ROOT / "task_engine" / filename + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + modules = [alias.name for alias in node.names] + elif isinstance(node, ast.ImportFrom): + modules = [node.module or ""] + else: + continue + if any( + module == prefix or module.startswith(prefix + ".") + for module in modules + for prefix in forbidden + ): + offenders.append(filename) + break + assert offenders == [] + + +def test_cross_engine_owners_are_explicit() -> None: + assert TaskAgent.__module__ == "embodichain.gen_sim.task_engine.agent" + assert ActionAgent.__module__ == "embodichain.gen_sim.action_engine.agent" + assert SceneAdapter.__module__.startswith( + "embodichain.gen_sim.task_engine.orchestration" + ) + assert TaskEngineCoordinator.__module__.startswith( + "embodichain.gen_sim.task_engine.orchestration" + ) + + +def test_task_engine_owns_its_module_entry_point() -> None: + assert task_engine_main.main is task_engine_cli.main + + +def test_legacy_cross_engine_packages_are_deleted() -> None: + assert not (_GEN_SIM_ROOT / "scene_bridge").exists() + assert not (_GEN_SIM_ROOT / "collaboration").exists() + assert not (_GEN_SIM_ROOT / "action_engine" / "collaboration").exists() diff --git a/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py b/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py new file mode 100644 index 000000000..053b3e93f --- /dev/null +++ b/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py @@ -0,0 +1,1102 @@ +# ---------------------------------------------------------------------------- +# 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 json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from embodichain.gen_sim.task_engine import cli +from embodichain.gen_sim.task_engine import _bundle_runner as bundle_runner +from embodichain.gen_sim.task_engine.orchestration.artifacts import ( + ArtifactTransaction, +) +from embodichain.gen_sim.task_engine.orchestration.contracts import ( + BINDING_REPORT_SCHEMA, + ROLE_BINDINGS_SCHEMA, + SCENE_MANIFEST_SCHEMA, + SCENE_REQUEST_SCHEMA, + SUCCESS_SPEC_SCHEMA, + TASK_CANDIDATE_SET_SCHEMA, + TASK_DRAFT_SCHEMA, + canonical_hash, +) +from embodichain.gen_sim.task_engine.orchestration.coordinator import ( + TaskEngineCoordinator, +) +from embodichain.gen_sim.task_engine.orchestration.scene_adapter import ( + SceneAdaptation, +) +from embodichain.gen_sim.task_engine.orchestration.scene_source import SceneSourceRef +from embodichain.gen_sim.action_engine.generation.artifacts import artifact_paths +from embodichain.gen_sim.action_engine.generation.models import PreparedScene +from embodichain.gen_sim.action_engine.protocol import ( + AGENT_CONFIG_FILENAME, + FAST_GYM_CONFIG_FILENAME, +) +from embodichain.gen_sim.action_engine.runtime import ( + ExecutionReport, + build_execution_provenance, +) +from embodichain.gen_sim.task_engine.scene import SceneEngineV1Adapter + +_UPRIGHT_CAN_INSTRUCTION = "test-instruction" + + +def _candidate_set() -> dict: + selector = { + "kind": "scene_ref", + "step_id": "", + "reference": "red can", + "quantifier": "one", + "count": 0, + } + none_selector = { + "kind": "none", + "step_id": "", + "reference": "", + "quantifier": "one", + "count": 0, + } + step = { + "id": "upright", + "task_type": "E2", + "object": selector, + "target": none_selector, + "relation": "none", + "required_arm": "auto", + "transfer_arm": "none", + "receive_arm": "none", + "orientation_goal": "upright", + "target_state": "none", + "target_setting": 0, + "layout": "none", + "axis": "none", + "direction": "none", + "terminal_behavior": "none", + "depends_on": [], + } + draft = { + "schema_version": TASK_DRAFT_SCHEMA, + "task_id": "upright_can", + "instruction": _UPRIGHT_CAN_INSTRUCTION, + "steps": [step], + } + candidate = { + "candidate_id": "candidate_01", + "draft": draft, + "scene_request": { + "schema_version": SCENE_REQUEST_SCHEMA, + "task_id": "upright_can", + "references": [ + { + "reference_id": "upright.object", + "step_id": "upright", + "role": "object", + "reference": "red can", + "quantifier": "one", + "count": 0, + "source_structure": "rigid_object", + "affordances": ["graspable", "orientable"], + "initial_state": {"orientation": "fallen"}, + "attributes": {}, + } + ], + }, + "success_spec": { + "schema_version": SUCCESS_SPEC_SCHEMA, + "task_id": "upright_can", + "op": "all", + "terms": [{"step_id": "upright", "type": "object_upright"}], + }, + "semantic_hash": canonical_hash([step]), + "vote_count": 1, + "attempts": 1, + "normalizations": [], + } + return { + "schema_version": TASK_CANDIDATE_SET_SCHEMA, + "task_id": "upright_can", + "instruction": _UPRIGHT_CAN_INSTRUCTION, + "candidates": [candidate], + "requested_candidate_count": 1, + "valid_response_count": 1, + "errors": [], + } + + +def _candidate_set_with_alternative() -> dict: + candidates = _candidate_set() + alternative = deepcopy(candidates["candidates"][0]) + alternative["candidate_id"] = "candidate_02" + alternative["draft"]["steps"][0]["required_arm"] = "left_arm" + alternative["semantic_hash"] = canonical_hash(alternative["draft"]["steps"]) + candidates["candidates"].append(alternative) + candidates["requested_candidate_count"] = 2 + candidates["valid_response_count"] = 2 + return candidates + + +def _prepared_scene(tmp_path: Path) -> PreparedScene: + scene_path = tmp_path / "scene_config.json" + scene_path.write_text("{}", encoding="utf-8") + scene_object = { + "uid": "red_can", + "source_uid": "red_can", + "role": "rigid_object", + "name": "red can", + "description": "A red can.", + "category": "can", + "color": "red", + "position": [0.0, 0.0, 0.5], + "affordances": ["graspable", "orientable"], + "initial_state": {"orientation": "fallen"}, + "attributes": {}, + } + return PreparedScene( + source_config_path=scene_path, + scene_dir=tmp_path, + planner_objects=(scene_object,), + background=(), + rigid_objects=(), + articulations=(), + uid_map={"red_can": "red_can"}, + table_top_z=None, + z_rotation_degrees=0.0, + body_scale_policy="preserve", + body_scale=(1.0, 1.0, 1.0), + asset_hashes={}, + ) + + +def _adaptation(tmp_path: Path, *, status: str = "bound") -> SceneAdaptation: + candidates = _candidate_set() + candidate = candidates["candidates"][0] + selected_id = candidate["candidate_id"] if status == "bound" else "" + role_bindings = ( + { + "schema_version": ROLE_BINDINGS_SCHEMA, + "task_id": "upright_can", + "candidate_id": "candidate_01", + "reference_bindings": {"upright.object": ["red_can"]}, + "role_bindings": {}, + } + if status == "bound" + else None + ) + return SceneAdaptation( + scene_manifest={ + "schema_version": SCENE_MANIFEST_SCHEMA, + "scene_id": "scene", + "source_format": "test", + "robot_profile": "dual_franka", + "objects": [ + { + "uid": "red_can", + "role": "rigid_object", + "name": "red can", + "description": "A red can.", + "category": "can", + "color": "red", + "affordances": ["graspable", "orientable"], + "initial_state": {"orientation": "fallen"}, + "attributes": {}, + } + ], + }, + role_bindings=role_bindings, + binding_report={ + "schema_version": BINDING_REPORT_SCHEMA, + "task_id": "upright_can", + "status": status, + "selected_candidate_id": selected_id, + "selection_reason": "test", + "candidates": [ + { + "candidate_id": "candidate_01", + "semantic_hash": candidate["semantic_hash"], + "status": "resolved" if status == "bound" else status, + "references": [ + { + "reference_id": "upright.object", + "status": ( + "resolved" if status == "bound" else "ambiguous" + ), + "confidence": 1.0, + "candidate_uids": ["red_can"], + "selected_uids": (["red_can"] if status == "bound" else []), + "reasons": [], + } + ], + "reasons": [], + } + ], + }, + selected_candidate=deepcopy(candidate) if status == "bound" else None, + prepared_scene=_prepared_scene(tmp_path), + source_config_path=tmp_path / "scene_config.json", + conservative_scene_graph={ + "schema_version": "embodichain.conservative-scene-graph/v1", + "scene_id": "scene", + "nodes": [ + { + "uid": "red_can", + "parent_uid": "unknown", + "parent_relation": "unknown", + "orientation": "unknown", + "source": "test", + } + ], + "relations": [], + }, + ) + + +def _adaptation_with_alternative(tmp_path: Path) -> SceneAdaptation: + candidate_set = _candidate_set_with_alternative() + adaptation = _adaptation(tmp_path) + alternative = candidate_set["candidates"][1] + alternative_audit = deepcopy(adaptation.binding_report["candidates"][0]) + alternative_audit["candidate_id"] = "candidate_02" + alternative_audit["semantic_hash"] = alternative["semantic_hash"] + alternative_bindings = { + **deepcopy(adaptation.role_bindings), + "candidate_id": "candidate_02", + } + return replace( + adaptation, + binding_report={ + **deepcopy(adaptation.binding_report), + "candidates": [ + *deepcopy(adaptation.binding_report["candidates"]), + alternative_audit, + ], + }, + candidate_bindings={"candidate_02": alternative_bindings}, + ) + + +def test_artifact_transaction_rolls_back_and_preserves_existing_output( + tmp_path: Path, +) -> None: + output = tmp_path / "bundle" + output.mkdir() + (output / "kept.txt").write_text("old", encoding="utf-8") + + with pytest.raises(RuntimeError, match="fail"): + with ArtifactTransaction(output, overwrite=True) as transaction: + assert transaction.staging_dir is not None + (transaction.staging_dir / "partial.txt").write_text( + "partial", encoding="utf-8" + ) + raise RuntimeError("fail before commit") + + assert (output / "kept.txt").read_text(encoding="utf-8") == "old" + assert not (output / "partial.txt").exists() + + +def test_prepare_rejects_output_overlapping_read_only_source(tmp_path: Path) -> None: + source = tmp_path / "gym_project" + source.mkdir() + coordinator = TaskEngineCoordinator( + task_agent=object(), + scene_adapter=SimpleNamespace(robot_profile="franka"), + action_agent=object(), + feasibility_broker=object(), + ) + + with pytest.raises(ValueError, match="must not overlap"): + coordinator.prepare( + "task", + "Pick up the object.", + source, + source / "task_run", + overwrite=True, + ) + + +def test_unbound_prepare_publishes_only_audit_artifacts(tmp_path: Path) -> None: + candidates = _candidate_set() + adaptation = _adaptation(tmp_path, status="ambiguous") + task_agent = SimpleNamespace(generate=lambda *args, **kwargs: candidates) + scene_adapter = SimpleNamespace( + robot_profile="franka", + adapt=lambda *args, **kwargs: adaptation, + ) + action_agent = SimpleNamespace( + plan=lambda *_args, **_kwargs: pytest.fail("Action Agent must not run") + ) + coordinator = TaskEngineCoordinator( + task_agent=task_agent, + scene_adapter=scene_adapter, + action_agent=action_agent, + bundle_generator=lambda *_args, **_kwargs: pytest.fail( + "legacy generator must not run" + ), + ) + + result = coordinator.prepare( + "upright_can", + _UPRIGHT_CAN_INSTRUCTION, + tmp_path / "scene_config.json", + tmp_path / "bundle", + candidate_count=1, + ) + + assert result.status == "ambiguous" + assert (result.output_dir / "task_candidate_set.json").is_file() + assert (result.output_dir / "binding_report.json").is_file() + assert not (result.output_dir / "scene_manifest.json").exists() + assert not (result.output_dir / "role_bindings.json").exists() + assert not (result.output_dir / "grounded_task_plan.json").exists() + assert not (result.output_dir / FAST_GYM_CONFIG_FILENAME).exists() + + +def test_prepare_reuses_precomputed_candidates_without_rerunning_task_agent( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + adaptation = _adaptation(tmp_path, status="ambiguous") + task_agent = SimpleNamespace( + generate=lambda *args, **kwargs: pytest.fail("Task Agent must not rerun") + ) + coordinator = TaskEngineCoordinator( + task_agent=task_agent, + scene_adapter=SimpleNamespace( + robot_profile="franka", + adapt=lambda *args, **kwargs: adaptation, + ), + action_agent=object(), + ) + + result = coordinator.prepare( + "upright_can", + _UPRIGHT_CAN_INSTRUCTION, + tmp_path / "scene_config.json", + tmp_path / "candidate-reuse", + candidate_set=candidates, + force_most_likely=True, + ) + + assert result.status == "ambiguous" + assert result.candidate_set == candidates + + +def test_prepare_inherits_adapter_robot_profile_for_raw_scene_path( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + adaptation = _adaptation(tmp_path, status="ambiguous") + captured: dict[str, object] = {} + + def adapt(_candidates, source, **_kwargs): + captured["source"] = source + return adaptation + + coordinator = TaskEngineCoordinator( + task_agent=SimpleNamespace( + generate=lambda *args, **kwargs: pytest.fail("Task Agent must not rerun") + ), + scene_adapter=SimpleNamespace(robot_profile="ur10", adapt=adapt), + action_agent=object(), + ) + + result = coordinator.prepare( + "upright_can", + _UPRIGHT_CAN_INSTRUCTION, + tmp_path / "scene_config.json", + tmp_path / "ur10-bundle", + candidate_set=candidates, + ) + + assert result.status == "ambiguous" + assert isinstance(captured["source"], SceneSourceRef) + assert captured["source"].robot_profile == "ur10" + + +def test_contradicted_feasibility_publishes_audit_without_planning( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + adaptation = _adaptation(tmp_path) + static_manifest = SceneEngineV1Adapter().adapt_prepared_scene( + adaptation.prepared_scene, + source_format="test", + robot_profile="dual_franka", + ) + adaptation = replace( + adaptation, + static_scene_manifest=static_manifest, + ) + task_agent = SimpleNamespace(generate=lambda *args, **kwargs: candidates) + scene_adapter = SimpleNamespace( + robot_profile="franka", + adapt=lambda *args, **kwargs: adaptation, + ) + registry = SimpleNamespace( + catalog=lambda: { + name: { + "runtime_available": name != "PickUp", + "unavailable_reason": ( + "PickUp disabled for test." if name == "PickUp" else None + ), + } + for name in ("PickUp", "MoveHeldObject", "Place") + } + ) + action_agent = SimpleNamespace( + registry=registry, + plan=lambda *_args, **_kwargs: pytest.fail("Action Agent must not plan"), + ) + + result = TaskEngineCoordinator( + task_agent=task_agent, + scene_adapter=scene_adapter, + action_agent=action_agent, + bundle_generator=lambda *_args, **_kwargs: pytest.fail( + "legacy generator must not run" + ), + ).prepare( + "upright_can", + _UPRIGHT_CAN_INSTRUCTION, + tmp_path / "scene_config.json", + tmp_path / "infeasible-bundle", + candidate_count=1, + ) + + assert result.status == "infeasible" + assert result.feasibility_report is not None + assert result.feasibility_report["status"] == "contradicted" + assert result.feasibility_report["remediation_class"] == "action_capability" + assert result.artifacts.static_scene_manifest.is_file() + assert result.artifacts.feasibility_report.is_file() + assert not result.artifacts.grounded_task_plan.exists() + + +def test_feasibility_contradiction_falls_back_to_next_resolved_candidate( + tmp_path: Path, +) -> None: + candidate_set = _candidate_set() + first = candidate_set["candidates"][0] + second = deepcopy(first) + second["candidate_id"] = "candidate_02" + second["semantic_hash"] = "b" * 64 + candidate_set["candidates"].append(second) + adaptation = _adaptation(tmp_path) + second_audit = deepcopy(adaptation.binding_report["candidates"][0]) + second_audit["candidate_id"] = "candidate_02" + second_audit["semantic_hash"] = "b" * 64 + binding_report = deepcopy(adaptation.binding_report) + binding_report["candidates"].append(second_audit) + second_bindings = { + **deepcopy(adaptation.role_bindings), + "candidate_id": "candidate_02", + } + adaptation = replace( + adaptation, + binding_report=binding_report, + candidate_bindings={"candidate_02": second_bindings}, + static_scene_manifest={}, + ) + + class _Broker: + @staticmethod + def assess(candidate, *_args, **_kwargs): + return { + "status": ( + "runtime_probe" + if candidate["candidate_id"] == "candidate_02" + else "contradicted" + ) + } + + registry = SimpleNamespace(catalog=lambda: {}) + coordinator = TaskEngineCoordinator( + action_agent=SimpleNamespace(registry=registry), + feasibility_broker=_Broker(), + ) + + updated, selected, bindings, report = coordinator._fallback_feasible_candidate( + candidate_set, + adaptation, + first, + adaptation.role_bindings, + {"status": "contradicted"}, + ) + + assert selected["candidate_id"] == "candidate_02" + assert bindings["candidate_id"] == "candidate_02" + assert report["status"] == "runtime_probe" + assert updated.binding_report["selected_candidate_id"] == "candidate_02" + assert ( + "static feasibility contradicted candidate_01" + in updated.binding_report["selection_reason"] + ) + + +def test_bound_prepare_uses_sidecar_and_publishes_complete_bundle( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + adaptation = _adaptation(tmp_path) + task_agent = SimpleNamespace(generate=lambda *args, **kwargs: candidates) + scene_adapter = SimpleNamespace( + robot_profile="franka", + adapt=lambda *args, **kwargs: adaptation, + ) + graph = {"graph": "planned"} + action_agent = SimpleNamespace(plan=lambda _plan: deepcopy(graph)) + generator_calls = [] + + def generator(_scene, output, **kwargs): + generator_calls.append(kwargs) + task_spec_path = Path(kwargs["task_spec"]) + assert task_spec_path.is_file() + assert (task_spec_path.parent / "scene_requirements.json").is_file() + paths = artifact_paths(output) + for path in ( + paths.gym_config, + paths.agent_config, + paths.task_spec, + paths.scene_requirements, + paths.seed_task_graph, + ): + path.parent.mkdir(parents=True, exist_ok=True) + value = graph if path == paths.seed_task_graph else {} + path.write_text(json.dumps(value), encoding="utf-8") + paths.seed_task_graph_png.write_bytes(b"png") + return paths + + result = TaskEngineCoordinator( + task_agent=task_agent, + scene_adapter=scene_adapter, + action_agent=action_agent, + bundle_generator=generator, + ).prepare( + "upright_can", + _UPRIGHT_CAN_INSTRUCTION, + tmp_path / "scene_config.json", + tmp_path / "bundle", + candidate_count=1, + ) + + assert result.bound + assert generator_calls + assert not (result.output_dir / ".task_engine_input").exists() + grounded = json.loads( + (result.output_dir / "grounded_task_plan.json").read_text(encoding="utf-8") + ) + assert grounded["success_spec"]["terms"] == [ + {"step_id": "task_01", "type": "object_upright"} + ] + assert (result.output_dir / "seed_task_graph.json").is_file() + + +def test_prepare_falls_back_after_candidate_action_planning_failure( + tmp_path: Path, +) -> None: + candidates = _candidate_set_with_alternative() + adaptation = _adaptation_with_alternative(tmp_path) + planned_candidates: list[str] = [] + graph = {"graph": "planned"} + + def plan(grounded_plan): + candidate_id = grounded_plan["selected_candidate_id"] + planned_candidates.append(candidate_id) + if candidate_id == "candidate_01": + raise ValueError( + "SeedGraph TaskGroup 'task_04' requires unavailable state " + "{'predicate': 'arm_free', 'arm': 'right_arm'}." + ) + return deepcopy(graph) + + def generator(_scene, output, **_kwargs): + paths = artifact_paths(output) + for path in ( + paths.gym_config, + paths.agent_config, + paths.task_spec, + paths.scene_requirements, + paths.seed_task_graph, + ): + path.parent.mkdir(parents=True, exist_ok=True) + value = graph if path == paths.seed_task_graph else {} + path.write_text(json.dumps(value), encoding="utf-8") + paths.seed_task_graph_png.write_bytes(b"png") + return paths + + result = TaskEngineCoordinator( + task_agent=SimpleNamespace(generate=lambda *args, **kwargs: candidates), + scene_adapter=SimpleNamespace( + robot_profile="franka", + adapt=lambda *args, **kwargs: adaptation, + ), + action_agent=SimpleNamespace(plan=plan), + bundle_generator=generator, + ).prepare( + "upright_can", + _UPRIGHT_CAN_INSTRUCTION, + tmp_path / "scene_config.json", + tmp_path / "fallback-bundle", + candidate_count=2, + ) + + assert result.bound + assert result.selected_candidate_id == "candidate_02" + assert planned_candidates == ["candidate_01", "candidate_02"] + assert ( + "candidate_01 failed action_planning" + in result.adaptation.binding_report["selection_reason"] + ) + assert not result.artifacts.preparation_failure.exists() + + +def test_prepare_publishes_failure_context_when_all_candidates_fail_planning( + tmp_path: Path, +) -> None: + candidates = _candidate_set_with_alternative() + adaptation = _adaptation_with_alternative(tmp_path) + output = tmp_path / "failed-bundle" + output.mkdir() + (output / "stale.txt").write_text("old", encoding="utf-8") + + result = TaskEngineCoordinator( + task_agent=SimpleNamespace(generate=lambda *args, **kwargs: candidates), + scene_adapter=SimpleNamespace( + robot_profile="franka", + adapt=lambda *args, **kwargs: adaptation, + ), + action_agent=SimpleNamespace( + plan=lambda _plan: (_ for _ in ()).throw( + ValueError( + "SeedGraph TaskGroup 'task_04' requires unavailable state " + "{'predicate': 'arm_free', 'arm': 'right_arm'}." + ) + ) + ), + bundle_generator=lambda *_args, **_kwargs: pytest.fail( + "bundle generation must not run" + ), + ).prepare( + "upright_can", + _UPRIGHT_CAN_INSTRUCTION, + tmp_path / "scene_config.json", + output, + candidate_count=2, + overwrite=True, + ) + + assert result.status == "planning_failed" + assert not result.bound + assert result.artifacts.preparation_failure.is_file() + assert not (result.output_dir / "stale.txt").exists() + failure = json.loads( + result.artifacts.preparation_failure.read_text(encoding="utf-8") + ) + assert failure["schema_version"] == "action_engine_preparation_failure_v1" + assert failure["task_id"] == "upright_can" + assert failure["selected_candidate_id"] == "candidate_01" + assert [attempt["candidate_id"] for attempt in failure["attempts"]] == [ + "candidate_01", + "candidate_02", + ] + for index, attempt in enumerate(failure["attempts"]): + candidate_id = f"candidate_{index + 1:02d}" + assert attempt["stage"] == "action_planning" + assert attempt["draft"] == candidates["candidates"][index]["draft"] + assert attempt["bindings"]["candidate_id"] == candidate_id + assert attempt["grounded_task_plan"]["selected_candidate_id"] == candidate_id + assert "unbound_action_plan" in attempt + assert "action_graph" in attempt + assert attempt["error"]["type"] == "ValueError" + assert "arm_free" in attempt["error"]["message"] + + +def test_private_bundle_runner_forwards_arguments_without_leaking_sys_argv( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + bundle = tmp_path / "bundle" + bundle.mkdir() + (bundle / AGENT_CONFIG_FILENAME).write_text( + json.dumps({"task_name": "task"}), encoding="utf-8" + ) + (bundle / FAST_GYM_CONFIG_FILENAME).write_text("{}", encoding="utf-8") + captured = [] + + def fake_cli() -> None: + import sys + + captured.append(list(sys.argv)) + + import embodichain.gen_sim.action_engine.cli as legacy_cli + + monkeypatch.setattr( + legacy_cli, + "run_agent", + SimpleNamespace(cli=fake_cli), + raising=False, + ) + import sys + + original = sys.argv + assert bundle_runner.main(["--bundle", str(bundle), "--seed", "7"]) == 0 + + assert sys.argv is original + assert captured[0][-2:] == ["--seed", "7"] + assert str(bundle / AGENT_CONFIG_FILENAME) in captured[0] + + +@pytest.mark.parametrize( + ("mode", "image", "scene", "edit"), + [ + ("image", "input.png", None, None), + ("image-edit", "input.png", None, "move the cup left"), + ("scene", None, "gym_project", None), + ("scene-edit", None, "gym_project", "move the cup left"), + ], +) +def test_unified_cli_accepts_exactly_four_modes( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + mode: str, + image: str | None, + scene: str | None, + edit: str | None, +) -> None: + captured = {} + + class FakeWorkflow: + def __init__(self, **_kwargs) -> None: + pass + + def run(self, request, **kwargs): + captured["request"] = request + captured["kwargs"] = kwargs + output = Path(request["output_dir"]) + return SimpleNamespace( + status="succeeded", + succeeded=True, + failure_class=None, + output_dir=output, + manifest_path=output / "run_manifest.json", + final_bundle=output / "final" / "bundle", + ) + + monkeypatch.setattr(cli, "SceneAdapter", lambda **_kwargs: object()) + monkeypatch.setattr(cli, "TaskEngineWorkflow", FakeWorkflow) + arguments = [ + "--mode", + mode, + "--task-id", + "task", + "--instruction", + "place the cup", + "--output-root", + str(tmp_path / "history"), + "--base-seed", + "9", + ] + if image is not None: + arguments.extend(["--image", str(tmp_path / image)]) + if scene is not None: + arguments.extend(["--scene", str(tmp_path / scene)]) + if edit is not None: + arguments.extend(["--scene-edit", edit]) + if mode == "image": + arguments.append("--dataset_saving") + + assert cli.main(arguments) == 0 + + request = captured["request"] + assert request["image_path"] == (None if image is None else str(tmp_path / image)) + assert request["gym_project"] == (None if scene is None else str(tmp_path / scene)) + assert request["scene_edit_prompt"] == edit + assert captured["kwargs"]["base_seed"] == 9 + assert captured["kwargs"]["dataset_saving"] is (mode == "image") + assert captured["kwargs"]["execute"] is True + payload = json.loads(capsys.readouterr().out) + assert payload["status"] == "succeeded" + assert payload["run_id"].replace("_", "").isdigit() + assert len(payload["run_id"]) == 15 + assert Path(payload["output_dir"]).parent == tmp_path / "history" + + +def test_unified_cli_reuses_history_root_without_modifying_prior_scene( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + history = tmp_path / "task1008" + source = ( + history + / "20260820_105939" + / "attempts" + / "scene_0001" + / "scene_revision" + / "scene_export" + ) + source.mkdir(parents=True) + marker = source / "scene_config.json" + marker.write_text('{"source": "unchanged"}\n', encoding="utf-8") + captured = {} + + class FakeWorkflow: + def __init__(self, **_kwargs) -> None: + pass + + def run(self, request, **_kwargs): + captured["request"] = request + output = Path(request["output_dir"]) + return SimpleNamespace( + status="succeeded", + succeeded=True, + failure_class=None, + output_dir=output, + manifest_path=output / "run_manifest.json", + final_bundle=output / "final" / "bundle", + ) + + monkeypatch.setattr(cli, "SceneAdapter", lambda **_kwargs: object()) + monkeypatch.setattr(cli, "TaskEngineWorkflow", FakeWorkflow) + + assert ( + cli.main( + [ + "--mode", + "scene", + "--task-id", + "task1008", + "--scene", + str(source), + "--instruction", + "place the cup on the book", + "--output-root", + str(history), + ] + ) + == 0 + ) + + output_dir = Path(captured["request"]["output_dir"]) + assert output_dir.parent == history + assert output_dir != source + assert marker.read_text(encoding="utf-8") == '{"source": "unchanged"}\n' + assert list(history.glob(".*.reserve")) == [] + + +def test_unified_cli_rejects_history_root_inside_source_before_reservation( + tmp_path: Path, +) -> None: + source = tmp_path / "scene_export" + source.mkdir() + output_root = source / "new_runs" + + with pytest.raises(ValueError, match="read-only source"): + cli.main( + [ + "--mode", + "scene", + "--task-id", + "task", + "--scene", + str(source), + "--instruction", + "place the cup", + "--output-root", + str(output_root), + ] + ) + + assert not output_root.exists() + + +def test_unified_cli_rejects_mode_input_mismatch(tmp_path: Path) -> None: + with pytest.raises(SystemExit, match="2"): + cli.main( + [ + "--mode", + "image", + "--task-id", + "task", + "--instruction", + "place the cup", + "--image", + str(tmp_path / "input.png"), + "--scene", + str(tmp_path / "scene"), + "--output-root", + str(tmp_path / "history"), + ] + ) + + +def test_public_cli_exposes_prepare_run_and_run_all_modes() -> None: + parser = cli.build_parser() + help_text = parser.format_help() + assert "prepare" in help_text + assert "run-all" in help_text + assert "run" in help_text + assert "--overwrite" not in help_text + assert "--run-after-prepare" not in help_text + arguments = parser.parse_args( + [ + "prepare", + "--mode", + "image", + "--task-id", + "task", + "--instruction", + "place the cup", + "--image", + "input.png", + "--output-root", + "history", + "--dataset_saving", + ] + ) + assert arguments.command == "prepare" + assert arguments.dataset_saving is True + + +def test_prepare_cli_stops_before_simulator_execution( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured = {} + + class FakeWorkflow: + def __init__(self, **_kwargs) -> None: + pass + + def run(self, request, **kwargs): + captured.update(kwargs) + output = Path(request["output_dir"]) + return SimpleNamespace( + status="prepared", + succeeded=False, + failure_class=None, + output_dir=output, + manifest_path=output / "run_manifest.json", + final_bundle=output / "final" / "bundle", + ) + + monkeypatch.setattr(cli, "SceneAdapter", lambda **_kwargs: object()) + monkeypatch.setattr(cli, "TaskEngineWorkflow", FakeWorkflow) + + result = cli.main( + [ + "prepare", + "--mode", + "image", + "--task-id", + "task", + "--instruction", + "place the cup", + "--image", + str(tmp_path / "input.png"), + "--output-root", + str(tmp_path / "history"), + ] + ) + + assert result == 0 + assert captured["execute"] is False + + +def test_run_cli_executes_an_existing_bundle( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + bundle = tmp_path / "bundle" + bundle.mkdir() + + class Executor: + def __call__(self, _bundle, output, **kwargs): + Path(output).mkdir() + assert kwargs["num_envs"] == 2 + return { + "status": "failed", + "environments": [ + {"success": True}, + {"success": False}, + ], + } + + monkeypatch.setattr(cli, "SubprocessActionExecutor", Executor) + + result = cli.main( + [ + "run", + "--bundle", + str(bundle), + "--output-root", + str(tmp_path / "history"), + "--num-envs", + "2", + ] + ) + + assert result == 0 + + +def test_private_bundle_runner_publishes_rejected_preflight_report( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + bundle = tmp_path / "bundle" + bundle.mkdir() + (bundle / AGENT_CONFIG_FILENAME).write_text( + json.dumps({"task_name": "task"}), encoding="utf-8" + ) + (bundle / FAST_GYM_CONFIG_FILENAME).write_text("{}", encoding="utf-8") + report = ExecutionReport( + task_id="task", + plan_hash="0" * 64, + action_graph_hash="1" * 64, + status="rejected", + run_id="preflight", + episode_id="0", + provenance=build_execution_provenance(), + environments=( + { + "env_id": "0", + "success": False, + "semantic_success": {}, + "action_count": 0, + "retry_count": 0, + "recovery_count": 0, + "revision_count": 0, + "failures": [], + }, + ), + error="ValueError: planning-only action", + ) + monkeypatch.setattr( + bundle_runner, + "_preflight_bundle", + lambda *args, **kwargs: report, + ) + + assert bundle_runner.main(["--bundle", str(bundle)]) == 2 + payload = json.loads((bundle / "execution_report.json").read_text(encoding="utf-8")) + assert payload["status"] == "rejected" + assert payload["action_count"] == 0 diff --git a/tests/gen_sim/task_engine/orchestration/test_legacy_scene.py b/tests/gen_sim/task_engine/orchestration/test_legacy_scene.py new file mode 100644 index 000000000..286441fb0 --- /dev/null +++ b/tests/gen_sim/task_engine/orchestration/test_legacy_scene.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. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import json +from pathlib import Path + +import trimesh + +from embodichain.gen_sim.action_engine.generation.source_scene import prepare_scene +from embodichain.gen_sim.task_engine.orchestration.legacy_scene import ( + convert_legacy_gym_project, + restore_locked_scene_entities, +) +from embodichain.gen_sim.task_engine.orchestration.scene_source import ( + fingerprint_scene_source, + scene_revision_id, +) +from embodichain.gen_sim.task_engine.scene import build_conservative_scene_graph + + +def _legacy_project(tmp_path: Path) -> Path: + project = tmp_path / "legacy" + assets = project / "assets" + assets.mkdir(parents=True) + trimesh.creation.box(extents=[1.0, 1.0, 0.1]).export( + assets / "table.glb", file_type="glb" + ) + trimesh.creation.cylinder(radius=0.03, height=0.12).export( + assets / "can.glb", file_type="glb" + ) + trimesh.creation.box(extents=[0.3, 0.2, 0.4]).export( + assets / "cabinet.glb", file_type="glb" + ) + (assets / "cabinet.urdf").write_text( + '' + '\n', + encoding="utf-8", + ) + config = { + "background": [ + { + "uid": "table_0", + "name": "table", + "description": "A work table.", + "category": "table", + "shape": {"shape_type": "Mesh", "fpath": "assets/table.glb"}, + "init_pos": [0.0, 0.0, 0.0], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + } + ], + "rigid_object": [ + { + "uid": "can_0", + "name": "red can", + "description": "A red can.", + "category": "can", + "shape": {"shape_type": "Mesh", "fpath": "assets/can.glb"}, + "init_pos": [0.0, 0.1, 0.2], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + } + ], + "articulation": [ + { + "uid": "cabinet_0", + "name": "cabinet", + "description": "A fixed cabinet.", + "category": "cabinet", + "fpath": "assets/cabinet.urdf", + "init_pos": [0.5, 0.0, 0.0], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + } + ], + } + (project / "gym_config.json").write_text(json.dumps(config), encoding="utf-8") + return project + + +def test_legacy_conversion_is_read_only_and_restores_locked_articulation( + tmp_path: Path, +) -> None: + project = _legacy_project(tmp_path) + original = fingerprint_scene_source(project) + + revision = convert_legacy_gym_project(project, tmp_path / "revision") + converted = json.loads(revision.scene_config_path.read_text(encoding="utf-8")) + manifest = json.loads(revision.manifest_path.read_text(encoding="utf-8")) + + assert fingerprint_scene_source(project) == original + assert converted["format"] == "embodichain.scene-export/v1" + assert converted["background"][0]["uid"] == "table" + assert converted["rigid_object"][0]["uid"] == "can" + assert converted["articulation"][0]["uid"] == "cabinet" + assert manifest["audit_hierarchy"] == "unknown" + assert manifest["operational_hierarchy"] == "assumed_on_table" + assert set(revision.locked_entity_uids) == {"table", "cabinet"} + + converted["articulation"] = [] + revision.scene_config_path.write_text(json.dumps(converted), encoding="utf-8") + restore_locked_scene_entities(revision.output_root) + restored = json.loads(revision.scene_config_path.read_text(encoding="utf-8")) + + assert restored["articulation"][0]["uid"] == "cabinet" + assert Path(restored["articulation"][0]["fpath"]).is_file() + assert fingerprint_scene_source(project) == original + + +def test_legacy_conversion_separates_audit_and_operational_hierarchy( + tmp_path: Path, +) -> None: + revision = convert_legacy_gym_project( + _legacy_project(tmp_path), + tmp_path / "revision", + ) + + operational = json.loads(revision.scene_graph_path.read_text(encoding="utf-8")) + conservative = build_conservative_scene_graph( + prepare_scene(revision.scene_config_path), + scene_id="legacy-scene", + ) + + operational_can = next( + node for node in operational["nodes"] if node["object_id"] == "can" + ) + conservative_can = next( + node for node in conservative["nodes"] if node["uid"] == "can" + ) + assert operational_can["parent_id"] == "table" + assert operational_can["parent_relation"] == "on" + assert conservative_can["parent_uid"] == "unknown" + assert conservative_can["parent_relation"] == "unknown" + assert conservative_can["source"] == "conservative_import" + + +def test_scene_identity_covers_transitive_urdf_meshes(tmp_path: Path) -> None: + project = _legacy_project(tmp_path) + original_fingerprint = fingerprint_scene_source(project) + original_revision = scene_revision_id(project) + + trimesh.creation.box(extents=[0.6, 0.2, 0.4]).export( + project / "assets" / "cabinet.glb", + file_type="glb", + ) + + changed_fingerprint = fingerprint_scene_source(project) + assert changed_fingerprint.asset_sha256 != original_fingerprint.asset_sha256 + assert scene_revision_id(project) != original_revision diff --git a/tests/gen_sim/task_engine/orchestration/test_scene_adapter.py b/tests/gen_sim/task_engine/orchestration/test_scene_adapter.py new file mode 100644 index 000000000..2ae3b749f --- /dev/null +++ b/tests/gen_sim/task_engine/orchestration/test_scene_adapter.py @@ -0,0 +1,780 @@ +# ---------------------------------------------------------------------------- +# 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 pathlib import Path + +import pytest + +import embodichain.gen_sim.task_engine.orchestration.scene_adapter as scene_adapter_module +from embodichain.gen_sim.task_engine.contracts import ( + SCENE_REQUEST_SCHEMA, + SUCCESS_SPEC_SCHEMA, + TASK_CANDIDATE_SET_SCHEMA, + TASK_DRAFT_SCHEMA, + canonical_hash, +) +from embodichain.gen_sim.task_engine.orchestration.scene_adapter import ( + SceneAdapter, + SceneAdapterProtocolError, +) +from embodichain.gen_sim.task_engine.orchestration.scene_source import ( + SceneSourceRef, + fingerprint_scene_source, + verify_scene_source_fingerprint, +) +from embodichain.gen_sim.task_engine.agent import ( + derive_scene_request, + derive_success_spec, +) + +_UPRIGHT_CAN_INSTRUCTION = "test-instruction" + + +@pytest.fixture +def scene_export(tmp_path: Path) -> Path: + export = tmp_path / "scene_export" + assets = export / "meshes" + assets.mkdir(parents=True) + for name in ("table", "red_can", "blue_can"): + (assets / f"{name}.glb").write_bytes(f"mesh:{name}".encode()) + config = { + "format": "embodichain.scene-export/v1", + "scene_id": "2026-03-18T10:20:30Z", + "background": [ + { + "uid": "table", + "name": "table", + "description": "A work table.", + "category": "table", + "affordances": ["support_surface"], + "shape": {"shape_type": "Mesh", "fpath": "meshes/table.glb"}, + "init_pos": [0.0, 0.0, 0.0], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + } + ], + "rigid_object": [ + { + "uid": f"{color}_can", + "name": f"{color} can", + "description": f"A {color} soda can.", + "category": "can", + "attributes": { + "color": color, + "geometry": {"position": [1.0, 2.0, 3.0]}, + }, + "affordances": ["graspable", "orientable", "placeable"], + "initial_state": {"orientation": "fallen"}, + "shape": { + "shape_type": "Mesh", + "fpath": f"meshes/{color}_can.glb", + }, + "init_pos": [0.0, offset, 0.7], + "init_rot": [0.0, 0.0, 90.0], + "body_scale": [1.0, 1.0, 1.0], + } + for color, offset in (("red", 0.2), ("blue", -0.2)) + ], + } + (export / "scene_config.json").write_text(json.dumps(config), encoding="utf-8") + return export + + +def _legacy_gym_project(tmp_path: Path, filename: str) -> Path: + project = tmp_path / filename.removesuffix(".json") + assets = project / "assets" + assets.mkdir(parents=True) + for name in ("table.glb", "red_can.glb", "cabinet.urdf"): + (assets / name).write_bytes(f"asset:{name}".encode()) + config = { + "background": [ + { + "uid": "table_0", + "name": "table", + "description": "A work table.", + "category": "table", + "shape": {"shape_type": "Mesh", "fpath": "assets/table.glb"}, + "init_pos": [0.0, 0.0, 0.0], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + } + ], + "rigid_object": [ + { + "uid": "red_can_0", + "name": "red can", + "description": "A red soda can.", + "category": "can", + "affordances": ["graspable", "orientable", "placeable"], + "initial_state": {"orientation": "fallen"}, + "shape": { + "shape_type": "Mesh", + "fpath": "assets/red_can.glb", + }, + "init_pos": [0.0, 0.2, 0.7], + "init_rot": [0.0, 0.0, 90.0], + "body_scale": [1.0, 1.0, 1.0], + } + ], + "articulation": [ + { + "uid": "cabinet_0", + "name": "cabinet", + "description": "A fixed articulated cabinet.", + "category": "cabinet", + "fpath": "assets/cabinet.urdf", + "init_pos": [0.4, 0.0, 0.0], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + } + ], + } + (project / filename).write_text(json.dumps(config), encoding="utf-8") + return project + + +def _selector(reference: str) -> dict: + return { + "kind": "scene_ref", + "step_id": "", + "reference": reference, + "quantifier": "one", + "count": 0, + } + + +def _none_selector() -> dict: + return { + "kind": "none", + "step_id": "", + "reference": "", + "quantifier": "one", + "count": 0, + } + + +def _candidate(candidate_id: str, reference: str, *, votes: int = 1) -> dict: + step = { + "id": "upright", + "task_type": "E2", + "object": _selector(reference), + "target": _none_selector(), + "relation": "none", + "required_arm": "auto", + "transfer_arm": "none", + "receive_arm": "none", + "orientation_goal": "upright", + "target_state": "none", + "target_setting": 0, + "layout": "none", + "axis": "none", + "direction": "none", + "terminal_behavior": "none", + "depends_on": [], + } + draft = { + "schema_version": TASK_DRAFT_SCHEMA, + "task_id": "upright_can", + "instruction": _UPRIGHT_CAN_INSTRUCTION, + "steps": [step], + } + return { + "candidate_id": candidate_id, + "draft": draft, + "scene_request": { + "schema_version": SCENE_REQUEST_SCHEMA, + "task_id": "upright_can", + "references": [ + { + "reference_id": "upright.object", + "step_id": "upright", + "role": "object", + "reference": reference, + "quantifier": "one", + "count": 0, + "source_structure": "rigid_object", + "affordances": ["graspable", "orientable"], + "initial_state": {"orientation": "fallen"}, + "attributes": {}, + } + ], + }, + "success_spec": { + "schema_version": SUCCESS_SPEC_SCHEMA, + "task_id": "upright_can", + "op": "all", + "terms": [{"step_id": "upright", "type": "object_upright"}], + }, + "semantic_hash": canonical_hash([step]), + "vote_count": votes, + "attempts": 1, + "normalizations": [], + } + + +def _candidate_set(candidates: list[dict]) -> dict: + return { + "schema_version": TASK_CANDIDATE_SET_SCHEMA, + "task_id": "upright_can", + "instruction": _UPRIGHT_CAN_INSTRUCTION, + "candidates": candidates, + "requested_candidate_count": sum(item["vote_count"] for item in candidates), + "valid_response_count": sum(item["vote_count"] for item in candidates), + "errors": [], + } + + +def _placement_candidate(candidate_id: str = "place") -> dict: + candidate = _candidate(candidate_id, "red can") + step = candidate["draft"]["steps"][0] + step.update( + { + "task_type": "E1", + "target": _selector("table"), + "relation": "on", + "orientation_goal": "preserve", + } + ) + candidate["scene_request"]["references"] = [ + { + "reference_id": "upright.object", + "step_id": "upright", + "role": "object", + "reference": "red can", + "quantifier": "one", + "count": 0, + "source_structure": "rigid_object", + "affordances": ["graspable", "placeable"], + "initial_state": {}, + "attributes": {}, + }, + { + "reference_id": "upright.target", + "step_id": "upright", + "role": "target", + "reference": "table", + "quantifier": "one", + "count": 0, + "source_structure": "physical_entity", + "affordances": [], + "initial_state": {}, + "attributes": {}, + }, + ] + candidate["success_spec"]["terms"] = [ + {"step_id": "upright", "type": "semantic_goal"} + ] + candidate["semantic_hash"] = canonical_hash([step]) + return candidate + + +def _grounder(**kwargs) -> dict: + prompt = kwargs["prompt"] + uid = "blue_can" if '"reference": "blue can"' in prompt else "red_can" + return { + "bindings": [ + { + "reference_id": "upright.object", + "status": "resolved", + "uids": [uid], + "confidence": 0.95, + } + ] + } + + +def test_scene_source_fingerprint_reads_without_copying(scene_export: Path) -> None: + before = sorted(path.relative_to(scene_export) for path in scene_export.rglob("*")) + fingerprint = fingerprint_scene_source(SceneSourceRef(scene_export)) + after = sorted(path.relative_to(scene_export) for path in scene_export.rglob("*")) + + assert fingerprint.config_path == scene_export / "scene_config.json" + assert len(fingerprint.config_sha256) == 64 + assert len(fingerprint.asset_sha256) == 3 + assert after == before + + +@pytest.mark.parametrize("filename", ["gym_config.json", "gym_config_merged.json"]) +def test_scene_adapter_supports_legacy_gym_configs( + tmp_path: Path, + filename: str, +) -> None: + project = _legacy_gym_project(tmp_path, filename) + result = SceneAdapter(grounding_caller=_grounder).adapt( + _candidate_set([_candidate("legacy", "red can")]), + project, + ) + + assert result.selected_candidate_id == "legacy" + assert result.static_scene_manifest["source_format"] == "legacy_gym_config" + assert any( + item["role"] == "articulation" + for item in result.static_scene_manifest["objects"] + ) + assert ( + result.prepared_scene.articulations[0]["fpath"] + == (project / "assets" / "cabinet.urdf").resolve().as_posix() + ) + + +def test_scene_source_fingerprint_covers_articulation_fpath(tmp_path: Path) -> None: + project = _legacy_gym_project(tmp_path, "gym_config.json") + original = fingerprint_scene_source(project) + articulation_path = project / "assets" / "cabinet.urdf" + + articulation_path.write_bytes(b"changed articulation") + changed = fingerprint_scene_source(project) + + assert articulation_path.resolve().as_posix() in original.asset_sha256 + assert changed.asset_sha256 != original.asset_sha256 + with pytest.raises(RuntimeError, match="changed after Task Engine preparation"): + verify_scene_source_fingerprint(original.to_dict()) + + +def test_scene_adapter_selects_bindable_majority_and_redacts_manifest( + scene_export: Path, +) -> None: + red = _candidate("red-majority", "red can", votes=2) + blue = _candidate("blue-minority", "blue can") + result = SceneAdapter(grounding_caller=_grounder).adapt( + _candidate_set([red, blue]), + scene_export, + ) + + hierarchy_by_uid = { + node["uid"]: node for node in result.conservative_scene_graph["nodes"] + } + assert hierarchy_by_uid["red_can"]["parent_uid"] == "unknown" + assert hierarchy_by_uid["red_can"]["parent_relation"] == "unknown" + + assert result.binding_report["status"] == "bound" + assert result.binding_report["candidates"][0]["status"] == "resolved" + assert result.selected_candidate_id == "red-majority" + assert result.reference_bindings == {"upright.object": ["red_can"]} + assert result.role_bindings["role_bindings"] == {} + red_manifest = next( + item for item in result.scene_manifest["objects"] if item["uid"] == "red_can" + ) + assert "position" not in json.dumps(red_manifest) + assert ( + result.prepared_scene.source_config_path == scene_export / "scene_config.json" + ) + assert result.static_scene_manifest is not None + static_by_uid = { + item["uid"]: item for item in result.static_scene_manifest["objects"] + } + assert static_by_uid["red_can"]["physics"]["body_type"] == "dynamic" + assert static_by_uid["red_can"]["geometry"]["asset_sha256"] + + +def test_scene_adapter_returns_report_for_business_level_non_binding( + scene_export: Path, +) -> None: + candidate = _candidate("missing", "green can") + + def not_found(**_kwargs): + return { + "bindings": [ + { + "reference_id": "upright.object", + "status": "not_found", + "uids": [], + "confidence": 0.0, + } + ] + } + + result = SceneAdapter(grounding_caller=not_found).adapt( + _candidate_set([candidate]), + scene_export, + ) + + assert result.selected_candidate is None + assert result.role_bindings is None + assert result.binding_report["status"] == "unsatisfied" + assert ( + result.binding_report["candidates"][0]["references"][0]["status"] == "not_found" + ) + + +def test_semantic_blueprint_selection_forces_ranked_low_confidence_uid() -> None: + candidate = _candidate("likely", "the can") + scene_objects = [ + { + "uid": "table", + "role": "table", + "name": "table", + "description": "A work table.", + "category": "table", + "init_pos": [0.0, 0.0, 0.0], + "affordances": ["support_surface"], + "initial_state": {}, + "attributes": {}, + }, + *[ + { + "uid": f"{color}_can", + "role": "rigid_object", + "name": f"{color} can", + "description": f"A {color} can.", + "category": "can", + "init_pos": [0.0, offset, 0.1], + "affordances": ["graspable", "orientable"], + "initial_state": {"orientation": "fallen"}, + "attributes": {"color": color}, + } + for color, offset in (("red", -0.1), ("blue", 0.1)) + ], + ] + + def ambiguous(**_kwargs): + return { + "bindings": [ + { + "reference_id": "upright.object", + "status": "ambiguous", + "uids": ["red_can", "blue_can"], + "confidence": 0.2, + } + ] + } + + result = SceneAdapter(grounding_caller=ambiguous).select_objects( + _candidate_set([candidate]), + scene_objects, + force_most_likely=True, + ) + + assert result.selected_candidate_id == "likely" + assert result.role_bindings["reference_bindings"] == {"upright.object": ["red_can"]} + reference = result.binding_report["candidates"][0]["references"][0] + assert reference["confidence"] == 0.2 + assert reference["candidate_uids"] == ["red_can", "blue_can"] + assert reference["selected_uids"] == ["red_can"] + assert reference["reasons"] == [ + "Forced the highest-ranked structurally compatible UID from an " + "ambiguous low-confidence response." + ] + + +def test_scene_adapter_uses_unique_bindable_then_injected_adjudication( + scene_export: Path, +) -> None: + red = _candidate("red", "red can") + blue = _candidate("blue", "blue can") + + def one_missing(**kwargs): + if '"reference": "red can"' in kwargs["prompt"]: + return { + "bindings": [ + { + "reference_id": "upright.object", + "status": "not_found", + "uids": [], + "confidence": 0.0, + } + ] + } + return _grounder(**kwargs) + + unique = SceneAdapter(grounding_caller=one_missing).adapt( + _candidate_set([red, blue]), + scene_export, + ) + assert unique.selected_candidate_id == "blue" + assert unique.binding_report["selection_reason"] == "unique_bindable" + + ambiguous = SceneAdapter(grounding_caller=_grounder).adapt( + _candidate_set([red, blue]), + scene_export, + ) + assert ambiguous.binding_report["status"] == "ambiguous" + + adjudicated = SceneAdapter( + grounding_caller=_grounder, + adjudicator=lambda **_kwargs: {"candidate_id": "blue"}, + ).adapt(_candidate_set([red, blue]), scene_export) + assert adjudicated.selected_candidate_id == "blue" + assert adjudicated.binding_report["selection_reason"] == "adjudicated_bindable" + + +def test_scene_adapter_runs_one_default_structured_adjudication( + scene_export: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + adjudications = 0 + + def caller(**kwargs): + nonlocal adjudications + if kwargs["schema"]["title"] == "ActionEngineTaskAdjudication": + adjudications += 1 + return {"candidate_id": "blue"} + return _grounder(**kwargs) + + monkeypatch.setattr( + scene_adapter_module, "_default_grounding_caller", lambda: caller + ) + result = SceneAdapter().adapt( + _candidate_set([_candidate("red", "red can"), _candidate("blue", "blue can")]), + scene_export, + ) + + assert result.selected_candidate_id == "blue" + assert result.binding_report["selection_reason"] == "adjudicated_bindable" + assert adjudications == 1 + + +def test_scene_adapter_accepts_direct_source_and_rejects_bad_protocol( + scene_export: Path, +) -> None: + candidate = _candidate("red", "red can") + direct = SceneAdapter(grounding_caller=_grounder).adapt( + _candidate_set([candidate]), + scene_export, + ) + result = SceneAdapter(grounding_caller=_grounder).adapt( + _candidate_set([candidate]), + SceneSourceRef(scene_export), + ) + assert result.scene_manifest == direct.scene_manifest + assert result.role_bindings == direct.role_bindings + + with pytest.raises(SceneAdapterProtocolError, match="unsupported fields"): + SceneAdapter( + grounding_caller=lambda **_kwargs: { + "bindings": [ + { + "reference_id": "upright.object", + "status": "not_found", + "uids": [], + "confidence": 0.0, + "invented": True, + } + ] + } + ).adapt(_candidate_set([candidate]), scene_export) + + +def test_explicit_scene_semantic_conflict_is_incompatible( + scene_export: Path, +) -> None: + config_path = scene_export / "scene_config.json" + config = json.loads(config_path.read_text(encoding="utf-8")) + config["rigid_object"][0]["initial_state"]["orientation"] = "upright" + config_path.write_text(json.dumps(config), encoding="utf-8") + + result = SceneAdapter(grounding_caller=_grounder).adapt( + _candidate_set([_candidate("red", "red can")]), + scene_export, + ) + reference = result.binding_report["candidates"][0]["references"][0] + assert result.binding_report["status"] == "unsatisfied" + assert reference["status"] == "incompatible" + assert result.binding_report["candidates"][0]["status"] == "incompatible" + assert "state 'orientation' conflicts" in reference["reasons"][0] + + +def test_scene_adapter_accepts_passive_support_target_and_rejects_self_reference( + scene_export: Path, +) -> None: + candidate = _placement_candidate() + + def place_on_table(**_kwargs): + return { + "bindings": [ + { + "reference_id": "upright.object", + "status": "resolved", + "uids": ["red_can"], + "confidence": 0.95, + }, + { + "reference_id": "upright.target", + "status": "resolved", + "uids": ["table"], + "confidence": 0.95, + }, + ] + } + + bound = SceneAdapter(grounding_caller=place_on_table).adapt( + _candidate_set([candidate]), scene_export + ) + assert bound.binding_report["status"] == "bound" + assert bound.reference_bindings["upright.target"] == ["table"] + + def self_reference(**_kwargs): + response = place_on_table() + response["bindings"][1]["uids"] = ["red_can"] + return response + + incompatible = SceneAdapter(grounding_caller=self_reference).adapt( + _candidate_set([candidate]), scene_export + ) + assert incompatible.binding_report["status"] == "unsatisfied" + assert incompatible.binding_report["candidates"][0]["status"] == "incompatible" + + +def test_scene_adapter_enforces_count_cardinality_in_audit( + scene_export: Path, +) -> None: + candidate = _candidate("two", "cans") + selector = candidate["draft"]["steps"][0]["object"] + selector.update(quantifier="count", count=2) + request = candidate["scene_request"]["references"][0] + request.update(reference="cans", quantifier="count", count=2) + candidate["semantic_hash"] = canonical_hash(candidate["draft"]["steps"]) + + def one_only(**_kwargs): + return { + "bindings": [ + { + "reference_id": "upright.object", + "status": "resolved", + "uids": ["red_can"], + "confidence": 0.95, + } + ] + } + + result = SceneAdapter(grounding_caller=one_only).adapt( + _candidate_set([candidate]), scene_export + ) + + assert result.binding_report["status"] == "unsatisfied" + audit = result.binding_report["candidates"][0] + assert audit["status"] == "incompatible" + assert "requires exactly 2 UIDs" in audit["references"][0]["reasons"][0] + + +def test_scene_adapter_binds_all_matching_uids( + scene_export: Path, +) -> None: + candidate = _candidate("all", "all cans") + candidate["draft"]["steps"][0]["object"].update(quantifier="all") + candidate["scene_request"] = derive_scene_request(candidate["draft"]) + candidate["semantic_hash"] = canonical_hash(candidate["draft"]["steps"]) + + def all_cans(**_kwargs): + return { + "bindings": [ + { + "reference_id": "upright.object", + "status": "resolved", + "uids": ["red_can", "blue_can"], + "confidence": 0.95, + } + ] + } + + result = SceneAdapter(grounding_caller=all_cans).adapt( + _candidate_set([candidate]), scene_export + ) + + assert result.binding_report["status"] == "bound" + assert result.reference_bindings == {"upright.object": ["red_can", "blue_can"]} + assert result.candidate_bindings[candidate["candidate_id"]][ + "reference_bindings" + ] == {"upright.object": ["red_can", "blue_can"]} + + +def test_scene_adapter_rejects_step_result_object_matching_same_step_target( + scene_export: Path, +) -> None: + config_path = scene_export / "scene_config.json" + config = json.loads(config_path.read_text(encoding="utf-8")) + config["rigid_object"][0]["affordances"].append("support_surface") + config_path.write_text(json.dumps(config), encoding="utf-8") + + candidate = _candidate("self-reference", "red can") + second = deepcopy(candidate["draft"]["steps"][0]) + second.update( + { + "id": "place_again", + "task_type": "E1", + "object": { + "kind": "step_result", + "step_id": "upright", + "reference": "", + "quantifier": "one", + "count": 0, + }, + "target": _selector("red can"), + "relation": "on", + "orientation_goal": "preserve", + "depends_on": ["upright"], + } + ) + candidate["draft"]["steps"].append(second) + candidate["scene_request"] = derive_scene_request(candidate["draft"]) + candidate["success_spec"] = derive_success_spec(candidate["draft"]) + candidate["semantic_hash"] = canonical_hash(candidate["draft"]["steps"]) + + def same_uid(**_kwargs): + return { + "bindings": [ + { + "reference_id": "upright.object", + "status": "resolved", + "uids": ["red_can"], + "confidence": 0.95, + }, + { + "reference_id": "place_again.target", + "status": "resolved", + "uids": ["red_can"], + "confidence": 0.95, + }, + ] + } + + result = SceneAdapter(grounding_caller=same_uid).adapt( + _candidate_set([candidate]), scene_export + ) + + assert result.binding_report["status"] == "unsatisfied" + target_audit = result.binding_report["candidates"][0]["references"][1] + assert target_audit["status"] == "incompatible" + assert "same UID as object and target" in target_audit["reasons"][0] + + +def test_scene_source_fingerprint_covers_assets_and_config(scene_export: Path) -> None: + original = fingerprint_scene_source(scene_export) + + asset_path = scene_export / "meshes" / "red_can.glb" + asset_path.write_bytes(b"changed asset") + changed_asset = fingerprint_scene_source(scene_export) + assert changed_asset.asset_sha256 != original.asset_sha256 + + config_path = scene_export / "scene_config.json" + config = json.loads(config_path.read_text(encoding="utf-8")) + config["rigid_object"][0]["body_scale"] = [1.1, 1.0, 1.0] + config["rigid_object"][0]["physics"] = {"mass": 0.25} + config_path.write_text(json.dumps(config), encoding="utf-8") + changed_config = fingerprint_scene_source(scene_export) + assert changed_config.config_sha256 != changed_asset.config_sha256 + + +def test_scene_source_verification_rejects_later_mutation(scene_export: Path) -> None: + expected = fingerprint_scene_source(scene_export).to_dict() + (scene_export / "meshes" / "red_can.glb").write_bytes(b"changed later") + + with pytest.raises(RuntimeError, match="changed after Task Engine preparation"): + verify_scene_source_fingerprint(expected) diff --git a/tests/gen_sim/task_engine/scene/__init__.py b/tests/gen_sim/task_engine/scene/__init__.py new file mode 100644 index 000000000..b1ffd4df2 --- /dev/null +++ b/tests/gen_sim/task_engine/scene/__init__.py @@ -0,0 +1,17 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Tests for Task Engine scene adaptation boundaries.""" diff --git a/tests/gen_sim/task_engine/scene/test_final_inspection.py b/tests/gen_sim/task_engine/scene/test_final_inspection.py new file mode 100644 index 000000000..937a075d9 --- /dev/null +++ b/tests/gen_sim/task_engine/scene/test_final_inspection.py @@ -0,0 +1,118 @@ +# ---------------------------------------------------------------------------- +# 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 + +import numpy as np +import trimesh + +from embodichain.gen_sim.task_engine.orchestration.scene_source import ( + scene_revision_id, +) +from embodichain.gen_sim.task_engine.scene.final_inspection import ( + inspect_final_scene, +) + + +def _scene_export(root: Path, *, scene_id: str, can_rotation: list[float]) -> Path: + export = root / "scene_export" + assets = export / "mesh_assets" + assets.mkdir(parents=True) + trimesh.creation.box(extents=[1.0, 0.1, 1.0]).export( + assets / "table.glb", file_type="glb" + ) + can = trimesh.creation.cylinder(radius=0.04, height=0.2) + can.apply_transform( + trimesh.transformations.rotation_matrix(np.pi / 2.0, [1.0, 0.0, 0.0]) + ) + can.export(assets / "can.glb", file_type="glb") + config = { + "format": "embodichain.scene-export/v1", + "scene_id": scene_id, + "background": [ + { + "uid": "table", + "name": "table", + "description": "A support table.", + "category": "table", + "shape": {"shape_type": "Mesh", "fpath": "mesh_assets/table.glb"}, + "init_pos": [0.0, 0.0, 0.0], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + } + ], + "rigid_object": [ + { + "uid": "can", + "name": "red can", + "description": "A red can.", + "category": "can", + "shape": {"shape_type": "Mesh", "fpath": "mesh_assets/can.glb"}, + "init_pos": [0.0, 0.0, 0.15], + "init_rot": can_rotation, + "body_scale": [1.0, 1.0, 1.0], + } + ], + } + path = export / "scene_config.json" + path.write_text(json.dumps(config), encoding="utf-8") + return path + + +def test_scene_revision_id_ignores_exporter_timestamp_and_location( + tmp_path: Path, +) -> None: + first = _scene_export( + tmp_path / "first", scene_id="scene-100", can_rotation=[0, 0, 0] + ) + second = _scene_export( + tmp_path / "second", scene_id="scene-200", can_rotation=[0, 0, 0] + ) + + assert scene_revision_id(first) == scene_revision_id(second) + + value = json.loads(second.read_text(encoding="utf-8")) + value["rigid_object"][0]["init_pos"][0] = 0.25 + second.write_text(json.dumps(value), encoding="utf-8") + assert scene_revision_id(first) != scene_revision_id(second) + + +def test_final_inspection_recomputes_support_and_orientation(tmp_path: Path) -> None: + source = _scene_export( + tmp_path / "standing", scene_id="scene", can_rotation=[0.0, 0.0, 0.0] + ) + + inspection = inspect_final_scene(source, revision_id=scene_revision_id(source)) + + can = next(item for item in inspection["objects"] if item["uid"] == "can") + assert can["orientation"] == "standing" + assert can["support"]["parent_uid"] == "table" + assert can["support"]["relation"] == "on" + assert can["support"]["xy_overlap_ratio"] > 0.9 + + +def test_final_inspection_detects_lying_rotation(tmp_path: Path) -> None: + source = _scene_export( + tmp_path / "lying", scene_id="scene", can_rotation=[90.0, 0.0, 0.0] + ) + + inspection = inspect_final_scene(source, revision_id=scene_revision_id(source)) + + can = next(item for item in inspection["objects"] if item["uid"] == "can") + assert can["orientation"] == "lying" diff --git a/tests/gen_sim/task_engine/scene/test_scene_boundary.py b/tests/gen_sim/task_engine/scene/test_scene_boundary.py new file mode 100644 index 000000000..21b1e0273 --- /dev/null +++ b/tests/gen_sim/task_engine/scene/test_scene_boundary.py @@ -0,0 +1,534 @@ +# ---------------------------------------------------------------------------- +# 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 pathlib import Path +from types import SimpleNamespace + +import pytest + +from embodichain.gen_sim.task_engine.scene import ( + FeasibilityBroker, + SceneEngineV1Adapter, +) +from embodichain.gen_sim.task_engine.agent import derive_scene_request +from embodichain.gen_sim.task_engine.contracts import TASK_DRAFT_SCHEMA + + +def _prepared_scene(tmp_path: Path) -> SimpleNamespace: + table = { + "uid": "table", + "source_uid": "table_0", + "role": "background", + "name": "table", + "description": "A support table.", + "category": "table", + "color": "brown", + "shape": {"shape_type": "Mesh", "fpath": "/assets/table.glb"}, + "init_pos": [0.0, 0.0, 0.0], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + "attributes": {}, + "initial_state": {}, + "affordances": [], + } + can = { + "uid": "red_can", + "source_uid": "red_can_0", + "role": "rigid_object", + "name": "red can", + "description": "A fallen red can.", + "category": "can", + "color": "red", + "shape": {"shape_type": "Mesh", "fpath": "/assets/can.glb"}, + "init_pos": [0.1, 0.0, 0.7], + "init_rot": [90.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + "attributes": {}, + "initial_state": {"orientation": "fallen"}, + "affordances": ["graspable", "orientable", "placeable"], + } + runtime_table = { + "uid": "table", + "shape": table["shape"], + "attrs": {"mass": 10.0}, + "body_type": "kinematic", + } + runtime_can = { + "uid": "red_can", + "shape": can["shape"], + "attrs": {"mass": 0.1}, + "body_type": "dynamic", + } + return SimpleNamespace( + source_config_path=tmp_path / "scene_config.json", + planner_objects=(table, can), + background=(runtime_table,), + rigid_objects=(runtime_can,), + articulations=(), + asset_hashes={"table": "a" * 64, "red_can": "b" * 64}, + ) + + +def _candidate(task_type: str, affordances: list[str]) -> dict: + return { + "candidate_id": "candidate_01", + "draft": { + "task_id": "task", + "steps": [{"id": "step_01", "task_type": task_type}], + }, + "scene_request": { + "references": [ + { + "reference_id": "step_01.object", + "role": "object", + "source_structure": "rigid_object", + "affordances": affordances, + "initial_state": ( + {"orientation": "fallen"} if task_type == "E2" else {} + ), + "attributes": {}, + } + ] + }, + } + + +def _catalog(*, pour_available: bool = False) -> dict[str, dict]: + return { + name: {"runtime_available": True, "unavailable_reason": None} + for name in ("PickUp", "MoveHeldObject", "Place") + } | { + "Pour": { + "runtime_available": pour_available, + "unavailable_reason": None if pour_available else "Pour is planning-only.", + } + } + + +def _selector(kind: str, *, reference: str = "") -> dict[str, object]: + return { + "kind": kind, + "step_id": "", + "reference": reference, + "quantifier": "one", + "count": 0, + } + + +def _relation_candidate(relation: str) -> dict: + draft = { + "schema_version": TASK_DRAFT_SCHEMA, + "task_id": "place_relative", + "instruction": "place the can relative to the target", + "steps": [ + { + "id": "step_01", + "task_type": "E1", + "object": _selector("scene_ref", reference="red can"), + "target": _selector("scene_ref", reference="target"), + "relation": relation, + "required_arm": "auto", + "transfer_arm": "none", + "receive_arm": "none", + "orientation_goal": "preserve", + "target_state": "none", + "target_setting": 0, + "layout": "none", + "axis": "none", + "direction": "none", + "terminal_behavior": "none", + "depends_on": [], + } + ], + } + return { + "candidate_id": "candidate_01", + "draft": draft, + "scene_request": derive_scene_request(draft), + } + + +def _manifest_with_target_kinds(tmp_path: Path) -> dict: + manifest = SceneEngineV1Adapter().adapt_prepared_scene( + _prepared_scene(tmp_path), + source_format="test", + robot_profile="dual_franka", + ) + by_uid = {item["uid"]: item for item in manifest["objects"]} + by_uid["red_can"]["affordances"].append( + { + "type": "container", + "status": "declared", + "confidence": None, + "source": "test", + "link_uid": "", + "frame": {}, + "parameters": {}, + } + ) + articulation = deepcopy(by_uid["red_can"]) + articulation.update( + uid="cabinet", + source_uid="cabinet_0", + role="articulation", + name="cabinet", + category="cabinet", + physics={}, + articulation={"runtime_uid": "cabinet"}, + affordances=[], + ) + manifest["objects"].append(articulation) + return manifest + + +def test_scene_engine_v1_adapter_preserves_static_execution_evidence( + tmp_path: Path, +) -> None: + manifest = SceneEngineV1Adapter().adapt_prepared_scene( + _prepared_scene(tmp_path), + source_format="embodichain.scene-export/v1", + robot_profile="dual_franka", + ) + + by_uid = {item["uid"]: item for item in manifest["objects"]} + assert manifest["adapter_capabilities"]["task_conditioned_generation"] is False + assert by_uid["red_can"]["geometry"]["asset_sha256"] == "b" * 64 + assert by_uid["red_can"]["physics"]["body_type"] == "dynamic" + assert {item["type"] for item in by_uid["table"]["affordances"]} == { + "support_surface" + } + assert ( + next( + item + for item in by_uid["red_can"]["affordances"] + if item["type"] == "graspable" + )["status"] + == "declared" + ) + + +def test_e2_feasibility_requires_runtime_probe_for_geometry(tmp_path: Path) -> None: + manifest = SceneEngineV1Adapter().adapt_prepared_scene( + _prepared_scene(tmp_path), + source_format="test", + robot_profile="dual_franka", + ) + report = FeasibilityBroker().assess( + _candidate("E2", ["graspable", "orientable"]), + {"step_01.object": ["red_can"]}, + manifest, + capability_catalog=_catalog(), + task_actions={"E2": ("PickUp", "MoveHeldObject", "Place")}, + ) + + assert report["status"] == "runtime_probe" + assert report["remediation_class"] == "none" + assert report["blockers"] == [] + assert report["summary"]["proven"] > 0 + assert report["summary"]["runtime_probe"] > 0 + + +def test_planning_only_action_is_reported_as_contradicted(tmp_path: Path) -> None: + manifest = SceneEngineV1Adapter().adapt_prepared_scene( + _prepared_scene(tmp_path), + source_format="test", + robot_profile="dual_franka", + ) + report = FeasibilityBroker().assess( + _candidate("E3", ["graspable", "pourable"]), + {"step_01.object": ["red_can"]}, + manifest, + capability_catalog=_catalog(), + task_actions={"E3": ("Pour",)}, + ) + + assert report["status"] == "contradicted" + assert report["remediation_class"] == "action_capability" + assert any("planning-only" in blocker for blocker in report["blockers"]) + + +def test_final_orientation_conflict_is_scene_remediable(tmp_path: Path) -> None: + manifest = SceneEngineV1Adapter().adapt_prepared_scene( + _prepared_scene(tmp_path), + source_format="test", + robot_profile="dual_franka", + ) + can = next(item for item in manifest["objects"] if item["uid"] == "red_can") + can["initial_state"]["orientation"] = "upright" + + report = FeasibilityBroker().assess( + _candidate("E2", ["graspable", "orientable"]), + {"step_01.object": ["red_can"]}, + manifest, + capability_catalog=_catalog(), + task_actions={"E2": ("PickUp", "MoveHeldObject", "Place")}, + ) + + assert report["status"] == "contradicted" + assert report["remediation_class"] == "scene_remediable" + + +def test_missing_affordance_remains_unknown_instead_of_becoming_supported( + tmp_path: Path, +) -> None: + manifest = SceneEngineV1Adapter().adapt_prepared_scene( + _prepared_scene(tmp_path), + source_format="test", + robot_profile="dual_franka", + ) + report = FeasibilityBroker().assess( + _candidate("E1", ["graspable", "liquid_safe"]), + {"step_01.object": ["red_can"]}, + manifest, + capability_catalog=_catalog(), + task_actions={"E1": ("PickUp", "MoveHeldObject", "Place")}, + ) + + assert report["status"] == "unknown" + assert any( + check["status"] == "unknown" and "liquid_safe" in check["reason"] + for check in report["checks"] + ) + + +def test_physical_object_can_be_a_runtime_support_without_support_affordance( + tmp_path: Path, +) -> None: + manifest = SceneEngineV1Adapter().adapt_prepared_scene( + _prepared_scene(tmp_path), + source_format="test", + robot_profile="dual_franka", + ) + candidate = _candidate("E1", ["graspable", "placeable"]) + candidate["draft"]["steps"][0].update( + target={"kind": "scene_ref"}, + relation="on", + ) + candidate["scene_request"]["references"].append( + { + "reference_id": "step_01.target", + "role": "target", + "source_structure": "physical_entity", + "affordances": [], + "initial_state": {}, + "attributes": {}, + } + ) + + report = FeasibilityBroker().assess( + candidate, + { + "step_01.object": ["red_can"], + "step_01.target": ["red_can"], + }, + manifest, + capability_catalog=_catalog(), + task_actions={"E1": ("PickUp", "MoveHeldObject", "Place")}, + ) + + structure = next( + check + for check in report["checks"] + if check["kind"] == "structure" and check["subject"] == "step_01.target:red_can" + ) + assert structure["status"] == "proven" + support_probe = next( + check for check in report["checks"] if check["kind"] == "placement_support" + ) + assert support_probe["status"] == "runtime_probe" + assert support_probe["evidence"]["runtime_obligations"] == [ + "placement_candidates", + "object_supported_by", + "stable_for", + "final_support_revalidation", + ] + assert report["blockers"] == [] + + +@pytest.mark.parametrize( + ("relation", "target_uid", "expected_structure", "expected_status"), + [ + ("on", "red_can", "physical_entity", "proven"), + ("on", "table", "physical_entity", "proven"), + ("on", "cabinet", "physical_entity", "contradicted"), + ("inside", "red_can", "rigid_object", "proven"), + ("inside", "table", "rigid_object", "contradicted"), + ("inside", "cabinet", "rigid_object", "contradicted"), + ("behind", "red_can", "spatial_reference", "proven"), + ("behind", "table", "spatial_reference", "proven"), + ("behind", "cabinet", "spatial_reference", "runtime_probe"), + ("front_of", "red_can", "spatial_reference", "proven"), + ("front_of", "table", "spatial_reference", "proven"), + ("front_of", "cabinet", "spatial_reference", "runtime_probe"), + ("left_of", "red_can", "spatial_reference", "proven"), + ("left_of", "table", "spatial_reference", "proven"), + ("left_of", "cabinet", "spatial_reference", "runtime_probe"), + ("right_of", "red_can", "spatial_reference", "proven"), + ("right_of", "table", "spatial_reference", "proven"), + ("right_of", "cabinet", "spatial_reference", "runtime_probe"), + ], +) +def test_relation_target_structure_matrix_uses_capability_semantics( + tmp_path: Path, + relation: str, + target_uid: str, + expected_structure: str, + expected_status: str, +) -> None: + candidate = _relation_candidate(relation) + target_request = next( + item + for item in candidate["scene_request"]["references"] + if item["role"] == "target" + ) + report = FeasibilityBroker().assess( + candidate, + { + "step_01.object": ["red_can"], + "step_01.target": [target_uid], + }, + _manifest_with_target_kinds(tmp_path), + capability_catalog=_catalog(), + task_actions={"E1": ("PickUp", "MoveHeldObject", "Place")}, + ) + + structure = next( + check + for check in report["checks"] + if check["kind"] == "structure" + and check["subject"] == f"step_01.target:{target_uid}" + ) + assert target_request["source_structure"] == expected_structure + assert structure["status"] == expected_status + + +def test_legacy_scene_entity_target_is_treated_as_an_abstract_structure( + tmp_path: Path, +) -> None: + candidate = _relation_candidate("behind") + target_request = next( + item + for item in candidate["scene_request"]["references"] + if item["role"] == "target" + ) + target_request["source_structure"] = "scene_entity" + + report = FeasibilityBroker().assess( + candidate, + { + "step_01.object": ["red_can"], + "step_01.target": ["red_can"], + }, + _manifest_with_target_kinds(tmp_path), + capability_catalog=_catalog(), + task_actions={"E1": ("PickUp", "MoveHeldObject", "Place")}, + ) + + structure = next( + check + for check in report["checks"] + if check["kind"] == "structure" and check["subject"] == "step_01.target:red_can" + ) + assert structure["status"] == "proven" + assert report["blockers"] == [] + + +def test_unknown_structure_contract_is_not_a_scene_contradiction( + tmp_path: Path, +) -> None: + candidate = _relation_candidate("behind") + target_request = next( + item + for item in candidate["scene_request"]["references"] + if item["role"] == "target" + ) + target_request["source_structure"] = "future_spatial_capability" + + report = FeasibilityBroker().assess( + candidate, + { + "step_01.object": ["red_can"], + "step_01.target": ["red_can"], + }, + _manifest_with_target_kinds(tmp_path), + capability_catalog=_catalog(), + task_actions={"E1": ("PickUp", "MoveHeldObject", "Place")}, + ) + + structure = next( + check + for check in report["checks"] + if check["kind"] == "structure" and check["subject"] == "step_01.target:red_can" + ) + assert structure["status"] == "unknown" + assert not any("future_spatial_capability" in item for item in report["blockers"]) + + +def test_required_arm_side_requires_the_live_robot_frame( + tmp_path: Path, +) -> None: + manifest = SceneEngineV1Adapter().adapt_prepared_scene( + _prepared_scene(tmp_path), + source_format="test", + robot_profile="dual_franka", + ) + red_can = next(item for item in manifest["objects"] if item["uid"] == "red_can") + red_can["initial_pose"]["position"][1] = -0.20 + candidate = _candidate("E2", ["graspable", "orientable"]) + candidate["draft"]["steps"][0]["required_arm"] = "right_arm" + + report = FeasibilityBroker().assess( + candidate, + {"step_01.object": ["red_can"]}, + manifest, + capability_catalog=_catalog(), + task_actions={"E2": ("PickUp", "MoveHeldObject", "Place")}, + ) + + probe = next( + check for check in report["checks"] if check["kind"] == "arm_layout_risk" + ) + assert probe["status"] == "runtime_probe" + assert probe["evidence"]["arm_side_frame"] == "live_robot" + assert probe["evidence"]["mismatch_risk"] is None + assert "expected_arm" not in probe["evidence"] + assert probe["evidence"]["geometry_certificate"] is False + assert report["blockers"] == [] + + +def test_workspace_report_covers_complete_task_phases(tmp_path: Path) -> None: + manifest = SceneEngineV1Adapter().adapt_prepared_scene( + _prepared_scene(tmp_path), + source_format="test", + robot_profile="dual_franka", + ) + report = FeasibilityBroker().assess( + _candidate("E2", ["graspable", "orientable"]), + {"step_01.object": ["red_can"]}, + manifest, + capability_catalog=_catalog(), + task_actions={"E2": ("PickUp", "MoveHeldObject", "Place")}, + ) + + workflow = next( + check for check in report["checks"] if check["kind"] == "task_workspace" + ) + phases = {item["phase"] for item in workflow["evidence"]["phases"]} + assert phases == {"pickup", "safety_clearance"} + assert workflow["status"] == "runtime_probe" diff --git a/tests/gen_sim/task_engine/test_agent.py b/tests/gen_sim/task_engine/test_agent.py new file mode 100644 index 000000000..cb6ab7cb2 --- /dev/null +++ b/tests/gen_sim/task_engine/test_agent.py @@ -0,0 +1,271 @@ +# ---------------------------------------------------------------------------- +# 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 threading +from time import sleep + +import pytest + +from embodichain.gen_sim.task_engine.contracts import ( + SUCCESS_SPEC_SCHEMA, + TASK_DRAFT_SCHEMA, + validate_success_spec, + validate_task_candidate, + validate_task_draft, +) +from embodichain.gen_sim.task_engine.agent import ( + TaskAgent, + TaskGenerationError, + derive_scene_request, + derive_success_spec, +) +from embodichain.gen_sim.task_engine.interpretation import InstructionDraftResult +from embodichain.gen_sim.task_engine.orchestration.coordinator import ( + lower_task_candidate, +) + +_TEST_INSTRUCTION = "test-instruction" + + +def _selector(kind="none", *, step_id="", reference="", quantifier="one", count=0): + return { + "kind": kind, + "step_id": step_id, + "reference": reference, + "quantifier": quantifier, + "count": count, + } + + +def _step(step_id="orient", reference="purple can"): + return { + "id": step_id, + "task_type": "E2", + "object": _selector("scene_ref", reference=reference), + "target": _selector(), + "relation": "none", + "required_arm": "auto", + "transfer_arm": "none", + "receive_arm": "none", + "orientation_goal": "upright", + "target_state": "none", + "target_setting": 0, + "layout": "none", + "axis": "none", + "direction": "none", + "terminal_behavior": "none", + "depends_on": [], + } + + +def _result(step): + return InstructionDraftResult( + intent={"steps": [deepcopy(step)]}, + model="injected_caller", + attempts=1, + latency_seconds=0.01, + normalizations=(), + ) + + +def test_task_agent_generates_concurrently_deduplicates_and_counts_votes(): + barrier = threading.Barrier(3) + lock = threading.Lock() + assigned = 0 + + def interpreter(_instruction, **_kwargs): + nonlocal assigned + with lock: + index = assigned + assigned += 1 + barrier.wait(timeout=2) + sleep(0.01) + if index < 2: + return _result(_step(step_id=f"arbitrary_{index}")) + return _result(_step(step_id="different", reference="orange can")) + + result = TaskAgent(interpreter=interpreter).generate("task", _TEST_INSTRUCTION) + + assert result["requested_candidate_count"] == 3 + assert result["valid_response_count"] == 3 + assert len(result["candidates"]) == 2 + assert sorted(item["vote_count"] for item in result["candidates"]) == [1, 2] + assert {item["draft"]["steps"][0]["id"] for item in result["candidates"]} == { + "step_01" + } + + +def test_scene_request_and_success_are_deterministic_contract_derivations(): + draft = { + "schema_version": TASK_DRAFT_SCHEMA, + "task_id": "upright", + "instruction": _TEST_INSTRUCTION, + "steps": [_step(reference="all cans")], + } + draft["steps"][0]["object"].update(quantifier="all") + + request = derive_scene_request(draft) + success = derive_success_spec(draft) + + assert request["references"] == [ + { + "reference_id": "orient.object", + "step_id": "orient", + "role": "object", + "reference": "all cans", + "quantifier": "all", + "count": 0, + "source_structure": "rigid_object", + "affordances": ["graspable", "orientable"], + "initial_state": {"orientation": "fallen"}, + "attributes": {}, + } + ] + assert success["terms"] == [{"step_id": "orient", "type": "object_upright"}] + + +@pytest.mark.parametrize( + ("relation", "expected_structure", "expected_affordances"), + [ + ("on", "physical_entity", []), + ("inside", "rigid_object", ["container"]), + ("behind", "spatial_reference", []), + ("front_of", "spatial_reference", []), + ("left_of", "spatial_reference", []), + ("right_of", "spatial_reference", []), + ], +) +def test_target_requirements_describe_capabilities_not_concrete_roles( + relation: str, + expected_structure: str, + expected_affordances: list[str], +) -> None: + step = _step(step_id="place", reference="green can") + step.update( + task_type="E1", + target=_selector("scene_ref", reference="red can"), + relation=relation, + orientation_goal="preserve", + ) + draft = { + "schema_version": TASK_DRAFT_SCHEMA, + "task_id": "stack", + "instruction": _TEST_INSTRUCTION, + "steps": [step], + } + + request = derive_scene_request(draft) + + target = next( + reference + for reference in request["references"] + if reference["role"] == "target" + ) + assert target["source_structure"] == expected_structure + assert target["affordances"] == expected_affordances + + +def test_lower_task_candidate_expands_success_for_all_binding(): + def interpreter(_instruction, **_kwargs): + step = _step(reference="all cans") + step["object"].update(quantifier="all") + return _result(step) + + candidate = TaskAgent(interpreter=interpreter).generate( + "upright", _TEST_INSTRUCTION, candidate_count=1 + )["candidates"][0] + grounded = lower_task_candidate( + candidate, + {"step_01.object": ["can_a", "can_b"]}, + [ + {"uid": "can_a", "role": "rigid_object", "description": "A can."}, + {"uid": "can_b", "role": "rigid_object", "description": "A can."}, + ], + "dual_franka", + ) + + assert grounded.task_spec["level"] == "L2" + assert [term["type"] for term in grounded.task_spec["success"]["terms"]] == [ + "object_upright", + "object_upright", + ] + + +def test_draft_rejects_grounded_fields_and_task_agent_fails_closed(): + draft = { + "schema_version": TASK_DRAFT_SCHEMA, + "task_id": "bad", + "instruction": "bad", + "steps": [_step()], + } + draft["steps"][0]["object"]["uid"] = "scene_uid" + with pytest.raises(ValueError, match="forbidden|exactly fields"): + validate_task_draft(draft) + + def invalid(_instruction, **_kwargs): + raise ValueError("invalid draft after repair") + + with pytest.raises(TaskGenerationError, match="All Task Agent candidates"): + TaskAgent(interpreter=invalid).generate("bad", "bad") + + +def test_task_candidate_rejects_scene_constraints_not_derived_from_draft(): + candidate = TaskAgent( + interpreter=lambda *_args, **_kwargs: _result(_step()) + ).generate("upright", _TEST_INSTRUCTION, candidate_count=1)["candidates"][0] + candidate["scene_request"]["references"][0]["affordances"] = [] + + with pytest.raises(ValueError, match="derived exactly"): + validate_task_candidate(candidate) + + +def test_success_spec_rejects_types_outside_task_ontology(): + with pytest.raises(ValueError, match="must be one of"): + validate_success_spec( + { + "schema_version": SUCCESS_SPEC_SCHEMA, + "task_id": "bad_success", + "op": "all", + "terms": [{"step_id": "step_01", "type": "looks_good"}], + } + ) + + +def test_task_agent_isolates_invalid_interpreter_results(): + lock = threading.Lock() + calls = 0 + + def interpreter(_instruction, **_kwargs): + nonlocal calls + with lock: + index = calls + calls += 1 + if index == 0: + invalid = _step() + invalid["object"]["uid"] = "red_can" + return _result(invalid) + return _result(_step()) + + result = TaskAgent(interpreter=interpreter).generate( + "upright", _TEST_INSTRUCTION, candidate_count=2 + ) + + assert result["valid_response_count"] == 1 + assert len(result["errors"]) == 1 + assert len(result["candidates"]) == 1 diff --git a/tests/gen_sim/task_engine/test_interpretation.py b/tests/gen_sim/task_engine/test_interpretation.py new file mode 100644 index 000000000..009362e87 --- /dev/null +++ b/tests/gen_sim/task_engine/test_interpretation.py @@ -0,0 +1,96 @@ +# ---------------------------------------------------------------------------- +# 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 pathlib import Path + +import pytest + +from embodichain.gen_sim.task_engine import interpretation as interpretation_module + + +def _write_dotenv(path: Path) -> None: + path.write_text( + "\n".join( + ( + "OPENAI_API_KEY=dotenv-key", + "OPENAI_BASE_URL=https://dotenv.example/v1", + "OPENAI_MODEL=dotenv-model", + ) + ), + encoding="utf-8", + ) + + +def _clear_process_provider(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ( + "OPENAI_API_KEY", + "OPENAI_BASE_URL", + "OPENAI_API_BASE", + "LLM_URL", + "TASK_ENGINE_LLM_MODEL", + "ACTION_ENGINE_LLM_MODEL", + "OPENAI_MODEL", + "LLM_MODEL", + ): + monkeypatch.delenv(name, raising=False) + + +def test_partial_process_transport_does_not_mix_with_dotenv( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + env_path = tmp_path / ".env" + _write_dotenv(env_path) + monkeypatch.setattr(interpretation_module, "_GEN_SIM_ENV_PATH", env_path) + monkeypatch.setattr( + interpretation_module, + "_GEN_CONFIG_PATH", + tmp_path / "missing.json", + ) + _clear_process_provider(monkeypatch) + monkeypatch.setenv("OPENAI_API_KEY", "unrelated-process-key") + + settings = interpretation_module._load_llm_settings(model=None) + + assert settings["api_key"] == "dotenv-key" + assert settings["base_url"] == "https://dotenv.example/v1" + assert settings["model"] == "dotenv-model" + + +def test_complete_process_transport_overrides_dotenv( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + env_path = tmp_path / ".env" + _write_dotenv(env_path) + monkeypatch.setattr(interpretation_module, "_GEN_SIM_ENV_PATH", env_path) + monkeypatch.setattr( + interpretation_module, + "_GEN_CONFIG_PATH", + tmp_path / "missing.json", + ) + _clear_process_provider(monkeypatch) + monkeypatch.setenv("OPENAI_API_KEY", "process-key") + monkeypatch.setenv("OPENAI_BASE_URL", "https://process.example/v1/") + monkeypatch.setenv("TASK_ENGINE_LLM_MODEL", "process-model") + + settings = interpretation_module._load_llm_settings(model=None) + + assert settings["api_key"] == "process-key" + assert settings["base_url"] == "https://process.example/v1" + assert settings["model"] == "process-model" diff --git a/tests/gen_sim/task_engine/test_parallel_workflow.py b/tests/gen_sim/task_engine/test_parallel_workflow.py new file mode 100644 index 000000000..23aaab655 --- /dev/null +++ b/tests/gen_sim/task_engine/test_parallel_workflow.py @@ -0,0 +1,912 @@ +# ---------------------------------------------------------------------------- +# 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 collections.abc import Mapping +from copy import deepcopy +import json +from pathlib import Path +import sys +from types import SimpleNamespace +from threading import Barrier + +import pytest + +from embodichain.gen_sim.action_engine.unbound import build_unbound_action_plan +from embodichain.gen_sim.action_engine.unbound import ActionCapabilityError +from embodichain.gen_sim.action_engine.runtime import ( + ExecutionReport, + build_execution_provenance, +) +from embodichain.gen_sim.scene_engine.errors import SceneServiceError +from embodichain.gen_sim.task_engine.config import ( + TaskEngineExecutionCfg, + TaskEnginePlanningCfg, + TaskEngineWorkflowCfg, +) +from embodichain.gen_sim.task_engine.orchestration.scene_adapter import ( + CandidateSelection, +) +from embodichain.gen_sim.task_engine.orchestration.scene_source import SceneSourceRef +from embodichain.gen_sim.task_engine.scene_backend import SceneAnalysis, SceneRevision +from embodichain.gen_sim.task_engine.workflow import ( + SubprocessActionExecutor, + TaskEngineWorkflow, + _environment_successes, + _run_streaming_process, +) +from embodichain.gen_sim.task_engine.workflow_contracts import TASK_RUN_REQUEST_SCHEMA + + +def _candidate_set() -> dict: + candidate = { + "candidate_id": "candidate_01", + "draft": { + "task_id": "place_can", + "instruction": "Place the can on the table.", + "steps": [ + { + "id": "place", + "task_type": "E1", + "object": { + "kind": "scene_ref", + "step_id": "", + "reference": "the can", + "quantifier": "one", + "count": 0, + }, + "target": { + "kind": "scene_ref", + "step_id": "", + "reference": "the table", + "quantifier": "one", + "count": 0, + }, + "depends_on": [], + } + ], + }, + } + return { + "task_id": "place_can", + "instruction": "Place the can on the table.", + "candidates": [candidate], + } + + +def _selection(candidate_set: Mapping[str, object]) -> CandidateSelection: + candidate = candidate_set["candidates"][0] + return CandidateSelection( + scene_manifest={}, + role_bindings={}, + binding_report={ + "status": "bound", + "selection_reason": "test", + "candidates": [{"candidate_id": "candidate_01", "status": "resolved"}], + }, + selected_candidate=candidate, + candidate_bindings={"candidate_01": {}}, + ) + + +def _request(tmp_path: Path, *, existing: bool = False, edit: bool = False) -> dict: + return { + "schema_version": TASK_RUN_REQUEST_SCHEMA, + "task_id": "place_can", + "task_instruction": "Place the can on the table.", + "image_path": None if existing else str(tmp_path / "input.png"), + "gym_project": str(tmp_path / "project") if existing else None, + "scene_edit_prompt": "Move the can left." if edit else None, + "output_dir": str(tmp_path / "run"), + } + + +class _TaskAgent: + def __init__(self, candidates: dict, barrier: Barrier | None = None) -> None: + self.candidates = candidates + self.barrier = barrier + + def generate(self, *_args, **_kwargs) -> dict: + if self.barrier is not None: + self.barrier.wait(timeout=2) + return self.candidates + + +class _ActionAgent: + def __init__(self, barrier: Barrier | None = None) -> None: + self.barrier = barrier + + def draft(self, candidate: Mapping[str, object]) -> dict: + if self.barrier is not None: + self.barrier.wait(timeout=2) + return build_unbound_action_plan(candidate) + + +class _FailingActionAgent: + def draft(self, _candidate: Mapping[str, object]) -> dict: + raise ActionCapabilityError("missing AtomicAction") + + +class _SceneBackend: + def __init__( + self, + selection: CandidateSelection, + *, + input_kind: str = "image", + input_barrier: Barrier | None = None, + materialize_barrier: Barrier | None = None, + materialize_failures: int = 0, + ) -> None: + self.selection = selection + self.input_kind = input_kind + self.input_barrier = input_barrier + self.materialize_barrier = materialize_barrier + self.materialize_failures = materialize_failures + self.seeds: list[int] = [] + + def analyze(self, request, output_root) -> SceneAnalysis: + if self.input_barrier is not None: + self.input_barrier.wait(timeout=2) + return SceneAnalysis( + input_kind=self.input_kind, + source=Path(request["image_path"] or request["gym_project"]), + blueprint=None, + source_fingerprint=None, + ) + + def select(self, *_args, **_kwargs) -> CandidateSelection: + return self.selection + + def materialize( + self, _analysis, _request, output_root, *, seed: int + ) -> SceneRevision: + if self.materialize_barrier is not None: + self.materialize_barrier.wait(timeout=2) + root = Path(output_root) + root.mkdir(parents=True) + self.seeds.append(seed) + if len(self.seeds) <= self.materialize_failures: + raise SceneServiceError("scene service failed") + source = root / "scene_config.json" + source.write_text("{}\n", encoding="utf-8") + return SceneRevision( + source=source, + output_root=root, + revision_id="0" * 64, + seed=seed, + edit_plan=None, + source_fingerprint=None, + ) + + def inspect(self, revision, output_path): + value = { + "schema_version": "embodichain.final-scene-inspection/v1", + "scene_revision_id": revision.revision_id, + "source_config_path": revision.source.as_posix(), + "contact_tolerance_m": 0.03, + "objects": [], + } + path = Path(output_path) + path.write_text(json.dumps(value), encoding="utf-8") + return value + + +class _Coordinator: + def __init__( + self, + statuses: list[str], + *, + infeasible_remediation: str = "scene_remediable", + ) -> None: + self.statuses = list(statuses) + self.infeasible_remediation = infeasible_remediation + self.calls = 0 + self.kwargs: list[dict] = [] + self.sources: list[object] = [] + + def prepare(self, _task_id, _instruction, _source, output_dir, **_kwargs): + status = self.statuses[min(self.calls, len(self.statuses) - 1)] + self.calls += 1 + self.kwargs.append(dict(_kwargs)) + self.sources.append(_source) + root = Path(output_dir) + root.mkdir(parents=True) + for name in ( + "conservative_scene_graph.json", + "seed_task_graph.json", + "grounded_task_plan.json", + ): + (root / name).write_text("{}\n", encoding="utf-8") + return SimpleNamespace( + status=status, + output_dir=root, + planning_attempts=(), + feasibility_report=( + {"remediation_class": self.infeasible_remediation} + if status == "infeasible" + else None + ), + selected_candidate_id="candidate_01" if status == "bound" else None, + ) + + +class _FailingCoordinator: + def prepare(self, *_args, **_kwargs): + raise RuntimeError("grounding service unavailable") + + +class _RebindingCoordinator(_Coordinator): + def __init__(self, final_candidate: Mapping[str, object]) -> None: + super().__init__(["bound"]) + self.final_candidate = final_candidate + + def prepare(self, *args, **kwargs): + result = super().prepare(*args, **kwargs) + result.selected_candidate_id = str(self.final_candidate["candidate_id"]) + result.unbound_action_plan = build_unbound_action_plan(self.final_candidate) + return result + + +class _InvalidSceneBackend(_SceneBackend): + def materialize(self, *_args, **kwargs): + self.seeds.append(int(kwargs["seed"])) + raise ValueError("invalid deterministic scene input") + + +class _Executor: + def __init__( + self, + successes: list[list[bool]], + *, + expected_dataset_saving: bool = False, + ) -> None: + self.successes = successes + self.expected_dataset_saving = expected_dataset_saving + self.calls = 0 + + def __call__( + self, + _bundle, + _output_root, + *, + seed: int, + num_envs: int, + dataset_saving: bool = False, + ): + values = self.successes[min(self.calls, len(self.successes) - 1)] + self.calls += 1 + assert len(values) == num_envs + assert dataset_saving is self.expected_dataset_saving + return { + "status": "succeeded" if all(values) else "failed", + "seed": seed, + "environments": [ + {"env_id": str(index), "success": success} + for index, success in enumerate(values) + ], + } + + +@pytest.mark.parametrize("existing", [False, True]) +@pytest.mark.parametrize("edit", [False, True]) +def test_parallel_workflow_supports_all_four_scene_inputs( + tmp_path: Path, + *, + existing: bool, + edit: bool, +) -> None: + candidates = _candidate_set() + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=_SceneBackend( + _selection(candidates), + input_kind="gym_project" if existing else "image", + ), + action_agent=_ActionAgent(), + coordinator=_Coordinator(["bound"]), + action_executor=_Executor( + [[True, False, False, False]], + expected_dataset_saving=True, + ), + ) + + result = workflow.run( + _request(tmp_path, existing=existing, edit=edit), + workflow_cfg=TaskEngineWorkflowCfg(), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + dataset_saving=True, + ) + + assert result.succeeded + + +def test_parallel_workflow_preserves_requested_robot_profile(tmp_path: Path) -> None: + candidates = _candidate_set() + coordinator = _Coordinator(["bound"]) + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_adapter=SimpleNamespace(robot_profile="ur10"), + scene_backend=_SceneBackend(_selection(candidates)), + action_agent=_ActionAgent(), + coordinator=coordinator, + action_executor=_Executor([[True]]), + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(), + execution_cfg=TaskEngineExecutionCfg(num_envs=1), + ) + + assert result.succeeded + assert isinstance(coordinator.sources[0], SceneSourceRef) + assert coordinator.sources[0].robot_profile == "ur10" + + +def test_parallel_workflow_accepts_one_success_and_publishes_all_graphs( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + input_barrier = Barrier(2) + work_barrier = Barrier(2) + scene = _SceneBackend( + _selection(candidates), + input_barrier=input_barrier, + materialize_barrier=work_barrier, + ) + coordinator = _Coordinator(["bound"]) + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates, input_barrier), + scene_backend=scene, + action_agent=_ActionAgent(work_barrier), + coordinator=coordinator, + action_executor=_Executor([[False, True, False, False]]), + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(), + planning_cfg=TaskEnginePlanningCfg( + candidate_count=3, + planning_mode="offline", + max_episodes=1, + max_episode_steps=4000, + ), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + base_seed=11, + run_id="20260820_072436", + ) + + assert result.succeeded + assert scene.seeds == [11] + assert result.final_bundle is not None + assert (result.final_bundle / "conservative_scene_graph.json").is_file() + assert (result.final_bundle / "seed_task_graph.json").is_file() + assert (result.final_bundle / "grounded_task_plan.json").is_file() + manifest = json.loads(result.manifest_path.read_text(encoding="utf-8")) + assert manifest["run_id"] == "20260820_072436" + assert manifest["configuration"]["planning"] == { + "candidate_count": 3, + "planning_mode": "offline", + "max_episodes": 1, + "max_episode_steps": 4000, + } + assert manifest["configuration"]["execution"]["dataset_saving"] is False + assert coordinator.kwargs[0]["max_episode_steps"] == 4000 + assert coordinator.kwargs[0]["final_inspection"]["scene_revision_id"] == "0" * 64 + assert ( + coordinator.kwargs[0]["unbound_action_plan"]["candidate_id"] == "candidate_01" + ) + assert manifest["attempts"][0]["action_attempts"][0]["status"] == "succeeded" + assert manifest["attempts"][0]["final_unbound_action_plan"]["candidate_id"] == ( + "candidate_01" + ) + assert manifest["attempts"][0]["unbound_transition"]["changed"] is False + state = json.loads(result.state_path.read_text(encoding="utf-8")) + succeeded = [ + event["stage"] for event in state["events"] if event["to"] == "succeeded" + ] + assert ( + succeeded.index("scene_finalization") + < succeeded.index("final_inspection") + < succeeded.index("final_binding") + ) + + +def test_prepare_only_publishes_bundle_without_action_execution( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + + def fail_execution(*_args, **_kwargs): + pytest.fail("prepare-only workflow must not execute Action Engine") + + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=_SceneBackend(_selection(candidates)), + action_agent=_ActionAgent(), + coordinator=_Coordinator(["bound"]), + action_executor=fail_execution, + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(), + execution_cfg=TaskEngineExecutionCfg(), + execute=False, + ) + + assert result.status == "prepared" + assert result.final_bundle is not None + assert result.final_bundle.is_dir() + + +def test_final_candidate_rebinding_updates_attempt_unbound_audit( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + final_candidate = deepcopy(candidates["candidates"][0]) + final_candidate["candidate_id"] = "candidate_02" + candidates["candidates"].append(final_candidate) + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=_SceneBackend(_selection(candidates)), + action_agent=_ActionAgent(), + coordinator=_RebindingCoordinator(final_candidate), + action_executor=_Executor([[True, False, False, False]]), + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + ) + + manifest = json.loads(result.manifest_path.read_text(encoding="utf-8")) + attempt = manifest["attempts"][0] + assert attempt["unbound_action_plan"]["candidate_id"] == "candidate_01" + assert attempt["final_unbound_action_plan"]["candidate_id"] == "candidate_02" + assert attempt["unbound_transition"]["changed"] is True + + +@pytest.mark.parametrize( + ("dataset_saving", "expects_filter"), + [(False, True), (True, False)], +) +def test_subprocess_executor_controls_dataset_saving_and_copies_trajectory( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + dataset_saving: bool, + expects_filter: bool, +) -> None: + bundle = tmp_path / "bundle" + bundle.mkdir() + trajectory = tmp_path / "trajectory-source" + trajectory.mkdir() + (trajectory / "episode.json").write_text("{}\n", encoding="utf-8") + captured = {} + provenance = build_execution_provenance(episode_seed=7) + + def fake_run(command, log_path): + captured["command"] = command + captured["log_path"] = Path(log_path) + Path(log_path).write_text("child output\n", encoding="utf-8") + report = ExecutionReport( + task_id="place_can", + plan_hash="0" * 64, + action_graph_hash="1" * 64, + status="succeeded", + run_id="run", + episode_id="0", + provenance=provenance, + environments=tuple( + { + "env_id": str(index), + "success": True, + "semantic_success": {}, + "action_count": 1, + "retry_count": 0, + "recovery_count": 0, + "revision_count": 0, + "failures": [], + } + for index in range(4) + ), + action_count=4, + record_dir=trajectory.as_posix(), + ) + (bundle / "execution_report.json").write_text( + json.dumps(report.as_mapping()), encoding="utf-8" + ) + return SimpleNamespace(returncode=0, stdout="ok", stderr="") + + monkeypatch.setattr( + "embodichain.gen_sim.task_engine.workflow._run_streaming_process", + fake_run, + ) + attempt = tmp_path / "attempt" + + report = SubprocessActionExecutor()( + bundle, + attempt, + seed=7, + num_envs=4, + dataset_saving=dataset_saving, + ) + + assert report["status"] == "succeeded" + assert captured["command"][1:5] == [ + "-u", + "-m", + "embodichain.gen_sim.task_engine._bundle_runner", + "--bundle", + ] + assert " prepare" not in " ".join(captured["command"]) + assert " workflow" not in " ".join(captured["command"]) + assert ("--filter_dataset_saving" in captured["command"]) is expects_filter + assert captured["log_path"] == attempt / "action.log" + assert (attempt / "action.log").read_text(encoding="utf-8") == "child output\n" + assert (attempt / "trajectory" / "episode.json").is_file() + process = json.loads((attempt / "process.json").read_text(encoding="utf-8")) + assert process["combined_log"] == "action.log" + assert process["stdout"] == "ok" + assert process["stderr"] == "" + + +def test_streaming_process_tees_combined_binary_output( + tmp_path: Path, + capfd: pytest.CaptureFixture[str], +) -> None: + log_path = tmp_path / "action.log" + script = ( + "import os; " + "os.write(1, b'stdout\\x00'); " + "os.write(2, b'stderr\\rprogress\\n'); " + "raise SystemExit(7)" + ) + + completed = _run_streaming_process( + [sys.executable, "-c", script], + log_path, + ) + + expected = b"stdout\x00stderr\rprogress\n" + assert completed.returncode == 7 + assert completed.stdout.encode("utf-8") == expected + assert completed.stderr == "" + assert log_path.read_bytes() == expected + terminal = capfd.readouterr().out + assert "stdout\x00" in terminal + assert "stderr\rprogress" in terminal + + +def test_scene_remediation_changes_seed_before_action_execution(tmp_path: Path) -> None: + candidates = _candidate_set() + scene = _SceneBackend(_selection(candidates)) + coordinator = _Coordinator(["infeasible", "bound"]) + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=scene, + action_agent=_ActionAgent(), + coordinator=coordinator, + action_executor=_Executor([[True, False, False, False]]), + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(max_scene_attempts=2), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + base_seed=20, + ) + + assert result.succeeded + assert scene.seeds == [20, 21] + assert coordinator.calls == 2 + manifest = json.loads(result.manifest_path.read_text(encoding="utf-8")) + assert [item["status"] for item in manifest["attempts"]] == [ + "preparation_failed", + "succeeded", + ] + + +def test_input_conflict_feasibility_does_not_regenerate_scene(tmp_path: Path) -> None: + candidates = _candidate_set() + scene = _SceneBackend(_selection(candidates)) + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=scene, + action_agent=_ActionAgent(), + coordinator=_Coordinator( + ["infeasible", "bound"], + infeasible_remediation="input_conflict", + ), + action_executor=_Executor([[True, False, False, False]]), + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(max_scene_attempts=2), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + ) + + assert result.failure_class == "input_conflict" + assert scene.seeds == [0] + + +def test_scene_service_retry_keeps_completed_unbound_plan(tmp_path: Path) -> None: + candidates = _candidate_set() + scene = _SceneBackend(_selection(candidates), materialize_failures=1) + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=scene, + action_agent=_ActionAgent(), + coordinator=_Coordinator(["bound"]), + action_executor=_Executor([[True, False, False, False]]), + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(max_scene_attempts=2), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + base_seed=30, + ) + + assert result.succeeded + assert scene.seeds == [30, 31] + manifest = json.loads(result.manifest_path.read_text(encoding="utf-8")) + assert manifest["attempts"][0]["unbound_action_plan"] is not None + + +def test_nonremediable_scene_error_does_not_change_scene_attempt_seed( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + scene = _InvalidSceneBackend(_selection(candidates)) + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=scene, + action_agent=_ActionAgent(), + coordinator=_Coordinator(["bound"]), + action_executor=_Executor([[True, False, False, False]]), + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(max_scene_attempts=3), + execution_cfg=TaskEngineExecutionCfg(), + base_seed=9, + ) + + assert not result.succeeded + assert scene.seeds == [9] + + +def test_execution_acceptance_requires_every_success_spec_term() -> None: + report = { + "environments": [ + { + "success": True, + "semantic_success": {"step_01": True, "step_02": False}, + }, + { + "success": True, + "semantic_success": {"step_01": True, "step_02": True}, + }, + ] + } + + assert _environment_successes( + report, + required_semantic_steps=("step_01", "step_02"), + ) == [False, True] + + +def test_unbound_failure_retains_completed_parallel_scene(tmp_path: Path) -> None: + candidates = _candidate_set() + scene = _SceneBackend(_selection(candidates)) + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=scene, + action_agent=_FailingActionAgent(), + coordinator=_Coordinator(["bound"]), + action_executor=_Executor([[True, False, False, False]]), + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + ) + + assert not result.succeeded + assert result.failure_class == "action_capability" + assert scene.seeds == [0] + manifest = json.loads(result.manifest_path.read_text(encoding="utf-8")) + assert manifest["attempts"][0]["scene_revision"] is not None + state = json.loads(result.state_path.read_text(encoding="utf-8")) + assert state["stages"]["scene_finalization"] == "succeeded" + assert state["stages"]["unbound_action"] == "failed" + + +def test_preparation_exception_is_published_as_audited_failure(tmp_path: Path) -> None: + candidates = _candidate_set() + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=_SceneBackend(_selection(candidates)), + action_agent=_ActionAgent(), + coordinator=_FailingCoordinator(), + action_executor=_Executor([[True, False, False, False]]), + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + ) + + assert not result.succeeded + assert result.failure_class == "preparation_error" + manifest = json.loads(result.manifest_path.read_text(encoding="utf-8")) + assert manifest["attempts"][0]["status"] == "preparation_error" + assert manifest["attempts"][0]["error"]["type"] == "RuntimeError" + + +def test_explicit_edit_may_materialize_initially_missing_reference( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + unresolved = CandidateSelection( + scene_manifest={}, + role_bindings={}, + binding_report={ + "status": "unsatisfied", + "selection_reason": "the can is not visible before the explicit edit", + "candidates": [{"candidate_id": "candidate_01", "status": "unsatisfied"}], + }, + selected_candidate=None, + candidate_bindings={"candidate_01": {}}, + ) + scene = _SceneBackend(unresolved, input_kind="gym_project") + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=scene, + action_agent=_ActionAgent(), + coordinator=_Coordinator(["bound"]), + action_executor=_Executor([[True, False, False, False]]), + ) + + result = workflow.run( + _request(tmp_path, existing=True, edit=True), + workflow_cfg=TaskEngineWorkflowCfg(), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + ) + + assert result.succeeded + provisional = json.loads( + (result.output_dir / "provisional_candidate.json").read_text(encoding="utf-8") + ) + assert provisional == { + "binding_status": "unsatisfied", + "candidate_id": "candidate_01", + "reason": "explicit_scene_edit_may_materialize_missing_reference", + } + + +def test_action_failure_retries_action_only_and_retains_attempts( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + scene = _SceneBackend(_selection(candidates)) + executor = _Executor([[False, False, False, False]]) + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=scene, + action_agent=_ActionAgent(), + coordinator=_Coordinator(["bound"]), + action_executor=executor, + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(max_action_attempts=3), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + ) + + assert not result.succeeded + assert result.failure_class == "action_execution" + assert scene.seeds == [0] + assert executor.calls == 3 + manifest = json.loads(result.manifest_path.read_text(encoding="utf-8")) + assert len(manifest["attempts"][0]["action_attempts"]) == 3 + + +def test_action_retry_stops_after_first_success(tmp_path: Path) -> None: + candidates = _candidate_set() + executor = _Executor( + [ + [False, False, False, False], + [True, True, True, True], + [True, True, True, True], + ] + ) + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=_SceneBackend(_selection(candidates)), + action_agent=_ActionAgent(), + coordinator=_Coordinator(["bound"]), + action_executor=executor, + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(max_action_attempts=3), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + ) + + assert result.succeeded + assert executor.calls == 2 + manifest = json.loads(result.manifest_path.read_text(encoding="utf-8")) + assert [item["status"] for item in manifest["attempts"][0]["action_attempts"]] == [ + "failed", + "succeeded", + ] + + +def test_existing_edit_binding_conflict_does_not_invent_scene_repair( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + scene = _SceneBackend(_selection(candidates), input_kind="gym_project") + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=scene, + action_agent=_ActionAgent(), + coordinator=_Coordinator(["unsatisfied"]), + action_executor=_Executor([[True, False, False, False]]), + ) + + result = workflow.run( + _request(tmp_path, existing=True, edit=True), + workflow_cfg=TaskEngineWorkflowCfg(max_scene_attempts=2), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + ) + + assert result.status == "input_conflict" + assert result.failure_class == "input_conflict" + assert scene.seeds == [0] + + +def test_image_binding_conflict_does_not_regenerate_scene(tmp_path: Path) -> None: + candidates = _candidate_set() + scene = _SceneBackend(_selection(candidates)) + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=scene, + action_agent=_ActionAgent(), + coordinator=_Coordinator(["unsatisfied"]), + action_executor=_Executor([[True, False, False, False]]), + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(max_scene_attempts=2), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + ) + + assert result.status == "input_conflict" + assert result.failure_class == "input_conflict" + assert scene.seeds == [0] diff --git a/tests/gen_sim/task_engine/test_run_directory.py b/tests/gen_sim/task_engine/test_run_directory.py new file mode 100644 index 000000000..e0305d59e --- /dev/null +++ b/tests/gen_sim/task_engine/test_run_directory.py @@ -0,0 +1,58 @@ +# ---------------------------------------------------------------------------- +# 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 datetime import datetime, timedelta, timezone +from pathlib import Path + +import pytest + +from embodichain.gen_sim.task_engine.run_directory import reserve_run_directory + +_NOW = datetime(2026, 8, 20, 7, 24, 36, tzinfo=timezone(timedelta(hours=8))) + + +def test_run_directory_uses_local_second_timestamp(tmp_path: Path) -> None: + root = tmp_path / "task2_2" + + with reserve_run_directory(root, now=_NOW) as allocation: + assert allocation.run_id == "20260820_072436" + assert allocation.path == root / "20260820_072436" + assert not allocation.path.exists() + allocation.path.mkdir() + + assert allocation.path.is_dir() + assert not (root / ".20260820_072436.reserve").exists() + + +def test_run_directory_adds_suffix_for_same_second_runs(tmp_path: Path) -> None: + root = tmp_path / "task2_2" + (root / "20260820_072436").mkdir(parents=True) + + with reserve_run_directory(root, now=_NOW) as first: + with reserve_run_directory(root, now=_NOW) as second: + assert first.run_id == "20260820_072436_01" + assert second.run_id == "20260820_072436_02" + + +def test_run_directory_rejects_naive_timestamp(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="timezone"): + with reserve_run_directory( + tmp_path, + now=datetime(2026, 8, 20, 7, 24, 36), + ): + pass diff --git a/tests/gen_sim/task_engine/test_scene_backend.py b/tests/gen_sim/task_engine/test_scene_backend.py new file mode 100644 index 000000000..cc6f85bd9 --- /dev/null +++ b/tests/gen_sim/task_engine/test_scene_backend.py @@ -0,0 +1,214 @@ +# ---------------------------------------------------------------------------- +# 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.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, +) +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.api import ( + SceneBlueprintPackage, + SceneMaterialization, +) +import embodichain.gen_sim.task_engine.scene_backend as scene_backend_module +from embodichain.gen_sim.task_engine.scene_backend import ( + SceneEngineBackend, + scene_blueprint_objects, +) +from embodichain.gen_sim.task_engine.workflow_contracts import TASK_RUN_REQUEST_SCHEMA + + +def _request(tmp_path: Path, project: Path, *, edit: str | None) -> dict: + return { + "schema_version": TASK_RUN_REQUEST_SCHEMA, + "task_id": "task", + "task_instruction": "Move the cup.", + "image_path": None, + "gym_project": project.as_posix(), + "scene_edit_prompt": edit, + "output_dir": (tmp_path / "run").as_posix(), + } + + +def _scene_export(tmp_path: Path) -> Path: + export = tmp_path / "project" / "scene_export" + assets = export / "mesh_assets" + assets.mkdir(parents=True) + (assets / "table.glb").write_bytes(b"glTF-table") + (assets / "cup.glb").write_bytes(b"glTF-cup") + (export / "scene_config.json").write_text( + json.dumps( + { + "format": "embodichain.scene-export/v1", + "scene_id": "scene", + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "mesh_assets/table.glb", + }, + } + ], + "rigid_object": [ + { + "uid": "cup", + "shape": { + "shape_type": "Mesh", + "fpath": "mesh_assets/cup.glb", + }, + } + ], + } + ), + encoding="utf-8", + ) + return export.parent + + +def test_blueprint_objects_preserve_semantics_without_geometry(tmp_path: Path) -> None: + scene = Scene( + objects=[ + SceneObject("table", "table", "table", "table", "A table."), + SceneObject("cup", "asset", "cup", "red cup", "A red cup."), + ] + ) + graph = SceneGraph( + nodes=[ + SceneGraphNode("table", None), + SceneGraphNode("cup", "table", "on", orientation_state="lying"), + ] + ) + package = SceneBlueprintPackage( + blueprint_id="blueprint", + image_path=tmp_path / "input.png", + output_root=tmp_path, + manifest_path=tmp_path / "scene_blueprint.json", + scene=scene, + scene_graph=graph, + ) + + objects = scene_blueprint_objects(package) + + cup = next(item for item in objects if item["uid"] == "cup") + assert cup["description"] == "A red cup." + assert cup["initial_state"] == {"orientation": "fallen"} + assert cup["affordances"] == [] + assert cup["init_pos"] == [0.0, 0.0, 0.0] + + +def test_existing_scene_edit_creates_revision_and_never_writes_source( + tmp_path: Path, + monkeypatch, +) -> None: + project = _scene_export(tmp_path) + source_config = project / "scene_export" / "scene_config.json" + source_value = json.loads(source_config.read_text(encoding="utf-8")) + articulation_path = project / "scene_export" / "cabinet.urdf" + articulation_path.write_text( + '\n', + encoding="utf-8", + ) + source_value["articulation"] = [ + { + "uid": "cabinet", + "name": "cabinet", + "description": "A fixed cabinet.", + "category": "cabinet", + "fpath": "cabinet.urdf", + } + ] + source_config.write_text(json.dumps(source_value), encoding="utf-8") + original = source_config.read_bytes() + prompts: list[str] = [] + + def fake_analyze_edit(*, output_root, edit_prompt): + prompts.append(edit_prompt) + return SimpleNamespace( + output_root=Path(output_root), + scene_edit_plan=SimpleNamespace( + to_dict=lambda: {"operations": [{"op": "move", "object_id": "cup"}]} + ), + ) + + def fake_materialize_edit(blueprint, *, seed=None): + assert seed == 7 + return SceneMaterialization( + scene=Scene(), + scene_graph=SceneGraph(nodes=[SceneGraphNode("table", None)]), + output_root=blueprint.output_root, + scene_config_path=blueprint.output_root + / "scene_export" + / "scene_config.json", + ) + + monkeypatch.setattr(scene_backend_module, "analyze_edit", fake_analyze_edit) + monkeypatch.setattr(scene_backend_module, "materialize_edit", fake_materialize_edit) + backend = SceneEngineBackend() + request = _request(tmp_path, project, edit="Move the cup left.") + analysis = backend.analyze(request, tmp_path / "analysis") + + revision = backend.materialize( + analysis, + request, + tmp_path / "revision", + seed=7, + ) + + assert prompts == ["Move the cup left."] + assert revision.source != source_config + assert revision.source.is_file() + assert len(revision.revision_id) == 64 + assert revision.edit_plan == {"operations": [{"op": "move", "object_id": "cup"}]} + assert source_config.read_bytes() == original + revision_config = json.loads(revision.source.read_text(encoding="utf-8")) + assert revision_config["articulation"][0]["uid"] == "cabinet" + audit = json.loads( + (tmp_path / "revision" / "scene_revision_attempt.json").read_text( + encoding="utf-8" + ) + ) + assert audit["seed"] == 7 + assert audit["revision_id"] == revision.revision_id + assert audit["edit_plan"] == revision.edit_plan + + +def test_final_inspection_rejects_scene_changed_after_revision(tmp_path: Path) -> None: + project = _scene_export(tmp_path) + backend = SceneEngineBackend() + request = _request(tmp_path, project, edit=None) + revision = backend.materialize( + backend.analyze(request, tmp_path / "analysis"), + request, + tmp_path / "unused", + seed=0, + ) + config_path = project / "scene_export" / "scene_config.json" + config = json.loads(config_path.read_text(encoding="utf-8")) + config["rigid_object"][0]["init_pos"] = [0.25, 0.0, 0.0] + config_path.write_text(json.dumps(config), encoding="utf-8") + + with pytest.raises(RuntimeError, match="changed before geometry inspection"): + backend.inspect(revision, tmp_path / "inspection.json") diff --git a/tests/gen_sim/task_engine/test_workflow.py b/tests/gen_sim/task_engine/test_workflow.py new file mode 100644 index 000000000..4a1564606 --- /dev/null +++ b/tests/gen_sim/task_engine/test_workflow.py @@ -0,0 +1,337 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from embodichain.gen_sim.task_engine.config import ( + TaskEngineExecutionCfg, + TaskEnginePlanningCfg, + TaskEngineWorkflowCfg, + load_task_engine_config, +) +from embodichain.gen_sim.task_engine.state_machine import ( + StageStatus, + WorkflowStage, + complete_stage, + fail_stage, + initial_state, + replay_events, + start_stage, + skip_stage, +) +from embodichain.gen_sim.task_engine.workflow_contracts import ( + TASK_RUN_REQUEST_SCHEMA, + scene_input_kind, + validate_scene_history_root, + validate_task_run_request, +) + + +def _request(tmp_path: Path, *, image: bool, edit: bool) -> dict[str, object]: + return { + "schema_version": TASK_RUN_REQUEST_SCHEMA, + "task_id": "pick-cup", + "task_instruction": "Pick up the red cup.", + "image_path": str(tmp_path / "input.png") if image else None, + "gym_project": None if image else str(tmp_path / "gym_project"), + "scene_edit_prompt": "Add a tray." if edit else None, + "output_dir": str(tmp_path / "output"), + } + + +@pytest.mark.parametrize("image", [False, True]) +@pytest.mark.parametrize("edit", [False, True]) +def test_run_request_accepts_all_four_input_combinations( + tmp_path: Path, + image: bool, + edit: bool, +) -> None: + request = validate_task_run_request(_request(tmp_path, image=image, edit=edit)) + assert scene_input_kind(request) == ("image" if image else "gym_project") + assert request["scene_edit_prompt"] == ("Add a tray." if edit else None) + + +def test_run_request_rejects_two_scene_inputs(tmp_path: Path) -> None: + request = _request(tmp_path, image=True, edit=False) + request["gym_project"] = str(tmp_path / "gym_project") + with pytest.raises(ValueError, match="exactly one"): + validate_task_run_request(request) + + +def test_run_request_rejects_scene_generation_prompt(tmp_path: Path) -> None: + request = _request(tmp_path, image=True, edit=False) + request["scene_generation_prompt"] = "Make a kitchen." + with pytest.raises(ValueError, match="fields differ"): + validate_task_run_request(request) + + +def test_run_request_rejects_output_inside_gym_project(tmp_path: Path) -> None: + request = _request(tmp_path, image=False, edit=False) + request["output_dir"] = str(tmp_path / "gym_project" / "task_run") + + with pytest.raises(ValueError, match="must not overlap"): + validate_task_run_request(request) + + +def test_run_request_rejects_output_containing_explicit_gym_config( + tmp_path: Path, +) -> None: + project = tmp_path / "gym_project" + project.mkdir() + config_path = project / "gym_config.json" + config_path.write_text("{}", encoding="utf-8") + request = _request(tmp_path, image=False, edit=False) + request["gym_project"] = str(config_path) + request["output_dir"] = str(project) + + with pytest.raises(ValueError, match="must not overlap"): + validate_task_run_request(request) + + +def test_scene_history_root_allows_a_source_from_a_prior_run( + tmp_path: Path, +) -> None: + history = tmp_path / "task_history" + source = history / "20260820_105939" / "attempts" / "scene_export" + source.mkdir(parents=True) + + validate_scene_history_root(source, history) + + request = _request(tmp_path, image=False, edit=False) + request["gym_project"] = str(source) + request["output_dir"] = str(history / "20260820_130000") + assert validate_task_run_request(request)["gym_project"] == source.as_posix() + + +@pytest.mark.parametrize("relative_output", [".", "new_runs", "new_runs/task"]) +def test_scene_history_root_rejects_writes_into_source_project( + tmp_path: Path, + relative_output: str, +) -> None: + source = tmp_path / "scene_export" + source.mkdir() + output_root = source / relative_output + + with pytest.raises(ValueError, match="read-only source"): + validate_scene_history_root(source, output_root) + + +def test_scene_history_root_resolves_symlinks_before_comparison( + tmp_path: Path, +) -> None: + source = tmp_path / "scene_export" + source.mkdir() + source_link = tmp_path / "scene_link" + source_link.symlink_to(source, target_is_directory=True) + + with pytest.raises(ValueError, match="read-only source"): + validate_scene_history_root(source_link, source / "new_runs") + + +def test_scene_history_root_protects_explicit_config_parent(tmp_path: Path) -> None: + source = tmp_path / "scene_export" + source.mkdir() + config = source / "scene_config.json" + config.write_text("{}\n", encoding="utf-8") + + with pytest.raises(ValueError, match="read-only source"): + validate_scene_history_root(config, source) + + +def test_task_and_scene_stages_can_run_concurrently(tmp_path: Path) -> None: + state = initial_state(_request(tmp_path, image=True, edit=False)) + state = start_stage(state, WorkflowStage.TASK_CANDIDATES) + state = start_stage(state, WorkflowStage.SCENE_PREPARATION) + assert state.stages[WorkflowStage.TASK_CANDIDATES] == StageStatus.RUNNING + assert state.stages[WorkflowStage.SCENE_PREPARATION] == StageStatus.RUNNING + assert state.stages[WorkflowStage.SCENE_EDIT] == StageStatus.SKIPPED + + +def test_candidate_selection_waits_for_both_branches(tmp_path: Path) -> None: + state = initial_state(_request(tmp_path, image=False, edit=False)) + state = start_stage(state, WorkflowStage.TASK_CANDIDATES) + state = complete_stage(state, WorkflowStage.TASK_CANDIDATES) + with pytest.raises(ValueError, match="incomplete dependencies"): + start_stage(state, WorkflowStage.CANDIDATE_SELECTION) + + +def test_unbound_action_can_run_while_user_scene_edit_is_running( + tmp_path: Path, +) -> None: + state = initial_state(_request(tmp_path, image=True, edit=True)) + for stage in (WorkflowStage.TASK_CANDIDATES, WorkflowStage.SCENE_PREPARATION): + state = start_stage(state, stage) + state = complete_stage(state, stage) + state = start_stage(state, WorkflowStage.CANDIDATE_SELECTION) + state = complete_stage(state, WorkflowStage.CANDIDATE_SELECTION) + state = start_stage(state, WorkflowStage.SCENE_EDIT) + state = start_stage(state, WorkflowStage.UNBOUND_ACTION) + + assert state.stages[WorkflowStage.SCENE_EDIT] == StageStatus.RUNNING + assert state.stages[WorkflowStage.UNBOUND_ACTION] == StageStatus.RUNNING + with pytest.raises(ValueError, match="incomplete dependencies"): + start_stage(state, WorkflowStage.SCENE_FINALIZATION) + + +def test_only_scene_edit_can_be_skipped(tmp_path: Path) -> None: + state = initial_state(_request(tmp_path, image=True, edit=True)) + + with pytest.raises(ValueError, match="Only the optional scene_edit stage"): + skip_stage(state, WorkflowStage.FINAL_BINDING) + + +def test_state_events_replay_to_the_same_snapshot(tmp_path: Path) -> None: + request = _request(tmp_path, image=True, edit=False) + state = initial_state(request) + state = start_stage(state, WorkflowStage.TASK_CANDIDATES) + state = start_stage(state, WorkflowStage.SCENE_PREPARATION) + state = complete_stage(state, WorkflowStage.TASK_CANDIDATES) + state = complete_stage(state, WorkflowStage.SCENE_PREPARATION) + state = start_stage(state, WorkflowStage.CANDIDATE_SELECTION) + state = complete_stage(state, WorkflowStage.CANDIDATE_SELECTION) + + replayed = replay_events(request, state.events) + + assert replayed.to_dict() == state.to_dict() + + +def test_state_replay_rejects_tampered_transition(tmp_path: Path) -> None: + request = _request(tmp_path, image=True, edit=False) + state = initial_state(request) + events = [dict(event) for event in state.events] + events[-1]["stage"] = WorkflowStage.FINAL_BINDING.value + + with pytest.raises(ValueError, match="event does not match"): + replay_events(request, events) + + +def test_state_replay_preserves_failure_reason(tmp_path: Path) -> None: + request = _request(tmp_path, image=True, edit=False) + state = initial_state(request) + state = start_stage(state, WorkflowStage.TASK_CANDIDATES) + state = fail_stage(state, WorkflowStage.TASK_CANDIDATES, reason="model timeout") + + replayed = replay_events(request, state.events) + + assert replayed.terminal + assert replayed.to_dict() == state.to_dict() + + +def test_later_retry_can_fail_a_previously_successful_stage(tmp_path: Path) -> None: + request = _request(tmp_path, image=True, edit=False) + state = initial_state(request) + state = start_stage(state, WorkflowStage.SCENE_PREPARATION) + state = complete_stage(state, WorkflowStage.SCENE_PREPARATION) + state = fail_stage( + state, + WorkflowStage.SCENE_PREPARATION, + reason="later scene attempt failed", + ) + + assert state.terminal + assert replay_events(request, state.events).to_dict() == state.to_dict() + + +def test_state_snapshot_mappings_are_immutable(tmp_path: Path) -> None: + state = initial_state(_request(tmp_path, image=True, edit=False)) + + with pytest.raises(TypeError): + state.stages[WorkflowStage.TASK_CANDIDATES] = StageStatus.SUCCEEDED + with pytest.raises(TypeError): + state.request["task_id"] = "changed" + with pytest.raises(TypeError): + state.events[0]["to"] = StageStatus.FAILED.value + + +def test_workflow_configuration_rejects_non_positive_limits() -> None: + with pytest.raises(ValueError, match="max_scene_attempts"): + TaskEngineWorkflowCfg(max_scene_attempts=0) + + +def test_packaged_workflow_configuration_uses_recovery_defaults() -> None: + workflow, planning, execution = load_task_engine_config() + + assert workflow.max_scene_attempts == 2 + assert workflow.max_action_attempts == 3 + assert planning.candidate_count == 3 + assert planning.planning_mode == "offline" + assert planning.max_episodes == 1 + assert planning.max_episode_steps == 4000 + assert execution.num_envs == 1 + assert execution.required_successes == 1 + + +def test_workflow_configuration_can_be_tuned_from_yaml(tmp_path: Path) -> None: + config = tmp_path / "task_engine.yaml" + config.write_text( + """\ +schema_version: embodichain.task-engine-defaults/v1 +workflow: + max_parallel_workers: 3 + max_scene_attempts: 4 + max_action_attempts: 5 +planning: + candidate_count: 7 + planning_mode: offline + max_episodes: 2 + max_episode_steps: 5000 +execution: + num_envs: 6 + success_policy: at_least + min_successful_envs: 2 +""", + encoding="utf-8", + ) + + workflow, planning, execution = load_task_engine_config(config) + + assert workflow.max_parallel_workers == 3 + assert workflow.max_scene_attempts == 4 + assert workflow.max_action_attempts == 5 + assert planning.candidate_count == 7 + assert planning.max_episodes == 2 + assert planning.max_episode_steps == 5000 + assert execution.num_envs == 6 + assert execution.required_successes == 2 + + +def test_execution_configuration_validates_success_policy() -> None: + assert TaskEngineExecutionCfg().num_envs == 1 + assert ( + TaskEngineExecutionCfg( + num_envs=4, + success_policy="at_least", + min_successful_envs=2, + ).required_successes + == 2 + ) + with pytest.raises(ValueError, match="success_policy=all"): + TaskEngineExecutionCfg( + num_envs=4, + success_policy="all", + min_successful_envs=1, + ) + + +def test_planning_configuration_rejects_invalid_values() -> None: + with pytest.raises(ValueError, match="candidate_count"): + TaskEnginePlanningCfg(candidate_count=0) + with pytest.raises(ValueError, match="planning_mode"): + TaskEnginePlanningCfg(planning_mode="unsupported") diff --git a/tests/sim/atomic_actions/test_affordance.py b/tests/sim/atomic_actions/test_affordance.py index 0e123eefe..bee13fff2 100644 --- a/tests/sim/atomic_actions/test_affordance.py +++ b/tests/sim/atomic_actions/test_affordance.py @@ -120,6 +120,29 @@ def test_best_grasp_poses_casts_approach_direction_to_generator_device(self): assert approach_direction.dtype == torch.float32 assert approach_direction.device == generator.device + def test_valid_grasp_poses_applies_object_aware_cost_callback(self): + aff = AntipodalAffordance() + generator = Mock() + generator.device = torch.device("cpu") + object_pose = torch.eye(4) + object_pose[0, 3] = 0.25 + grasp_poses = torch.eye(4).repeat(2, 1, 1) + costs = torch.tensor([0.2, 0.4]) + + def get_valid_grasp_poses(**kwargs): + adjusted = kwargs["pose_cost_fn"](grasp_poses, costs) + return True, grasp_poses, 0.0, adjusted + + generator.get_valid_grasp_poses.side_effect = get_valid_grasp_poses + aff._generator = generator + + results = aff.get_valid_grasp_poses( + object_pose.unsqueeze(0), + grasp_cost_fn=lambda obj, _grasps, current: current + obj[0, 3], + ) + + assert torch.allclose(results[0][1], torch.tensor([0.45, 0.65])) + class TestTwistAffordance: def test_requires_explicit_grasp_position_and_axis_origin(self): diff --git a/tests/sim/atomic_actions/test_primitives_helpers.py b/tests/sim/atomic_actions/test_primitives_helpers.py index 2ff754fe0..f5f297738 100644 --- a/tests/sim/atomic_actions/test_primitives_helpers.py +++ b/tests/sim/atomic_actions/test_primitives_helpers.py @@ -27,6 +27,10 @@ resolve_batched_pose, resolve_object_target, ) +from embodichain.lab.sim.atomic_actions.primitives.pick_up import ( + PickUpOptions, + _upright_yaw_pose_variants, +) BATCH_SIZE = 2 ROBOT_DOF = 6 @@ -83,3 +87,19 @@ def test_resolve_object_target_uses_custom_name_in_shape_error() -> None: device=torch.device("cpu"), name="placing_object_target_pose", ) + + +def test_upright_yaw_pose_variants_preserve_translation() -> None: + pose = torch.eye(4).repeat(2, 1, 1) + pose[:, :3, 3] = torch.tensor([[0.2, -0.1, 0.8], [-0.3, 0.4, 0.7]]) + + variants = _upright_yaw_pose_variants(pose, 4) + + assert variants.shape == (2, 4, 4, 4) + assert torch.allclose(variants[:, :, :3, 3], pose[:, None, :3, 3].expand(-1, 4, -1)) + assert torch.allclose(variants[:, 0], pose) + + +def test_upright_yaw_samples_must_be_positive() -> None: + with pytest.raises(ValueError, match="upright_yaw_samples"): + PickUpOptions(upright_yaw_samples=0) diff --git a/tests/test_main.py b/tests/test_main.py index 5466e7c11..0777d549c 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -32,6 +32,7 @@ "preview_lerobot_data", "run-env", "scene-engine", + "task-engine", "simready", "train-rl", "preview-scene",