From ef4449c3d0667c7e8cdb9ede68db24a398958d4f Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Fri, 21 Aug 2026 00:21:19 +0800 Subject: [PATCH] feat(action-engine): add SeedGraph planning and bundle generation --- .../gen_sim/action_engine/cli/__init__.py | 21 + .../cli/generate_action_agent_config.py | 240 +++ .../action_engine/compiler/__init__.py | 33 + .../gen_sim/action_engine/compiler/core.py | 590 ++++++ .../gen_sim/action_engine/compiler/v2.py | 424 ++++ .../action_engine/generation/__init__.py | 33 + .../action_engine/generation/artifacts.py | 141 ++ .../action_engine/generation/assets.py | 158 ++ .../action_engine/generation/generator.py | 866 ++++++++ .../action_engine/generation/source_scene.py | 600 ++++++ .../generation/templates/default_lights.json | 3 + .../generation/templates/default_sensors.json | 14 + .../templates/dual_franka_robot.json | 185 ++ .../generation/templates/dual_ur_robot.json | 126 ++ .../generation/templates/vlm_sensors.json | 58 + .../action_engine/graph_visualization.py | 938 ++++++++ .../action_engine/planning/__init__.py | 62 + .../gen_sim/action_engine/planning/dual.py | 218 ++ .../gen_sim/action_engine/planning/linker.py | 1018 +++++++++ .../gen_sim/action_engine/planning/online.py | 371 ++++ .../gen_sim/action_engine/planning/planner.py | 821 +++++++ .../action_engine/planning/selection.py | 364 ++++ .../planning/task_planner_prompt.py | 159 ++ .../gen_sim/action_engine/planning/vision.py | 808 +++++++ .../gen_sim/action_engine/runtime/__init__.py | 32 + .../gen_sim/action_engine/runtime/loader.py | 292 +++ .../gen_sim/action_engine/runtime/models.py | 314 +++ .../gen_sim/action_engine/runtime/state.py | 84 + .../gen_sim/action_engine/tasks/recipes.py | 997 +++++++++ tests/__init__.py | 21 + tests/gen_sim/__init__.py | 21 + tests/gen_sim/action_engine/__init__.py | 19 + .../action_engine/compiler/__init__.py | 19 + .../action_engine/compiler/test_compiler.py | 572 +++++ .../gen_sim/action_engine/compiler/test_v2.py | 172 ++ .../generation/test_generation.py | 1533 ++++++++++++++ .../action_engine/planning/__init__.py | 21 + .../action_engine/planning/test_online_v2.py | 480 +++++ .../action_engine/planning/test_planner.py | 730 +++++++ .../action_engine/tasks/test_factory.py | 486 +++++ .../tasks/test_interpretation.py | 1885 +++++++++++++++++ .../tasks/test_language_decoupling.py | 420 ++++ .../action_engine/test_graph_visualization.py | 362 ++++ .../gen_sim/action_engine/test_orientation.py | 145 ++ 44 files changed, 16856 insertions(+) create mode 100644 embodichain/gen_sim/action_engine/cli/__init__.py create mode 100644 embodichain/gen_sim/action_engine/cli/generate_action_agent_config.py create mode 100644 embodichain/gen_sim/action_engine/compiler/__init__.py create mode 100644 embodichain/gen_sim/action_engine/compiler/core.py create mode 100644 embodichain/gen_sim/action_engine/compiler/v2.py create mode 100644 embodichain/gen_sim/action_engine/generation/__init__.py create mode 100644 embodichain/gen_sim/action_engine/generation/artifacts.py create mode 100644 embodichain/gen_sim/action_engine/generation/assets.py create mode 100644 embodichain/gen_sim/action_engine/generation/generator.py create mode 100644 embodichain/gen_sim/action_engine/generation/source_scene.py create mode 100644 embodichain/gen_sim/action_engine/generation/templates/default_lights.json create mode 100644 embodichain/gen_sim/action_engine/generation/templates/default_sensors.json create mode 100644 embodichain/gen_sim/action_engine/generation/templates/dual_franka_robot.json create mode 100644 embodichain/gen_sim/action_engine/generation/templates/dual_ur_robot.json create mode 100644 embodichain/gen_sim/action_engine/generation/templates/vlm_sensors.json create mode 100644 embodichain/gen_sim/action_engine/graph_visualization.py create mode 100644 embodichain/gen_sim/action_engine/planning/__init__.py create mode 100644 embodichain/gen_sim/action_engine/planning/dual.py create mode 100644 embodichain/gen_sim/action_engine/planning/linker.py create mode 100644 embodichain/gen_sim/action_engine/planning/online.py create mode 100644 embodichain/gen_sim/action_engine/planning/planner.py create mode 100644 embodichain/gen_sim/action_engine/planning/selection.py create mode 100644 embodichain/gen_sim/action_engine/planning/task_planner_prompt.py create mode 100644 embodichain/gen_sim/action_engine/planning/vision.py create mode 100644 embodichain/gen_sim/action_engine/runtime/__init__.py create mode 100644 embodichain/gen_sim/action_engine/runtime/loader.py create mode 100644 embodichain/gen_sim/action_engine/runtime/models.py create mode 100644 embodichain/gen_sim/action_engine/runtime/state.py create mode 100644 embodichain/gen_sim/action_engine/tasks/recipes.py create mode 100644 tests/__init__.py create mode 100644 tests/gen_sim/__init__.py create mode 100644 tests/gen_sim/action_engine/__init__.py create mode 100644 tests/gen_sim/action_engine/compiler/__init__.py create mode 100644 tests/gen_sim/action_engine/compiler/test_compiler.py create mode 100644 tests/gen_sim/action_engine/compiler/test_v2.py create mode 100644 tests/gen_sim/action_engine/generation/test_generation.py create mode 100644 tests/gen_sim/action_engine/planning/__init__.py create mode 100644 tests/gen_sim/action_engine/planning/test_online_v2.py create mode 100644 tests/gen_sim/action_engine/planning/test_planner.py create mode 100644 tests/gen_sim/action_engine/tasks/test_factory.py create mode 100644 tests/gen_sim/action_engine/tasks/test_interpretation.py create mode 100644 tests/gen_sim/action_engine/tasks/test_language_decoupling.py create mode 100644 tests/gen_sim/action_engine/test_graph_visualization.py create mode 100644 tests/gen_sim/action_engine/test_orientation.py 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/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/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/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/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/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/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/runtime/__init__.py b/embodichain/gen_sim/action_engine/runtime/__init__.py new file mode 100644 index 000000000..a91770d6e --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/__init__.py @@ -0,0 +1,32 @@ +# ---------------------------------------------------------------------------- +# 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 .loader import load_agent_execution_program, load_execution_program +from .models import ExecutionProgram, ExecutionReport, ExecutionResult +from .state import ExecutionState + +__all__ = [ + "ExecutionProgram", + "ExecutionState", + "ExecutionResult", + "ExecutionReport", + "load_agent_execution_program", + "load_execution_program", +] 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/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/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/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/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/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/generation/test_generation.py b/tests/gen_sim/action_engine/generation/test_generation.py new file mode 100644 index 000000000..4a520357d --- /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"] == 1 + 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"] == 1 + + +@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_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/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_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_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_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"